###Model selection in neural networks is hard.

Before digging into the Bayesian perspective, we’re going to give some context for the model selection problem in deep learning. In most machine learning applications, our goal is to produce some function approximator that fits a target function on some unknown data-generating distribution, from which we have samples in the form of a training set. Typically, some points from this training set is hidden from the model during training, and an estimate of the model’s expected risk (its average error on the data generating distribution) is obtained by computing the validation loss on this held-out data. The model with the lowest validation loss is assumed to also have the lowest expected risk. While the validation loss is great as an estimator of empirical risk, it’s not ideal for hyperparameter tuning because a) it’s expensive to optimize and b) it doesn’t explain why a model will generalize well.

Generalization bounds can be used to address both of these issues. A generalization bound gives a high-probability guarantee that the test set error will not exceed a certain value, and can typically be decomposed into an empirical risk estimate (i.e. the training loss) and a model complexity penalty. The model complexity penalty can be interpreted as giving a causal hypothesis for what makes a model generalize well, addressing point a) above. When both of these terms are differentiable, the bound can be optimized directly, addressing point b) above. While generalization bounds can get non-vacuous values on neural networks, they’re generally not tight enough to be useful for model selection. We conjecture that part of the reason these bounds might struggle so much in the deep learning regime is that the model complexity term used typically only depends on the final value of the model’s parameters after training, rather than taking into account properties of the training scheme. Some recent work has developed PAC-Bayesian generalization bounds based on a similar idea, but there’s likely much more that could be explored.

The idea that models which train faster should generalize better is not a new one: generalization bounds based on the stability of gradient descent have existed for years. These bounds explicitly depend on the number of training steps taken to reach a minimum. However, they also require some assumptions that aren’t obvious to show in the generic deep learning setting. There have also been some more oblique connections in the literature. For example, Arora et al. propose a generalization bound with a data complexity term that can be related to an upper bound on the convergence rate of gradient descent on a convex function.

Instead of trying to directly solve the model selection problem for neural networks, we’re going to take a digression into an easier setting, Bayesian models, and develop some theoretically grounded intuition on the model selection there. Don’t forget about this section though – we’ll come back to it in a bit.

###Bayesian Inference in 2 Minutes

Being a Bayesian is all about using Bayes’ rule to update prior beliefs based on evidence. For example, given a model \(M\) and some data \(D\), our belief that the model generated the data can be expressed in terms of the likelihood the model assigns the data, and our prior probability of the data and the model.

\[P(M|D) = \frac{P(D|M)P(D)}{P(M)}\]

To quantify terms like ‘evidence’ and ‘prior belief’, we translate some standard quantities from the risk minimization framework into probabilistic analogues. Instead of optimizing the parameters of a function using its loss on the data, we have a model \(\mathcal{M}\) which defines a probability distribution over the data \(P(\mathcal{D}|\theta)\) given some parameters \(\theta \in \Theta\), along with a probability distribution over the parameters \(P(\theta)\).

One quantity of interest is to Bayesians is the posterior distribution over parameters \(P(\theta|\mathcal{D})\): finding this distribution is known as parameter inference. While this computation is straightforward in some models like Gaussian Processes, it’s often not tractable to compute exactly for complex models like Bayesian neural networks. However, it’s still possible to get samples from this distribution via techniques such as ensemble methods, where we randomly draw a number of parameter samples from the prior and then optimize each parameter sample independently to get a sample from the posterior.

The quantity that we will be particularly interested in is the marginal likelihood \(P(\mathcal{D})\) (we’ll abbreviate this as ML from here on out). This tells us how likely the data is under the model, and can be written as follows.

\[P(\mathcal{D}) = \int_{\theta} P(\mathcal{D}|\theta) P(\theta) d\theta \]

If we have a collection of models, picking the one with the highest marginal likelihood is what MacKay calls Type II maximum likelihood estimation and is a popular approach to Bayesian model selection. Maximizing the marginal likelihood is less prone to overfitting than maximum likelihood estimation at the parameter level, as the marginal likelihood can be viewed as having a built-in model complexity term that helps to prevent overfitting.

In linear models, we can get the posterior samples that these estimators need by the ensemble methods described previously. Specifically: in Bayesian linear regression with noiseless Gaussian likelihoods, we can sample parameters from the prior and then optimize the \(\ell_2\) regression loss. The minimum of this optimization problem will correspond to a sample from the posterior of the model conditioned on the data used for optimization. This gives us a way of translating Bayesian updating to gradient optimization with ensembles, which we can adapt to models with noisy likelihoods (assuming we know the ground-truth targets) by adding noise to the targets and optimizing a penalized regression loss which encourages parameters to stay close to their initialized values.

Optimizing a linear regression objective lets us sample from the posterior.

• randomly initialize some collection of parameters \((\theta_0^k)\)
• for each \(D_i \in D\), run GD on \(D_{<i}\) to convergence for each parameter \(\theta^k_{i-1}\)
• take this parameter \(\theta^k_i \sim p(\theta|D_{<i})\) as a posterior sample to estimate \(\log P(D_i|D_{<i})\)
• and then add this estimate to a running sum.

In the worst case, this procedure can be enormously expensive, requiring us to run gradient descent to convergence \(n\) times. In practice we found that this wasn’t as expensive as it might as first seem, because if the model fits the data well then the optimal parameters for \(D_{\le i}\) should be similar to those for \(D_{<i}\), so it won’t take long for gradient descent to reach them.

We can use gradient descent to estimate \(L(D)\), but it also turns out that \(L(D)\) can tell us something about why gradient descent might prefer models with a high (lower bound on their) marginal likelihood. We’ll focus on the setting where we train linear combinations of models while these models are themselves performing Bayesian updating, in a loose sense mimicking what happens when we optimize a hierarchical model like a neural network.

