Variational Deep Learning via Implicit Regularization
Authors: Jonathan Wenger
Beau Coker
Juraj Marusic
John Patrick Cunningham
Venue: ICLR 2026 Poster (2026)
Keywords: Implicit Regularization
Bayesian Deep Learning
Generalized Variational Inference
Implicit Bias of SGD
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We demonstrate theoretically and empirically that one can exploit the implicit bias of SGD for variational inference in Bayesian neural networks.

Abstract: Modern deep learning models generalize remarkably well in-distribution, despite being overparametrized and trained with little to no explicit regularization. Instead, current theory credits implicit regularization imposed by the choice of architecture, hyperparameters, and optimization procedure. However, deep neural networks can be surprisingly non-robust, resulting in overconfident predictions and poor out-of-distribution generalization. Bayesian deep learning addresses this via model averaging, but typically requires significant computational resources as well as carefully elicited priors to avoid overriding the benefits of implicit regularization. Instead, in this work, we propose to regularize variational neural networks solely by relying on the implicit bias of (stochastic) gradient descent. We theoretically characterize this inductive bias in overparametrized linear models as generalized variational inference and demonstrate the importance of the choice of parametrization. Empirically, our approach demonstrates strong in- and out-of-distribution performance without additional hyperparameter tuning and with minimal computational overhead.

Published as a conference paper at ICLR 2026
VARIATIONAL DEEP LEARNING VIA IMPLICIT REGULARIZATION
Jonathan Wengerf
Beau Coker Columbia University
Juraj Marusic
John P. Cunningham Columbia University
Columbia University
Columbia University
ABSTRACT
Modern deep learning models generalize remarkably well in-distribution, despite being overparametrized and trained with little to no explicit regularization. In- stead, current theory credits implicit regularization imposed by the choice of ar- chitecture, hyperparameters, and optimization procedure. However, deep neural networks can be surprisingly non-robust, resulting in overconfident predictions and poor out-of--dstribel d-tistriibution generalization. Bayesian deep learning addresses this via model averaging. but typically requires significant computational resources as well carefully elicited priors to avoid overriding the benefits of implicit regu- larization. Instead, in this work, we propose to regularize variational neural net- works solely by relying the implicit bias of (stochastic) gradient descent. We theoretically characterize this inductive bias in overparametrized linear models as generalized variational inference and demonstrate the importance of the choice of parametrization. Empirically, our approach demonstrates strong in- and out- of ff-distribution performance without additional hyperparameter tuning and with minimal computational overhead.
INTRODUCTION
The success of deep learning across many application domains is, on the surface, remarkable, given that deep neural networks are usually overparameterized and trained with little to no explicit regular- ization. The generalization properties observed in practice have been explained by implicit regular- ization instead, resulting from the choice of architecture [1], hyperparameters [2, 3], and optimizer [4-10]. Notably, the corresponding inductive biases often require no additional computation, in contrast to enforcing a desired inductive bias through explicit regularization.
In the last two decades, there has been increasing focus improving the reliability and robust- ness of deep learning models via (approximately) Bayesian approaches [11] to improve performance on out-of-distributtin data [12], in continual learning [13], and in sequential decision-making [14]. However, despite its promise, in practice, Bayesian deep learning can suffer from issues with prior elicitation [15], can be challenging scale [16], and explicit regularization via prior combined with approximate inference may result in pathological inductive biases and uncertainty [17-20].
In this work, we demonstrate both theoretically and empirically how exploit the implicit bias of optimization for approximate inference in probabilistic neural networks, thus regularizing training implicitly rather than explicitly via the prior. This not only narrows the gap to how standard neural networks are trained, but also reduces the computational overhead of training compared to varia- tional inference. More specifically, we propose learn a variational distribution over the weights of deep neural network by maximizing the expected log- likeliheehood in analogy to training via max- imum likelihood in the standard case. However, in contrast to variational Bayes, there is no explicit regularization via Kullback-Leibler divergence to the prior. Surprisingly, we show theoretically and empirically that training this way does not cause uncertainty to collapse away from the training data, if initialized and parametrized correctly. More so, for overparametrized lincar models rigor- ously characterize the implicit bias of SGD as generalized variational inference with a 2- Wassersstein regularizer penalizing deviations from the prior. Figure 1 illustrates our approach toy example.
1 Equal contribution.

---

