Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning

1. Introduction

• Marginal likelihood uses training data alone, but its estimation is challenging and it is rarely used for deep-learning model selection.Many Bayesian deep-learning approaches instead use validation data, while standard Laplace estimation requires Hessian computations that do not scale to large problems.Source: §1 · p. 1
• Differentiable hyperparameters can be optimized online during training, whereas discrete choices such as architectures can be selected afterward.The paper describes this as a scalable, online version of MacKay’s proposal.Source: §1 · pp. 1–2
• Across models of varying size and architecture, higher training marginal likelihood tends to accompany higher test accuracy, while smaller models tend to score higher at similar performance.Figure 1 also reports a higher marginal likelihood for a 3-layer network than a 1-layer network, whose deeper fit better explains the sinusoidal trend.Source: Figure 1 · p. 1, Figure 2 · p. 2
• The method applies scalable Laplace estimation with Generalized Gauss-Newton or Empirical Fisher Hessian approximations, further reduced through diagonal or block-diagonal approximations.These approximations had been used for inference but not previously applied to model selection in the described work.Source: §1 · p. 2
• Approximate marginal likelihoods select reasonable models and perform on par with or better than cross-validation across regression and classification benchmarks.The best architecture selected by marginal likelihood aligns with test performance.Source: §1 · p. 2, Figure 2 · p. 2

2. Background

• A supervised dataset consists of input-output pairs, with real-valued network outputs and real or categorical outcomes.The model is specified by a network architecture and hyperparameters, and a Bayesian posterior combines likelihood and prior.Source: §2 · p. 2
• Marginal likelihood balances model fit and complexity: models that are too simple or too complex receive lower probability for the training data.Selecting the model with the highest training-data probability is described as a simple comparison procedure.Source: §2 · pp. 2–3
• Laplace’s method approximates the posterior locally around a maximum with a Gaussian and uses it to approximate the marginal likelihood.Computing the full P × P Hessian and its determinant is infeasible in general.Source: §2 · p. 3
• The marginal-likelihood log-determinant favors Hessians with small eigenvalues, linking the objective to low curvature and flatness near a solution.The passage says flat minima have been argued to generalize better and notes empirical and PAC-Bayesian support for that claim.Source: §2 · p. 3
• Scalable Generalized Gauss-Newton and Empirical Fisher Hessian approximations had been used for inference, but their application to marginal-likelihood estimation was largely unexplored.The cited exception is Khan et al. (2019), on small regression tasks using a full GGN.Source: §2 · p. 3
• The GGN approximation uses network Jacobians and likelihood curvature, while the Empirical Fisher uses outer products of likelihood gradients.The Empirical Fisher has stated complexity O(P^2N), compared with O(P^2NC + PNC^2) for GGN; the passage describes it as O(C) cheaper when C is typically much smaller than P.Source: §2 · p. 3
• Laplace-GGN and Laplace-EF denote Laplace’s method with the Hessian approximated by GGN or Empirical Fisher, respectively.These are the names used for the two Hessian-approximation variants.Source: §2 · p. 3

3. Model Selection with Laplace-GGN and -EF

• The method uses Laplace’s method with GGN or EF Hessian approximations, whose determinant computations can be structured for scalability.For small datasets, the Woodbury identity rewrites a parameter-sized determinant in terms of matrices sized by the number of data points.Source: §3 · p. 3, §3.3 · p. 4

3.1. Step 1: Online model selection during training

• During training, the algorithm alternates network-parameter optimization with gradient updates to differentiable hyperparameters such as prior variances, observation noise, or softmax temperature.Hyperparameter updates use the marginal-likelihood estimate and can also support early stopping without validation data.Source: §3.1 · pp. 3–4

3.2. Step 2: Model selection after training

• After training, compare marginal-likelihood estimates and choose the model with the higher value to select between discrete alternatives such as architectures.The method treats this comparison as a hypothesis test between models.Source: §3.2 · p. 4

3.3. Scalable Laplace approximations

• The Kronecker-factored GGN approximates each layer’s GGN with a Kronecker product and avoids damping, which would distort the Laplace-GGN.Its matrix decompositions are quadratic in the number of neurons per layer.Source: §3.3 · p. 5
• The diagonal approximations provide a cost-effective alternative to block-diagonal approximations in the online algorithm, despite often weaker predictive performance in other settings.The authors report that these approximations work well in practice, especially for convolutional neural networks.Source: §3.3 · p. 5

3.4. Computational considerations

• The algorithm’s runtime depends on estimation frequency, consecutive hyperparameter updates, burn-in, and the chosen Hessian approximation.Marginal-likelihood estimation requires a full training-data pass, but expensive computations can be reused across consecutive update steps.Source: §3.4 · p. 5
• The online procedure is reported robust to settings of K, F, B, and γ, and can optimize around 400 hyperparameters for a ResNet.The authors report that it typically needs no more epochs to converge than standard training.Source: §3.4 · p. 5