We’ll assume some collection of Bayesian linear regression models \(M_1, \dots, M_k\) and a regression data set of the form \(D = (X_i, Y_i)_{i=1}^N\); each of these models \(M_j\) will output predictions \(\hat{Y^j_i}\) sampled from its posterior \(P(Y_i|D_{<i}, X_i)\). On top of these models we are training a linear predictor \(w\) to minimize the objective \(\|w^\top(\widehat{Y}_i) - Y_i\|^2\). We’ll assume that these models make predictions of a similar scale (i.e. there isn’t a model that consistently outputs \(0.0000001 \times Y_i\) whose optimal weight is then 10000000, while another model which consistently outputs \(Y_i\) would get weight 1), and that they don’t have complementary errors (i.e. it’s not the case that model 1 consistently underestimates the target \(Y\) while model 2 consistently overestimates it in such a way that the optimal solution is to assign a high weight to both so that they cancel each other out). Under these two assumptions, we can show that under the following optimization problem

\[ w^* = \text{arg min}_w \sum \mathbb{E}_{\hat{Y}_i \sim P(\cdot | D_{<i})} \bigg [ \|w^\top \widehat{Y}_i - Y_i\|^2 \bigg ]\]

the optimal solution \(w^*\) will assign the highest weight to the model with the highest \(\mathcal{L}(D)\). In other words, running this concurrent optimization procedure and then doing magnitude pruning on the linear combination is going to be equivalent (in expectation) to picking the model with the highest lower bound on its marginal likelihood.

We compared all three estimators on a feature selection task and found that they provided a similar model ranking to that given by the log ML. This is a nice existence proof that the lower bounds give sensible model rankings, and we provide a few more example problems in the paper. The results shown here are for the task of feature selection – i.e., how many (and which) features should we use in our model? We simulate a dataset inspired by of the form \((\textbf{X}, \textbf{y})\), where \(x_i = (y_i + \epsilon_1, y_i + \dots, y_i + \epsilon_{15}, \epsilon_{16}, \dots, \epsilon_{30})\), and consider a set of models \(\{\mathcal{M}_k\}\) with feature embeddings \(\phi_k(x_i) = x_i[1, \dots, k]\). The optimal model selects the first 15 features. This example is typical of a range of model selection tasks that we looked at where the hyperparameter formed a nice “Occam’s Hill”. The estimators tended to agree around the optimal hyperparameter/model, with the estimator for \(L(D)\) penalizing models whose posteriors were too peaked more than the log ML did in models that didn’t fit the data well.

We find similar results when we compare the weight assigned to a Bayesian model in a linear model combination trained online as the Bayesian model is being iteratively updated, as in the procedure described in the previous section. We compare against two other linear model combination approaches: one where we train the linear model combination on samples from the already-fitted posterior \(P(D_i|D)\), and another where we fit the linear model combination on samples from the prior \(P(D_i|\emptyset)\). We again see that \(L(D)\) and \(\log P(D)\) agree close to the optimum. We also see a similar ranking of models between \(L(D)\) and the linear model combination weights given by concurrent sampling – the ‘posterior sampling’ and ‘prior sampling’ approaches do not exhibit this trend, as expected.

How quickly a model’s training loss decreases while training depends on a number of factors including the optimizer, hyperparameters, model architecture, and dataset. However, all else being equal, a model’s loss will decrease faster if the gradient update based on minibatch \(i\) of the dataset also decreases the loss on minibatch \(i+1\). This relates to the notion of stiffness discussed by Fort et al., who argue that models for which gradients from minibatches all look similar should generalize better than those whose minibatch updates are orthogonal. Summing over the minibatch training losses then gives us a straightforward way to measure how well the gradients are generalizing. Models which attain a low loss with fewer training steps will have a lower sum over training losses (SOTL).

To illustrate how SOTL might measure within-minibatch gradient generalization, consider a first-order Taylor approximation of the training loss. The change in loss for minibatch \(i\) after taking a gradient step computed on that minibatch looks as follows

\[\begin{align} \ell(D_i ; \theta_{t+1}) - \alpha g_{t}) &\approx \ell(D_i; \theta_{t}) - \alpha \nabla_\theta \ell (D_i; \theta_{t}) ^\top g_{t}\\ & = \ell(D_i; \theta_{t_i}) - \alpha \|g_{t}\|^2 \end{align}\]

Meanwhile, the change due to gradients tep \(g_{t_i}\) on a disjoint minibatch \(D_j\) can be written \[\begin{align} \ell(D_j; \theta_{t+1}) = \ell(D_j; \theta_t - \alpha g_{t}) &\approx \ell(D_j; \theta_t) - \nabla_\theta \ell (D_j; \theta)^\top g_t \\ &= \ell(D_j; \theta_t) - \nabla_\theta \ell (D_j; \theta)^\top \nabla_\theta \ell(D_i; \theta) \end{align}\]

This dot product term measures how correlated gradients based on disjoint minibatches are. For the first epoch of training, this gives an unbiased estimate of the improvement to the expected risk obtained by the gradient step. After the first epoch, because the parameters will depend on all data points in the training dataset, we can’t get nice guarantees on this being a good estimate of the change in the test set loss. Nonetheless, for a large training set we can argue that because the parameters depend minimally on any particular minibatch, we should still get a pretty good estimate of the change in the test set error.

Concretely, we hypothesize that under fixed-stepsize SGD, models that obtain a lower sum over training losses should get the best test set performance. We evaluated a handful of different neural network architectures trained on image datasets and found that this correlation was fairly consistent. A more in-depth empirical analysis in the setting of neural architecture serach was performed by Ru et al.