Amortising Inference and Meta-Learning Priors in Neural Networks

AI summary

Overview

Research area: Bayesian deep learning and probabilistic meta-learning, sitting at the intersection of Bayesian neural networks (BNNs) and neural processes (NPs).

Technical level: Advanced. The paper relies on Bayesian linear regression, variational inference, layerwise conditional posteriors, and the neural process literature.

Scope: The paper introduces the Bayesian neural network process (BNNP), a latent-variable neural process whose latent variable is the full set of BNN weights and whose decoder is the neural network parameterised by a sample of that latent variable, trained with a new objective called PP-AVI.

What This Paper Is About

Bayesian deep learning promises principled uncertainty quantification, but it is unclear how to choose sensible prior distributions over neural network weights, since weights are uninterpretable and the convenient priors used in practice collapse large BNNs into "simple smoothing devices" that behave like Gaussian processes. The authors ask whether meta-learning can be used to design well-specified BNN priors from a collection of related datasets, by devising a scheme for per-dataset amortised inference in BNNs. The resulting BNNP also serves as a new kind of neural process, adding capabilities such as within-task minibatching and tunable prior flexibility.

Key Contributions

1. The amortised linear layer. A method for closed-form approximate inference over the weights of a linear layer situated arbitrarily inside a neural network, conditional on the weights of previous layers. Layerwise pseudo-likelihood terms (a pseudo-observation and a noise level for every datapoint) are produced by per-layer inference networks, and the resulting Gaussian posterior mean and covariance are computed exactly via Bayesian linear regression. Stacking these across layers yields the variational posterior (q(\mathbf{W}|\mathcal{D}) = \prod_{l=1}^{L} p(\mathbf{W}^{l} | \mathbf{X}^{l-1}, \mathbf{Y}^{l})).

2. The Bayesian neural network process (BNNP). A latent-variable neural process in which the latent variable is the set of BNN weights and the decoder is the BNN itself, enabling per-dataset amortised inference and meta-learning of BNN weight priors.

3. The PP-AVI training objective. A new objective combining the log posterior-predictive density of a target set with the standard ELBO on the context set, with a proposition (Proposition 1) stating that as the number of meta-datasets tends to infinity, maximising it directly targets three desiderata: accurate approximate posteriors, a faithful prior, and high-quality predictions. It can be estimated unbiasedly from a minibatch of tasks.

4. New neural-process capabilities. Within-task minibatching via sequential Bayesian inference for scalability to large context sets (exactly matching the full-batch approximate posterior at prediction time), an online-learning scheme, and a knob for adjusting prior flexibility by fixing a subset of prior parameters so overfitting can be avoided when few tasks are observed. Two architectural extensions are also outlined: the attentive BNNP (AttBNNP) with a transformer-based encoder, and the Bayesian neural attentive machine (BNAM) with amortised attention blocks in the decoder, which is not a valid stochastic process because it lacks consistency.

Main Findings

• The BNNP achieves high-quality approximate inference. Under a BNN architecture with fixed hyperparameters, data generated by sampling a function from the BNN's prior, uniformly sampling inputs, and adding Gaussian noise with standard deviation 0.1, the BNNP's ELBO gap to the log marginal likelihood (the KL divergence from the approximate to the true posterior) is small across a range of likelihood noise settings, relative to mean-field VI (MFVI), global inducing-point VI (GIVI), and increasingly high-rank correlated variants (UCVI, LCVI, FCVI). Results are averaged over four repeat runs. Approximate inference quality decreases for all methods with smaller likelihood noise, more so for unstructured approximations, which favour overly large noise settings even though the true noise level is (10^{-1}).

• Meaningful BNN priors exist, and the BNNP can find them. Across random sawtooth functions, Heaviside (binary) functions, functions from a standard BNN prior, synthetic ECG signals, and MNIST (cast as pixelwise meta-regression with random pixel masking, using the AttnBNNP), the learned BNNP priors are almost indistinguishable from the true data-generating processes, even though the priors are simple Gaussians. Posterior predictive samples show increased uncertainty away from observations while preserving underlying functional structure.

• Super-resolution arises natively from a learned prior. The AttnBNNP's prior encoded the functional behaviour of handwritten digits, so functions queried on a 100x100 grid could be generated after training on the 28x28 grid, without further training. The authors state that the AttnBNNP's image generation is not state-of-the-art, but many samples are clearly recognisable digits.

• A good prior is not all you need. In one real-world setting (ERA5 precipitation prediction over Europe centred on Switzerland, using longitude, latitude and temperature, with no context points from Switzerland) and three synthetic settings (squared-exponential GP, Heaviside, sawtooth), a well-specified learned prior almost always improved performance for SWAG, MFVI, Langevin Monte Carlo, GIVI, the BNNP and HMC. However, considerable variation remains between methods under learned priors, so high-quality approximate inference is still necessary. Performance was measured by per-datapoint log posterior predictive density (LPPD) and mean absolute error (MAE) over 16 test datasets; for the real-world setting, LMC and HMC were replaced by their stochastic-gradient counterparts SGLD and SGHMC.

• Restricting prior flexibility improves meta-level data efficiency. On Abalone age prediction recast as meta-regression (three sex classes used to split the data; male and female for meta-training, infant for testing; seven input features) and on the Paul15 single-cell RNA sequencing task (predicting cell specialisation from 3451 gene expressions, reduced by PCA to 100 features and split into 19 subsets, with ten used for meta-training and one for testing), the best model was a BNNP with a partially trainable prior: proportions 0.8 and 0.5 respectively. Fully flexible NPs, including the BNNP trained at 1.0, overfit in the data-scarce Abalone problem. Baseline comparators were MFVI and GIVI with standard priors, the original latent-variable NP, the Bayesian NP (BNP), and a transformer NP under an autoregressive sampling scheme (AR-TNP). Results were averaged over four trials.

