Controlled Langevin Dynamics for Sampling of Feedforward Neural Networks Trained with Minibatches

Alessandro Zambon; Francesca Caruso; Riccardo Zecchina; Guido Tiana

cond-mat.dis-nnMar 16, 2026

Controlled Langevin Dynamics for Sampling of Feedforward Neural Networks Trained with Minibatches

Authors: Alessandro Zambon, Francesca Caruso, Riccardo Zecchina, Guido Tiana

Organizations: Department of Physics, Università degli Studi di Milano and INFN, via Celoria 16, 20133 Milano, Italy · Department of Computing Sciences and Bocconi Institute for Data Science and Analytics (BIDSA), Bocconi University, 20136 Milano, Italy

Abstract

Sampling the parameter space of artificial neural networks according to a Boltzmann distribution provides insight into the geometry of low-loss solutions and offers an alternative to conventional loss minimization for training. However, exact sampling methods such as hybrid Monte Carlo (hMC), while formally correct, become computationally prohibitive for realistic datasets because they require repeated evaluation of full-batch gradients. We introduce a pseudo-Langevin (pL) dynamics that enables efficient Boltzmann sampling of feed-forward neural networks trained with large datasets by using minibatches in a controlled manner. The method exploits the statistical properties of minibatch gradient noise and adjusts fictitious masses and friction coefficients to ensure that the induced stochastic process samples efficiently the desired equilibrium distribution. We validate numerically the approach by comparing its equilibrium statistics with those obtained from exact hMC sampling. Performance benchmarks demonstrate that, while hMC rapidly becomes inefficient as network size increases, the pL scheme maintains high computational diffusion and scales favorably to networks with over one million parameters. Finally, we show that sampling at intermediate temperatures yields optimal generalization performance, comparable to SGD, without requiring a validation set or early stopping procedure. These results establish controlled minibatch Langevin dynamics as a practical and scalable tool for exploring and exploiting the solution space of large neural networks.

Figures & tables

Figure 1: Sketch of the feed–forward ANN used in this work, with (from left to right) $L_{1}$ , $L_{2}$ and $L_{3}$ neurons for the first, second and third layer respectively. On the left and on the right of the NN, a possible couple of input and output vectors respectively, that is a $d$ -dimensional spin vector and a normalized probability distribution on the $K=L_{3}$ classes.

Figure 2: Scheme of the pseudo–Langevin (pL) algorithm.

Figure 3: a) The distribution $\rho$ of the components of the mini-batch noise $\boldsymbol{\mathcal{R}}^{(b)}_{\tau}\left(\boldsymbol{w}(t)\right)$ during a pL simulation as a function of time $t$ (blue curves), compared to expected distribution $\mathcal{N}(0,1)$ (orange curves). The gray planes indicate the times at which the values of the mini-batch noise variance matrix $\mathcal{V}_{\tau}$ are updated. b) The distribution $\rho$ of the elements of the correlation matrix $C_{ij}$ during a pL simulation as a function of time $t$ , computed among the components of the mini-batch noise $\boldsymbol{\mathcal{R}}^{(b)}_{\tau}\left(\boldsymbol{w}(t)\right)$ (blue curves), of the white noise $\boldsymbol{\mathcal{R}}(t)$ (orange curves) and of the weighted sum of the two (green curves). The curves are slightly shifted along the z-axis to make them distinguishable. c) The average of the component–wise noise autocorrelation $\overline{\chi}(\Delta t)$ for the mini-batch noise (upper panel, blue curves), for the white noise (middle panel, orange curves) and for their weighted sum (lower panel, green curves), computed for small time intervals during a pL simulation parametrized by the time $t$ . d) The average of the mean cross-entropy $\mathcal{L}$ (upper plot), the squared norm $|\boldsymbol{w}|^{2}$ (center plot) and the training error $\epsilon$ (lower plot) as a function of the temperature $T$ obtained from hMC (blue curves) and pL (orange curves) for the smallest model $N=11160$ . The relative error between the estimates is also reported (gray dotted line).

Figure 4: In the upper plot, the diffusion coefficient $D^{*}$ scaled by the network dimension $N$ as a function of $N$ , for the three studied sizes $N=11160$ , $N=101610$ and $N=1006110$ and for both sampling methods hMC (blue curve) and pL (orange curve). In the lower plots, the mean squared distance $\overline{d^{2}}$ scaled by the network dimension $N$ as a function of wall-clock time interval $\Delta t_{W}$ between sampled vectors, for the three studied sizes $N=11160$ , $N=101610$ and $N=1006110$ and for both sampling methods hMC (blue curve) and pL (orange curve). The mean squared distances are obtained by averaging the values from independent simulations starting from different equilibrated weight vectors.

