As a representative continuous-depth neural network approach, stochastic differential equation (SDE)-based Bayesian neural networks (BNNs) have attracted considerable attention due to their solid theoretical foundations and strong potential for real-world applications. However, their reliance on numerical SDE solvers inevitably incurs a large number of function evaluations (NFEs), resulting in high computational cost and occasional convergence instability. To address these challenges, we propose a Nesterov-accelerated gradient (NAG) enhanced SDE-BNN model. By integrating NAG into the SDE-BNN framework along with an NFE-dependent residual skip connection, our method accelerates convergence and substantially reduces NFEs during both training and testing. Extensive empirical results show that our model consistently outperforms conventional SDE-BNNs across various tasks, including image classification and sequence modeling, achieving lower NFEs and improved predictive accuracy.
Figures & tables
Fig. 1 : Comparison of mentioned residual mechanisms: (a) Direct residual connection for mt in eq ( 8 ). (b) Proposed NFEf -dependent residual skip connection for mt in eq ( 11 ).
Experiments
Model
Hyper-parameter
1D Regression
MNIST [ 30 ]
CIFAR-10 [ 31 ]
Walker2D [ 32 ]
SDE-BNN (fixed step)
and Nesterov-SDE-BNN (fixed step)
Augment dim.
2
2
2
0
# blocks
1
1
1-1-1
1
Diffusion σ
0.2
0.1
0.1
0.1
KL coef.
0.0
1e-5
100
1
TABLE I : The following hyperparameters are used for each evaluation method corresponding to the experimental results reported in the paper.
Fig. 2 : Predictive prior and posterior of Nesterov-SDEBNN on a non-monotonic toy dataset. The blue and red shaded regions denote the 95% confidence intervals of the prior and posterior, respectively, while the solid lines represent their corresponding mean predictions.
MNIST
CIFAR-10
Model
Accuracy (%) ↑
AUC (%) ↑
NLL ( ×10−2 ) ↓
Accuracy (%) ↑
AUC (%) ↑
NLL ( ×10−1 ) ↓
SDE-BNN(fixed step)
98.90 ± 0.07
97.59 ± 0.04
14.37 ± 0.89
87.56 ± 0.36
83.05 ± 0.29
9.89 ± 0.22
Nesterov-SDEBNN(fixed step)
99.04 ± 0.12
97.87 ± 0.06
7.37 ± 0.65
88.36 ± 0.24
85.61 ± 0.45
5.97 ± 0.11
SDE-BNN(adaptive step)
98.87 ± 0.12
97.57 ± 0.03
13.85 ± 1.26
85.87 ± 0.28
78.42 ± 0.60
6.77 ± 0.09
Nesterov-SDEBNN(adaptive step)
99.04 ± 0.03
97.88 ± 0.04
7.15 ± 0.34
86.99 ± 0.23
83.94 ± 0.32
5.32 ± 0.10
TABLE II : Quantitative evaluation on MNIST and CIFAR-10. Comparison of SDE-BNN and Nesterov-SDEBNN performance (Accuracy, AUC, and NLL) using fixed-step and adaptive-step integration.
Fig. 3 : Comparison of test accuracy between SDE-BNN and Nesterov-SDEBNN: (Left) MNIST (Right) CIFAR-10.
Fig. 4 : Comparison of test NFEs between SDE-BNN and Nesterov-SDEBNN: (Left) MNIST (Right) CIFAR-10.
Fig. 5 : Walker2D test loss performance. Nesterov-SDEBNN vs. SDE-BNN under fixed-step (Left) and adaptive-step (Right) solver configurations.
Fig. 6 : Forward-pass NFEs on the Walker2D kinematic simulation. Comparison between SDE-BNN and Nesterov-SDEBNN during (Left) training and (Right) testing using adaptive-step solver.
Lookahead-based acceleration methods, such as Nesterov's momentum, are widely used in optimization, but they often become unreliable in deep learning training mainly due to stochastic gradient noise and non-convex loss landscapes. In particular, standard lookahead relies on short-horizon update signals (e.g., differences between consecutive iterates), which are inherently noisy and can lead to unstable extrapolation directions. This work revisits Nesterov's acceleration from a trajectory perspective and argues that effective acceleration in deep learning should harness the low-frequency trends of optimization trajectories rather than extrapolating noisy one-step updates. Leveraging this insight, we propose EMA-Nesterov, a simple modification that replaces the standard Nesterov's lookahead direction with an exponential moving average (EMA) of parameter updates. This yields a stabilized lookahead direction that captures and harnesses the evolving trend of the training trajectory through a low-pass filter, while remaining adaptive to progressive changes via the geometric weighting structure of EMA. We show that EMA-Nesterov retains a theoretical accelerated convergence rate in convex problems that is analogous to Nesterov's accelerated gradient method. Furthermore, we provide empirical evidence on language model pre-training to verify that EMA-Nesterov is broadly applicable across a range of fine-tuned base optimizers, including Adam, SOAP, Muon, as well as complex optimizers that achieve state-of-the-art performance on optimization benchmarks (NanoGPT). Compared to prior lookahead methods, EMA-Nesterov achieves better performance by avoiding the instability of short-horizon lookahead and the non-adaptivity of long-horizon lookahead.
Chung-Yiu Yau, Dawei Li, Athanasios Glentis +3
University of Minnesota · 2Amazon AGI · 3The Chinese University of Hong Kong
The practical adoption of sampling-based inference (SAI) in Bayesian neural networks (BNNs) remains limited, partly due to persistent misconceptions about the feasibility and efficiency of sampling. This position paper argues that SAI has achieved computational parity with optimization-based methods and is at the verge of superseding such methods for effective and efficient inference in BNNs. This development should be in the interest of the whole community, promoting BNNs as a principled paradigm with its long-standing yet unfulfilled promise of providing principled uncertainty quantification for neural networks. SAI can even do more -- yielding superior prediction performance through model averaging, serving as the foundation for a plethora of possible downstream tasks, and providing crucial insights into the landscape of BNNs. In order to make such a change happen and unfold the potential of sampling, overcoming current misconceptions is a necessary first step. The next step is to realign research efforts toward addressing remaining challenges in SAI. In particular, the community must focus on two core problems: sufficient exploration of the posterior landscape and high-fidelity distillation of posterior samples for efficient downstream inference. By addressing conceptual and practical obstacles, we can unlock the full potential of SAI and establish it as a central tool in Bayesian deep learning.
Emanuel Sommer, David Rügamer
Department of Statistics, LMU Munich, Munich, Germany · Munich Center for Machine Learning, Munich, Germany.
Bayesian neural networks are typically trained against the evidence lower bound (ELBO), whose Jensen gap closes only when the variational posterior is exact. We instead train by local consistency: gradient descent on the Bethe free energy, driving the belief at every factor toward agreement with its neighbours rather than placing a loss on the output. The resulting objective scores each observation by its own predictive density: a strictly proper rule whose optimum is the true conditional, for any likelihood with a tractable predictive convolution. Instantiated with a Gaussian last layer over a deterministic backbone, exact inference appears as one known corner: the neural-linear marginal likelihood. That corner is evidence-optimal; the shared-cavity, free-routed interior is predictive-optimal, improving NLL and calibration over it. This instance, SCROLL (Shared-Cavity fRee-rOuting Last-Layer), is a single-pass Bayesian neural network: batchable, any-likelihood, and implicitly empirical-Bayes-prior precision, observation noise, covariance, and backbone fit in one gradient pass. Prior work enters the interior only through the ELBO and its Jensen gap, even in this conjugate setting. At a single training run and forward pass per architecture-where the validation-tuned conventional references cross-validate λ and ensembles pay 5-50× at inference-a fixed SCROLL variant is best-or-tied on NLL and calibration on 7/8 UCI regression benchmarks, and best on 4/5 across three large tabular datasets (up to 515k examples) and two frozen text/vision embeddings.