PaperlayerView exact revision on arXiv

Source-linked AI summary

Subspace Inference for Bayesian Deep Learning

Pavel Izmailov, Wesley J. Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, Andrew Gordon Wilson

arXiv:1907.07504v1cs.LGstat.ML

TL;DR

Bayesian inference for deep neural networks is difficult because their parameter spaces are extremely high dimensional. The paper constructs low-dimensional subspaces containing diverse high-quality models, performs posterior inference there, and finds strong prediction and uncertainty performance across regression and image classification.

  • Problem

    Extremely high-dimensional neural-network parameter spaces make efficient and practically useful Bayesian inference challenging.

  • Method

    The paper constructs low-dimensional subspaces, performs approximate posterior inference within them, and forms Bayesian model averages.

  • Results

    The approach produces strong practical results across regression and image classification, including notable gains from 5-dimensional subspaces for a 36-million-dimensional WideResNet on CIFAR-100.

  • Takeaways & Limitations

    Low-dimensional subspaces can contain enough variability for practically effective Bayesian model averaging despite the enormous size of modern neural-network weight spaces.

  • Takeaways & Limitations

    The approach assumes that useful subspaces can be constructed around high-quality regions, and some richer low-loss-curve subspaces require additional computation.

Abstract

from arXiv · show

Bayesian inference was once a gold standard for learning with neural networks, providing accurate full predictive distributions and well calibrated uncertainty. However, scaling Bayesian inference techniques to deep neural networks is challenging due to the high dimensionality of the parameter space. In this paper, we construct low-dimensional subspaces of parameter space, such as the first principal components of the stochastic gradient descent (SGD) trajectory, which contain diverse sets of high performing models. In these subspaces, we are able to apply elliptical slice sampling and variational inference, which struggle in the full parameter space. We show that Bayesian model averaging over the induced posterior in these subspaces produces accurate predictions and well calibrated predictive uncertainty for both regression and image classification.

1 INTRODUCTION

Subspace Inference performs Bayesian inference in designed low-dimensional regions of neural-network weight space, aiming to retain diverse high-quality models while avoiding full-space dimensionality challenges. The paper studies several subspace constructions and reports strong accuracy and likelihood across regression and classification tasks.

  • Modern neural networks have extremely high-dimensional parameter spaces that challenge standard Bayesian inference procedures.
  • Subspace Inference performs posterior inference over parameters within a designed low-dimensional subspace S of the weight space.
  • The proposed subspaces are intended to contain diverse representations corresponding to different high-quality predictions, supporting Bayesian model averaging.
  • Candidate subspaces include random subspaces, PCA subspaces from SGD trajectories, and subspaces containing low-loss curves between independently trained solutions.
  • The paper evaluates uncertainty visualization, UCI regression, and CIFAR-10/CIFAR-100 classification, reporting consistently strong test accuracy and likelihood.
  • 5-dimensional subspaces can yield notable Bayesian model-averaging gains for a 36-million-dimensional WideResNet trained on CIFAR-100.

2 RELATED WORK

Related work places Subspace Inference among projected, variational, constrained-posterior, and scalable Bayesian approaches. Its distinguishing focus is parameter-space inference in subspaces selected after SGD-based unconstrained inference.

  • SWAG forms an approximate Gaussian posterior whose low-rank covariance spans a low-dimensional subspace using principal components of SGD iterates.
  • Projected Bayesian inference has also been studied for constrained latent-variable models, including probabilistic PCA.
  • Autoencoder-based variational inference in a subspace formed from fast geometric ensembling requires several models and autoencoder fitting, limiting scalability.
  • Metapriors and Bayesian hyper-networks use low-dimensional spaces or stored meta-models to represent parameters for Bayesian neural networks.
  • Constrained-posterior methods may choose constraints a priori, whereas Subspace Inference chooses constraints after unconstrained inference via SGD.
  • Bayesian coresets and compressed regression reduce inference cost in data space, unlike Subspace Inference, which operates in parameter space.

3 INFERENCE WITHIN A SUBSPACE