• Specific numeric values for the plots are not reported in the text. The LPPD and MAE values behind Figures 6 and 7, and the ELBO and KL values behind Figure 2, are presented graphically; exact numbers are not given in the paper text.

• Why the approximate posterior is good. Each layer's posterior is conditioned on the previous layers' weights, modelling inter-layer correlations, but the authors note that FCVI also models such correlations and performs poorly. They suggest instead that the method is unaffected by posterior multimodality induced by weight-space symmetries, because each layer only "sees" one mode of the posterior.

Methodology in Plain English

Instead of trying to put one giant prior over all network weights and then inferring them at once, the authors build the BNN posterior layer by layer. For each layer, small "inference networks" look at individual input-output pairs and produce a pseudo-observation and a pseudo-noise level for that layer's weights — essentially a set of soft, invented targets. Combined with a Gaussian prior over that layer's weights, these pseudo-likelihoods admit an exact closed-form Gaussian posterior, obtained using standard Bayesian linear regression. The twist is that each layer's posterior is computed conditionally on the already-sampled weights of the layers before it, so sampling starts at the first layer and proceeds forward. Stacking all these layers gives the full approximate BNN posterior.

Because the model is now a neural process, training uses a meta-dataset: many tasks drawn from a shared data-generating process, each split into a context set (with labels) and a target set (points to predict). The training objective has two terms — the log posterior predictive density on target points, and the usual ELBO on the context set — so the model is pushed both to predict well and to infer accurately. The inference-network parameters act as variational parameters while the prior parameters are model parameters, and separating these two roles lets the authors freeze part of the prior (for example, fixing the last layer's prior to a zero-centred unit-variance diagonal Gaussian) to control how flexible the learned prior is. At prediction time, large context sets can be handled by streaming minibatches of data through sequential Bayesian updates, layer by layer, without keeping all pseudo-observations in memory.

Why This Matters

This work reframes the hardest practical problem in Bayesian deep learning — where do priors come from? — as a meta-learning problem, and gives evidence that Gaussian priors over weights can encode genuinely complex, even multimodal, stochastic processes. It also supplies a scientific tool: with a well-specified prior, researchers can finally isolate how much approximate inference quality actually matters, and the answer here is that it still matters a lot.

Real-world applications named in the paper (domains where neural processes have already been applied):

• Weather and climate modelling, including the ERA5 precipitation prediction task tested here.
• Causal machine learning.
• Bayesian optimisation.
• Cosmological applications.

Industry relevance: Practitioners who have many related datasets but want calibrated uncertainty — for example in climate and geospatial services, scientific discovery pipelines, and any setting where a small context set must be extrapolated — could use the BNNP as a drop-in neural process that additionally exposes prior flexibility as a tuning knob. The within-task minibatching property is directly relevant to deployment on large or high-dimensional context sets under memory constraints, and the authors' closing argument is that practitioners should spend today's abundant data learning powerful priors rather than squeezing posteriors out of dubious ones.

Future Directions

• Scaling the BNNP as a general-purpose neural process. Inference cost scales unfavourably with architecture width, so the authors explicitly leave investigation of wider and larger-scale use to future work.
• The BNAM as a decoder-side model. Amortised attention blocks could enable more complex tasks, but the authors note the resulting model breaks consistency; whether a consistent version can be built is open.
• Single-dataset prior learning. The BNNP solves the prior problem when multiple datasets are available, but the authors state the problem remains unsolved for the single-dataset case.
• Attention in the encoder. Processing the whole context set jointly could give better pseudo-likelihood parameters than treating points independently, at an extra (\mathcal{O}(n_c^2)) cost, but it would break the exactness of the minibatching scheme with respect to the full-batch forward pass; this trade-off is left unexplored.
• Relation to concurrent work. The authors note that a concurrent scheme for meta-learning BNN priors by others uses in-context network linearisation rather than amortisation and proposes a mixture of priors, which their results suggest is unnecessary — a claim future work could test directly.

Target Audience

Researchers and graduate students in Bayesian deep learning, probabilistic meta-learning, and neural processes who want to understand how priors over network weights can be learned rather than assumed. It is also relevant to applied scientists working with collections of related tasks who need calibrated uncertainty — climate and geospatial analysts, computational biologists working with single-cell data, and practitioners building probabilistic models where only a handful of datasets, or a handful of context points, are available. A reader needs prior familiarity with variational inference and neural processes to follow the derivations, though the empirical sections are accessible to anyone comfortable with predictive log-likelihood and error metrics.

Authors’ abstract

One of the core facets of Bayesianism is in the updating of prior beliefs in light of new evidence$\text{ -- }$so how can we maintain a Bayesian approach if we have no prior beliefs in the first place? This is one of the central challenges in the field of Bayesian deep learning, where it is not clear how to represent beliefs about a prediction task by prior distributions over model parameters. Bridging the fields of Bayesian deep learning and probabilistic meta-learning, we introduce a way to $\textit{learn}$ a weights prior from a collection of datasets by introducing a way to perform per-dataset amortised variational inference. The model we develop can be viewed as a neural process whose latent variable is the set of weights of a BNN and whose decoder is the neural network parameterised by a sample of the latent variable itself. This unique model allows us to study the behaviour of Bayesian neural networks under well-specified priors, use Bayesian neural networks as flexible generative models, and perform desirable but previously elusive feats in neural processes such as within-task minibatching or meta-learning under extreme data-starvation.