PaperlayerView exact revision on arXiv

Source-linked AI summary

Functional Variational Bayesian Neural Networks

Shengyang Sun, Guodong Zhang, Jiaxin Shi, Roger Grosse

arXiv:1903.05779v1cs.LGstat.ML

TL;DR

BNNs struggle to express meaningful function-level priors and approximate posteriors in weight space. fBNNs instead optimize a functional ELBO using finite measurement sets and SSGE, enabling structured priors; experiments report strong extrapolation, uncertainty, and scaling results, with measurement-set selection limiting generalization.

  • Problem

    BNNs make meaningful priors and accurate posterior inference difficult because weight-space distributions poorly reflect the functions computed by networks.

  • Method

    fBNNs perform variational inference over stochastic processes by optimizing a functional ELBO approximated with finite measurement sets and SSGE.

  • Results

    fBNNs extrapolate well with structured priors, provide reliable uncertainty estimates, and outperform competing approaches across small- and large-scale settings.

  • Takeaways & Limitations

    Function-space inference supports richly structured explicit and implicit priors while allowing explicit posterior function samples for downstream optimization.

  • Takeaways & Limitations

    Measurement sets selected to maximize the KL term may concentrate near training data, weakening encouragement to exploit structured priors for extrapolation.

Abstract

from arXiv · show

Variational Bayesian neural networks (BNNs) perform variational inference over weights, but it is difficult to specify meaningful priors and approximate posteriors in a high-dimensional weight space. We introduce functional variational Bayesian neural networks (fBNNs), which maximize an Evidence Lower BOund (ELBO) defined directly on stochastic processes, i.e. distributions over functions. We prove that the KL divergence between stochastic processes equals the supremum of marginal KL divergences over all finite sets of inputs. Based on this, we introduce a practical training objective which approximates the functional ELBO using finite measurement sets and the spectral Stein gradient estimator. With fBNNs, we can specify priors entailing rich structures, including Gaussian processes and implicit stochastic processes. Empirically, we find fBNNs extrapolate well using various structured priors, provide reliable uncertainty estimates, and scale to large datasets.

1 INTRODUCTION

fBNNs move variational inference from weights to functions, enabling structured stochastic-process priors and more stable behavior as network capacity grows. Experiments report sensible extrapolation, competitive performance, and reliable uncertainty estimates.

  • 1 INTRODUCTION: Weight-space BNNs make meaningful priors and accurate posterior inference difficult because weights relate ambiguously to computed functions and have complex dependencies.
  • 1 INTRODUCTION: fBNNs perform variational inference directly over distributions of functions rather than high-dimensional network weights.The approach defines a functional ELBO and uses finite measurement sets with SSGE-based gradient estimates.
  • 1 INTRODUCTION: On the toy y = x3 task, larger fBNNs produce consistent predictions across network sizes, unlike BBB, while a 1 × 100 factorized Gaussian fBNN lacks expressive diversity.
  • 1 INTRODUCTION: fBNNs support richly structured stochastic-process priors, including Gaussian processes and implicit distributions over piecewise functions.These priors can encode dependencies such as smoothness and periodicity.
  • 1 INTRODUCTION: As network capacity increases, fBNN predictions remain well behaved, whereas ordinary BNN priors and posterior inference can degenerate in larger networks.The functional prior allows the underlying BNN to grow without changing the functional variational inference problem.
  • 1 INTRODUCTION: fBNNs generate sensible extrapolations with structured priors and achieve strong results on regression and contextual-bandit benchmarks.The authors report reliable uncertainty estimates and state-of-the-art contextual-bandit performance.

2 BACKGROUND