Figure 5: The mean value and the standard deviation of the generalization error $\epsilon_{g}$ (a) , the mean cross-entropy function $\mathcal{L}$ (b) and the average of all the variances of the mini–batch gradient components in the first two layers $\overline{\mathcal{V}_{\tau}}$ (c) sampled at equilibrium at different temperatures $T$ using the pL scheme. In the inserted plot, the generalization error $\epsilon_{g}$ as function of the wall-clock time $t_{W}$ during two simulations starting from initialized models at two different temperatures, $T=1.0\cdot 10^{-7}$ (orange curve) and $T=3.0\cdot 10^{-7}$ (green curve), and during an Adam training beyond early–stopping (red dotted curve). The gray straight lines reported in the upper plot and in the inserted one represent the best mean generalization error found with Adam training. All values of $\epsilon_{g}$ have been computed on the same test dataset.

Figure 6: a) Sample images $\left(28\times 28\right)$ of the FashionMNIST training set and their projection into $d=100$ binary vectors used during training and pL samplings. Each $d$ -vector is represented as a $\left(10\times 10\right)$ matrix. b) The generalization error $\epsilon_{g}$ computed on the test set of the FashionMNIST dataset as a function of the wall–clock time $t_{W}$ during trainings with Adam as optimizer (upper left) and pL samplings performed at different temperatures, i.e. $T=5\cdot 10^{-8}$ (upper right), $T=1\cdot 10^{-7}$ (lower left) and $T=5\cdot 10^{-7}$ (lower right). Curves with the same color in different plots start from the same stochastically generated weight vector.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

$N$	hMC $[s]$	pL $[s]$
$11160$	$2.7\cdot 10^{0}$	$\left(2.0+1.5\right)\cdot 10^{-3}$
$10160$	$5.1\cdot 10^{1}$	$\left(3.5+3.2\right)\cdot 10^{-3}$
$1006110$	$3.2\cdot 10^{3}$	$\left(5.2+5.1\right)\cdot 10^{-2}$

Appendix

Table A1: The minimum wall-clock time interval between sampled vectors

\Delta t_{W}^{\mathrm{min}}

for all model sizes

N

, and for both sampling schemes, hMC and pL. For hMC, we report the times relative to a single move with the optimal number of integration steps, that is

10^{3}

for all model sizes. For pL, the first value within the parenthesis represents the time needed for a single integration, while the second one is the time spent to extract

100

mini–batches for the estimate of

\mathcal{V}_{\tau}

. Also, all pL values are relative to a mini–batch size

S=10^{-2}P

Figure A.1: The normalized mean squared distance $\overline{d^{2}}/N$ as a function of wall-clock time interval $\Delta t_{W}$ , for the biggest model size $N=1006110$ sampled with hMC (thick blue curve), mb-hMC (dotted blue curve), pL (thick orange curve) and naive-pL (dotted orange curve). The mean squared distances are obtained by averaging the values from independent simulations starting from different equilibrated weight vectors. All curves have been shifted so that $\langle d^{2}\rangle(0)=10^{-9}$ , since the y-axis is in log-scale.

Figure A.2: The normalized mean squared distance $\overline{d^{2}}/N$ as a function of wall-clock time interval $\Delta t_{W}$ , for the biggest model size $N=1006110$ using pL with $S=10^{-2}P$ (upper orange curve), pL with $S=5\cdot 10^{-3}P$ (intermediate dark-orange curve), pL with $S=10^{-3}P$ (lower red curve) and hMC (blue curve). The mean squared distances are obtained by averaging the values from independent simulations starting from different equilibrated weight vectors.

Figure A.3: The average of the mean cross-entropy $\mathcal{L}$ (upper plot), the squared norm $|\boldsymbol{w}|^{2}$ (center plot) and the training error $\epsilon$ (lower plot) as a function of the size $N$ of the model at $T=10^{-6}$ , obtained from hMC (blue curves) and pL (orange curves) simulations at equilibrium. The orange and blue curves have been slightly shifted along the y-axis to make them distinguishable. The relative error with respect to the hMC prediction between the estimates is also reported (gray dotted line).