The method defines a neural-network model on a K-dimensional affine subspace, performs posterior inference over its free coordinates, and averages predictions by mapping samples back to the original weights. Low dimensionality enables flexible approximate inference, while temperature controls posterior concentration.

  • 3.1 MODEL DEFINITION: A K-dimensional subspace uses a fixed shift vector and projection matrix, leaving z ∈ R^K as the model’s free parameters.
  • 3.1 MODEL DEFINITION: The subspace model evaluates the original network at weights ˆw + Pz, then performs Bayesian inference over z.
  • 3.1 MODEL DEFINITION: Because the mapping from full weights to the subspace is not invertible, the subspace model is treated as a different model sharing functional properties with the original.
  • 3.2 BAYESIAN MODEL AVERAGING: Posterior samples in z are transformed into weight samples ˜w = ˆw + P˜z and used for Monte Carlo Bayesian model averaging on new data.
  • 3.3 APPROXIMATE INFERENCE PROCEDURES: With K much smaller than the full parameter dimension, inference becomes more tractable and can use methods such as HMC, elliptical slice sampling, Laplace, variational inference, or RealNVP.
  • 3.4 PRIOR CHOICE: For sufficiently diffuse priors, performance depends less on the prior alone than on how it interacts with the network’s induced function-space distribution.
  • 3.5 PREVENTING POSTERIOR CONCENTRATION WITH FIXED TEMPERATURE POSTERIORS: When K ≪ N, the posterior can concentrate around the MLE and produce overconfident uncertainty estimates; temperature scaling is introduced to address this issue.
  • 3.5 PREVENTING POSTERIOR CONCENTRATION WITH FIXED TEMPERATURE POSTERIORS: Near T = 1, inference can fail to improve over maximum-likelihood training, while large T produces less confident predictions; good performance is possible across a broad range of T.

4 SUBSPACE CONSTRUCTION

The paper constructs low-dimensional inference subspaces using random directions, PCA components of SGD trajectories, or low-loss curves between solutions. These choices trade construction cost against the ability to capture diverse high-performing models.

  • Subspace construction precedes posterior inference and Bayesian model averaging, with random, PCA, and mode-connected subspaces as design choices.The procedure separates subspace construction from inference and model averaging.
  • Random subspaces: Random subspaces use normalized Gaussian directions and are quick to generate, but contain little information about the model.They require drawing Kp random normal numbers.
  • PCA subspaces: The PCA prior is chosen as a centered Gaussian because the subspace is centered on a good solution and the log likelihood appears approximately quadratic.The shift vector is the SWA solution, and PCA vectors may be scaled using singular values.
  • PCA subspaces: The SGD-trajectory subspace is motivated by containing diverse models while remaining cheap to construct.Prior work argues that the trajectory subspace satisfies both diversity and construction-efficiency goals.
  • PCA subspaces: PCA subspaces use the first principal components of deviations from an SWA solution along a high-learning-rate SGD trajectory.Randomized SVD is applied to deviations from snapshots collected during the final M epochs.
  • Computational cost: PCA construction is a one-time computation that took 4 minutes on the largest reported WideResNet example, compared with approximately 8 hours for training.Random subspaces are virtually free, while construction generally adds minimal overhead relative to standard training.
  • Curve subspaces: Curve subspaces contain diverse high-performing models along low-loss paths between independently trained solutions, but require roughly 3× the cost of training one network.They are the most expensive construction because they require two endpoint solutions and a curve-finding procedure.

5 EXPERIMENTS

The experiments evaluate subspace inference across regression uncertainty visualization, UCI regression, and CIFAR image classification. Low-dimensional PCA and curve subspaces support effective posterior inference, accurate predictions, and improved uncertainty calibration relative to several baselines.

  • Experimental setup: Subspace inference combines subspace construction, posterior inference, and Bayesian model averaging to evaluate random, PCA, and curve subspaces with ESS and variational methods.The experiments use these modular choices across regression and image-classification tasks.
  • Regression uncertainty: PCA and curve subspaces produce uncertainty that grows away from regression data, whereas random subspaces and full-space variational inference are overconfident in extrapolation.The curve subspace is described as the most adaptive, while PCA captures a diverse collection of models.
  • UCI regression: On UCI regression datasets, subspace inference outperforms SGD on several tasks, is competitive on others, and typically improves slightly over SWAG.It also outperforms SGD and DVI on naval, concrete, and yacht, while remaining competitive elsewhere.
  • UCI regression: Subspace inference provides at least as good predictive-interval coverage as SGD and substantially better calibration on elevators and protein.The calibration comparison uses 95% predictive intervals.
  • Posterior approximation: ESS captures posterior shape across the evaluated subspaces, while variational methods struggle with the more complex posterior in the mode-connecting curve subspace.The text attributes this difficulty partly to the constraints of Gaussian variational representations and notes that ESS can navigate such distributions.
  • Subspace choice: Across subspace choices, results improve from random to PCA to curve subspaces, while PCA is selected for later experiments because it offers good performance at lower computational cost than curve subspaces.The reported explanation links the progression to increasing diversity and quality of models within each subspace.
  • Image classification: On CIFAR-10 and CIFAR-100, PCA-subspace inference is competitive with SWAG and consistently outperforms most other tested baselines, including MC-dropout, temperature scaling, and KFAC-Laplace.The evaluated architectures include VGG-16, PreResNet-164, and WideResNet28x10.

