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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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) L1L_{1} , L2L_{2} and L3L_{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 dd -dimensional spin vector and a normalized probability distribution on the K=L3K=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 Rτ(b)(w(t))\boldsymbol{\mathcal{R}}^{(b)}_{\tau}\left(\boldsymbol{w}(t)\right) during a pL simulation as a function of time tt (blue curves), compared to expected distribution N(0,1)\mathcal{N}(0,1) (orange curves). The gray planes indicate the times at which the values of the mini-batch noise variance matrix Vτ\mathcal{V}_{\tau} are updated. b) The distribution ρ\rho of the elements of the correlation matrix CijC_{ij} during a pL simulation as a function of time tt , computed among the components of the mini-batch noise Rτ(b)(w(t))\boldsymbol{\mathcal{R}}^{(b)}_{\tau}\left(\boldsymbol{w}(t)\right) (blue curves), of the white noise R(t)\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 χ‾(Δt)\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 tt . d) The average of the mean cross-entropy L\mathcal{L} (upper plot), the squared norm ∣w∣2|\boldsymbol{w}|^{2} (center plot) and the training error ϵ\epsilon (lower plot) as a function of the temperature TT obtained from hMC (blue curves) and pL (orange curves) for the smallest model N=11160N=11160 . The relative error between the estimates is also reported (gray dotted line).
Figure 4: In the upper plot, the diffusion coefficient D∗D^{*} scaled by the network dimension NN as a function of NN , for the three studied sizes N=11160N=11160 , N=101610N=101610 and N=1006110N=1006110 and for both sampling methods hMC (blue curve) and pL (orange curve). In the lower plots, the mean squared distance d2‾\overline{d^{2}} scaled by the network dimension NN as a function of wall-clock time interval ΔtW\Delta t_{W} between sampled vectors, for the three studied sizes N=11160N=11160 , N=101610N=101610 and N=1006110N=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 ϵg\epsilon_{g} (a) , the mean cross-entropy function L\mathcal{L} (b) and the average of all the variances of the mini–batch gradient components in the first two layers Vτ‾\overline{\mathcal{V}_{\tau}} (c) sampled at equilibrium at different temperatures TT using the pL scheme. In the inserted plot, the generalization error ϵg\epsilon_{g} as function of the wall-clock time tWt_{W} during two simulations starting from initialized models at two different temperatures, T=1.0⋅10−7T=1.0\cdot 10^{-7} (orange curve) and T=3.0⋅10−7T=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 ϵg\epsilon_{g} have been computed on the same test dataset.
Figure 6: a) Sample images (28×28)\left(28\times 28\right) of the FashionMNIST training set and their projection into d=100d=100 binary vectors used during training and pL samplings. Each dd -vector is represented as a (10×10)\left(10\times 10\right) matrix. b) The generalization error ϵg\epsilon_{g} computed on the test set of the FashionMNIST dataset as a function of the wall–clock time tWt_{W} during trainings with Adam as optimizer (upper left) and pL samplings performed at different temperatures, i.e. T=5⋅10−8T=5\cdot 10^{-8} (upper right), T=1⋅10−7T=1\cdot 10^{-7} (lower left) and T=5⋅10−7T=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

