Simone Rossi
← All publications

On permutation symmetries in Bayesian neural network posteriors: a variational perspective

  • Bayesian Deep Learning
  • Variational Inference
  • Deep Learning Theory

We align independently trained variational neural-network posteriors through neuron permutations and observe low predictive-loss barriers along the resulting distribution interpolations.

Paper NeurIPS 2023

Simone Rossi1, Ankit Singh2, Thomas Hannagan2

  1. EURECOM
  2. Stellantis
PublicationPDFPreprint

Independent variational inference runs can produce Bayesian neural networks with similar predictions and distant weight distributions. We align entire Gaussian weight distributions through neuron permutations, including their uncertainty, before interpolating between them. On MNIST and CIFAR-10, this removes the visible predictive-loss barrier for multilayer perceptrons and substantially reduces it for ResNet-20. Neuron ordering therefore accounts for part of the apparent separation between these variational solutions. A substantial barrier remains on CIFAR-100.

Two Bayesian regression solutions and their naive and aligned interpolations on the Snelson dataset.
Two solutions fit the Snelson regression data, with predictive means and two-standard-deviation bands. Naive interpolation loses the fit, while aligned interpolation retains it, illustrating why neuron correspondence matters when combining distributions.

Why neuron ordering changes posterior geometry

A hidden unit has an arbitrary index. Reordering a layer and compensating in the following layer preserves every network output. For a multilayer perceptron with elementwise activations, a permutation matrix PP acts on weights WlW_l and biases blb_l as

Wl′=PWl,bl′=Pbl,Wl+1′=Wl+1P⊤.W_l'=PW_l,\qquad b_l'=Pb_l,\qquad W_{l+1}'=W_{l+1}P^\top.

Applying this transformation to every draw from a weight distribution q(θ)q(\theta) gives the push-forward distribution P#qP_\#q. Both distributions induce the same random function. Consequently, comparing independently trained posteriors in their original coordinates can obscure their relationship. We search within this class of equivalent representations for a version of q1q_1 that aligns with q0q_0.

A CIFAR-10 MLP log-posterior slice containing two variational solutions and a permuted copy of one solution.
A log-posterior slice places q0q_0, q1q_1, and the permuted distribution P#q1P_\#q_1 in the same view. Alignment moves one representation toward a connection through higher-density regions while preserving its predictions.

Interpolating probability mass and measuring predictive quality

For NN evaluation inputs xi⋆x_i^\star and labels yi⋆y_i^\star, we assess each distribution using the marginalized predictive log-likelihood

L(q)=∑i=1Nlog⁡∫p(yi⋆∣θ,xi⋆) dq(θ).\mathcal L(q)=\sum_{i=1}^{N}\log\int p(y_i^\star\mid\theta,x_i^\star)\,\mathrm dq(\theta).

The integral averages predictions over weights before taking the logarithm. We examine how far this score falls below the interpolation of the endpoint scores, equivalently measuring a barrier in negative predictive log-likelihood.

The path should move probability mass through parameter space. We therefore use a Wasserstein geodesic. For diagonal Gaussian endpoints qi=N(mi,diag⁡(si2))q_i=\mathcal N(m_i,\operatorname{diag}(s_i^2)), with means mim_i and positive standard-deviation vectors sis_i, it becomes

qτ=N(mτ,diag⁡(sτ2)),mτ=(1−τ)m0+τm1,sτ=(1−τ)s0+τs1,τ∈[0,1].\begin{aligned} q_\tau&=\mathcal N(m_\tau,\operatorname{diag}(s_\tau^2)),\\ m_\tau&=(1-\tau)m_0+\tau m_1,\\ s_\tau&=(1-\tau)s_0+\tau s_1,\qquad \tau\in[0,1]. \end{aligned}

This interpolates means and standard deviations. A mixture instead reweights the endpoints, so its predictions can remain strong while leaving the intervening parameter region unexplored.

Mixture and optimal-transport interpolation, with distribution sketches and predictive likelihood along each path.
For CIFAR-10 MLP solutions, mixture interpolation keeps predictive likelihood nearly constant, whereas transport interpolation reveals a large dip. The distribution sketches explain how reweighting endpoint modes differs from moving their mass.

Matching means and uncertainty together

For diagonal Gaussians, squared Wasserstein distance reduces to two Euclidean terms:

W22(P#q1,q0)=∥m0−Pm1∥22+∥s0−Ps1∥22.\mathcal W_2^2(P_\#q_1,q_0) =\lVert m_0-Pm_1\rVert_2^2 +\lVert s_0-Ps_1\rVert_2^2.

The matching criterion therefore includes the location and spread of each weight distribution. For a single vector, minimizing this expression is a linear assignment problem. In a network, each hidden-layer permutation acts on adjacent weight matrices, coupling their assignments. We alternate between layers, fixing the other permutations and solving the remaining assignment exactly. These updates improve a combinatorial matching objective; the overall search remains a heuristic for the coupled problem. The interpolated path then connects q0q_0 to the aligned endpoint P#q1P_\#q_1.

Predictive barriers across architectures and training settings

We train independent mean-field Gaussian posteriors by maximizing the evidence lower bound. The classification experiments use MLPs and ResNet-20, Gaussian priors, and filter response normalization for the convolutional networks. They omit data augmentation. Evaluation uses 128 weight samples and 25 evenly spaced interpolation coefficients.

Predictive likelihood along unaligned and aligned paths for MNIST and CIFAR-10 networks.
Training and test predictive likelihoods along the paths show no noticeable MLP barrier after alignment and a much smaller residual ResNet-20 barrier. The gray curves expose the degradation caused by interpolating the original coordinates.

Across these comparisons, alignment substantially improves the full interpolation path on training and test data. Increasing width further reduces aligned barriers. For ResNet-20, the corresponding unaligned barriers remain large, making the interaction between width and neuron correspondence particularly clear.

Training and test predictive-likelihood barriers against model width for an MLP and ResNet-20.
Wider aligned MLPs on MNIST and ResNet-20 models on CIFAR-10 have lower barriers. ResNet-20 retains a large barrier without alignment across the tested widths, showing how width and neuron correspondence jointly affect these paths.

The prior-variance and posterior-temperature sweeps probe the sensitivity of this behavior. Aligned MLP barriers stay close to zero over the tested prior variances, while ResNet-20 retains architecture-dependent residual barriers. A temperature T<1T<1 sharpens the posterior target by raising its unnormalized density to the power 1/T1/T; these cold settings can reduce the remaining barrier.

Predictive-likelihood barriers over prior variance for MNIST, Fashion-MNIST, and CIFAR-10 models.
Changing Gaussian prior variance leaves aligned MLP barriers small, while the ResNet-20 residual varies more. These controls distinguish the effect of alignment from sensitivity to prior regularization.
Predictive-likelihood barriers as posterior temperature varies for MLP and ResNet-20 models.
The temperature sweep compares cold and warm posterior targets. Alignment lowers barriers throughout the sweep, with architecture-dependent residuals and smaller values at some cold settings.

A remaining barrier on CIFAR-100

On CIFAR-100, our matching procedure leaves a large predictive-loss barrier for ResNet-20. The improvement observed on MNIST and CIFAR-10 therefore depends on the task and the permutation found by the layerwise search. The evaluations characterize predictive quality at a finite set of points along paths between mean-field Gaussian approximations.

ResNet-20 predictive likelihood along CIFAR-100 interpolations, showing a large remaining barrier after alignment.
On CIFAR-100, aligned and unaligned paths both pass through substantially worse predictive distributions. This case exposes a limit of the layerwise matching procedure on a more demanding classification task.