4. Experiments

• Experiments compare with cross-validation and manually selected models, using full GGN and EF or Kronecker-factored and diagonal approximations according to problem size.The paper also compares MAP predictions with a Laplace-GGN posterior predictive based on a linearized network.Source: §4 · p. 6

4.1. Illustrative examples

• In toy regression, the larger model fits without overfitting after regularization optimization; on banana classification, test accuracy rises from 86% for weak regularization to 89%.The marginal likelihood identifies the better regression model.Source: §4.1 · pp. 6–7
• For prior precision, the online algorithm converges from three initial values to the grid-search optimum within the same number of steps, and marginal likelihood identifies the best-generalizing value.The examples show predictive distributions for over- and underfitting choices as well as the optimum.Source: Figure 4 · p. 6, §4.1 · p. 6

4.2. UCI regression

• On UCI regression, MAP predictions from the online method outperform or match Bayesian predictions with cross-validated hyperparameters on 6 of 8 datasets; full GGN performs best.The comparisons cover CV MAP, VI, and Laplace predictives and marginal-likelihood methods using GGN, EF, and GGN-KFAC.Source: Table 1 · p. 7, §4.2 · p. 7

4.3. Image classification

• Across image-classification architectures, marginal-likelihood selection generally improves performance over cross-validated models, especially on CIFAR-10, and favors CNNs over MLPs and ResNets over CNNs.The reported comparisons also show higher accuracies corresponding to higher marginal likelihoods.Source: Table 2 · p. 8, §4.3 · p. 8
• Without data augmentation, the method achieves better test log-likelihood, ECE, and OOD-AUC than the augmented baseline; with augmentation, its accuracy matches the baseline.The OOD-AUC uses SVHN as the out-of-distribution data.Source: Figure 5 · p. 8, §4.3 · p. 8
• On CIFAR-10, online marginal-likelihood optimization reduces overfitting and the generalization gap versus weight decay; training takes around twice one run, while cross-validation could test two settings.The Kronecker-factored Laplace-GGN approximation accounts for the reported twofold training time.Source: §4.3 · p. 8, Figure 6 · p. 9
• Across CIFAR-10/100, ResNets achieve higher marginal likelihoods than CNNs, and width tends to improve marginal likelihood more than depth; test-accuracy rank correlation is 97%.On FMNIST, the rank correlation is 69%, likely lower because marginal likelihood often selects significantly smaller rather than slightly better-performing models.Source: §4.3 · p. 9, Figure 7 · p. 9

5. Conclusion

• Marginal likelihood can reflect model quality using training data alone, with a mild increase in computation and support for comparing models and selecting hyperparameters during training.The authors link this approach to using the neighborhood of a solution to improve robustness.Source: §5 · p. 9
• The authors see training-data-based selection as potentially useful where little is known about deployment and future usage.They suggest removing dependence on validation data could widen deep learning’s impact and open new application areas.Source: §5 · p. 9
• Further work includes improving Hessian approximations and computation, online selection of discrete hyperparameters, and extending the approach beyond supervised learning.Examples named include bandits, Bayesian optimization, active learning, and continual learning.Source: §5 · p. 9

A.1. Laplace approximation to the marginal likelihood

• The local Laplace approximation expands the log joint around an arbitrary parameter value; at a MAP point, the gradient correction vanishes.The standard approximation uses the zero-gradient condition, while the local form includes an additional correction term.Source: §A.1 · p. 12
• The online method omits the correction term because it did not help empirically and its implied parameter update can typically lead to divergence.The authors follow the common deep-learning simplification of assuming the gradient is zero.Source: §A.1 · p. 12

A.2. Efficient determinant computation

• Woodbury’s identity lets determinant and correction-term computations shift complexity from the number of parameters to the number of data points.The Hessian in these computations is approximated by either the GGN or EF.Source: §A.2 · p. 13

A.3. Kronecker-factored Laplace-GGN without damping

• The undamped Kronecker approximation represents each layer’s GGN with two factors and computes its determinant from their eigenvalues across the block-diagonal layers.The determinant computation uses an outer product of the factor eigenvalues; the authors note damping distorts the Laplace approximation when prior parameters are optimized.Source: §A.3 · p. 13
• Damping yields negative likelihoods up to two times worse and accuracy decreases of up to 4 percentage points in the online algorithm.The ablation concludes that damping should be avoided in this context.Source: §A.3 · p. 13, Table A1 · p. 13

B. Computational details

• The implementation uses per-layer isotropic Gaussian priors; regression optimizes observation-noise variance, while the illustrative classification example optimizes softmax temperature.The determinant is the main part where the Hessian approximations differ, and update frequency and burn-in affect how often it is computed.Source: B. Computational details · p. 14

B.1. Full Laplace-GGN

• Full Laplace-GGN computation depends on whether P is smaller or larger than NC, with determinant costs of O(P^3) or O(N^3C^3), respectively.With an isotropic prior, subsequent updates can be performed in O(P) or O(N).Source: §B.1 · p. 14