6 CONCLUSION

The paper addresses high-dimensional Bayesian inference by performing inference in low-dimensional subspaces with sufficient neural-network variability. These subspaces support practical Bayesian model averaging, strong predictive performance, and uncertainty representation while introducing a measure-theoretic modeling limitation.

  • Conclusion: Subspace inference addresses the dimensionality challenge by constructing variable low-dimensional subspaces and applying approximate Bayesian inference within them.The approach combines subspace construction, posterior inference, and Bayesian model averaging.
  • Conclusion: Low-loss curve subspaces can further improve performance by connecting independently trained solutions, though they require additional computation.These subspaces complement the simpler affine PCA construction.
  • Conclusion: Subspace inference represents uncertainty that grows away from data and shrinks toward data, addressing a challenge for Bayesian deep-learning methods.The conclusion specifically highlights this behavior in regression settings.
  • Conclusion: Mapping into a lower-dimensional subspace loses distribution volume and does not define an invertible reparameterization of the fully parameterized model.The subspace model should therefore be understood as a different model sharing functional properties with the full model.
  • Conclusion: Reducing dimensionality can accelerate MCMC exploration, but very low-dimensional subspaces may cause posteriors to concentrate around a single point.The paper identifies faster mixing as a benefit and missing degeneracies as a possible concern.

B APPROXIMATE INFERENCE METHODS

Approximate inference in the subspace can use either MCMC sampling or deterministic variational approximations. The paper considers slice sampling, NUTS, factorized Gaussian VI, and RealNVP flows, each offering different computational or distributional properties.

  • Inference choices: Posterior inference in the subspace can use MCMC sampling or deterministic approximations such as Laplace or variational inference.The inference procedure is treated as an experimental design choice.
  • Slice Sampling: Elliptical slice sampling is suitable for low-dimensional subspaces because it is gradient-free, has no tuning parameters, and assumes a Gaussian prior.It samples from the projected posterior and can support likelihood accumulation when all data do not fit in memory.
  • NUTS: NUTS dynamically tunes HMC step size and leapfrog counts, requiring a joint likelihood and gradients but incurring full gradient calls.This makes NUTS nearly black-box while potentially increasing computational cost.
  • Simple Variational Inference: Factorized Gaussian VI enables minibatch training and Bayesian model averaging but is constrained in the posterior distributions it can represent.It is one of the simplest and most common variational families.
  • RealNVP: RealNVP uses invertible neural networks to parameterize flexible non-Gaussian variational posterior families.Normalizing flows provide a more expressive alternative to factorized Gaussian approximations.

C EIGEN-GAPS OF THE FISHER AND HESSIAN MATRICES

The Fisher and Hessian spectra exhibit a separated group of large eigenvalues, motivating the use of a low-dimensional subspace for dominant training directions. The paper relates this spectral structure to SGD dynamics and examines trajectory-based predictive distributions.

  • Eigenvalue gaps: The Fisher and Hessian spectra show a clear separation between roughly 20 top eigenvalues and the remaining crowded eigenvalues.Figure 6 plots 300 eigenvalues for PreResNet164 on CIFAR100.
  • Eigenvalue gaps: The spectral gap suggests that gradient-descent updates primarily occupy the subspace spanned by the leading eigenvectors.Prior empirical and theoretical work is cited as support for this interpretation of SGD dynamics.
  • Predictive analysis: The paper visualizes predictive distributions for multiple inference methods and subspaces, including posterior-density samples in a separate figure.These visualizations are used to compare the resulting inference behavior.