NNhMC [s][s]pL [s][s]
11160111602.7⋅1002.7\cdot 10^{0}(2.0+1.5)⋅10−3\left(2.0+1.5\right)\cdot 10^{-3}
10160101605.1⋅1015.1\cdot 10^{1}(3.5+3.2)⋅10−3\left(3.5+3.2\right)\cdot 10^{-3}
100611010061103.2⋅1033.2\cdot 10^{3}(5.2+5.1)⋅10−2\left(5.2+5.1\right)\cdot 10^{-2}
Appendix
Table A1: The minimum wall-clock time interval between sampled vectors ΔtWmin\Delta t_{W}^{\mathrm{min}} for all model sizes NN , 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 10310^{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 100100 mini–batches for the estimate of Vτ\mathcal{V}_{\tau} . Also, all pL values are relative to a mini–batch size S=10−2PS=10^{-2}P .
Figure A.1: The normalized mean squared distance d2‾/N\overline{d^{2}}/N as a function of wall-clock time interval ΔtW\Delta t_{W} , for the biggest model size N=1006110N=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 ⟨d2⟩(0)=10−9\langle d^{2}\rangle(0)=10^{-9} , since the y-axis is in log-scale.
Figure A.2: The normalized mean squared distance d2‾/N\overline{d^{2}}/N as a function of wall-clock time interval ΔtW\Delta t_{W} , for the biggest model size N=1006110N=1006110 using pL with S=10−2PS=10^{-2}P (upper orange curve), pL with S=5⋅10−3PS=5\cdot 10^{-3}P (intermediate dark-orange curve), pL with S=10−3PS=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 L\mathcal{L} (upper plot), the squared norm ∣w∣2|\boldsymbol{w}|^{2} (center plot) and the training error ϵ\epsilon (lower plot) as a function of the size NN of the model at T=10−6T=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 d2‾/N\overline{d^{2}}/N as a function of wall-clock time interval ΔtW\Delta t_{W} between sampled vectors, for the three studied sizes N=11160N=11160 , N=101610N=101610 and N=1006110N=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 TT using the pL scheme.
More in this topic

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

Explore similar work

CardsList
Oct 17, 2026

Global Convergence in Deep Networks via the Second Law of Thermodynamics

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 L2L_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.

Matan Tsipory, Ofir Gaash, Itamar Harel +2

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 DynamicsStochastic Optimization ConvergenceStochastic Gradient Langevin DynamicsNeural Network StabilityNeural Tangent Kernel
AskPDF
May 14, 2026

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

The minima of modern neural network loss functions are typically not isolated, rather they form connected components of reparameterization invariant solutions on the training data. Analytically characterizing these solutions is a hard problem, but sampling approaches are feasible. By construction, existing methods either spread over low-loss regions, and thus do not sample reparameterization invariant solutions exactly, or are inherently local, which limits exploration of other minima valleys. We propose sampling such reparameterization invariant models using a dynamical system based on kinetic energy, subject to a gravitational pull and a friction term that dissipates energy from the system. Our proposed sampler, DiMS, is guaranteed to sample exactly from the minimum level sets and depends on physically motivated hyperparameters which allows control over the exploration capabilities of the sampler. We consider uncertainty quantification in Bayesian inference as the motivating problem and observe improved performance compared to previously proposed approaches.

Albert Kjøller Jacobsen, Leo Uhre Jakobsen, Johanna Marie Gegenfurtner +1

Section for Cognitive Systems, Technical University of Denmark · Center for Quantum Information Physics, New York University

Bayesian Neural NetworksDynamical SystemsBayesian Inference
AskPDF
Oct 9, 2026

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

Existing convergence guarantees for non-log-concave sampling with stochastic gradient Langevin Monte Carlo (SG-LMC) typically assume globally Lipschitz gradients and uniformly bounded noise variance. Both assumptions can fail even in elementary sampling problems, and relaxing them better reflects practical applications such as Bayesian fine-tuning of large language models (LLMs). We relax both assumptions through generalized smoothness and our proposed state-dependent heavy-tailed oracle with finite pp-th moments, 1<p≤21<p\leq 2. The oracle permits infinite variance when p<2p<2 and includes the Blum-Gladyshev (BG-0) condition. We study these relaxations separately and jointly, proposing three methods. Under globally Lipschitz gradients, robust LMC (R-LMC) handles the proposed oracle through geometric-median centering and residual clipping. With exact gradients, adaptive LMC (A-LMC) controls discretization error under generalized smoothness. Robust adaptive LMC (RA-LMC) addresses both relaxations through robust estimation and a noise-aware step size. For all three methods, we establish non-asymptotic convergence guarantees in Fisher information (FI) without assuming dissipativity or a functional inequality. Synthetic experiments demonstrate FI convergence and the benefit of each mechanism. Bayesian fine-tuning experiments with LLMs and vision-language models support our oracle model and show that RA-LMC matches or outperforms state-of-the-art baselines in accuracy and uncertainty quantification.

M. Berk Sahin, Ege C. Kaya, Arda Fazla +2

Elmore Family School of Electrical and Computer Engineering, Purdue University · Weldon School of Biomedical Engineering, Purdue University

Stochastic Gradient Langevin Dynamics
AskPDF