The background frames BNNs as weight-based Bayesian models and introduces stochastic processes through their finite-dimensional marginals. It also motivates implicit-distribution methods such as SSGE for settings where densities are unavailable.

  • 2 BACKGROUND: A Bayesian neural network combines a weight prior p(w) with a likelihood p(D|w), while variational methods optimize an ELBO over an approximate weight posterior.
  • 2 BACKGROUND: Common weight priors are chosen for computational convenience, but their relationship to the functions computed by networks is difficult to characterize.
  • 2 BACKGROUND: A stochastic process is a collection of random variables indexed by X, with each outcome defining a sample function from X to Y.
  • 2 BACKGROUND: Finite index sets induce joint marginal distributions over function values; for Gaussian processes, these marginals are multivariate Gaussians.
  • 2 BACKGROUND: The Kolmogorov Extension Theorem characterizes a stochastic process through marginals over all finite index sets subject to consistency conditions such as exchangeability.
  • 2 BACKGROUND: Implicit distributions have tractable sampling processes despite intractable densities, including neural-network output distributions and neural samplers.
  • 2 BACKGROUND: SSGE estimates the derivative of an implicit distribution’s log density using samples and a positive-definite kernel.
  • 2 BACKGROUND: SSGE approximates its estimator through kernel eigenfunctions, Nyström approximations, truncation, and Monte Carlo estimates.

3 FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS

fBNNs perform variational inference over distributions of functions, using finite measurement sets to make the functional ELBO practical. The method supports explicit and implicit stochastic-process priors while addressing intractable function-space KL terms with sampling and score estimation.

  • 3.1 FUNCTIONAL EVIDENCE LOWER BOUND (fELBO): fBNNs define variational inference over functions rather than weights, with stochastic-process priors and neural-network variational posteriors.The prior may be a Gaussian process or an implicit process such as a distribution over piecewise linear functions.
  • 3.1 FUNCTIONAL EVIDENCE LOWER BOUND (fELBO): The KL divergence between stochastic processes equals the supremum of marginal KL divergences over finite measurement sets.This reduces an infinite-dimensional divergence to KL divergences over function values evaluated at finite input sets.
  • 3.1 FUNCTIONAL EVIDENCE LOWER BOUND (fELBO): Maximizing the functional ELBO becomes a two-player game in which the network chooses the posterior and an adversary chooses measurement points.Restricting measurement sets to a fixed finite size yields a practically useful objective, although the infimum may not be attainable in general.
  • 3.2 CHOOSING THE MEASUREMENT SET: Sampling-based measurement sets combine training inputs with random domain points so the network matches prior structure where predictions and extrapolations are needed.The adversarial approach generalized poorly because maximizing the KL term tends to select points near the training data.
  • 3.2 CHOOSING THE MEASUREMENT SET: Finite measurement sets provide an upper bound on the true functional ELBO, while including all training inputs restores a lower bound on log p(D).Under Gaussian-process posteriors and variational families, finite-point objectives also admit the stated consistency results.
  • 3.3 KL DIVERGENCE GRADIENTS: The spectral Stein gradient estimator approximates intractable score terms for both explicit and implicit stochastic-process priors.The algorithm samples measurement points and function values, estimates KL gradients, and updates variational parameters using the resulting objective.

4 RELATED WORK

Related work spans weight-space variational BNNs, functional-prior approaches, and scalable Gaussian processes. fBNNs differ by performing inference directly in function space and by supporting richer stochastic-process priors and explicit function samples.

  • Bayesian neural networks: Bayes By Backprop uses a fully factorized Gaussian posterior with reparameterization-based gradients for variational inference in neural networks.Earlier approaches used biased gradient estimators, while Blundell et al. proposed an unbiased estimator.
  • Functional Priors: Earlier functional-prior methods still performed variational inference in weight space or added prior regularization to a weight-space ELBO.Noise Contrastive Priors use a random function-space noise prior but are not rich enough to express the structured priors targeted here.
  • Functional Priors: In periodic extrapolation, fBNNs matched exact Gaussian-process predictions, whereas standard BBB failed to fit training data and KL-scaled BBB failed to extrapolate sensibly.The comparison uses RBF and PER + RBF kernels, with the latter encoding periodic structure.
  • Functional Priors: VIP reverses the fBNN arrangement by specifying BNN priors and using Gaussian processes to approximate the posterior.Consequently, VIP cannot exploit richly structured Gaussian-process or other stochastic-process priors.
  • Scalable Gaussian Processes: Exact Gaussian processes scale as O(N^3) time and O(N^2) memory, motivating sparse and random-feature approximations for large datasets.Sparse methods remain difficult for very large datasets, while random-feature and KISS-GP methods require kernel-specific tailoring.