E UCI REGRESSION EXPERIMENTAL DETAILS

The UCI regression experiments use repeated train-test splits and evaluate subspace inference against deterministic variational and deep Gaussian-process baselines. On small datasets, PCA-subspace ESS and VI outperform DVI on two of five datasets.

  • Experimental setup: Experiments replicate over 20 trials, reserving 90% of each dataset for training and 10% for testing.The setup follows prior UCI regression evaluation protocols.
  • Experimental setup: Trajectory covariance eigenvalues are estimated with randomized SVD across three architectures and two CIFAR datasets, decaying toward 0 around 10-20 steps.The eigenvalues are plotted as explained-variance proportions on a log scale.
  • Gaussian test likelihood: Bayesian model averaging estimates predictive mean and variance from samples drawn from the approximate posterior.The predictive variance combines averaged observation variance with variation in predictive means.
  • Small Regression: ESS and VI in the PCA subspace outperform deterministic VI on two of five small UCI regression datasets.The experiments use one-hidden-layer networks with 50 units and heteroscedastic predictive uncertainty.

E.1.3 Large-Scale Regression

Large-scale UCI regression experiments use heteroscedastic neural networks and compare subspace inference with SGD, approximate Gaussian processes, and Bayesian final layers. The evaluation uses repeated train/test splits and standardized inputs.

  • Evaluation: The regression figures visualize predictive means, sampled predictive functions, and predictive-distribution uncertainty across inference methods and subspaces.The shaded region represents ±3 standard deviations at each input.
  • Model setup: Large UCI tasks use deep networks with widths such as [1000, 1000, 500, 50, 2] and train with specified learning-rate, momentum, batch-size, and weight-decay settings.Skillcraft uses a smaller architecture and a shorter training schedule.
  • Model setup: The likelihood combines input-dependent variance with a global variance parameter optimized alongside the network.The global parameter is analogous to Gaussian-process jitter; omitting it often produced over-confident predictions.
  • Evaluation: Subspace inference is evaluated against approximate Gaussian processes and Bayesian final-layer methods using RMSE comparisons.The comparisons include DKL, OrthVGP, FastFood kernel GPs, and Bayesian final layers.
  • Evaluation: Experiments repeat each model over 10 random train/test splits, with each test set containing 10% of the dataset and standardized features.Inputs are pre-processed to mean zero and variance one.

F IMAGE CLASSIFICATION RESULTS

The image-classification experiments follow the established CIFAR evaluation framework and report negative log-likelihood and accuracy for the proposed method and baselines.

  • Experimental framework: CIFAR experiments follow the framework of Maddox et al. (2019).The reported evaluation uses the same general experimental framework as that work.
  • Metrics: The experiments report negative log-likelihood and accuracy for subspace inference and baseline methods.Results are presented in Tables 8 and 9.

F.1 EFFECT OF TEMPERATURE

Temperature affects subspace-inference uncertainty and likelihood more strongly than predictive performance in the reported CIFAR-100 and UCI Elevators experiments. In both settings, a broad temperature range gives near-optimal performance.

  • CIFAR-100: 10^3 ≤ T ≤ 10^4 yields strong Bayesian model-averaging results for PreResNet-164 on CIFAR-100.The experiment uses elliptical slice sampling in a 5-dimensional PCA subspace.
  • Cross-task pattern: Temperature changes affect uncertainty estimates and negative log-likelihood more than predictive accuracy in the CIFAR-100 experiment.Figure 11 reports test negative log-likelihood and accuracy as functions of temperature.
  • UCI Elevators: On UCI Elevators, performance remains nearly constant and close to optimal across a certain temperature range.The effect of temperature on likelihood is larger than its effect on RMSE.
  • Evaluation context: The broader evaluations include prediction coverage, RMSE, log-likelihood, and accuracy across UCI and image-classification settings.The reported tables compare subspace inference with SGD, SWAG, approximate Gaussian processes, Bayesian final layers, temperature scaling, and dropout.
  • Experimental design: The temperature study uses three independent runs and reports mean ±1 standard deviation for the plotted metrics.The CIFAR-100 and UCI Elevators plots use a 5-dimensional PCA subspace with elliptical slice sampling.
Loading 1907.07504v1…