B.2. Full Laplace-EF

• The full EF alternative scales independently of output count C and uses the cheaper dimension of P and N for determinant computation.Its costs are O(P^3) when P < N and O(N^3) when N < P; isotropic priors permit O(P) or O(N) updates.Source: §B.2 · p. 14

B.3. Kronecker-factored Laplace-GGN or EF

• Kronecker-factored approximations store per-layer factors and make prior-precision and noise-variance updates O(P).Their factor computations and eigendecompositions have an up-front cost that multiple hyperparameter updates can amortize.Source: §B.3 · p. 15

B.4. Diagonal Laplace-GGN

• Diagonal Laplace-GGN determinant computation is O(P), but the method is not significantly faster than Kronecker-factored Laplace and still scales with output count C.It makes a different approximation and may work better in some circumstances.Source: §B.4 · p. 15

B.5. Diagonal Laplace-EF

• Diagonal Laplace-EF needs a pass through training data to sum squared individual gradients, roughly O(NP), followed by O(P) hyperparameter updates.The passage characterizes this as cheaper and simpler than diagonal Laplace-GGN.Source: §B.5 · p. 15

Bayesian regression predictive

• In regression, predictive variance is the sum of model uncertainty Var[f] and observation noise σ2, which the method learns online.The example decomposes the predictive into epistemic, model-dependent uncertainty and aleatoric observation noise.Source: §C.1 · pp. 15–16
• The posterior predictive does not significantly improve UCI regression predictions except on the small yacht dataset, which may require better uncertainty quantification.Table C2 is described as additional posterior-predictive performance on the UCI regression benchmark.Source: Table C2 · p. 16, Table 1 · p. 17

Additional UCI regression results

• Diagonal EF and GGN online optimization slightly reduces performance versus full GGN but is not significantly worse, and still performs quite well against cross-validated MAP.Source: Table C3 · p. 16

Bayesian classification predictive

• For classification, the Bayesian predictive with better marginal likelihood is certain in the data region and increasingly uncertain away from it, without post-training parameter tuning.The posterior predictive is estimated using 1000 samples on network outputs.Source: §C.1 · p. 16

Regression grid search

• The online grid-search example finds that the best marginal likelihood corresponds to the best held-out test negative log likelihood, and the online algorithm finds the optimum without test data.The result is reported for a grid search over one prior parameter with different initializations.Source: §C.1 · p. 16

C.2. UCI regression

• Figure C.2 plots test negative log likelihood against prior precision δ and shows the proposed method reliably converging to optimal predictive performance.Source: Figure C.2 · p. 16
• The MAP predictive generally performs better than the posterior predictive, except on the small yacht dataset.Source: §C.2 · p. 16
• The negative test log-likelihood comparison contrasts diagonal GGN and EF approximations with full GGN.Source: §C.2 · p. 17

C.3. UCI classification

• UCI classification performance is comparable to cross-validation; full EF leads GGN and KFAC, while diagonal variants are weaker, especially for Bayesian prediction.Source: Table C4
• UCI classification methods show no obvious negative-log-likelihood differences, while cross-validated MAP and full EF-MAP perform better than most methods in accuracy and expected calibration error.Source: §C.3 · p. 17

• Diagonal EF performs slightly better on all image-classification datasets than diagonal GGN and is cheaper to compute.Source: Table C6 · p. 18, §C.4 · p. 18
• On CIFAR-10, cheaper diagonal EF typically performs slightly better than the KFAC variant.Source: §C.4 · p. 18, Table C7 · p. 19
• Without data augmentation, the smallest CIFAR-10 ResNet has the highest marginal likelihood; with augmentation, ResNet-44 seems best, while parameter growth barely affects marginal likelihood.Source: §C.4 · pp. 18–19, Table C8 · p. 19

• The experiments compare FMNIST MLPs and CNNs and CIFAR CNNs and ResNets, varying network depth and width across the listed architectures.Source: §C.5 · p. 19
• On CIFAR, increasing width tends to raise marginal likelihood and test accuracy, while increasing depth raises marginal likelihood but saturates earlier than width.Source: §C.5 · p. 19
• On FMNIST, less-wide MLPs achieve better marginal likelihood and test accuracy; increasing MLP depth tends to reduce both, while deeper CNNs benefit.Source: §C.5 · p. 19

• The step size γ need only prevent divergence; γ = 1 can reduce the hyperparameter steps needed, and practitioners can monitor the objective or parameters for divergence or oscillation.Source: §C.6 · p. 20
• Burn-in B reduces early runtime while network parameters are noisy; B = 0 is optimal for performance if affordable, and B = 50 works equivalently well.Source: §C.6 · p. 20
• Although computing the Laplace approximation is expensive, hyperparameter steps are amortized; F between 5 and 10 and K ≥100 suffice in practice, with no gain from K = 1000.Source: §C.6 · p. 20