5 EXPERIMENTS

Experiments test whether fBNNs exploit explicit and implicit structural priors for extrapolation, regression, scalability, and contextual bandits. Across these settings, fBNNs generally produced sensible structured predictions, competitive regression performance, and reliable uncertainty estimates.

  • The experiments target structured-prior extrapolation and competitiveness against other BNNs on regression and contextual-bandit benchmarks.
  • 5.1.1 LEARNING PERIODIC STRUCTURES: For periodic regression, fBNNs made predictions similar to Gaussian processes with matching kernels, while BBB failed to recover the periodic pattern.The periodic structure was encoded by the PER + RBF kernel.
  • 5.1.1 LEARNING PERIODIC STRUCTURES: fBNN posterior samples reproduced roughly piecewise constant or piecewise linear structure, although they did not capture the full diversity of possible explanations.Sharp transitions appeared despite using tanh activations.
  • 5.2.1 SMALL SCALE DATASETS: On most small-scale regression datasets, fBNNs outperformed BBB and Noisy K-FAC, sometimes by a significant margin.The benchmarks report averaged test RMSE and log-likelihood.
  • 5.2.2 LARGE SCALE DATASETS: fBNNs scale to large datasets through stochastic estimation of the expected log-likelihood term.Large-scale experiments compared fBNNs with BBB on UCI datasets using validation-based hyperparameter selection and early stopping.
  • 5.3 CONTEXTUAL BANDITS: In contextual bandits, fBNNs outperformed other methods by a wide margin and maintained consistent performance as networks became deeper and wider.BBB suffered significant degradation when hidden size increased from 50 to 500.

6 CONCLUSIONS

The paper formalizes variational inference over stochastic processes by relating process-level KL divergence to finite-dimensional marginals. It uses this foundation to motivate practical functional variational inference and fBNNs.

  • The paper’s conclusion relies on finite measurement sets as the basis for relating stochastic-process divergences to marginal divergences.This finite-set perspective is consistent with the measure construction described by the Kolmogorov Extension Theorem.
  • Process-level measures are connected to finite-dimensional distributions through canonical projections and the cylindrical σ-algebra.Cylinder sets impose restrictions on function values at finitely many inputs.
  • The Kolmogorov Extension Theorem constructs a unique stochastic-process measure from compatible probability measures on all finite subsets.In Gaussian processes, these finite-set measures are marginal Gaussian distributions.

A.2 FUNCTIONAL KL DIVERGENCE

The paper establishes that stochastic-process KL divergence is determined by finite-dimensional marginals, supporting practical functional variational inference and its conditional-process interpretation.

  • Every measurable set in the cylindrical function space depends on at most countably many restricted indices.The proof establishes this by showing the relevant sets form a sigma-algebra containing the cylinder generators.
  • KL divergence between stochastic processes equals the supremum of marginal KL divergences over all finite index subsets.The finite-set marginals arise through canonical projection maps.
  • The theorem follows by reducing finite partitions to countable index sets, relating partition and KL divergences, and approximating countable-index KL by finite-index KL.The proof uses three explicit steps before showing arbitrary finite index sets are covered.
  • For conditional stochastic processes, the KL divergence reduces to the marginal KL divergence at observed locations.When the conditioning dataset is empty, this becomes the marginal KL between posterior and prior processes on observed locations.
  • This reduction justifies using finite measurement points in adversarial and sampling-based functional variational inference.The practical methods therefore evaluate process-level divergence through finite measurement sets.