Figure A.4: The relative mean squared distance $\overline{d^{2}}/N$ as a function of wall-clock time interval $\Delta t_{W}$ between sampled vectors, for the three studied sizes $N=11160$ , $N=101610$ and $N=1006110$ and for both hMC and pL (left and right columns, respectively). For each size and for each method, the optimal mean squared distance curve (brighter color) is shown along with a suboptimal one (darker color).

Figure A.5: Mean and standard deviation of the average of all the variances in the first layer (upper plot), in the second layer (middle plot) and in the third layer (lower plot) for the weight vectors sampled at equilibrium at different temperatures $T$ using the pL scheme.

Stochastic Gradient Langevin Dynamics

Stochastic Gradient Langevin Dynamics (SGLD) uses noisy gradient updates to scale Bayesian sampling and uncertainty estimation to large datasets.

See all 16 papers

Global Convergence in Deep Networks via the Second Law of Thermodynamics
Oct 17, 2026Matan Tsipory, Ofir Gaash, Itamar Harel +2Neural Network Training Dynamics Stochastic Optimization Convergence
Technion - Israel Institute of Technology, The Andrew and Erna Viterbi Faculty of Electrical & Computer Engineering, Technion City, Haifa 3200003, Israel. · Technion - Israel Institute of Technology, Faculty of Data and Decision Sciences, Technion City, Haifa 3200003, Israel. · Toyota Technological Institute at Chicago, Chicago, IL, USA.
Don't Stop Me Yet: Sampling Loss Minima via Dissipative Riemannian Mechanics
May 14, 2026Albert Kjøller Jacobsen, Leo Uhre Jakobsen, Johanna Marie Gegenfurtner +1Bayesian Neural Networks Dynamical Systems
Section for Cognitive Systems, Technical University of Denmark · Center for Quantum Information Physics, New York University
Robust Langevin Sampling under State-Dependent Heavy-Tailed Noise with Applications to LLM Fine-Tuning
Oct 9, 2026M. Berk Sahin, Ege C. Kaya, Arda Fazla +2Stochastic Gradient Langevin Dynamics
Elmore Family School of Electrical and Computer Engineering, Purdue University · Weldon School of Biomedical Engineering, Purdue University

Oct 17, 2026J/K move · Enter open · S save

Global Convergence in Deep Networks via the Second Law of Thermodynamics

Matan Tsipory, Ofir Gaash, Itamar Harel, Nathan Srebro+1

Technion - Israel Institute of Technology, The Andrew and Erna Viterbi Faculty of Electrical & Computer Engineering, Technion City, Haifa 3200003, Israel. · Technion - Israel Institute of Technology, Faculty of Data and Decision Sciences, Technion City, Haifa 3200003, Israel. · Toyota Technological Institute at Chicago, Chicago, IL, USA.

Neural Network Training Dynamics Stochastic Optimization Convergence Stochastic Gradient Langevin Dynamics Neural Network Stability Neural Tangent Kernel

Open paper AskPDF arXiv

Existing convergence analyses of finite-width overparameterized deep neural networks under standard parameterization often rely on parameters remaining close to their initialization and typically neglect the role of noise, despite empirical evidence that stochasticity can improve optimization and generalization. To bridge this gap, we study

L_2

-regularized continuous-time Langevin dynamics (CLD) with isotropic noise, for training feedforward networks with smooth activations. We introduce a change-of-measure inequality, inspired by the second law of thermodynamics, that enables control of key properties of the training dynamics under the time-evolving parameter distribution. Leveraging a sharpened lower bound on the minimum eigenvalue of the Neural Tangent Kernel (NTK) at initialization, we show that injected noise helps preserve NTK stability with high probability without confining the weights to a small neighborhood of their initialization. As a result, we derive a non-asymptotic upper bound on the expected training loss. Our analysis yields explicit conditions on width and noise temperature that guarantee convergence to a finite loss floor. Finally, we experimentally show the network Jacobian can change from initialization, departing the commonly analyzed lazy regime.

Abstract

Stochastic Gradient Langevin Dynamics

Explore similar work

Global Convergence in Deep Networks via the Second Law of Thermodynamics

Don't Stop Me Yet: Sampling Loss Minima via Dissipative Riemannian Mechanics

Robust Langevin Sampling under State-Dependent Heavy-Tailed Noise with Applications to LLM Fine-Tuning