Improving Infinitely Deep Bayesian Neural Networks with Nesterov's Accelerated Gradient Method

Chenxu Yu; Wenqi Fang

stat.MLMar 26, 2026

Improving Infinitely Deep Bayesian Neural Networks with Nesterov's Accelerated Gradient Method

Authors: Chenxu Yu, Wenqi Fang

Organizations: Shenzhen Institutes of Advanced Technology Chinese Academy of Sciences Shenzhen, P.R.China

Abstract

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

Panel of Figure 1 — Fig. 1 : Comparison of mentioned residual mechanisms: $(a)$ Direct residual connection for $m_{t}$ in eq ( 8 ). $(b)$ Proposed $\text{NFE}_{f}$ -dependent residual skip connection for $m_{t}$ 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 $\sigma$	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 (%) $\boldsymbol{\uparrow}$	AUC (%) $\boldsymbol{\uparrow}$	NLL ( $\times 10^{-2}$ ) $\boldsymbol{\downarrow}$	Accuracy (%) $\boldsymbol{\uparrow}$	AUC (%) $\boldsymbol{\uparrow}$	NLL ( $\times 10^{-1}$ ) $\boldsymbol{\downarrow}$
SDE-BNN(fixed step)	98.90 $\pm$ 0.07	97.59 $\pm$ 0.04	14.37 $\pm$ 0.89	87.56 $\pm$ 0.36	83.05 $\pm$ 0.29	9.89 $\pm$ 0.22
Nesterov-SDEBNN(fixed step)	99.04 $\pm$ 0.12	97.87 $\pm$ 0.06	7.37 $\pm$ 0.65	88.36 $\pm$ 0.24	85.61 $\pm$ 0.45	5.97 $\pm$ 0.11
SDE-BNN(adaptive step)	98.87 $\pm$ 0.12	97.57 $\pm$ 0.03	13.85 $\pm$ 1.26	85.87 $\pm$ 0.28	78.42 $\pm$ 0.60	6.77 $\pm$ 0.09
Nesterov-SDEBNN(adaptive step)	99.04 $\pm$ 0.03	97.88 $\pm$ 0.04	7.15 $\pm$ 0.34	86.99 $\pm$ 0.23	83.94 $\pm$ 0.32	5.32 $\pm$ 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.

Explore similar work

EMA-Nesterov: Stabilizing Nesterov's Lookahead for Accelerated Deep Learning Optimization
May 25, 2026Chung-Yiu Yau, Dawei Li, Athanasios Glentis +3Deep Learning Optimization Gradient Descent
University of Minnesota · 2Amazon AGI · 3The Chinese University of Hong Kong
Position: The Time for Sampling Is Now! Charting a New Course for Bayesian Deep Learning
May 20, 2026Emanuel Sommer, David RügamerBayesian Neural Networks Efficient Neural Network Inference
Department of Statistics, LMU Munich, Munich, Germany · Munich Center for Machine Learning, Munich, Germany.
Direct Bethe Free Energy Minimization for Bayesian Neural Networks
May 8, 2026Pavel ProchazkaBayesian Neural Networks Variational Inference
Cisco Inc.

May 25, 2026J/K move · Enter open · S save

EMA-Nesterov: Stabilizing Nesterov's Lookahead for Accelerated Deep Learning Optimization

Chung-Yiu Yau, Dawei Li, Athanasios Glentis, Valentyn Boreiko+2

University of Minnesota · 2Amazon AGI · 3The Chinese University of Hong Kong

Deep Learning Optimization Gradient Descent Accelerated Gradient Methods

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

EMA-Nesterov: Stabilizing Nesterov's Lookahead for Accelerated Deep Learning Optimization

Position: The Time for Sampling Is Now! Charting a New Course for Bayesian Deep Learning

Direct Bethe Free Energy Minimization for Bayesian Neural Networks