B.1 PROOF FOR EVIDENCE LOWER BOUND

The consistency proof connects functional ELBO optimization with equality of Gaussian-process posteriors by evaluating KL on training and measurement points.

  • The proof considers measurement points XM outside the training data and decomposes the functional objective over training and measurement function values.The construction sets XM = X\XD and uses joint values at the two point sets.
  • The analysis assumes both the variational posterior and target posterior are Gaussian processes with respective mean and covariance functions.These are represented as GP(mp, kp) and GP(mq, kq).
  • The functional objective remains bounded above by log p(D), with equality requiring zero KL between the joint variational and posterior distributions.The argument assumes measurement points cover all training locations.
  • Zero marginal KL at arbitrary measurement sets implies matching posterior and variational means and covariances at those points.For M > 1, the proof obtains equality of means at each point and covariances for each pair.
  • Because Gaussian processes are uniquely determined by their mean and covariance functions, matching those functions establishes consistency.The conclusion follows after the pointwise equalities hold for arbitrary measurement locations.

C.3 BAYESIAN OPTIMIZATION

The Bayesian optimization experiments test whether fBNNs provide usable function minima samples and structured posterior behavior across several Gaussian-process priors.

  • The experiment uses Max Value Entropy Search, which maximizes information gain about the minimum value y⋆.The acquisition function depends on posterior mean, standard deviation, and samples of function minima.
  • Gaussian-process priors use RBF, Order-1 ArcCosine, and Matern12 kernels, comparing fBNN, BBB, and Random Feature methods.fBNN uses the true kernel, whereas Random Feature uses RBF features for all three kernels.
  • fBNN and Random Feature outperform BBB by a large margin on all three functions.The comparison is based on minimal value found over optimization iterations and averages 10 independent runs.
  • fBNN slightly trails Random Feature for RBF priors but outperforms it for ArcCosine and Matern12 functions.The paper attributes this pattern to Random Feature’s RBF mismatch with the latter two kernels and fBNN’s use of true kernels.
  • The experiment highlights that fBNNs can learn parametric function posteriors for different structured priors.This capability supports the observed optimization performance across mismatched kernel structures.

C.5 LARGE SCALE REGRESSION WITH DEEPER NETWORKS

On large-scale regression with deeper networks, fBNNs are compared with BBB and SVGP to assess scaling as dataset size increases.

  • The experiment uses five-hidden-layer networks with 100 units per layer for BBB and fBNNs, while comparing against SVGP.SVGP uses 1000 inducing points initialized by k-means.
  • SVGP performs better than BBB and fBNNs on the smallest naval dataset.This is the reported exception within the large-scale regression comparison.
  • As dataset size increases, SVGP performs worse than BBB and fBNNs by a large margin.The paper links this to the limited capacity of its 1000 inducing points as sufficient statistics for large datasets.
  • BBB and fBNNs can use larger networks without the intractable computational cost associated with scaling sparse Gaussian-process representations.The paper presents this as the contrast with SVGP’s fixed inducing-point capacity.

D IMPLEMENTATION DETAILS

The implementation handles explicit and implicit stochastic-process priors, stabilizing Gaussian-process gradient estimates with injected noise and sampling structured piecewise functions.

  • Gaussian-process priors permit analytical computation of ∇f log pφ(f X) because p(f X) has an explicit multivariate Gaussian density.
  • Injecting Gaussian noise into function values stabilizes gradient computation by effectively replacing the kernel matrix K with K + γ2I.The noisy objective estimates gradients for distributions convolved with the noise distribution.
  • The injected-noise approach preserves the reparameterization trick for the parametric variational posterior, so prior estimation formulas remain applicable.
  • The method supports implicit priors through piecewise constant and piecewise linear stochastic-process constructions.Functions are generated on [0, 1] by sampling Poisson-distributed change points and values defining the pieces.
Loading 1903.05779v1…