Published as a conference paper at ICLR 2026
Standard NN
Implicit Bias VI (ours)
Mean-field VI (KL)
Generalized VI (W3)
-2
-2
(a) Implicir regularization.
(b) Explicit regularization.
Figure 1: Variational deep learning via implicit regularization. Neural networks generalize well without explicit regularization due to implicit regularization from the architecture and optimization. We exploit this implicit bias for variational deep learning, removing the computational overhead of explicit regularization and narrowing the gap to deep learning practice. As illustrated for a two- hidden layer MLP and proven rigorously for overparametrized linear models in Theorems and 2, the implicit bias of (SYGD in variational networks (see (a)) can be understood as generalized vari- ational inference with Wasserstein regularizer (see (b)). This differs from the standard ELBO objective with a KL divergence to the prior as used for example in mean-field VI (see (b)).
Contributions In this work, we propose a new approach to Bayesian deep learning that gener- alizes robustly by exploiting the implicit regularization of (stochastic) gradient descent. We fully characterize this implicit bias for regression (Theorem 1) and binary classification (Theorem 2) in overparameterized linear models, generalizing results for non-probabilistic models and drawing rigorous connection to generalized Bayesian inference. We also demonstrate the importance of the parametrization for the inductive bias and its impact hyperparameter choice. In several bench- marks, we demonstrate competitive performance to state-of-the baselines for Bayesian deep learning, at minimal computational overhead compared standard neural networks. Finally, we provide an open-source implementation of our approach as standalone library: inferno.
2 BACKGROUND
Given training dataset [X, y) = ((.,, ) of input-output pairs, supervised learning seeks function fw to predict the corresponding output y(x) for a test input æ. The parameters E R.P of the function are typically trained via empirical risk minimization, i.e.
w. E min (w)
with
e, (ww = f(y fw (X)) + Ar(w),
(1)
where the loss Ely. fw (X) encourages fitting the training data and the regularizer r(w), given some À >> O, discourages overfitting, which can lead poor generalization on test data.
Implicit Bias of Optimization One remarkable observation in deep learning is that training over- parametrized neural networks (P N) with gradient descent without explicit regularization can nonetheless lead to effective in-distribution) generalization, despite there being many global min- ima of the loss corresponding to functions fw. which achieve zero training error [21]. This can be explained by the optimizer, initialization, and parametrization implicilly regularizing the optimiza- tion problem arg min... E(y, fw (X)), thereby preferring certain global minima [c.g. 7, 23}. Nonetheless, deep neural networks can be surprisingly brittle when predicting out-of-distribution. often displaying overconfidence and significant drop in generalization performance.
Bayesian Deep Learning Approximate Bayesian techniques like the Laplace approximation [24- 26]. stochastic weight averaging [27, 28), deep ensembles [29], and variational approaches [30-33] attempt to address the aforementioned shortcomings of deep learning by learning a distribution over functions as opposed to merely point estimate. The idea being that weighted combination of models, all of which achieve low training error, generalizes more robustly while at the same time providing uncertainty quantification.
Variational Inference In Bayesian inference this weighted combination is defined by the posterior distribution p(w X, y) x ply X, ( ()))(u) over weights, induced by likelihood p(y wr) and
2

---

Published as a conference paper at ICLR 2026
choice of prior p(w) that expresses explicit preference for some models over others. Approxi- mating the posterior with (o(w) z p(w X, y) by maximizing a lower bound to the log-evidence leads to the following variational optimization problem [34]:
0, E arg min (0)
s.t.
(0) Eg w)) l- p(y w)) + KL(go(w) ((w))
(2)
0
Equation (2) is an instance of the empirical risk minimization objective in Equation (I), with the key difference that one optimizes over variational parameters of family of distributions go(w) E Q- If that family includes the posterior, 4o 40(w)) (w) = p(w X, y) is the unique global minimum. In the case of a potentially misspecified prior or likelihood, the variational formulation (2) can be generalized to arbitrary loss functions l and statistical distances D to the prior [35-37], such that
(+) Bg (w)) (L(y, (x))) + A p).
(3)
3 VARIATIONAL DEEP LEARNING VIA IMPLICIT REGULARIZATION
Our overarching goal is to enable deep neural networks to generalize robustly out-of-distribution without sacrificing their in-distribution performance, minimal computational overhead. We ap- proach this goal within the framework of Bayesian deep learning, by learning distribution over neural networks sw. induced by parametrized variational distribution qo(w) (w) over its weights. However, rather than approximating the Bayesian posterior, which trades off training error against an explicit, a priori preference for certain models, we enforce that all models have zero training error while using implicit regularization to weight them. Doing so preserves the implicit regular- ization of the optimizer, which determines the generalization performance neural networks to a substantial degree, rather than purely relying on explicit regularization induced by the prior. Impor- tantly, this approach leads to robust out-of-distribution generalization, while providing uncertainty quantification at small computational overhead over standard deep learning.
3.1 TRAINING VIA THE EXPECTED Loss
We propose to train variational neural network defined by architecture fw and variational distribution over weights go (w) by minimizing the expected loss {()) in analogy to how deep neural networks are usually trained. In other words, the optimal variational parameters given by
0. E arg min (w) (f(y, Ja ())) D(
(4)
2mm (())
At first glance removing the divergence term from the variational objective in Eq. (3) seems prob- lematic because the new objective is clearly minimized when the variational distribution is a point mass at the minimum loss solution, i.e. Qo, (w) = 6ur, (w) where w. E arg min,, E(y, (x) (X)). This seemingly defeats the point of Bayesian deep learning framework, given that there is no variability in predictions on test data. Moreover, the new objective no longer involves prior dis- tribution, ostensibly removing the ability to manually favor some models over others entirely, The key to understanding our approach is that, in the overparameterized setting, point mass is only one of many optima corresponding to distributions qo, (w), and is the implicit bias of the optimiza- tion procedure that chooses among them. As we will see, if one trains an overparametrized linear model via the expected loss using (stochastic) gradient descent, this implicit bias can be explicitly characterized to depend on the initialization.
3.2 IMPLICIT BIAS OF SGD AS GENERALIZED VARIATIONAL INFERENCE
Assume we train an overparametrized linear model with a Gaussian variational family via the ex- pected loss. For an appropriate learning rate sequence, (stochastic) gradient descent converges to a global minimum g E arg min ((0) of the training objective. As we show in Section 4, if SGD initialized to the prior, i.e. qa, (o,(w) (w) p( w(w), its implicit bias can be understood as selecting the distribution over models with zero training error that is closest to the prior in 2- Wasserstein distance:
aoc
arg min To
W2(ge. .)).
s.t. 0 Éaarg min E(O)
3

---

Published as a conference paper at ICLR 2026
Therefore, we can interpret the implicit bias of (S)GD when training a variational linear model as performing generalized variational inference. More precisely, the above is equivalent to yooo mini- mizing the objective in Equation (3) for certain regularization strength, but with regularizer that is not a KI. divergence as would be for standard variational inference, but rather a 2- 2- distance to the prior. This characterization directly generalizes results for mnon-poobabilistio mod- els, where the implicit bias of SGD selects minima that are close to the initialization in Euclidean distance [5, 21J. We therefore call our method Implicit Bias Variational Inference (IBVI), From practical perspective, by exploiting the implicit regularization of SGD, rather than performing gen- eralized variational inference directly, we no longer need to compute the regularizer explicitly or allocate memory for the prior hyperparameters.
Section 4 provides detailed version of the regression result introduced here and proves a similar result for binary classification, Our experiments in Section 5 focus on the application to deep neural networks, where we generally expect the implicit regularization be more complex.
3.3 COMPUTATIONAL EFFICIENCY
In practice, minibatch the expected loss both over training data and parameter samples um drawn from the variational distribution ge (w) such that
{())= (v
(w) (l(g) for (X))) N.M
No
EM llyn. fw. (xx))
(5)
The training cost is primarily determined by two factors. The number of parameter samples M we draw for each evaluation of the objective, and the variational family, which determines the number of additional parameters of the model and the cost for sampling a set of parameters in each forward pass. We wish to keep the overhead compared to vanilla deep neural network small as possible.
Training With A Single Parameter Sample (M = 1) When drawing fewer parameter samples wmm the training objective in Eq. (5) becomes noisier similar to using a smaller batch size. This concerning since the optimization procedure may not converge given this additional noise. However, one train
SGD
0.10
Test Ervor
0.05
with a single parameter sample only, simply by reducing the learning rate appropriately, we show experimentally in Fig-
0.00
ure and Section S3.2. Therefore, given a set of sampled pa- rameters, the cost of forward and backward pass is identical to a standard neural network (up to the overhead of the covari- parameters). When using fewer parameter samples in the expected loss, training is unstable unless the learning rate is cho- sen sufficiently small. For fixed number of optimizer steps this decreases performance, but either training for more steps, or us- ing momentum closes this gap.
SGD Momentum
0.10
Test Error
0.05
0.00
-4
100
10a
10a
Parameter Samples
Variational Family and Covariance Structure We choose Gaussian variational distribution go (ww) over (a subset of the) weights of the neural network. While at first glance this may seem restrictive, there is ample evidence that variational fami- lies in deep neural networks do not need be complex be expressive [38, 39]. In fact, in analogy to deep feedforward NNs with ReLU activations being universal approximators [40]. one can show that Bayesian neural networks with ReLU activations and at least one Gaussian hidden layer are universal conditional distribution approximators, meaning they can approximate any continuous conditional distribution arbitrarily well [39]. As we show in Section 4, training an overparametrized linear model with SGD via the expected loss amounts to generalized variational inference if the covariance is factorized, i.c. Σ SS where S E R is a dense matrix with rank R < P. The implicit bias SGD for arbitrary parametrizations of the covariance matrix remains open problem. Throughout our experiments we use Gaussian layers with factorized covariances for all architectures.
Figure 2: Training with single parameter sample given small enough learning rate. Lighter color shades correspond to smaller learning rates. See also Section S3.2.
We only need them to initialize the optimizer after which we can free up the allocated memory.
4

---

Published as a conference paper at ICLR 2026
Standard Param. (SP)
Maximal Update Param. (juP)
Relative Test Accuracy
1.0
1.5
2.0
Train
1.0
0.8
1,5
0.5
0.6
2
2
2
2
2-7
210
Learning Rate
Learning Rate
Hidden Size
Hidden Size
Purametrization
LR Selection Method
128
512
2048
Standard
Grid Search
256
1024
Maximal Update
Transferred Grid Search
When scaling the size of neural network, one has to re-tune the hyperparameters, such as the learning rate, when using the standard parametrization (SP). The same is true for probabilistic networks as we show here on CIFAR-I0 (left). However, when using our proposed extension of the maximal update parametrization (μP) [41] to probabilistic networks, one can tune the learning rate on small model and achieve optimal generalization for larger models by "transferring" the optimal learning rate from smaller model (center and right).
Figure 3: Hyperparameter Transfer.
3.4 PARAMETRIZATION, FEATURE LEARNING AND HYPERPARAMETER TRANSFER
The inductive bias of SGD depends the initialization and choice of parametrizaaion, bijective map p O' + θ reparametrizing a (variational) model such that fe = Jp(en)- When training deep neural networks, it is not unusual to use layer-specific learning rates. These can be absorbed into the weights of the model and the initialization, meaning they effectively just define a different parametrization [Lemma J.1, 41]. While parameterization well-studied for non- con-proaabilistic deep learning, it has been identified as one of the "grand challenges of Bayesian computation" [42].
In deep learning, the "standard parameterizaaton" (SP) initializes the weights of a neural network randomly from a distribution with variance xx 1/fan in (e.g., in Kaiming initialization, the PyTorch default) and makes no further adjustments the forward pass or learning rate. In contrast, the maximal update parametrization (uP) [43] ensures feature learning even the width of the network tends to infinity. In addition, under μP, hyperparameters like the learning rate, can be tuned on small model and transferred to large-scale model [41].
Given our interpretation training via the expected loss as generalized variational inference with a prior implied by the parametrization and initialization, natural question is whether μPP to the variational setting and thus inherit its inductive bias. In the probabilistic setting, feature learning occurs when the distribution over hidden units changes from initialization. At any point during training, the ith hidden unit layer is a function of four random variables: the variational mean and covariance parameters (1,, S), Gaussian noise z, and the previous layer hidden units:
extend
① (x) = W,h W,h (-)) = (m, + (.
(6)
The parameters are random because of the stochasticity in the initialization and/or optimization procedure, while the noise is randomly drawn during each forward pass. Since the S₁z term is sum over terms, where R the rank of S PxR applying the central limit theorem we propose scaling this term by R-1/2 and then applying μP to the mean and covariance parameters. In practice, we implement the scaling via an adjustment to the covariance initialization and learning rate. Section S2 in the supplement provides empirical investigating of this scaling, demonstrating feature learning in the last hidden layer as the width is increased.
Figure 3 demonstrates that proposed maximal update parametrization enables hyperparameter transfer a probabilistic model. We train two-hidden- layer MLPs on CIFAR1O, using a low rank covariance in the final two layers. Under standard parametrization (left panel), the learning rate that results in the smallest training loss decreases with hidden size. In contrast, under jPP (middle panel), it remains the same across hidden sizes. The right panel of Fig. 3 demonstrates the practical
5

---

Published as a conference paper at ICLR 2026
implications for model selection. For each parametrization and each hidden size D, select the learning rate based on a grid search. In "transferred grid search" we do a grid search using the smallest model (hidden size 128) and transfer the best validating learning rate to the hidden size D model, whereas in "grid search" we perform the grid search on the hidden size model. Relative to the test accuracy of the best performing model across learning rate and parametrization, that (a) μP outperforms SP, though the gap decreases with hidden size, and (b) the transfer strategy works well for uP but poorly for SP once the hidden size exceeds 256.
3.5 RELATED WORK
Variational inference in the context of Bayesian deep learning has seen rapid development in cent years [30-33, 44, 45]. Using a Wasserstein regularizer [37] in the context of generalized VI [36] is arguably most related to our work, given our theoretical results. Structure in the variational parameters has always played an important role for computational reasons [38, 46, 47] and often only a few layers are treated probabilistically [39]. with some methods only considering the last layer, effectively treating the neural network as feature extractor [48, 49]. The Laplace approxi- mation if applied in the last-layer also falls under this category, which has the advantage that it can be applied post-hoe [13, 24-26, 50-54]. Deep ensembles repeat the standard training process using multiple random initializations [29, 55] and have been linked to Bayesian methods [56, 57J with cer- tain caveats [58, 59]. While we use SGD only to optimize the variational parameters and arguably average over samples by using momentum, SGD has also been used widely to directly approximate samples from target distribution [27, 56, 60, 61]. a popular example being stochastic weight aver- aging (SWA) [27, 28]. Our theoretical analysis extends recent developments on the implicit bias of overparameterized linear models [4, 5, 7] to the probabilistic setting. For classification, works have focused on convergence rates [6]. SGD [7], SGD with momentum [8], and the multiclass setting [10]. Results on the implicit bias of neural network training [22] often assume large widths [9, 62- 65] allowing similar arguments as for linear models. The former is exemplified by the neural tangent parametrization, under which neural networks behave like kernel methods the infinite width limit [66]. Yang et al. [41, 43, 64, 65] developed an alternative parameterization that still admits feature learning in the infinite width limit, which we extended to the case of variational networks.
4 THEORETICAL ANALYSIS
Consider an overparameterized linear model with a Gaussian prior, trained via the expected loss using (stochastic) gradient descent. We show that, in both regression (Theorem 1) and binary clas- sification (Theorem 2), our approach can he understood as generalized variational inference with a 2- 2-Wasserstein regularizer, which penalizes deviation from the prior among models with zero training error. Theorems I and 2 recover analogous results for non-probabilistic models [4, 5, 21].
4.1 LINEAR REGRESSION
Theorem (Implicit Bias in Regression)
Let fmo x w be overparametrized linear model with P N. Define Gaussian prior p(w) N(w: Mo, SyST) and likelihood p(y w) = Nyy (X), o²?) and assume varia- tional family qo(w) N (w: μ., STT with 0 S1, such that μ TRP and S TR. PxR where R < P. If the learning rate sequence (Nt); is chosen such that the limit point 0 lim OCD identified by gradient descent, initialized Og (Ho, So). a (global) minimizer of the expected oogliielihood ((O), then
GG
arg min 0=(n.S) s.L. 0Earg min 8(0)
Further, this also holds in the case of stochastic gradient descent and when using momentum.
Proof. See Section SI.I.I.
Theorem states that, among those variational parameters which minimize the expected loss, SGD (with momentum) converges the unique variational distribution which is closest in 2- Wassersteii

---

Published as a conference paper at ICLR 2026
distance to the prior. This characterization of the implicit regularization of SGD generalized varia- tional inference differs from a standard ELBO objective (2) in VI via the choice of regularizer. Since the variational parameters minimize the expected loss in Equation (7), all samples from the predic- tive distribution interpolate the training data (see Figure 1 (b)) right panel), the same way a standard neural network would. In contrast, when training with a KL. regularizer, the uncertainty does not collapse at the training data (see Figure 1(b), left panel). In fact, a KL regularizer would diverge infinity for Gaussian with vanishing variance. Now, for test points that are increasingly out-of- distribution, i.e. less aligned with the span of the training data, the variational predictive matches the prior predictive more closely. Interestingly. g0 is equal the distribution over weights of an ensemble of linear models initialized from the prior and trained independently (see Section S1.1.3). Next, we prove similar result for binary classification.
4.2 BINARY CLASSIFICATION OF LINEARLY SEPARABLE DATA
Consider binary classification problem with labels yn E (-1, a linear model f=) x w, and a variational distribution qo (w) with variational parameters O. The expected empirical loss is {(0)- = Ene[N] [g (f(yn=",,)) w)) We assume without loss of generality? that all labels are positive, i.e. yn. for all n, and that the dataset is linearly separable.
Assumption The dataset is linearly separable: E RR such that Vn w x. 0.
For an overparametrized linear model, if X E be linearly separable." Define the solution to the hard margin SVM, the L2 max margin vector as
NN has full row rank the dataset is guaranteed to
arg min [ 1 wER P
w xn 1,
and the set of support vectors S = arg e x 2 indexing the data points on the margin. We make the following additional assumption which is satisfied with high probability under mild assumptions on the training data distribution and degree of overparametrization [67, 68].
Assumption 2 The SVM support vectors span the dataset: span( (zn INE(N)) = ssaan ( Ines).
We can now characterize the implicit bias in the case of binary classification.
Theorem 2 (Implicit Bias in Binary Classification)
Let Sw (x)) Teuu uu be an (overparumetrized) linear model and define a Gaussian prior p(w) N N(w; no; S- Assume variational distribution qo go(w)) =N (w: p. SST) over the weights E gd with variational parameters 0 (H, such that S E R and R < P. Assume we are using the exponential loss f(u) exp(-u) and optimize the expected empirical loss ((0) via gradient descent initialized the prior, i.e. 0, (H.. with sufficiently small learning rate n. Then for almost any dataset which is linearly separable (Assumption 1) and for which the support vectors span the data (Assumption 2), the rescaled gradient descent iterates (rGD)
@:GD
STGD)
GD p +
X SED
(9)
log(t]
converge to a limit point g: l se G for which it holds that
@DGD E arg min W W2(..))) p),
(10)
e= (u. s.t. O
where the feasible set O (1, S) P PaangeXXTij
v and Vn Var (fm( 0} consists of mean parameters which, if projected onto the training data, are equivalent to the L2 mux margin vector and covariance parameters such that there is no uncertainty at training data.
Proof. See Section S1.2.
Theorem 2 states that the mean parameters 14, converge to the 7-2 max-margin vector i in the span of the training data, i.c. the data manifold, and there uncertainty collapses to zero. This is analogous
"This is not a restriction since we can always absorb the sign into the inputs, such that æ', = y Wn. We can always choose x (XxT)-1 i.c. the weights linearly interpolating y (1,..., ... T)a.
7

---

Published as a conference paper at ICLR 2026
to the regression case, where zero training loss enforces interpolation of the training data. In the null space of the training data, i.c. off of the data manifold, the model falls back on the prior as enforced by the 2-Wasserstein distance. The assumption of an exponential loss is standard in the literature and we expect this to extend to (binary) cross-entropy in the same way it does in results for standard neural networks [4, 6-8, 10]. Similarly, to SGD with momentum [cf. 7, 8J. While Theorem 2 is similar to Theorem 1, there are some subtle differences. First, the feasible set for the minimization problem in Equation (10) is not the set of minima of the expected loss. This is because the exponential function does not have an optimum in contrast to quadratic function. However, the sequence of variational parameters identified by gradient descent still satisfies lim {(O,) 0. Second, without transformation of the mean parameters, the exponential loss results in the mean parameters being unbounded. This necessitates
conjecture that Theorem 2 can be extended
the transformation in Equation (9) as we explain in detail in Section $1.3.
5 EXPERIMENTS
We benchmark the generalization and robustness of approach, Implicit Bias VI (IBVI,, against standard neural networks and several baselines for uncertainty quantification, namely Temperature Scaling (TS) [69], Laplace approximation (LA-GS) & (LA-ML) [24-26], Weight-Space VI (WSVI) [30, 31], SWA-Gausiian (SWAG) [28], and Deep Ensembles (DE) [29], on set of standard bench- mark datasets for image classification and robustness to input corruptions. We use convolutional architectures (LeNet5 [70] or ResNet34 [71], which, for all datasets but MNIST, are initialized with pretrained weights except for the imput and output layer. All models are trained with SGD with momentum 7 = 0.9 and a batch size of N₂ 128 for 200 epochs in single precision on an NVIDIA GH200 GPU. Results shown are averaged across five random seeds. A detailed description of the datasets, metrics, models and training can be found in Section S3. An implementation of our method can be found at:
In-Distribution Generalization and Uncertainty Quantification In order to assess the in- distribution generalization, measure the test error, negative log-likelihood (NLL), and calibra- tion error (ECE) on MNIST, CIFARIO, 10, CIFARIO 100 and TinyImageNet. As Figure 4 shows for CI- FAR100, and Figure S10 for all datasets, the test error for post-hoc methods (TS, LA-GS, LA-ML) is unchanged. As expected. SWAG and IBVI perform similarly with only Ensembles providing an increase in accuracy, but at substantial memory overhead compared to most other approaches. Sim- ilarity of IBVI to Ensembles is perhaps expected in light of their equivalence for linear models (see Proposition S1). In-distribution uncertainty quantification measured terms of NLL is improved substantially by TS, DE, and IBVI, with only LA and WSVI showing occasional worsening of NLL compared to the base model. The full results in Figure S10 show that TS, DE, and IBVI consistently are also the best calibrated. As described in Section 3.3, for IBVI we train with a single sample only and a probabilistic input and output layer with low-rank covariance, reducing the memory overhead compared to standard neural network to as little as xe 10% with similar training time (see Figure 4). See Section S3.3.2 for the full experimental results including different parametrizations (SP vs jcP).
Robustness to Input Corruptions We evaluate the robustness of the different models on MNISTC [72], CIFARIOC, CIFAR100C, and TinylmageNetC [73]. These are corrupted versions of the original datasets, where the images are modified via a set of 15 corruptions, such as impulse noise. blur, pixelation, etc. We selected the maximum severity for each corruption and averaged the performance across all. As expected, the performance of all models drops compared to the in- distribution performance measured on the standard test sets as Figure 5 shows. Besides DE which consistently show lower test error, also IB shows improved accuracy on corrupted data compared to all other approaches. When using the maximal update parametrization, SWAG shows good ac- curacy on the two larger datasets (see Figure S12). TS, DE, and IB perform consistently well in terms of uncertainty quantification (both for NLL and ECE) across all datasets, with LA-ML being somewhat competitive in terms of NLL. However, compared to the in-distribution setting IBVI has better uncertainty quantification than Ensembles across all datasets.
8

---

Published as a conference paper at ICLR 2026
CIFARI00
CIFARI00
CIFARI00
Test Top-5 Error 4
Num. Parameters
108
Training Runtime
4h
0.15
Test NLL
0.10
107
MNIST CIFARIO CIFARIOD
Tiny ImageNet
Method
Standard
Laplace (Last-layer. GS)
Weight-space VI (Mean-field) field) Implicit Bias VI (Low-rank)
SWAG
Temperature Scaling
Laplace (Last-layer. ML)
Ensemble
Figure 4: In-distribution generalization and uncertainty quantification. Implicit Bias VI (IBVI) has similar test error to other Bayesian deep learning approaches and achieves competitive uncertainty quantification on in-distribution data. While ensembles have improved accuracy, they come at an additional memory overhead. Training probabilistic model via IBVI has only a minor computa- tional overhead during training, both time and memory, over standard deep learning.
MNISTC
CIFARIOC
Emor 4
CIFAR100C
TinyImageNoC
Test T T
0.35
0.65
Test Error
0.2
0.275
0.30
G0M
0.250
0.60
0.1
Test Test NLL 4
2
1.0
00O
0.5
Test ECE 4
0.15
0.4
0.50
0.3
0.10
0.2
0.2
0.25
Method
Standard
Laplace (Last-layer,
Weight-space VI (Meun-field)
SWAG Ensemble
Temperature Scaling
Laplace (Last-layes, MI.)
Implicit Bias VI (Low-rank)
Figure 5: Generalization robustness benchmark problems. When comparing different methods for Bayesian deep learning with regards to robustness to 15 different input corruptions, our approach, Implicit Bias VI, consistently has competitive uncertainty quantification across different datasets and metries without sacrificing accuracy compared non-probabillsttc network.

---

Published as a conference paper at ICLR 2026
Limitations Compared to standard neural networks, when training via Implicit Bias VI, ob- served that often lower learning rates were necessary due to the additional stochasticity in the ob- jective (see also Section 3.3). While this does not have significant impact generalization, sometimes requires slightly more epochs to achieve similar in-distribution performance to standard neural networks. Effectively, early in training it takes bit more time for IB VI become sufficiently certain about those features which are critical for in-distribution performance. This also means that folk knowledge on learning rate settings for specific architectures may not immediately transfer. In the experiments we train models with probabilistic in- and output layers with our approach, but we have so far not explored other covariance structures or where in the network probabilistic layers are most beneficial. While there is theoretical evidence that even just single probabilistic hidden layer may be sufficient [39], we believe there potential for improvement. Beyond the prior induced by the choice of parametrization, we did not experiment with more informative or learned priors, which could potentially give significant performance improvements on certain tasks [15].
CONCLUSION
In this paper, we demonstrated how to improve the robustness of deep neural networks while quanti- fying predictive uncertainty by exploiting the implicit regularization of (stochastic) gradient descent. We rigorously characterized this implicit bias for an overparametrized linear model and showed that approach is equivalent to generalized variational inference with 2- Wasserstein regularizer at reduced computational cost. We demonstrated the importance of parameterization and how it im- pacts the inductive bias via the initialization thus conferring desirable properties such learning rate transfer. Lastly, we empirically demonstrated competitive performance with state-of-the-ae art methods for Bayesian deep learning on a set of in- and dittt---iistibttionn benchmarks with minimal computational overhead over standard deep learning. In principle, our approach is not restricted to Gaussian variational families and should seamlessly extend location-scale families, which could further improve performance. Finally, it would be interesting to explore connections between Im- plicit Bias VI and Bayesian deep learning in function-space [e.g., 37, 54, 74-76).
ACKNOWLEDGMENTS
JW, BC, JM and JPC are supported by the Gatsby Charitable Foundation (GAT3708). the Simons Foundation (542963), the NSF AI Institute for Artificial and Natural Intelligence (ARNI: NSF DBI 2229929) and the Kavli Foundation. This work used the DeltaAl system the National Center for Supercomputing Applications through allocations CIS250340 and CIS250292 from the Advanced Cyherinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which supported by U.S. National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and 138299... The authors would like to thank Hanna Detki for valuable input, which significantly improved this paper.
REFERENCES
[I] M. Goldblum, M. Finzi, K. Rowan, and A. G. Wilson. "The No Free Lunch Theorem, Kol- mogorov Complexity, and the Role of Inductive Biases in Machine Learning.. In: Interna- tional Conference on Machine Learning (ICML). 2024. DOI: 10. 48550/ arxiv 2304. 05366 (cit. on p. 1).
[2] M. S. Nacson, Mulayoff, G. Ongie, Michaeli, and D. Soudry. "The Implicit Bias of Minima Stability in Multivariate Shallow ReLU Networks". In: International Conference on Learning Representations (ICLR). 2023. DOE: 10. 48550 / wat.ii.iiivv.... 2306 17499 (cit. on
p. 1).
[3] R. Mulayoff, T. Michaeli, and D. Soudry. "The Implicit Bias of Minima Stability: A View from Function Space". In: Advances in Neural Information Processing Systems (NeurIPS).
2021. URL: https / proceedings neurips cc paper / 2021 / hash (e4aaaee383ed555 wwtrddccee1bbbcbh7cc
Obwccbbbbbbooc (rbboccb9422222 7942=279 -bst-bbttrattt... html (cit. on p. 1).
[4] D. Soudry, E. Hoffer, M. S. Nacson, S. Gunasekar, and N. Srebro. "The Implicit Bias of Gradient Descent on Separable Data". In: Journal of Machine Learning Research (JMLR) (2018). DOI: 10. 48550/ /a00//axivvi 1 1710.110 10345 (cit. pp. 1, 2, 6, 8, 23, 24, 28, 32).
10

---

Published as a conference paper at ICLR 2026
[5] S. Gunasekar, J. Lee, D. Soudry, and N. Srebro. "Characterizing Implicit Bias in Terms of Optimization Geometry". In: International Conference on Machine Learning (ICML). 2018. DOL: 10. 48550/ aatoo/rriiv.... w8w.v1..2.... 0 (cit. pp. 1, 4, 6, 23).
[6] M. S. Nacson, D. Lee, S. Gunasekar, H. P. Savarese, N. Srebro, and D. Soudry. "Con- vergence of Gradient Descent on Separable Data". In: International Conference on Artificial Intelligence and Statisties (AISTATS). 2019. DOr: 10. 48550 arxiv 1803. 01905 (cit.
on pp. 1, 6, 8).
[7] M. S. Nacson, N. Srebro, and D. Soudry. "Stochastic Gradient Descent on Separable Data: Exact Convergence with a Fixed Learning Rate". In: International Conference on Artificial Intelligence and Statistics (AISTATS). 2019. DOE: 10 4.550/550// aawpvaaiivvv/..//iiii 18060000000000000000 01796 (cit.
on pp. 1, 2, 6, 8).
[8] B. Wang. Q. Meng. H. Zhang, R. Sun, W. Chen, Z.-M. Ma, and T.-Y. Liu. "Does Momentum Change the Implicit Regularization on Separable Data?" In: Advances in Neural Information Processing Systems (NeurIPS) (2022) (cit. on pp. 6, 8).
[9] H. Jin and G. Montáfar. Implicit Bias of Gradient Descent for Mean Squared Error Re- gression with Two-Layer Wide Neural Networks. arXiv:2006.07356 [stat]. May 2023. DOT:
10. 185550/ar330
2006. 07356 (cit. on pp. 1, 6).
[10] H. Ravi, C. Scott, D. Soudry, and Y. Wang- "The Implicit Bias of Gradient Descent on Sep- arable Multiclass Data". In: Advances in Neural Information Processing Systems (NeurIPS).
2024. DOI: 10. 485550000 arr/vi/r......//// 24 01350 (cit. on pp- 1, 6, 8).
[11] T. Papamarkou, M. Skoularidou, K. Palla, L. Aitchison, J. Arbel, D. Dunson, M. Filippone, V. Fortuin, P. Hennig. J. M. Hernández-Lobato A. Hubin, A. Immer, T. Karaletsos, M. E. Khan, A. Kristiadi, Y. Li, S. Mandt, C. Nemeth, M. A. Osborne, T. G. J. Rudner, D. Rügamerr Y. W. Teh, M. Welling, A. G. Wilson, and R. Zhang. "Position: Bayesian Deep Learning Needed in the Age of Large-Scale AI". In: International Conference on Machine Learning (ICML).
2024. DOT: 10. 4 / arvvv. 24022 00809 (cit. on p. 1).
[12]
D. Tran J. Liu, M. W. Dusenberry, D. Phan, M. Collier, J. Ren, K. Han, Z. Wang. Z. Mariet, H. Hu, N. Band, T. G. J. Rudner, K. Singhal, Z. Nado, J. Amersfoort, A. Kirsch, R. Jenatton, N. Thain, H. Yuan, K. Buchanan, K. Murphy, D. Sculley, Y. Gal, Z. Ghahramani, J. Snoek, and B. Lakshminarayanan. Plex: Towards Reliability using Pretrained Large Model Extensions. July 15, 2022. DOI: 10. 48550/ wiw.iiivva... 2ww.xiv..220.... 07411 1. arXiv: 2207 0 07411 [cs]. URL: http: /ttp//aavvvv///////v///eee/a.ava org/ abs/ 02227000423 (visited on 05/16/2025) (eit. on p. 1).
[13] H. Ritter, A. Botev, and D. Barber. "Online Structured Laplace Approximations For Over- coming Catastrophic Forgetting". ". In: Advances in Neural Information Processing Systems (NeurIPS). 2018. DOI: 10. 4850 arrii.iiv..../vv. www.v.105..0 00005.07810 07810 (cit. on pp. 1, 6).
[14] Y. L. Li, T. G. J. Rudner, and A. G. Wilson. "A Study of Bayesian Neural Network Surro- gates for Bayesian Optimization". In: International Conference on Learning Representations
(ICLR). 2024. DOI: 10. 48550/ arxiv. 2305 20028 8 (cit. on p. 1).
[15] V. Fortuin. "Priors in Bayesian Deep Learning: A Review". In: International Statistical Re-
view 90.3 (2022), pp. 563-591. DOE: 10. 11 11l1l11111 insr 12502 (cit. on pp. 1, 10).
[16] P. Izmailov, S. Vikram, M. D. Hoffman, and A. G. Wilson. "What Are Bayesian Neural Net- work Posteriors Really Like?" In: International Conference on Machine Learning (ICML). 2021. DOI: 10. 48550/ attpaaxxvvvvv.... 21....144221 14421 (cit. on p. 1).
[17] B. Adlam, J. Snoek, and S. L. Smith. Cold Posteriors and Aleatoric Uncertainty. July 31,
2020. DOE: 10. 48550 / arxiv 2008 00029. 9. arXiv: 2008 00029 [stat]. URL: http://a /ttp///axivv.v// orq/ abs///ssssss 00029 (visited on 05/15/2025) (cit. on p. 1).
[18] т. Cinquin, A. Immer, M. Horn, and V. Fortuin. "Pathologies in priors and inference for Bayesian transformers". In: NewrIPS Bayesian Deep Learning Workshop. 2021. DOI: 10.
4 48550/arX0 2110. 04020 (cit. on p. 1).
[19] B. Coker, W. Bruinsma, D. R. Burt, W. Pan, and F. Doshi-Velez "Wide Mean-Field Bayesian Neural Networks Ignore the Data". In: International Conference Artificial In- telligence and Statistics (AISTA7S). 2022. DOT: 10. 18550/aaxxa 2202. 11670 1670 (cit. on
p. 1).
[20] A. Y. K. Foong, D. R. Burt, Y. Li, and R. E. Turner. "On the Expressiveness of Approximate Inference in Bayesian Neural Networks". In: Advances in Neural Information Processing Systems (NeurIPS). 2020. DOL: 10. 1..4. 48550/arma arvviv 1909. 00719 9 (cit. on p. 1).

---

Published as a conference paper at ICLR 2026
[21] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. "Understanding deep learning quires rethinking generalization". In: International Conference on Learning Representations
(ICLR). 2017. DO1: 10. 48550/ we.xxii..611. 03530 (cit. on pp. 2, 4, 6).
[22] G. Vardi. "On the Implicit Bias in Deep-Leaming Algorithms", In: Commun. ACM 66.6 (May 2023), pp. 86-93. DOI: 10. 1 145 0045/3557070 (cit. on pp. 2, 6).
[23] B. Vasudeva, P. Deora, and C. Thrampoulidis. Implicit Bias and Fast Convergence Rates for Self-attention. 2024, DOE: 10. 48550/ aaro0//rxiv../ 2402. 05738 (cit. on p. 2).
[24] D.J. C. MacKay. "A Practical Bayesian Framework for Backpropagation Networks". In: Neu- ral Computation 4 (1992). ISSN: 0899- 1530-888X. DO1: 10. www.//eeccnoooeeeeeooooeeoeooeee 1992....99...... 4.
3. 448 (cit. on pp. 2, 6, 8). [25] H. Ritter, A. Botev, and D. Barber. "A Scalable Laplace Approximation for Neural Net- works". In: International Conference on Learning Representations (ICLR). 2018 (cit. on
pp. 2. 6, 8).
[26] E. Daxberger, A. Kristiadi, A. Immer, R. Eschenhagen, M. Bauer, and Hennig. 'Laplace Redux - Effortless Bayesian Deep Learning". In: Advances in Neural Information Processing Systems (NeurIPS). 2021. DOI: 10.
2106. 6 14806 (cit. on pp. 2, 6, 8, 41).
[27] P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. "Averaging Weights Leads to Wider Optima and Better Generalization" In: Conference Uncertainty in Arti- ficial Intelligence (UAI). 2018. URL: https: / arxiv. org/ / abs / /8003 05407v3 (cit. pp. 2, 6).
[28] W. Maddox, Garipov, P. Izmailov, D. Vetrov, and A. G. Wilson. "A Simple Baseline for Bayesian Uncertainty in Deep Learning". In: Advances in Neural Information Processing Systems (NeurIPS). 2019. DOT: 10. 4 25550/rxivv
902. 0247 7 6 (cit. on pp. 2, 6, 8, 42).
[29] B. Lakshminarayanan, A. Pritzel, and C. Blundell. "Simple and Scalable Predictive Uncer- tainty Estimation using Deep Ensembles". In: Advances in Neural Information Processing Systems (NeurIPS), 2017. DOE: 10 48550 / arxiv 1rxiv.1612. 1612. 01474. URL: http://
arxiv. www.v.org//bba.oo.oo 612. 474 (cit. on pp. 2, 6, 42).
[30] A. Graves. "Practical Variational Inference for Neural Networks". In: Advances in Neural Information Processing Systems (NeurIPS). 2011. URL: https: : /pwppp/pppesppp/ppp/eeeeeeeeessa nips. cc/ Saper_filest www.fiilssspppppeeee httpp/apprr///// 2011 http///211hhhhhhhehheaee wwwaah/eebbcceeeelleeeeeeeeeeeee
Abstract. html (cit. pp. 2, 6, 8, 42).
[31] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. "Weight Uncertainty in Neu- ral Networks". In: International Conference on Machine Learning (ICML). 2015. DO1: 10.
4 (8850/aaaXaa 1 1505. 05424 (cit. on pp. 2, 6, 8, 42).
[32] K. Osawa, S. Swaroop, A. Jain, R. Eschenhagen, R. E. Turner, R. Yokota, and M. E. Khan. "Practical Deep Learning with Bayesian Principles" ", In: Advances in Neural Information Pro- cessing Systems (NeurIPS). 2019. DOL: 10. 48550/ /arvivriivvvv./////v/ 1906. 02506 (cit. on pp. 2, 6).
[33] Y. Shen, N. Daheim, Cong, P. Nickl, G. M. Marconi, C. Bazan, R. Yokota, I. Gurevych, D. Cremers, M. E. Khan, and T. Möllenhoff. "Variational Learning is Effective for Large Deep Networks". In: International Conference on Machine Learning (ICMI.). 2024. DOI:
10. 485550/axxii 24 22.22... . 176641 (cit. on pp. 2, 6).
[34] Zellner. "Optimal Information Processing and Bayes's ' Theorem". In: The American Statis- tician 42.4 (1988), pp. 278-280. DO1: 10. 2307 2685143 (cit. on p. 3).
[35] P. G. Bissiri, C. Holmes, and S. Walker. "A General Framework for Updating Belief Distri- butions". In: Journal of the Royal Statistical Society: Series (Statistical Methodology) 78.5 (Nov. 2016), pp. 1103-1130 ISSN: 1369-7412, 1467-9868. DOI: 10. 11w1/1111/// htt.i//rrbb..... 12158 (cit. on p. 3).
[36] J. Knoblauch, J. Jewson, and T. Damoulas. "An Optimization--entric View on Bayes' Rule: Reviewing and Generalizing Variational Inference". In: Journal of Machine Learning Re- search (JMLR) 23.132 (2022), pp. 1-109. ISSN: 1533-7928. URL: http: : / jmlr org/ papers/ //13/223//99/// 9- 1047. .1mml (cit. on pp. 3, 6).
[37] V. D. Wild, R. Hu, and D. Sejdinovic. "Generalized Variational Inference in Function Spaces: Gaussian Measures meet Bayesian Deep Learning" In: Advances in Neural Information Pro- cessing Systems (NeurIPS). Oct. 2022. DOI: 10. 48550/ / arXiv. 2205 06342 (cit. on
pp. 3, 6, 10).
12

---

Published as a conference paper at ICLR 2026
[38] S. Farquhar, L. Smith, and Y. Gal. "Liberty or Depth: Deep Bayesian Neural Nets Do Not Need Complex Weight Posterior Approximations". In: Advances in Neural Information Pro- cessing Systems (NeurIPS). 2020. DOL: 10. 48550 http/aiiv..00// 2002. 03704.
URL: http:
/ //arxiv.. .rg/ abs / 2202 03704 (cit. on pp. 4, 6).
[39] M. Sharma, S. Farquhar, E. Nalisnick, and T. Rainforth. "Do Bayesian Neural Networks Need
To Be Fully Stochastic?" In: International Conference on Artificial Intelligence and Statistics www.i.ivvii....................e
(AISTATS). 2023, DOI: 10.
91 (cit. on pp. 4, 6, 10).
[40] B. Hanin and M. Sellke. Approximating Continuous Functions by ReLU Nets of Minimal Width, arXiv:1710.1 http:/ /////arss /arxivv..... /t.g///gg//ooo abs / 11100 11278 (cit. on p. 4.4).
[stat]. Mar. 2018, DOI: 10. 0..48555//22 /...0///xi..... 1710 11278. URL:
[41] G. Yang, E. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. "Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperpa- rameter Transfer". In: Advances in Neural Information Processing Systems (NeurIPS). 2021.
DOL: 10..1..48550 48540/555400 aarxiv...vv/ 2203 2003...3222 0346 (cit. on pp. 5, 6, 36, 37).
[42] A. Bhattacharya, A. Linero, and C. J. Oates. "Grand Challenges in Bayesian Computation".
In: Bulletin of the International Society for Bayesian Analysis (ISBA) 31.3 (Sept. 2024). DOI:
10. 48550/ a arxivv 2 2410. 004 96496 (cit. on 5)5
[43] G. Yang and E. J. Hu. "Tensor Programs IV: Feature Learning in Infinite-Width Neural Net- works". In: International Conference on Machine Learning (ICML). 2021. DOI: 10. 48550/
arXiv. 2011 111.11.1... 14522 (cit. on pp. 5, 6, 34).
[44] G. Zhang, S. Sun, D. Duvenaud, and R. Grosse. Noisy Natural Gradient as Variational Infer-
ence. Feb. 26, 2018. DOL: 10, 48550/ arxiv..iiivvvvvvi 1712. 02090, arXiv: 1712. 02390 [cs]. URI.: http: / aarii org. abs 7121. 02390 (visited on 05/15/2025) (cit. on p. 6).
[45] M.-N. Tran, N. Nguyen, D. Nott, and R. Kohn. Bayesian Deep Net GLM and GLMM. May 25,
2018. DOr: 10. 48550 /
1805 10157. arXiv: 1805 10157 [stat]. URL: http://arxils org ht///b///1100 1 10157 (visited on 05/15/2025) (cit. on p. 6).
[46] C. Louizos and M. Welling. Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors. June 23, 2016. DOr: 10. 48550 / arxiv 1603, 04733. arXiv:
1 1603. . 04733 [stat]. URL: http: : / aaxiv. org/ / abs 1603 04733 (visited on 05/15/2025) (cit. on P. 6).
[47] A. Mishkin, F. Kunstner, D. Nielsen, M. Schmidt, and M. E. Khan. SLANG: Fast Structured
Covariance Approximations for Bayesian Deep Learning with Natural Gradient. Jan. 12,
1811.111111 04504. arXiv: 1811. 04504 [cs]. URL: http: / / //axxiv.... org/ abs abbs/18111 04504 (visited on 05/15/2025) (cit. on p. 6). [48] J. Harrison, Willes, and J. Snoek. "Variational Bayesian Last Layers". In: International Conference on Learning Representations (ICLR). Apr. 2024. DOI: 10. 48550 / arxiv.
2019. DOL: 10. 10.48888888 008.4455555555
2 2404. 11599 (cit. on p. 6).
[49] J. Z. Liu, Z. Lin, S. Padhy, D. Tran, T. Bedrax- and B. Lakshminarayanan. "Sim- ple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness". In: Advances in Neural Information Processing Systems (NeurIPS). Oct. 2020.
DOI: 10. 13550/arXivv 2006. 10110000 108 (cit. on p. 6).
[50] M. E. Khan, A. Immer, E. Abedi, and M. Korzepa. "Approximate Inference Turns Deep Networks into Gaussian Processes".. In: Advances in Neural Information Processing Systems (NeurtPS). 2019. DOI: 10. 4 48550/ / ariiv 1906. 01930 (cit. on p. 6).
[51] A. Immer, M. Korzepa, and M. Bauer. "Improving predictions of Bayesian neural nets via lo- cal linearization". In: International Conference on Artificial Intelligence and Statistics (AIS- TATS). 2021 (cit. on p. 6).
[52] E. Daxberger, E. Nalisnick, J. U. Allingham, J. Antorán, and J. M. Hernández-Lobato. "Bayesian Deep Learning via Subnetwork Inference". In: International Conference on Ma- chine Learning (ICML). 2021. DOI: 10. / arxiv v. 2010. 1 14989 (cit. on p. 6).
[53] A. Kristiadi, A. Immer, R. Eschenhagen, and Fortuin. "Promises and Pitfalls of the Lin- earized Laplace in Bayesian Optimization". In: Advances in Approximate Bayesian Inference (AABI). 2023. DOE: 10. 48550/ 48550/arxii. arxiv. 2304. 08309 (cit. on P. 6).
[54] T. Cinquin, M. Pförtner, V. Fortuin, P. Hennig, and R. Bamler. "FSP-Laplace: Function- Space Priors for the Laplace Approximation in Bayesian Deep Learning". In: Advances in Neural Information Processing Systems (NeurIPS). Oct. 2024. DOI: 10. . 48550 / arxiv. 2407. 13711. URL: http::/arre wrg/ /207 13711 (cit. on pp. 6, 10).
13

---

Published as a conference paper at ICLR 2026
[55] S. Fort, H. Hu, and B. Lakshminarayanan. Deep Ensembles: A Loss Landscape Perspective.
June 25, 2020. DOL: 10. 48550 / artii. 1912. 02757. arXiv: 1 1912. 02757 [stat]. URL: http: http:://axxiv.rree httpg//bs//111// 02757 (visited on 05/15/2025) (cit. on p. 6).
[56] A. G. Wilson and P. Izmailov. "Bayesian Deep Learning and a Probabilistic Perspective of Generalization". In: Advances in Neural Information Processing Systems (NeurIPS), 2020.
DOI: 10. 48550/arxiv 2002 08791 (cit. on p. 6).
[57] V. D. Wild, S. Ghalebikesabi, D. Sejdinovic, and J. Knoblauch. "A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods". In: Advances in Neural Information Processing Systems (NeurIPS). 2023. DOE: 10. 48550 / arxiv 2305 15027 (cit. on
p. 6).
[58] T. Abe, F. K. Buchanan, G. Pleiss, R. Zemel, and J. Cunningham. "Deep Ensembles Work, But Are They Necessary?" In: Advances in Neural Information Processing Systems (NeurlPS). 2022. DOI: 10. 4 8550/ arriiv. 2202. 0985 (cit. on p. 6).
[59] N. Dem,, J. P. Cunningham, and G. Pleiss. Theoretical Limitations of Ensembles in the Age of Overparameterization. mxXiv:2410.1220 [stat]. Oct. 2024. DOT: 10. 48550 / arxiv. 2410. 1 6201 (cit. on p. 6).
[60] C. Mingard, G. Valle-Pérez, J. Skalse, and A. A. Louis. "Is SGD a Bayesian sampler? Well, almost." " In: Journal of Machine Learning Research (JMLR) (2020) (cit. on p. 6).
[61] J. A. Lin, J. Antorán, S. Padhy, D. Janz, J. M. Hernández-Lobare and A. Terenin. "Sampling from Gaussian Process Posteriors using Stochastic Gradient Descent". In: Advances in Neural Information Processing Systems (NeurIPS). 2023. DOI: 10. 20550//rrii..o 2306.11589
(cit. on p. 6).
[62] J. Lee, L.. Xiao, S. Schoenholz, Y. Bahri, R. Novak, Sohl-Dickstein, and J. Pennington. "Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent". In: Journal of Statistical Mechanics: Theory and Experiment 2020.12 (2020). DOE: 10. 1088/ 1742- 54 68/ / abc (cit. p. 6).
[63] J. Lai, M. Xu, R. Chen, and Q- Lin. Generalization Ability of Wide Neural Networks on R. Feb. 12. 2023. DO1: 10. 48550/ / arxiv. 2tl.iii...02...2 2302. 05933. arXiv: 2302. 05933 [ [stat]. URL: http: //ap:ii///ivvv//////a///a org/ abs / 2302. 05933 (visited on 05/15/2025) (cit. on p. 6).
[64] G. Yang. "Tensor Programs I: Wide Feedforward or Recurrent Neural Networks of Any Ar- chitecture are Gaussian Processes". In: Advances in Neural Information Processing Systems (NeurIPS). 2019. DOI: 10. 4 88500 / arxiv. 1910.12411 (cit. on p. 6).
[65] G. Yang. Tensor Programs II: Neural Tangent Kernel for Any Architecture. 2020. DOE: 10.
4 8 15500/axxiv 2006 6 14548 (cit. on P. 6).
[66] A. Jacot, Gabriel, and C. Hongler. "Neural Tangent Kernel: Convergence and Generaliza-
tion in Neural Networks". In: Advances in Neural Information Processing Systems (NeurIPS).
2018. DOI: 10. 585500 arxiv. 1806. 07572 (cit. on p. 6).
[67] V. Muthukumar, A. Narang, V. Subramanian, M. Belkin, D. Hsu, and A. Sahai. "Classification vs regression in overparameterized regimes: Does the loss function matter?" In: Journal of Machine Learning Research (JMLR) (Oct. 2021). DOr: 10.11.8800 48550/ /0155///r/iv..// 2005. 08054 (cit. on p- 7).
[68] D. Hsu, V. Muthukumar, and J. Xu. On the proliferation support vectors in high dimensions. 2022. DOI: 10. 48550/ arxiv..... 2009 10670. URL: http: : / arxiv. org abss / 2009. 10670 (cit. on p. 7).
[69] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger. "On Calibration of Modern Neural Net- works". In: International Conference on Machine Learning (ICML). 2017. DOI: 10. 48550/ arXiv. 1706. 04599 (cit. on pp. 8, 33, 41).
[70] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. "Gradient-based learning applied to doc- ument recognition". In: Proceddings of the IEEE 86.11 (Nov. 1998), pp. 2278-2324. ISSN: 1558-2256. DOI: 10. 1109/ / 726791 URI.: https: http//eeeeeeeeeeeeeee/e Seeeeieee feee. org/ document/7267 91 (cit. on pp. 8, 40).
[71] K. He, X. Zhang, S. Ren, and J. Sun. "Deep Residual Learning for Image Recognition". In: Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), IEBE, 2016, pp. 770-778. ISBN: 978-1-4673-8851-1.
Dor: 10- 1109 / CVPR.
/ttp...og//obcceeee////////000/ 2016. 90, URL: http: http://eeeeeporreeoe wwwllioreeeeeeee.e ieee http////ommmuncecc/e /7780499 (cit. on
pp. 8. 41).
14

---

Published as a conference paper at ICLR 2026
[72] N. Mu and J. Gilmer. "MNIST-C: A Robustness Benchmark for Computer Vision". In: ICML Workshop on Uncertainty and Robustness in Deep Learning June 2019. DOE 10. 48550/ arxiv.1906. 02337. URL: attp://arme org/ abs 1906 02337 (cit. pp. 8, 39).
[73] D. Hendryeks and T. Dietterich. "Benchmarking Neural Network Robustness to Common Corruptions and Perturbations". In: International Conference on Learning Representations (ICLR). 2019. DOT: 10 48550 arxiv.. 1903 12261. URL: http: / wtwpiv.ivvv.///////..e///.aa org/
abs/1 903. 12261 (cit. on pp. 8, 39).
[74] D. R. Burt, S. W. Ober, A. Garriga-Alonso, and M. van der Wilk. Understanding Variational Inference in Function-Space, Nov. 2020. DOE: 10. 48550 /tt//aaiiivvv..e.iiv 2011. 09421 (cit. on
p. 10).
[75] S. Qiu, T. J. Rudner, S. Kapoor, and A. G. Wilson. "Should We Learn Most Likely Func- tions or Parameters?"
Advances Neural Information Processing Systems (NeurIPS). 2023. DOI: 10. 48550/ aarviviivv//// 2312.31111222 11111.1550 15990 (cit. p. 10).
[76] T. G. Rudner, Z. Chen, Y. W. Teh, and Y. Gal. "Tractable Function-Space Variational Infer- ence in Bayesian Neural Networks". In: Advances in Neural Information Processing Systems (NeurIPS). 2023. DOI: 10. 85500 aarraiiiii../..////// 22122.232... 17199 (cit. on p. 10).
[77] S. P. Boyd and L. Vandenberghe. Conver optimization. Cambridge University Press, 2004.
ISBN: 978-0-521-833778 (cit. on p. 17).
[78] Y. Nesterov. "A method for solving the convex programming problem with convergence rate
O )" In: Dok! Akad Nauk SSSR 269 (1983), 543 (cit. on 19).
[79] B. T. Polyak. "Some methods of speeding up the convergence of iteration methods". In:
USSR Computational Mathematics and Mathematical Physics 4.5 (1964), pp. 1-17. DOI:
10. 101 00166/0055555 (64) 90137-5 (cit. on p. 19).
[80] A. Krizhevsky et al. Learning multiple layers of features from tiny images. Tech. rep. 2009 (cit. on p. 39).
[81] Y. Le and X. Yang. "Tiny ImageNet Visual Recognition Challenge", In: Stanford CS 23/N (2015). URL: http: : /ttp//cs..nnn/////// stanford edu edu/tiny-- www.iiiangeenaeeeeeeeeeeeeeeee 200 zip (cit. p. 39).
[82] P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L.. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He. Accurate, Large Minibatch SGD: Training ImageNet in Hour. Tech. rep. 2018.
URL: http: / / aarxiv.... org/ abs 1766. 02677 (cit. on p. 40).
[83] S. L. Smith and Q. Le. "A Bayesian Perspective on Generalization and Stochastic Gradient Descent". In: International Conference on Learning Representations (ICLR). 2018. DOI: 10. 48550/arXiv 1 1710 06451 (cit. on p. 40).
[84] S. L. Smith, P.-. Kindermans, C. Ying, and Q. Le. "Don't Decay the Learning Rate, In- crease the Batch Size". In: International Conference on Learning Representations (ICLR). 2018. DOI: 10. 48 3550/arXiv... Siiii.17111 1711.1111 00489 (cit. on p. p0).
[85] T. maintainers and contributors. Torch PyTorch's Computer Vision library. https: 2016 (cit. on p. 41).
/ / giihhuu. com wtt.cooo////.eeeeeec wwon/pyyoocc//iiiii
15

---

Published as a conference paper at ICLR 2026
SUPPLEMENTARY MATERIAL
This supplementary material contains additional results and proofs for all theoretical statements. References referring to sections, equations theorem-type environments within this document prefixed with 'S', while references to, or results from, the main paper are stated as is.
S1 Theoretical Results
16
$1.1
Overparametrized Linear Regression
17
$1.11
Characterization of Implicit Bias (Proof of Theorem
17
$1.1.2
Non-Asymptotic Error Analysis
20
S1.1.3
Connection to Ensembles
22
S1.2
Binary Classification of Linearly Separable Data
23
$1.2.1
Preliminaries
24
S1.2.2
Gradient Flow for the Expected Loss
24
$1.2.3
Complete Proof of Theorem 2
27
S1.3 NLL Overfitting and the Need for (Temperature) Scaling
32
82 Parametrization, Feature Learning and Hy Hyperparameter Transfer
33
$2.1
Definitions Stability and Feature Learning
................................
34
S2.2 Initialization Scaling for Linear Network
................................
34
S2.3 Proposed Scaling
................................
36
S2.4 Details on Hyperparameter Transfer Experiment
37
S3 Experiments
38
S3.1 Setup and Details
38
$3.1.1
Datasets
................................
39
S3.1.2 Metrics
39
$3.2
Time and Memory-Ffffiient Training...
40
S3.3
In- and Out-of-distriban Generalization
40
$3.3.1
Architectures, Training, Methods
40
$3.3.2
In-Distribution Generalization and Uncertainty Quantification
43
S3.3.3
Robustness to Input Corruptions
44
S3.3.4
Comparison to Generalized VI with 2- Wasserstein Regularization
44
S1 THEORETICAL RESULTS
Lemma S1
Let q(w) N(w: μ.. E), p(w)= N(w; nox Eo) such that μ. po ERP. E,En E RPP posi- tive semi-definite and let VA E PNN, Vn (P--N) be matrices with pairwise orthonormal columns that together define an orthonormal basis of R.,, i.e. for V = [VA VBl it holds that VVT VVT=VTV- = and span(V) = RR.. Assume further that
-VEV=
(S11)
then the squared 2- 2-Wasserstein distance is given by
W 1|V vo W3(N(V V.V VH + CC, (S12)
where the constant C' is independent of (p, E).
Proof. Consider the matrix
V EV = ONXN V
EVE
V EVB
Since VTVV is symmetric positive semi-definite, its off-diagonal block V ΣV, satisfies
(I 00 -00)
= VVVV= =0
16

---

Published as a conference paper at ICLR 2026
by Boyd and Vandenberghe [A5.5, 77]. Therefore, we have
VVEV=
ONXN
V
VVVV
ONXN 0 NP-NN
0 (PP-NN
(S13)
V EVA
V
EVB
V
The squared 2-Wasserstein distance between ((w) and p(w) is given by (2. p) = |1 - Hell + tr( - 2 En).
For the squared norm term holds by unitary invariance of ||-112 that
V ( _o) m VB (
|| 140 12
VV ( m
VV vimol
V
V 140 1
Now for the trace term have that
tr(VV - 2( Eo))
=tr(V EV) - 2tr
VA) tn = "tr + tr 2 [
-Ir(V[VV) + (VV
(S14)
where we used Eq. (S11) and = denotes equality up to constants independent of (n, E).
Now by Eq. (S13), we have that Σ = VBMVE for M = VEEVE and its unique principal square root given by E V VØJ since
(V (v (V.M. Va VeM- also holds that the unique principal square root
( V
since direct calculation gives
E VB //
V Eo V. 3
Therefore we have that
tr(v ".tp::";"""..nt...o....././/..., -
=
Putting all together we obtain
W)( = V μ
VI p0
V 14 -V3 Ir( IV VVn) +r 221
2
- 2tr(( V E. VB.)) ||V k moll W2( V.EVV), V1L.o
|
V
H4
1 + VITT 1 ViJ
+ + +
which completes the proof.
S1.1 OVERPARAMETRIIZE LINEAR REGRESSION
S1.1.1 CHARACTERIZATION OF IMPLICIT BIAS (PROOF OF THEOREM 1)
Theorem (Implicit Bias in Regression)
Let for x x w be an overparametrized linear model with P > N. Define Gaussian prior p(w) = (w:: 14n SoST) and likelihood p(y w) = NNy; Suu (X), a2 and assume a varia- tional family qo (w(w) = (ww; μ., SST) with 0 = (μ, S) such that E IR'' and S E R PxR where R < P. If the learning rate sequence (ne); s chosen such that the limit point @GD lim DGD identified by gradient descent, initialized at 0 = (Ho. So). is a (global) minimizer of the expected log- likelihood E(O), then
OGD
arg min 0=(u.S)
W2(
0Earg E(0)
Further, this also holds in the case of stochastic gradient descent and when using momentum.
17

---

Published as a conference paper at ICLR 2026
Training
Prediction
Training
Prediction
Expected Loss
Loss
Weights
Weight Distributions ((w)
X
Initialization 20
Initializution wa
Init. Prior 9A
Init. Prior qD
NN trained with explicit regularization.
(b) BNN trained with explicit regularization.
Figure S1: Implicit regularization in standard neural networks versus in probabilistic networks. Left panels: neural network trained without explicit regularization can converge to different global minima of the loss. Optimization of the weights will implicitly regularize towards one or the other. Right panels: Analogously, there are multiple distributions over neural networks that are global min- ima of the expected loss. Optimization of the distribution over the weights will implicitly regularize towards one or the other. Our approach uses this implicit regularization instead of an explicit regu- larization prior.
Proof. Let 0. (n., S.) be minimizer of O()). By assumption it holds that the expected negative log-likelihood is equal to the following non-negative loss function up to an additive constant:
O Eg (1m) (E(3) fw (X)) B (olog p(y w))
2x By (1ly Xwl12)
202 (lly Xul2 tx(xX 2 0,
where E = S and nen-negativity follows from E being symmetric positive semi-definite. There- fore any (global) minimizer 0, (u,, Σ.) necessarily satisfies
1 M 11 0,
(S15)
(r XXXXI1 0.
(S16)
Let V Viange Vmui E PxP be the orthonormal matrix of right singular vectors of X = UAVT, where Vaange E PxN and V E PPx(P-N), Since X E NNP and we are in the overparametrized regime, i.e. P > N, the optimal mean parameter decomposes into the least- squares solution and null space contribution
μ. Viange U. Vaull X y + V
(S17)
Furthermore, it holds for positive semi-definite E E RPxP that
0 tr( xr(EEx)) I Lr(AV
VT EV Vange range
ANNN
ir([ANXN
-tr(ANxN VT EV Vannge A range
EV Σ V range Virngelis
where X2 are the squared singular values X, which are strictly positive since rank (X) Therefore using Equation (S16) any global minimizer necessarily satisfies [Vringe Σ, Vrangelis for E {1...., N). Now since V E. Vmage is symmetric positive semi-definite and its diagonal zero, so is its trace and therefore the sum of its non-negative eigenvalues is necessarily zero. Thus all eigenvalues are zero and therefore
V EV range Vangge 0.
(S18)
18

---

Published as a conference paper at ICLR 2026
Now by Lemma SI have that the squared 2- 2-Wasserstein distance between ge, (w) N(w; p., E.) and the initialization p(w) NV NN: po En) is given up to a constant independent of (m,, E.) by
W W2( _ || | rangolt VT
Vinnge (N (AMM(VITTL (Vmnn V D. Vma)) V || V Mo + ( (V mill/h,, V E, Vnan), , (V mull/40, V Eo Vmai))
.N (V Dul/ V Eo Vann)
VVVVVVV W2 / /4+; V Σ, E, .N (V Mall V Eg Vnuti)
Therefore among variational distributions ge, with parameters 0, that minimize the expected loss ((0), any such 0, that minimizes the squared 2- Wasserstein distance to the prior satisfies
(VT
mulllbe, VT E, Vmn)) = (V VTT Zo Vmll)..
(S19)
11
M
(Stochastic) Gradient Descent It remains to show that (stochastic) gradient descent identifies a minimum of the expected loss {()), such that the above holds. By assumption have for the loss on a batch X5 of data that
{(0)= E (w))
w (l(yb: Se (x))) ( log w))
20 (1 X 2/2 + tr tr(x,EX)))
Therefore, at convergence of (stochastic) gradient descent the variational parameters 0xc (Moo: are given by
po
1/ (0:1)) Mo X (yb XoH(-1)
t=1
as well as
Sx = So
11 V se (0.
So Σ
1tt x X
0²
and therefore
00
200
Voul
ViT
44
17t VT XT
X
V
440
Erange( (X]
S S20 = VT Su
Tt
VT
x X
VT So
columns Crange (X!)
where we used continuity of linear maps between finite-dimensional spaces. It follows that
Mo. V Eno Viuil
V Vnall V So Vmull
V Eo Vnull-
Therefore any limit point of (stochastic) gradient descent that minimizes the expected log- likelihoooo also minimizes the 2- Wasserstein distance to the prior, since Ox satisfies Equation (S19).
Momentum In case we are using (stochastic) gradient descent with momentum, the updates are given by
P++1 14 + Ap: T/ V (0 + A0,
(S20)
S S, + n AS, 17p sl, + a,)),
where
AAμ, AS, = 0, - 0:-1,
A0, =
AO = 0.
for parameters Te: a, 2 0, which includes Nesterov's 's acceleration (Yr = a;)) [78] and heavy ball momentum (or: 0) [79].
19

---

Published as a conference paper at ICLR 2026
To prove that the updates of the variational parameters are always orthogonal to the null space of X,, we proceed by induction. The base case is trivial since A = 0. Assume now that V Apt = and V AS, = O. then by Equation (S20), we have
V
mall A = m ( - p4) = Vinll Apt Th nall V f + a,20,) = 0
VIu AS mill (S++1 St) = V null AS, The null V s l (0) -,. o, =
where we used the induction hypothesis and the fact that the gradients are orthogonal to the null space shown earlier.
Therefore by the same argument as above we have that Ox computed via (stochastic) gradient descent with momentum satisfies Equation (S19), which directly implies Theorem 1.
S1.1.2 NON-ASYMPTOTIC ERROR ANALYSIS
Theorem S3 (Non-Asymptotic Error of Gradient Flow)
Let fw (x) x_ be linear model. Define prior p(w) Nww; po: SoST) and assume p(w). Further, define a variational distribution go(w)= = ((w: SST) with 0 = (p. S) such that μ Ip and S E Pxx where R S P. Let 0 (n(1), (())) be the variational parameters at time i 2 given by the gradient flow of the
noise-free observations y(.) fw(() for
expected loss
(tt) = Vo((((()
(S21)
initialized at (()) = (1o0, So). Then the expected squared error of the mean prediction
TP (suus)
(Jeest [a(ee()Ftest))
VVrrwwggro
(S22)
any test point Feest R.P In other words, assuming the training and test data drawn from the prior predictive, the predictive error of f.(e(() at any time t 2 0 is exactly quantified by the predictive uncertainty of the variational distribution, not only initialization and in the limit t + 00.
Proof. The dynamics of the variational parameters as defined by the gradient flow in Equation (S21) given by
(t) = - -V,, v,L(((()) xT - = - - w) = d z ((((t) - w),
$(t) s(((((()) --x
Since these dynamics are matrix differential equations, the mean and covariance parameters function of time are given by
(tt) Stt) = = xix Sy-
X X w),
(S23)
(S24)
Thus the expected predictive error time step > is given by
E
(Iivess
Su(ce) (Ietet))1?)
E
x ) (1ettest x To (tt)l1)) Browe
Ev
x w
x (w + te- xxxxXX xx'xx(om (H5O - -w)) w))113
[ X x'x xx(((u w)ll
Emp
where we used Equation (S23). We have since E[w)) noo that the above
tr Cov(w m40 C
X x 2o e xxxi)
-Ir =Ir(xTest -x -x'XI SoS. -x Censtt
20

---

Published as a conference paper at ICLR 2026
tr (eees SiS; Drest)
Var (Jun (2rest)
where we used Equation (S24) in the second-to- equality. This completes the proof.
Theorem S4 (Non- -Nym-Asympootic Error SGD)
Let (x) = x w be a linear model. Define prior p(w) = =N(w;; po: SoST) and assume noise-free observations y(-) (.) for w ~ p(w). Further, define a variational distribution ( (ww; μ. S with
(n, S) such that μ R2" and S RPPR where R S P. Assume the expected loss is given by {(O) Ego(w) (21l
and let 0 (u(t),((t)
be the variational parameters at step t of (stochastic) gradient descent with learning rate sequence (ne),, initialized (0) (Ho, So). Then the expected squared error of the mean prediction
E yroes
(Weeest (3ees()rsst)) Var (Irest)))
(S25)
at any test point Teest E Rº, In other words, assuming the training and test data are drawn from the prior predictive, the predictive error of fu(e(() at any optimization step is exactly quantified by the predictive uncertainty of the variational distribution.
Further, if the learning rate n S Xmmax (X. x.)
for all steps t. then
tr tr(Cov CoVVwaquunn (w)) S tr(Cov (w)),
(S26)
i.e. uncertainty about the parameters decreases monotonically during optimization.
Proof. The expected loss is given up to an additive constant by
ε(0)= Ego (w) (c(y, fw (x)))
2 (lly x x))).
Now let (X,, y:) be the minibatch at step 2 1. Then it holds that
fatte (Erest) Vurst ITent (a(t) w).
(S27)
Further, the mean parameters identified by SGD are given by
(t) μu(t 1) w ni V ((((t - 1))
cu(t - 1) - w XxX (Xefe(t - yt)
= plt - 1) - w x, x.(((tt (tluutt-) -1) --- w))
(I x. (x)( 1) ww)
II - ni XJ X, (m)()) m))
Bi w)
where we defined B, 11 j (I 1/ X] X;). The covariance parameters are given by
S(1) = S(I 1) - T/ (
S(I 1) - nx x,,,', X,S(L 1)
(I - x xx)(X 1)
II (I xx)s(())
3=1
B.S.
Therefore the predictive error step & {0, (.... given by
((,ra)) (leeeese SL(()) (Zteet))) = E
Veeee (1)
X w Loce
21

---

Published as a conference paper at ICLR 2026
Ew || (H((t) - -w)) -w)l3)
Ew
[1 (140 200 20))12)
We have since E(p -w) = 0, that the above
= tr Treat 2 T B, Cov(w m0)
tr(wxlest B,S Bi Frest)
tr(acess ((1) Brest)
Vatw~quco (f (tess)))
This proves Equation (S25).
X, X2 has a spectrum in the interval To prove the second statement, we begin by showing that 17: Xi [0, 1]. We have by Weyl's theorem, since I and Ct+1 == -1/t+1 X 2+1 Xt+1 are hermitian, that
Ap(I + Ci+1) Ap(() + Amax (C41)) M mmxx (X 1 x S Ap (I + C(41) < mm11mmia (XT (141X1+41
Xx(I) Amin(C
(
Amax X(+1)
X(+1) Ap + C++1) 2
Amax
0 Ap + C.+1) S
where we used the assumption on the learning rate that Vt q I Xxxx)) Now by von Neumann's 's trace inequality, it holds that
Lr( Cov 90-0044))) (w)) = tr((I 11+1 X 1+1 X SIS - m111 X 41 X
ttr(S,s (I Mt X XI+I((I M111 X X(+1))
(p ((p)))))p((())) - X x
x
Σ Ap (.S.ss))pp((1 M - X,+1))²
Σ (S.S.)
tr(Cov worge (w)).
S1.1.3 CONNECTION TO ENSEMBLES
Proposition S1 (Connection Ensembles)
Consider an ensemble of overparametrized linear models fw((x) = xTw initialized with weights drawn from the prior wh") ~ N(w: Hoo S.ST). Assume each model is trained independently to convergence via (SJGD such that w. arg min fx))). Then the distribution over the weights of the trained ensemble qEn((W) is equal to the variational approximation qo, (w) learned via (SJGD initialized at the prior hyperparameters 0n- = (po, So), i.e.
PEns (w) = goco (w).
(S28)
Proof. The parameters wh of the (independently) trained ensemble members identified via (stochastic) gradient descent are given by
w arg minl|wr wn l|2 WEF
22

---

Published as a conference paper at ICLR 2026
where / = (w E IF (x( Xw= = y) is the set of interpolating solutions [5, Sec. 2.1]. Since can write equivalently via the minimum norm solution and arbitrary null space contribution, s.t. F ww = X + Waull Wnull E null(X), we have
X y + arg min |walll
(wo
(1)
x XTy)lla
(mnEnal(
- y + +pro]mal(())
(6) W
X
Crange(X"]
where we used the characterization of an orthogonal projection onto a linear subspace as the (unique) closest point in the subspace. Finally, we use that the minimum norm solution is in the range space of the data and rewrite the projection in matrix form, s.t.
x y + P (6)
(6) Therefore the distribution over the parameters w of the ensemble members computed via (S)GD with initial parameters wo ~ N N(w; noo ( is given by
qEns (w)) = = w; y + Po
S PT
mull
=
SE
Now the expected negative log-likelihood of the distribution over the parameters of the trained en- semble members qEns (w) with hyperparameters Ogns (Meiss SEin) is
((0mm.)
ie (lly X Mans 112 tr SEn Ens xT)
ST
20
and therefore Ons a minimizer of the expected log- log-likelihood. Further it holds that
V (oo) V
yu M V nall( (o) T V S.S.S. V.ail and thus by Equation (S19), the distribution of the trained ensemble parameters minimizes the 2- Wasserstein distance to the prior distribution, i.c.
V Eo V
arg min We(g(w))NNur po: Ea)). (w:;)=((
VEns
Combining this with the characterization of the variational posterior Theorem proves the claim.
S1.2 BINARY CLASSIFICATION OF LINEARLY SEPARABLE DATA
this subsection we provide proofs of claims from Section 4.2. We begin with presenting some preliminary results from Soudry et al. [4] which will be used throughout the proof. Next, we will analyze the gradient flow of the expected loss. We extend the results for the gradient flow to gradient descent and derive the characterization of the implicit bias, completing the proof of Theorem 2.
Theorem (Implicit Bias in Binary Classification)
Let fw (x) x w be an (overparametrizee) linear model and define Gaussian prior p(w) = N po; S. Assume variational distribution (o(w) N (un; μ., SST) over the weights w E R with variational parameters f (H,S) S) such that S E PxR and R < P. Assume we are using the exponential loss E(u) exp(-up u) and optimize the expected empirical loss ?(0) via gradient descent initialized at the prior, i.e. Oo (Ho. So), with sufficiently small learning rate n, Then for almost any dataset which linearly separable (Assumption 1) and for which the support vectors span the data (Assumption 2), the rescaled gradient descent iterates (rGD)
g
(μr GD
STGD)
GD + P Ho SCD
Togstt
converge to limit point @D lim o for which it holds that
gL E arg min W p))
(10)
0=(u.9) 6ce.
A0 and Yn Vax (x 0} consists of mean parameters which, if projected onto the training data, are equivalent to the L2 max margin vector and covariance parameters such that there is uncertainty at training data.
where the feasible set θ {1,, S)
P Prange
23

---

Published as a conference paper at ICLR 2026
S1.2.1 PRELIMINARIES
Recall that the expected loss is given by
{()) - -m= (go(w) (1)))
(S29)
and specifically, for the exponential loss, we have
0)) (1,, S) Σm exp (- f æ.))
(S30)
Throughout these proofs, for any mean parameter iterate μ4, define the residual
Pt 14 ib
log ib
(S31)
where v is the solution to the hard margin SVM, and w is the vector which satisfies
tn nexp where weights are defined through the KKT conditions on the hard margin SVM problem, i.c.
(- x. ii)
Onz
(S32)
i Σ On
(S33)
RES
In Lemma 12 (Appendix B) of Soudry et al. [4], is shown that, for almost any dataset, there are no more than P support vectors and on 0, Vn S. Furthermore, we denote the minimum margin to a non-support vector as:
min x;, J i nés
1.
(S34)
Finally, we define Ps PxP the orthogonal projection matrix the subspace spanned by the support vectors, and Ps = Ps as the complementary projection.
$1.2.2 GRADIENT FLOW FOR THE EXPECTED LOSS
Similar as in Soudry et al. [4], we begin by studying the gradient flow dynamics, i.e. taking the continuous time limit of gradient descent:
= =
(S35)
which can be written componentwise as:
N
j. = - pellpr S:) Σ
exp
M,
on + 2 2 on In
(S36)
n=1
N
5, -
Σ exp (-n In
T ar S.S.S.S xin ) a
(S37)
We begin by showing that the total uncertainty, as measured by the Frobenius norm of the covariance factor, is bounded during the gradient flow dynamics. To that end, derive the following dynamics:
d - = ir(ST = Σ di
1 S.ST x exp (1, on + 2
|. S 0,
(S38)
n=1
and therefore
| < ||
(S39)
Finally, by Cauchy-Schwarz inequality, have that
S. P < |. < S.
(S40)
We continue by studying the convergence behavior of the mean parameter p4-
24

---

Published as a conference paper at ICLR 2026
Mean parameter
Our goal is to show that |rr ll is bounded. Equation (S31) implies that
1 w.
i2 ih w -ul(He St)
(S41)
This in turn implies that
|r,1 = +[rr:
di
N
exp
M4 Un +
2 2,, Fin
F. 2 =r:r:
TT
1
n=1
Σ exp
- log(())o² 2n - w w
+
-
S.S. In r: a
2 r,
2
NES
(S42)
Σ exp
log(() on w Un -wxx
7 T S.S, Hn 2:/r, w
n
Σ exp
(-wo (-wTTn) xn)
[exp (-0 -2
+ x S.S. In
-1) x
NES
i to
exp -wid on + n. S.S. Bn
exp (-x x..
LES
where in last line we used the fact that w Fn Enes exp(-xx,,t) -,, s as in (S33). We begin by examining the first bracket, studying three possible cases for each of the summands. First, note that if w.r,, S O. then since $xTSS,STT x, 2 0, we have that
for n as in (S32), and that
(- rr + 2", "0, xn -1) aTr,, SO. exp
(S43)
Next, by defining B = ||Soll max, || 11,, if < x < 2. we have that
((
(exp
exp (- r. + + S,
B
B
exp
-
1
(S44)
2
and if x, 2 B
2. we have that
exp
--pTr + 1,,,,,,,,,0 S. xn -1 2.T s O.
(S45)
Finally, for arbitrary c 2 max B, 1}, (.Trel 2 c, we have that
exp (-x + 2 -t, S. xn -1 æTT s (exp) (
1
(exp 2 -1 E< 0. (S46)
B
-1))
Furthermore, let 7x = minnes u xn and 7-' maxnes u. J xn. Now, by taking & > max BB max(B,1) 1] large enugh such that
B -1 -1) 2 |S| exp(cxx) cxp 2
B
B z
exp(->7
exp
-
(S47)
2
if there exists support vector n E S such that [2] 2 e, then
Σ exp u xx)) (exp) WES
xITT (-x
+ +2+220 - x -1) N., S O.
(S48)
The idea of this is that if there exists a support vector such that x rel is sufficiently big, then the first bracket in Eq. (S42) is negative.
25

---

Published as a conference paper at ICLR 2026
On the other hand, for the second bracket in Eq. (S42), note that for n & S, we have that x/.00 ni and hence
tb
Σ
1
-wo on + " ænn exp (-xr x;,
exp
nÇS
(S49)
1
1
Soll max F. o7, 1
Σ ngS
exp (-wo xma) -0
exp
where in the last line we used that ze S 1, IR and fact that || F S | co.
We will now combine the results from above to show that the residual ry is bounded in the following way: there exists support vector n E S such that rer 2 e for big enough e > 0, then d ai |1r = "). If such a support vector does not exist at time t, we will show that vi is containted inside compact set. To that end, if |Srrll E1. we have that
max 20 WES
Σ nES
x Psr ?
| |
1 [
min (Xs)
(S50)
S
where the first inequality we used the fact that PJa, Dn for E Hence by choosing < such that o (Xs (x))f/(S] = e², where the is chosen in Eq. (S47), we have that
1 d
| 21
2 dt 1rol Ot
(S51)
On the other hand, if || c1, recall that
rt = (14 - 14) n log w
(S52)
and since all updates to the mean parameter are in the space spanned by the support vectors (As- sumption 2), we have that
Psr: = PsHo - Psw
(S53)
We can now conclude that
PSrt S e + |rrll S Psrr + Psrr S + ||Ps + Psr0 < co.
(S54)
Finally, combining the results from Eq. (S49) and Eq. (S54), recalling that is bounded for all <> 0. This completes the first part of the proof and shows that
have that ||r:||
Me u logt w
u log t O(1),
(S55)
and in particular
lim 1-000
(S56)
|
||0|
We proceed by showing that the limit covariance parameter vanishes in the span of the support vectors.
Covariance parameter We begin by substituting the definition of the residual T2 (S31) into the gradient flow dynamics for the covariance factor S;:
S, = V s el e14, St)
(S57)
exp( uT xn + [...S.SS xn) T
Next, we split the sum into contributions from support vectors and non-support vectors. For n S, use the property xx u 1; for & S, the margin is strictly larger than one, which introduces higher-ond order decay in
S, = Σ
exp(( - t0 xn - r xa) exp (f x:n)) x. St
t
nES
(S58)
exp
on rI xm) exp STT.S x,
R$S
26

---

Published as a conference paper at ICLR 2026
Since F. is bounded (from the previous part of the proof), the exponential prefactor is uniformly bounded away from zero. We formalize this by defining
= min min exp (-w0 - on - ri xnn) > 0.
(S59)
N 120
We also let min denote the smallest non eigenvalue of the matrix Enes Tn Finally, to measure the size of S, restricted to the support-vector subspace, we define
A = tr (P We now compute the derivative of A, over time. Differentiating and substituting the dynamics of S. yields
d 2 d
A. tr
PSS..S.
Ps
- - Σ
eexp(-on -w xn (x)
tr aP
(S60)
nES + (1.)
At this point use two facts: 1. from (S59), the exponential prefactor is bounded below by C' >> O, from the definition of minn we can control the quadratic form Enes T,,",, Applying both gives
Co A. + (
1 1 A <
(S61)
2 di
Finally, by Grönwall's lemma, there exists constant K > and starting time to > such that
2Camin
K
to
to
2Comin
+ k (x--1)
Vt to.
(S62)
Since both |S|Coamin 0 and 1, conclude that A, 0 as co. In words: the covariance factor vanishes when projected onto the span of the support vectors, i.e.
Vn E S lim x, S.S.S. 2, 0, t+000
(S63)
as claimed.
$1.2.3 COMPLETE PROOF OF THROREM
We will now extend the results for the gradient flow to gradient descent and then use these results to characterize the implicit bias of gradient descent as generalized variational inference.
Throughout this proof, let
N Σ
A
exp (11 aT +
2 S. an En
(S64)
be a positive definite matrix at iteration We begin the section with a few lemmata which will be used throughout the proof.
Lemma S2
Suppose that we start gradient descent from (nrr So). 1ʃ < Amax (Ao) then for the gradient descent iterates
S S. 17 S.),
(S65)
have that ||Se]] Sool for all 0.
Proof. First, note that the gradient descent update for the covariance factor is given by
S = S,(I nA,)),
(S66)
and hence we have that
| || n)l,)l & ||S|||]][ nA:1)12-
(S67)
27

---

Published as a conference paper at ICLR 2026
Now, since r7 & Amax (Ao) Amaxx (Ac) -1 for all 2 0 and noting that A, 2 0, we have that
II 17,,)) 12 S 1,
(S68)
and therefore
||S Σ ||
(S69)
Finally, conclude that ||Sel|r S |oll| r for all t 0, required.
Lemma S3
Suppose that we start gradient descent from So). 1f < Amax then for the gradient descent iterates
1141 u4 - v St), Consequently, we also have that
(S70)
we have that Ene IV Vullfuu Sa)l o0.
lim; p S.)1
Proof. Note that our loss function is not globally smooth in μ.. However, if we initialize at (no, So), the gradient descent iterates with 7) < Amax (A,) maintain bounded local smoothness. The state- ment now follows directly from Lemma 10 in Soudry et al. [4].
Lemma S4
By choosing 61 as Eq. (S51), if PSrrl 61. we have that
(r+t1 r.)) rı S 0 1 (1²
1
|rell-
(S71)
82
If Psrel < <1,, there exists constant C' such that
(r++1 -r) r (
(S72)
Proof. We follow similar steps as in the gradient flow case. It holds that
(ri+1 - rr) rr
(1 -t)- (log(t 1) log(t))) r
17 Σ exp (-u, o. nn S. Fn
a.T - u log(o + +1
(S73)
ú TT 1 + +)- + Σ exp
-Mufan + 1. Inn
n$S
Σ exp (-wTxx,) x.,) + exp RES
- In + x S.S. x
1
T.
where in the last equality used Equation (S33) to expand in TTTr, Furthermore, bound all four terms as follows, beginning with the first term:
1
w re((tt log(1 + +)))) < ||r
(S74)
p2)
where we used that log(1 t- t- ( For the second term, using the same argument as in Equation (S49), we derive that
Σ exp (-u] Hn + 1 a an Nt,,, S s ( 1.6 RÉS
(S75)
For the third item, from Eq. (S48) and Eq. (S50), we have that Psr 2 ei implies that
7
Σ
exp
(-w0 u xx))
exp (1r σn
1 J
S. xr
x r, S 0.
(S76)
WES
28

---

Published as a conference paper at ICLR 2026
The first result follows from combining the above three inequalities.
Next, if PSrell < c1, by defining == -Soll? following the steps in Eq. (S44), we have that
1 Σ exp néS
(11 D.n + 2 S.
xn ar,,r, S 1//S| exp
B
B
(S77)
2
2
and hence, combining this with Assumption 2 which implies that r3 is bounded as in Eq. (S54), one can find a constant such that
(ru+1 - ( S.
(S78)
Proof of Theorem
Proof. As in the simple version of the proof, begin by considering the convergence behavior of the mean parameter 14:-
Mean parameter Our goal is again to show that ||re|] is bounded. To that end, we will provide an upper bound to the following equation
|raal
||re+ rell + 2 re)" r, +reell
(S79)
First, consider the first term in the above equation:
re1 - rrl
|c+1 - w log(t + 1) w pr + v log(t) + r0|2
V Si) u log + |11
(S80)
1/ V v S.) 1 | log 2 + +
< 2
1 V E S.)
|
where in the first inequality we used the standard inequality that (x + y) < 2(xx + y²) and in the second inequality we used the fact that log(1 + z) I x for x 2 0. Now, from Lemma S3 and the fact that t is summable, conclude that there exists C1 such that
80
M++1 rell < C1 < oo.
(S81)
t=1
Next, for the second term, recall that in Lemma S4 we showed that if P 2 61, then, for some constants C2, C3 we have that, eventually
(r-11 - rz) r, S C2 4"
+ C3 +C Irall:
(S82)
and that if Psrrll < then there exists constant C4 < o0 such that
(r++1 - rz) rt S C4.
(S83)
We will show that when ||Pr < c1, the residual r, is contained in a compact set, and when [|Pttlll 2 61, the residual can't escape to infinity. We now formally show this claim.
Let S, be the frst time such that |Srr Ej, if such time does not exist, we are done since the support vectors span the data and hence |r:|] is bounded. Now, let Tj be the first time after S₁ such that |Prr 61, where we allow Ti oc if such time does not exist. Continuing in this manner, we define the sequences S < Ti < S2 < T2 ...., where we allow T; co for some i.
We prooced by showing that [|r:]] is uniformly bounded on each of the intervals [S,, T.). To that end, note that for /. E [St, T;), we have that
| ||r, 12 < 2C
+ 2Cs |ell + | - rell?
(S84)
29

---

Published as a conference paper at ICLR 2026
and hence, using the fact that 1, by the discrete version of Grönwall's lemma, that
max (1) 1|2 |rs, (1?) - K, (E(S;.T)
(S85)
for some constant K oc independent of Furthemore, we also know from Eq. (S83) that
|rs.l| 1 + 2C4 + ||rs. - - & 61 + 2C4 + max ||r41 - rell < co, 120
(S86)
showing that the first jump outisde the ci-ball is bounded. Combining the two results, we conclude that ||r=]]| is uniformly bounded on each the intervals [S,, T:).
Finally, by noting that the support vectors span the data, we have that ||rel is uniformly bounded on each of the intervals (T,,))1)) Combining the two results, we conclude that ||r/]] uniformly bounded for all 2 0 and hence we have that
14 peel
lim E-
(S87)
and the following lemma.
Lemma S5
For the mean parameter t4t, we have that
1r = log(()) O()).
(S88)
Proof. This follows immediately from the definition of the residual in Equation (S31):
H4 v log + re + Wr;
and the fact that r. and w. bounded showed above.
We continue with the analysis of the covariance parameter over optimization iterations.
Covariance parameter As before, let A, = tr(PsS Ps) be the trace of the projection of the covariance parameter the space of support vectors in S. By following the ideas from the gradient flow case, we have the following dynamics:
A
S.S. n Ps)
tr tr(Ps Ps) (000)-2nnr)) tn A,Ps) + 12 tr ASS. A.Ps)
< A: 27 Comin tr (r(PsS.s. 0
+
1
(S89)
to
t2
21 Ca A 0
A:
7
where we used the same arguments as in Equation (S60) to derive the last inequality, in addition to noting that {max(AA) < (+z) order bound the last term. Hence, we can write
2n Co A O
A411 A, <
(S90)
tcc
+2
Again, by the discrete version of Grönwall's 's lemma, we derive the equivalent result to Eq. (S62). Now, noting that E. diverges, the fact that and 1/Co O, we conclude that A, converges to zero. This implies that the covariance parameter converges to zero in the span of the support vectors, i.e.
Vn S lim x S.S.S. xn 0,
(S91)
desired.
30

---

Published as a conference paper at ICLR 2026
Characterization Generalized Variational Inference As a final step we need to show that the solution identified by gradient descent if appropriately transformed identifies the minimum 2- Wasserstein solution in the feasible set. Define the feasible set
+ S) Psμ = i and Vn E S: Var (fw (x.)) = 0}
(S92)
() Psμ i and Vn SES: ==,,ss nn 0}
(S93)
and the variational parameters identified by rescaled gradient descent as
1 log(() Mr + ( Mo; S.
0rG lim 0,00
lim t-000
(S94)
1000
holds by Lemma s5 that
Psp Ps
lim 1o log(t) 14 + Pşr = 1
(S95)
and additionally by Equation (S91) have for all S that
(S xnn lim S.S. xn, 0.
(S96)
t0000
Therefore, the limit point 0.00 of rescaled gradient descent is in the feasible set. It remains to show that it is also minimizer of the 2- Wassersstei distance the prior initialization. We will first show more general result that does not require Assumption 2.
To that end define (Vs VxIS )) E RRP'"" where Vs PxPs' is an orthonormal basis of the span of the support vectors range( XJ). VXIS E (Px(NN-P) an orthonormal basis of its orthogonal complement in range(XT) and V E IR Px(x-N) the corresponding orthonormal basis of the null space mull( X) of the data. Let V (Vs Vmn(x)) E (P- Ps) and define the projected variational distribution and prior onto the span of the support vectors and the null space of the data as
c p (w) Ni Pvu. Pv (w: (T. Σ
(S97)
piod (i) N (w;; Pv Pv Eg PT) = N (ii; iáo Eo
(S98)
where v E RN Now earlier we showed that the limit point of rescaled gradient descent is in the feasible set, defined in Equation (S94), and thus the same holds for the projected limit point of rescaled gradient descent, i.e.
(p GD STGD) E e.
(S99)
in particular
PSp..UD GDD = PS = w
(S100)
Vn S: ( on JT S (S:(D) an =0. O.
(
(S101)
Therefore, we have for all n E S that
0 = x S.
SOGD (5:00) T xn = [|(5::O)] T xn xn112 112
( T Φa = 0
(S102)
(D TV 0
(S103)
and thus VJ SIGD ( Vvs- = 0. Therefore by Lemma SI holds for the squared 2- Wasserstein distance between the projected limit point of rescaled gradient descent and the projected prior that
mnj w??
(0. pm)))
+ || V v3J 1012 w? (N ViT VVinnl ..
> (ViTIiso. ViT Eo Vinil)
[v3
V3li i V5 vo)
11 + +W3 +W3(NN Vinin VIT VVaau .N C.VVVVliiii (Viiiiio. Viti Eo Viull)
V it V5 140
||² + w3 (N(VituAa V- Vinss) ) .N( (ViT Viniiito, V Eo Vnall)
0
31

---

Published as a conference paper at ICLR 2026
(V(VVIL PsS Psto 112 W33 N V mil Vinit Vvnun .N V minkke V Zn Vauil)
-11200 - Psjuol12 122 + W3 N V mullt VT VVaul) .N (Viniioo V-T En Viuil) W3 (V((NVV ( /malfi,, Viain VV mall
(
mull
.V (OVVVV Va Vmull Eo Vanl)
where we used that Psp = ú for any (j, S) the feasible set e.. Therefore it suffices to show that the projected solution OGGD minimizes
W (V(V(vnnni,
Vi V EVain ,. V mil1/40, V Eo Vinit) )
2 0.
(S104)
We have using the definition of the iterates Equation (9) that
lim 1 1-000 log(t) pe + P
VT Gr full ft
VT Pv
mill
(S105)
mull i + P x = ViLli/to
(S106)
where we used vo E range ( Further, it holds for the gradient of the expected loss (S30) with respect to the covariance factor parameters that
VT SNGD null nul! Pv.
VT
null STOD
So
The v se sl S:)
(S107)
null
1=1
Erange(X"]
V-T Sv null S.
V V Sy-
(S108)
Therefore we have that
(V(N((
V i
GOD
SOOD mall
Vmull .N
VTT mili VTT En Vmnit =0
(S109)
and thus the projected variational parameters @:GD are both feasible (S99) and minimize the squared 2- Wasserstein distance the projected initialization prior (S104). This completes the proof for the generalized version of Theorem 2 without Assumption 2, which we state here for convenience. Lemma S6
Given the assumptions of Theorem 2, except for Assumption 2 meaning the support vectors Xs do necessarily span the data, holds for the limit point of rescaled gradient descent that
G E arg min W 9g ppro))
(S110)
(=(u.S)
AEO.
in addition Assumption 2 holds, i.c the support vectors span the training data X, such that span((mn) (n(N)) = span( ( Jnes),
(S111)
then the orthogonal complement of the support vectors range(X (T) has dimension Ps and thus the projection Pv = IPx is the identity and therefore
=
proj 4
ge
and
p = p.
(S112)
This completes the proof of Theorem 2.
S1.3 NLL OVERFITTING AND THE NEED FOR (TEMPERATURE) SCALING
In Theorem 2, we assume we rescale the mean parameters. This is because the exponential loss can be made arbitrarily small for a mean vector that is aligned with the T-2 max- max-mn vector simply by increasing its magnitude. In fact, the sequence of mean parameters identified by gradient descent diverges to infinity a logarithmic rate μ GD as log(t) (f)r as we show Lemma $5 and illustrate in Figure S2 (right panel).
"This has heen observed previously in the deterministic case (see Theorem 3 of Soudry et al. [4] and thus naturally also appears our probabilistic extension.
32

---

Published as a conference paper at ICLR 2026
Norm Mcan Params
Validation Error J
100
Validation NLL
10-1
10-1
5
5000
10000
5000
10000
5000
10000
Epoch
Epoch
Epoch
Implicit Bias
Implicit Bias VI Theoretical Scaling
Implicit Bias + Temperature Scaling
Figure S2: NLL overfitting in classification due to implicit bias of the mean parameters. As shown here for two-hidden layer neural network on synthetic data, when training with vanilla SGD the mean parameters diverge to infinity || (log(()) (right) and thus the classifier will eventually overfit in terms of negative log-likelihood (left and middle). Rescaling the GD iterates as in Theo- rem 2 or using temperature scaling [69] avoids overfitting.
This bias of the mean parameters towards the max-margin solution does not impact the train loss or validation error, but leads overfitting in terms of validation NLL (see Figure S2) long there at least one misclassified datapoint æ, since then the (average) validation NLL is given by
((0,00)) Begen (w) (exp( (exp(-yxx
GD exp(x 14
1 (
(S113)
exp(log(t)a* i _ St (S(D)) + xo
However, by rescaling the mean parameters as we do Theorem this can be prevented Fig- ure S2 (middle panel) illustrates for a two-hidden layer neural network on synthetic data. Such overfitting in terms of NLL has been studied extensively empirically with the perhaps most com- mon remedy being Temperature Scaling (TS) [69]. As we show empirically in Figure S2, instead of using the theoretical rescaling, using temperature scaling performs very well, especially in the non-asympiotic regime, which is why we also adopt it for our experiments in Section 5.
The aforementioned divergence of the mean parameters infinity also explains the need for the pro- jection of the prior mean parameters in Equation (9), since any bias from the initialization vanishes the limit of infinite training. At first glance the additional projection seems computationally pro- hibitive for anything but zero mean prior, but close inspection of the implicit bias of the covariance parameters in Theorem 2 shows that at convergence
Vn Vargo (fm (z.)) .tT.SSTT x, = = range S) C null(X)
(S114)
Meaning we can approximate basis of the null space of the training data by computing a QR decomposition of the covariance factor in (PRR) once at the end of training. For P the inclusion becomes equality and the projection can be computed exactly.
S2 PARAMETRIZATION, FEATURE LEARNING AND HYPERPARAMETER TRANSFER
Notation For this section need a more detailed neural network notation. Denote an L-hidden
layer, width-D D feedforward neural network by f(xx E D with inputs x E Rr, weights W pre-activations (0) E RD()) and post-activations (or "features") (())) (w) E RDD") That is.
(1)(())) (1) xr and, for 1,...,,
g 11 (x) = snn (x)). h (x+1)())=) = ((+1) (0) (x:),
and the network output is given by (xx W(
(L) (2x), where (o) is activation function.
For convenience, may abuse notation and write 00) (a) and h (z) f(z). Through- out we (1) to indicate the layer, subscript ,, indicate the training time (i.e., epoch), A = o, o to indicate the change since initialization, and [o]. [oly to indicate the compo- nent within vector matrix.
33

---

Published as a conference paper at ICLR 2026
S2.1 DEFINITIONS OF STABILITY AND FEATURE LEARNING
The following definitions extend those of Yang and Hu [43] to the variational setting. Definition S1 (be scaling)
In layer I, the variational parameters are initialized as
[, ~ N (O, 1s 1G ~ N (0. D- 28")) )-20%)
and the learning rates for the mean and covariance parameters, respectively, are set to (1) nD nD- e n nn 20)
The hyperparameter 7) represents a global learning rate that can be tuned, as for example in the hyperparameter transfer experiment from Section 3.4.
For the next two definitions, let m, (X) moment of a random variable X with respect to which represents all reparameterizatioo noise in the random variable X. All Landau notation Section S2 refers to asymptotic behavior width D in probability over reparameterization noise We say that vector sequence ( where each
denote the rth central moment
Un IR,, is O(D-a) if the scalar sequence (v b1
(RMSE(E))) is O(D-a).
Definition S2 (Stability of Moment r)
A neural network is stable in moment if all of the following hold for all x and 1 (1,....,1 L).
1. At initialization (t = 0):
(a) The pre- and post-activations are ((1):
my (Ng (1)
(z)), m, (g0 (x)) = 0(1)
(b) The function is O()):
0(1.
2. At any point during training >> 0:
(a) The change from initialization in the pre- and post-activations are O
Am, (hl (z)) (g. (1) (x)) = =0(1)
O().
(b) The function is O(1):
my (f(((())) O(1)
Definition S3 (Feature Learning of Moment r)
Feature learning occurs in moment r in layer if, for any t > O, the change from initialization is 2(1):
Amr
9 12 (a)) = 2(1.
As we will see later, Figure S5 and Figure S6 investigate feature learning for the first two moments.
S2.2 INITIALIZATION SCALING FOR A LINEAR NETWORK
In this section we illustrate how the initialization scaling ((60), [(0) can be chosen for stability. For simplicity, we consider a linear feedforward network of width D evaluated on a single input E R.. We assume Gaussian variational family that factorizes across layers. This implies the hidden units evolve as h (1-1) w (1+1)) h
(1) and the weights are linked to the variational parameters
by vec( (1()) ( = Pt (1) + S (.)..
z.
Therefore, the mean and variance of the ith component hidden units in layer E {1, L + 1),
where i E 1 1..., D(O, are given by
E. [n. [1)
[14 1 Eg 111 (
34

---

Published as a conference paper at ICLR 2026
Var: [ (n."11) = 1t4, (u/")Te (e
(( lu 1 + (r([s]])] I..A 1 (s
- lr,),
where 7 = (iD(I-))
(i+ (i+1) and the second moment of and covariance layer-1 hidden
units denoted by
4 = E.
h h (T)
c = (0) E.
h
Eg h (1)
Mean We start with the mean of the hidden units, which conveniently depends only on the mean variational parameters and the previous layer hidden units.
D
E2
(100
Σ
① 140
1 E. (1h0 (h0-11)
j=1
o D(l-1). D D-b 11)
o bl2) D-
1=1
lE f2..... 13. D (6(0) -))
Therefore, we require ((1) 2 and ((0)
2
2 for I L+11]
Variance Next we examine the variance of hidden units. Consider the first term, which represents the contribution of the mean parameters.
D(-1))
D
[40
Σ j=1
112 I00
li, 11 +
Σ
Ino"l, (c(l (iorlmo01 14,
j#8
(DDD-1) D 26 11) + (DLDI 1) D-b
D-26
69) . D- 600
=O(DU-)), D D-26 26(1)
S0(D 10(D-261) 26(1)
1=1
0 10(D-(2b1)) (266))
lele L
Therefore, we require ((1) and (() 2
/ for (2...., T..,, + 1}. Notice these are the same
requirements above for the mean of the hidden units. We summarize the scaling for the mean parameters as
(
0 1=1 1 IE 1 1)
>
(S115)
Now consider the second term in the variance of the hidden units. Assume the rank scales with the input and output dimension of layer as ( (Dil-)) 1) D(0Po
where p(1) E [0, 1].
R(
lr..)
Σ
A
( 1rr
(1-1)
=
r=1
R!
DC
D
Σ
(11 17, [A! (1 + Σ
11,
(4-1) 1r
1=1
o ((()) D(1-1) D 20
0-2600 + (O -1) D- 500
1 D.DD 50
o(R())
Dil-1) D(1-))(11)))))))))
)
26
35

---

Published as a conference paper at ICLR 2026
(28
-01))
1= le (2....,L. 1= L+1.
D (28 -100-1-2200 -1-201))
(2 -1---tt-
Therefore require ((0)
(0)
5 2 2 p for (2...., L}, and ((+++))
2 + pla
Notice we can write these conditions in terms of the mean scaling as
1=1
500 2 10
E 1.,
(S116)
= 1+1.
2
S2.3 PROPOSED SCALING
The previous section derives the necessary conditions for stability at initialization. Recall from Section 3.4 that propose scaling the contribution of the covariance parameters to the forward pass, i.e. the Sz term, by R- since each element in the term is sum over R random variables, where the rank of S. the more detailed notation this section, the proposed scaling implics the forward pass in linear layer is given by
[h (9) [W] 11-1) I 11 h (1-1)
(S117)
In practice, rather than scaling [S by R- 1/2 the forward pass, we apply Lemma J.1 from Yang et al. [41] to instead scale the initialization by R 1/2 and, in SGD, the learning rate by R R-1 Scaling by the rank allows treating the mean and covariance parameters if they were weights parameterized by pP in a non-probabilistic network, inheriting any scaling that has already been derived for that architecture.
From Table 3 of Yang et al. [41], we therefore scale the mean parameters as
0
1 -=1
1=
b0
1/2
E (2....,
and c
0
le (2....,
(S118)
1 1 1=1+1
1 1+
Assuming R() parameters as
(D(t-1) ()00)'"))) before, where p()) E [0, 1]. the scale the covariance
(())
1=1
p())
1=1
000 (())
+
E L} and 200 cl))
+
2p(0)
LE (2. L} (S119)
L++1
(()
1 1 L+1+1
By comparing to Equations $115 and S116, the output layer are initialized as large as possible while still maintaining stability. The output layer parameters scale to zero faster, since, in pP for the weights of non-probabilistic networks, we set B(L+1) to 1 instead of 1/2.
the mean and covariance parameters in all but
Note that in Section S2.2 we did not consider input and output dimensions that scaled with the width D for simplicity. For our experiments, take the exact μP initialization and learning rate scaling from Yang et al. [41] which includes, for example, a 1/ ffnninn scaling in the input layer - for the means and then make the rank adjustment for the covariance parameters as described above.
We investigate the proposed scaling in Figures S4 and S5. We train two-hidden-layer (L = 2) MLPs of hidden sizes 8, 16, 32, and 64 on single observation (z, y) (I, 1) using a squared error loss. We use SGD with a learning rate of 0.05. For the variational networks, we assume multivariate Gaussian variational family with a full rank covariance.
Figures S3 and S4 show the RMSE of the change in the hidden units from initialization, Ag. (x) (1) gi (x) Do (x)), function of the hidden size. The RMSE of the hidden units at initialization,
36

---

Published as a conference paper at ICLR 2026
(1) is also shown in blue. Each panel corresponds to layer of the network, the first two panels 9o
correspond to features 9: (x) and (2) (z)), respectively, while the third panel corresponds to the output of the network, 9: (3) (x) (x). The difference between the figures is the paramaterization. Figure S3 uses standard parameterization (SP) while Figure S4 uses maximal update parametrization (μP). We observe that (a) the features change more under pP than SP and (b) training more stable across hidden sizes under yP than SP, especially for smaller networks.
Figures S5 and S6 show the analogous results for a variational network. The top row shows the change in the mean of the hidden units, while the bottom row shows the change in the standard deviation. As the non-prohabillitic case, we observe that (a) both the mean and standard deviation of the features change more under μP than SP and (b) training is more stable across hidden sizes under μP than SP, especially for smaller networks.
1 1=1
I = 2
1 = 3
2
RMSE
RMSE (0)
1
23
23
26
23
25
26
2
2
Hidden Size
Hidden Size
Hidden Size
Epochs (t)
0
32
128
512
2048
16
64
256
1024
4096
Figure S3: MLP, Standard Parameterization. RMSE of the change in the hidden units and. in blue, their initial values. Shaded region represents 95% confidence interval over 5 random initializations. The MLP is trained under SP.
RMSE (9)
2
RMSE(g. (9)
23
24
25
26
25
23
25
Hidden Size
Hidden Size
Hidden Size
Epochs (f)
0
2
8
32
128
512
2048
16
256
1024
4096
Figure S4: MLP, Maximal Update Parameterization. RMSE of the change in the hidden units and, in blue, their initial values. Shaded region represents 95% confidence interval over 5 random initializations. The MLP is trained under pPP
S2.4 DETAILS ON HYPERPARAMETER TRANSFER EXPERIMENT
As discussed in Section 3.4 we train two-hidden-Hee layer MLPs of width 128, 256, 512, 1024, and 2048 on CIFAR-10. For comparability Figure in Tensor Programs V [41] we the same hyperpa-
37

---

Published as a conference paper at ICLR 2026
=
1=2
1= 3
(0)
-Big
2
(ISS(S RMSE(E(
RMSE(B(ge ())
RM ) SSd )
2
23
24
25
26
23
24
2
29
23
2
2
26
Hidden Size
Hidden Size
Hidden Size
Epochs (f)
128
512
2048
16
64
256
1024
4096
Figure S5: Variational MLP,, Standard Parameterization. RMSE of the change in the hidden units and, in blue, their initial values. Shaded region represents 95% confidence interval over random initializations. The variational MLP trained under SP with full rank covariance in each layer.
rameters but applied to the mean parameters." 5 For the input layer, scale the mean parameters initialization by factor of 16 and in the forward pass by a factor of 1/16. For the output layer, we scale the mean parameters by 0.0 initialization and by 32.0 in the forward pass. We use 20 epochs, batch size 64, and a grid of global learning rates ranging from 2-8 to 2 with cosine annealing dur- the grid search results shown the right panel Figure 3, we use validation NL.L for model selection and then evaluate the relative test error compared to the best performing model for that width across parameterizations and learning rates.
ing training.
S3 EXPERIMENTS
This section outlines in more detail the experimental setup, including datasets (Section S3.1.1), metrics (Section S3.1.2), architectures, the training setup and method details (Section S3.3.1). also contains additional experiments to the ones in the main paper (Sections S3.2, S3.3.2 and S3.3.3).
$3.1
SETUP AND DETAILS
In all of experiments used the following datasets and metrics.
'Specifically, used the hyperparameters indicated here: https::/ com/m coooooof mup/blob/ma eeeeampplee wemmplls/..lll /dmoo. pynb
38

---

Published as a conference paper at ICLR 2026
=
1=2
1 == 3
(0)
-Big
RMSE(E (9)
2
RMSE(E(ge (
R. ) SSd )
2
10
RMSE
23
24
25
26
23
24
25
28
23
2
25
2
Hidden Size
Hidden Size
Hidden Size
Epochs (f)
128
512
2048
16
64
256
1024
4096
Figure S6: Variational MLP, Maximal Update Parametrization. RMSE of the change in the hidden units and, in blue, their initial values. Shaded region represents 95% confidence interval over random initializations. The variational MLP trained under yP with a full rank covariance in each layer.
S3.1.1 DATASETS
Table S1: Benchmark datasets used in our experiments. All corrupted datasets are only intended for evaluation and thus only have test sets consisting of 15 different corruptions of the original test set.
Dataset
N
Ntest
Dm
C
Train / Validation Split
MNIST [70]
60 000
10000
28
x 28
10
(0.9, 0.1)
CIFAR-10 [80]
50000
10000
3 x 3
32
x 32
10
(0.9, 0.1)
CIFAR-100 [80]
50000
10000
32
x 32
100
(0.9, 0.1)
TinyImageNet [81]
100.000
10000
x 64
64
200
(0.9.0.1)
MNIST-C [72]
150000
28
28
10
CIFAR-10-C1 [73]
150000
32
CIFAR-100-C [73]
150000
x
32
x 32
100
TinylmageNet-C [73]
150000
3 X
64
x 64
200
S3.1.2 METRICS
Accuracy The (top-k) accuracy defined as
Nesst
Accuracy (v. ④)
1 E0)ik)-
(S120)
Niest
Negative Log-Likelihood (NLL) The (normalized) negative log likelihood for classification is
given by
Niss
NLL(y, 9)
log P
(S121)
N
39

---

Published as a conference paper at ICLR 2026
where Pg., is the probability model assigns to the predicted class Dn-
Expected Calibration Error (ECE) The expected calibration error measures how well a model is calibrated, i.c how closely the predicted class probability matches the accuracy of the model. Assume the predicted probabilities of the model on the test set are binned into a given binning of the unit interval. Compute the accuracy aj and average predicted probability ú3 of each bin, then the expected calibration error is given by
ECE
bylaj 431 j=1
Σ
(S122)
where by is the fraction of datapoints in bin j E (1,....,] J}.
S3.2 TIME AND MEMORY-EFFICIEN TRAINING
To keep the time and memory overhead low during training, we would like draw as few samples of the parameters as possible to evaluate the training objective f(O). Drawing M parameter samples for the loss increases the time and memory overhead of forward and backward pass M times (disregarding parallelism). Therefore it is paramount for efficiency to use as few parameter samples as possible, ideally M
When drawing fewer samples from the variational distribution, the variance in the training loss and gradients increases. practice this means one has to potentially choose smaller learning rate to still achieve good performance. This is analogous to the previously observed linear relationship N₂ X ?) between the optimal batch size N₂ and learning rate 7/ [e.g., 82-84]. Figure S7 shows this relationship between the number of parameter samples used for training and the learning rate MNIST for a two- hiddeidden layer MLP width 128.
0.10
0.4
Test N
Test Error
Test ECE
0.05
0.2
10-2
0.00
0.0
-1
10
10a
10a
10
10a
10a
10
10a
10a
Parameter Samples
Parameter Samples
Parameter Samples
Learning Rate
0.003
0.01
0.03
0.1
Figure S7: Generalization versus number of parameter samples. For a fixed number of epochs and batch size, fewer samples require smaller learning rate. For fixed learning rate, generalization performance quickly plateaus with more parameter samples.
As Figure $8 shows, when using momentum, generalization performance tends to increase, but only either the number of samples is increased, or the learning rate is decreased accordingly. similar relationship between noise in the objective and the use of momentum has previously been observed by Smith and Le [83], which propose and empirically verify a scaling law for the optimal batch size No 1-y as a function of the momentum parameter ? 0.
S3.3 IN- AND
GENERALIZATION
This section recounts details of the methods we benchmark in Section 5, how they are trained and additional experimental results.
$ ARCHITECTURES, TRAINING, AND METHODS
Architectures We use convolutional architectures for all experiments Section For MNIST, standard LeNet-5 [70] with ReL.U activations. For CIFAR-10, CIFAR-100 and TinyIma-
40

---

Published as a conference paper at ICLR 2026
0.10
0.4
Test Test NLL
Test Error
Test ECE
0.05
0.2
10
0.00
100
10
10a
10
10
10a
100
10
10a
Parameter Samples
Parameter Samples
Parameter Samples
Learning Rate
0.003
0.01
0.1
Figure S8: Generalization versus number of parameter samples when using momentum, Using mo- mentum improves generalization performance, but when using fewer parameter samples, smaller learning rate necessary than for vanilla SGD as predicted by ??.
0.6
Valid. Eror J.
Valid. NLL
0.4
Valid. ECE
0.4
2
0.2
0.2
0.0
0.0
10
103
102
10
102
10
Optimizer Step
Optimizer Step
Optimizer Step
Parameter Samples
64
16
2
128
128
Figure S9: Validation error during training for different numbers of parameter samples. The dif- ference in generalization error between different number of parameter samples vanishes with more optimization steps both for SGD (-) and when using momentum (-), if the learning rate is suffi- ciently small (in this example 17 = 0.003).
geNet we use a ResNet-34 [71] where the first layer is a 2D convolution with kernel.sise ze=3,, stride=1 and padding= 1g account for the image resolution of CIFAR and TinyImageN and the normalization layers are GroupNorm layers. We use pretrained weights from ImageNet for all but the first and last layer of the ResNets from torcavision [85] and fully finetune all parameters during training.
Training We train all models using SGD with momentum (7 0.9) with batch size N₂ = 128 and learning rate = 0.005 for 200 epochs. We do not use learning rate scheduler since we found that neither cosine annealing nor learning rate warm-up improved the results.
Temperature Scaling [69] For temperature scaling optimize the scalar temperature param- eter in the last layer on the validation set via the L-BFGS implementation torch with an initial learning rate yis 0.1, maximum number of 100 iterations per optimization step and
Laplace Approximation (Last-Layer, GS + ML) [26] As recommended by Daxberger et al. [26] we use post-hoe KFAC last-layer Laplace approximation with GGN approximation to the Hessian. We tune the hyperparameters post-hoc using type-ll maximum likelihood (ML), As an alternative we also do a grid search (GS) for the prior scale, which we found to be somewhat more robust in experiments. Finally, compute the predictive using (extended) probit approxi- mation. Our implementation the Laplace approximation is a thin wrapper of laplace [26] and use its default hyperparameters throughout.
41

---

Published as a conference paper at ICLR 2026
Weight-space VI (Mean-field) [30, 31] For variational inference, used a mean-field variational family and trained via an ELBO objective with a weighting of the Kullback- term the prior. We chose unit-variane, Gaussian prior with mean that set to the pretrained weights, except for the in- and output layer which had zero mean. We found that using a KI. weight and more than single sample (here M 8) was necessary achieve competitive performance. The KL weight was chosen to be inversely proportional to the number of parameters of the model, for which we observed better performance than a KL weight that was independent of the architecture. At test time we compute the predictive by averaging logits using 32 samples.
regularization
Implicit Bias VI [ours] For all architectures in Section we use Gaussian in- and output layer with low-rank covariance (R 10, 20). We train with single parameter sample M through- out and do temperature scaling at the end of training on the validation set with the same settings as when just performing temperature scaling. We do temperature scaling in classification due to the specific form of the implicit bias in classification as described in Section S1.3. Since IBVI trains by optimizing minibatch approximation of the expected negative log-likelihood (an average over ogg-robbabilities with respect to parameter samples), we also average log-prooe bogprobbbiilities at test-time compute the predictive distribution class probabilities. Although did not see significant difference between averaging log probaabbiiitiss probabilities or logits. Like for WSVI we use 32 samples at test time.
SWAG [28] We used slightly modified implementation of SWAG based on torch-uncertainty and the original implementation by Maddox et al. [28]. The begin- ning of the averaging cycle set to half the number of total epochs and cycle length of one. i.c. SWAG updates happen every epoch. For all other hyperparameters we use the default settings.
Deep Ensembles [29] We use five ensemble members initialized and trained independently. We compute the predictive by averaging the predicted probabilities of the ensemble members in line with standard practice [29]. We did not see a significant difference performance between averaging logits or averaging class probabilities.
42

---

Published as a conference paper at ICLR 2026
S3.3.2 IN-DISTRIBUTION GENERALIZATION AND UNCERTAINTY QUANTIHICATION
The full results from the in-distribution generalization experiment in Section 5 can be found in Figure $10. The Figure SII. When finetuning a pretrained model, we found that on some datasets (CIFAR-100, TinylmageNet) μP resulted in somewhat lower performance, contrary the results in Section 3.4, where we trained from scratch. This suggests that, when pretraining, there may be a modification to the parametrization that could improve generalization.
experiment but done in the Maximal Update parametrization is depicted in
MNIST
CIFARIO
CIFARI00
TinyImageNet
Test Error +
0.02
0.125
Test Test Top-5 Error
0.15
0.175
0.01
0.075
0.10
0.150
1.0
Test NLL
0.10
0.5
0.05
10-1
0.5
0.5
Test ECE
10-2
0.0
0.0
0.0
Method
Standard
Laplace (Last-layer, GS)
Weight-space VI (Mean-field)
SWAG
Temperature Scaling
Laplace (Last-layer. ML)
Implicit Bias VI (Low-rank)
Ensemble
Figure S10: In-distribution generalization and uncertainty quantification (Standard parametriza-
tion).
MNIST
CIFARIO
Test Top-5 Emrr 4
CIFARI00
Imageee
Tost Error 4
0.012
0.2
0.4
0.100
0.010
DOM
0.075
0.1
0.008
Test NLL 1
2.0
0.5
0.5
0.4
1.5
OOO
0.0
0.2
Test ECE
10
0.2
0.4
10
0.1
0.2
0.1
Method
Standard
Laplace llaestllyer, GS) Laplace layest- ML)
Weight-space VI (Meun-field)
SWAG
Temperature Scaling
Implicit (Low-rank)
Ensemble
Figure S11: In-distribution generalization and uncertainty quantification (Maximal Update parametrization).
43

---

Published as a conference paper at ICLR 2026
S3.3.3 ROBUSTNESS TO INPUT CORRUPTIONS
Besides the benchmark Figure S11, we also evaluated the models trained using the Maximal Update parametrization on the corrupted datasets. The results can be found in Figure S12.
MNISTC
CIFARIOC
Test Top-5 Emoor 4
CIFARI00C
TinyImageNeIC
Test Error 4
0.45
0.2
0.28
0.40
0.75
0.70
COM
0.1
Test NLL
0OM
0.5
Test ECE 4
0.4
0.50
10-1
0.2
0.25
Method
Standard
Luplace (Last-layer,
Weight-space VI (Meun-field)
SWAG
Temperature Scaling
Laplace (Last-layer, ML)
Implicit Bias VI (Low-rank)
Ensemble
Figure S12: Generalization on robustness benchmark probleess (Maximal Update parametrization)
S3.3.4 COMPARISON TO GENERALIZED VI WITH 2- 2-SSSEEST REGULARIZATION
Theorems and 2 characterize the implicit bias of gradient descent for overparametrized lin- car model as a preference for distributions minimizing the expected loss, which are closest in 2- Wasserstein distance the initialization. Given this characterization, by the KKT conditions there exists Lagrange multiplier À 2 0 such that the optimal variational parameters DGD define a sta- tionary point of the following unconstrained optimization objective:
(0) E(() W p)
(S123)
In other words, Implicit Bias VI is equivalent to Generalized VI (GVD) with 2- 2-Wasserstein regu- larizer and some regularization strength 2 0 for overparametrized linear models.
Experiment Results To understand the difference in performance between IBVI and Generalized VI with 2- 2-Wasserstein regularizer for deep neural networks, we trained models via the GVI ob- jective in Equation (S123) for different regularization strengths > with the same setup as in Section 5. The results on in-distribution test data can be found in Figure S13 and the results for corrupted test data are in Figure $14. Both on in- and out-of-distribution data GVI performs similar or worse than IBVI for all regularization strengths we tested in terms of test error. IBVI and GVI perform roughly similar in terms of uncertainty quantification with GVI only performing better for regularization strengths that harm accuracy.
44

---

Published as a conference paper at ICLR 2026
MNIST
CIFARIO
0.11
Test Error V
0.0150
0.10
0.0125
0.09
0.0100
Test NLL
0.6
0.06
0.4
0.04
0.075
Test BCB
0.010
0.050
0.005
0.025
Method
Standard
Generalized VI (Low-rank)
À = 0.001
X 0.1
Implicit VI (Low-rank)
0.0001
0.01
A
Figure S13: In-distribution generalization and uncertainty quantification of IBVI and GVI.
MNISTC
CIFARIOC
0.300
Test Error
0.16
0.275
0.14
0.12
0.250
Test NLL
1.0
0.5
Test ECE
0.10
0.2
0.05
Method
Standard
Generalized VI (Low-rank)
= 0.001
A 0.1
Implicit Bias VT (Low-rank)
0.0001
0.01
Figure S14: Out-of-distribbttin generaiizatonnnnnnnnnin and uncertainty quantification of IBVI and GVL.
45


Reviews:
[{"review_id": "XHrY1n6mA0", "reviewer": "Reviewer_QjSV", "summary": "he authors propose and analyze a variant of the so-called “generalized variational inference” objective (e.g., Eq. 1 of Knoblauch et al.) that drops that regularization term, arguing that the problem is already well-posed from implicit rather than explicit regularization. The well-posedness derives from the fact that in overparameterized models, the set of global optima is quite large (not just a point mass), and so SGD tends to converge to the member of the set of global optima which is nearest to the initialization. Because of this, the problem is implicitly regularized to the prior already (b/c the prior is used for initialization), so explicit regularization is unnecessary. The authors formalize this result for linear regression and binary classification by showing the implicit regularization takes the form of a Wasserstein distance from the the variational distribution to the prior.", "questions": "- It seems (cf. line 45 and line 2228) that the loss function used by the authors is the negative log likelihood. Considering that the proposed method only minimizes this, and nothing else, it’s surprising that other approaches can outperform on this exact metrics (Test NLL for MNISTC, for example, in Figure 5). Can the authors provide any insight on this phenomenon?\n- Taking implicit regularization as well-established, can the authors comment more on the desirability of the solutions found by this fact? Is the solution closest to the prior advantageous with respect to generalizability, etc.?\n- The authors may consider citing other works that use the negative log likelihood loss, perhaps to comment on differences. For example, the class of neural posterior estimation (NPE) methods comes to mind. Can you elaborate on similarities or differences to this class of methods?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The ideas of the paper are central, yet deep, and well articulated by the authors\n- To my knowledge, the contributions are novel in the literature and interestingly meld together the ideas of Bayesian ensembles of models with variational inference and optimization. I find the takeaway message of the paper to be quite “unifying”.\n- Extensibility to two different problem settings (regression and classification) augments the impact of the work.\n- The authors provide sufficient background to make the work accessible to a broader audience", "weaknesses": "- The selling point of the paper falls short of the results: deep learning is not included in the main results of S4, which assume overparameterized linear models. I still find the contribution interesting even in this context, but I’m unsure how much the results can be said to explain the success of “deep learning”. I understand the discussion on page 5 and the CIFAR experiment tie it more closely to practical deep learning, but the theoretical results do not explain the success of deep learning in the same way that the results of, for example, the neural tangent kernel (NTK) does. The lack of a clear connection to deep learning undermines the contribution somewhat. \n- The main experimental results (Fig. 4) do not demonstrate clear outperformance of the method relative to the competitors. However, I commend the breadth of the competing approaches considered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7EJMoA8C51", "reviewer": "Reviewer_NjWy", "summary": "- this paper exploits the implicit regularization effect of (stochastic) gradient descent (SGD) for Bayesian deep learning.\n- In particular, the authors characterize the implicit bias induced by SGD for regression and binary classification task in overparameterized linear models.\n- Experiments show competitive performance compared to state-of-the-art Bayesian deep learning baselines.", "questions": "See the weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written and structured.\n- The authors propose to training a Bayesian neural network using only the expected loss, removing the divergence term from the variational objective. They then show that among the variational parameters (assumed to follow a Gaussian distribution) that minimize the expected loss, SGD converges to the one closest to the prior in terms of Wasserstein-2 distance -- analogous to how SGD converges to minimum-norm solution in overparameterised linear network.\n- The experimental results demonstrate that this implicit bias achieves competitive test error and uncertainty quantification across different datasets, thereby empirically confirming the proposed effect.\n- Overall, I enjoyed reading this paper.", "weaknesses": "- The theoretical analysis is limited to overparameterized linear model (one layer). Extending the proof to more than 2 layers might not be straightforward.\n- The proof assumes Gaussian variational distribution, and the theorems do not generalize to other distributions which might be more effective for other problem settings/ datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gxqzJbwZB3", "reviewer": "Reviewer_uPWm", "summary": "This paper proposes a method to get the benefits of Bayesian deep learning—like improved robustness and uncertainty estimates—without the usual computational cost.\n\nInstead of adding an explicit regularization term (like KL divergence) to the loss function, they rely on the implicit regularization of SGD. They initialize a variational distribution over the network's weights to match a prior, and then train by minimizing only the data-fitting loss.\n\nThe key insight is that SGD naturally finds the solution that is closest to its starting point. This behavior acts as a powerful regularizer, achieving strong performance with minimal overhead. They provide proof that this technique actually yields the divergence minimiser for Wasserstein distance in the linear case for overparametrised models. \n\nThey provide empirical results showcasing that their proposed methodology offers competitive performance on a broad range of datasets.", "questions": "1. Can it possibly be generalised to other family of variational posteriors and not just Wasserstein distance for the divergence term?\n2. Did you try any other parameterisations of the mean field covariance term? \n3. Would appreciate an ablation study showing that this offers similar performance compared to using an explicit Wasserstein distance divergence term.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. Very sound approach and presentation, very much like the simplicity of just not having an extra divergence term to deal with the numerical issues of.\n2. Theorem 1 is very nice, gives an intuition as to why this could work for much larger models.\n3. Experiments convincing, comprehensive across diverse tasks and datasets", "weaknesses": "1. This method seems to be possibly limited to mean field variational posteriors, and only considers the Wasserstein distance as the divergence term. \n2. Missing baseline comparison, would be nice if you could compare against a variational method that uses Wasserstein distance explicitly in the \"ELBO\", possibly https://proceedings.neurips.cc/paper_files/paper/2022/file/18210aa6209b9adfc97b8c17c3741d95-Paper-Conference.pdf or any other paper that proposes Wasserstein distance explicitly as a divergence term.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)