Neural Thermodynamic Laws for Large Language Model Training

TL;DR

The paper introduces Neural Thermodynamic Laws (NTL) to connect large language model training dynamics with thermodynamics by modeling the loss landscape as a river-valley structure, decomposed into fast valley dynamics and slow river dynamics via $$\\ell=\\ell_f+\\ell_s$$. A tractable toy model $$\\ell(x,y)=c(y)+\\tfrac{1}{2}a(y)x^2$$ enables exact analysis of SGD/SignGD, revealing how an effective temperature $T\\sim\\eta$ and a heat capacity emerge, and deriving an optimal decay schedule $$\\eta_t\\approx \\frac{\\eta/2}{1+t/t_h}$$ with a characteristic time $t_h$. The framework further relates entropy in fast directions to entropic forces that influence slow dynamics, and draws analogies to Fourier conduction and the second/third laws to explain relaxation and entropic trapping. Empirical validation on GPT-2 small shows the key predictions hold in early training, and the results yield practical guidelines for LR warmup-stable-decay schemes. Overall, NTL provides a mechanistic, physics-inspired lens on LLM training that links loss decomposition, equilibrium/annealing behavior, and optimal LR schedules with thermodynamic principles, suggesting principled directions for future experimentation on larger models.

Abstract

Beyond neural scaling laws, little is known about the laws underlying large language models (LLMs). We introduce Neural Thermodynamic Laws (NTL) -- a new framework that offers fresh insights into LLM training dynamics. On the theoretical side, we demonstrate that key thermodynamic quantities (e.g., temperature, entropy, heat capacity, thermal conduction) and classical thermodynamic principles (e.g., the three laws of thermodynamics and the equipartition theorem) naturally emerge under river-valley loss landscape assumptions. On the practical side, this scientific perspective yields intuitive guidelines for designing learning rate schedules.

Figures (10)

• Figure 1: Connections between LLM training dynamics and thermodynamics.
• Figure 2: (a) LLM pretraining usually uses the WSD (warmup-stable-decay) learning rate schedule. $\eta_{\rm min}$ is the final learning rate. (b) validation loss is a linear function of $\eta_{\rm min}$ for large $\eta_{\rm min}$. (c) $\Delta\ell$ is a linear function of $\Delta\eta$ for small $\Delta\eta$.
• Figure 3: Annealing toy examples. (a)(b)(c) Isotropic loss $\ell = \sum_{i=1}^n \frac{1}{2}a\theta_i^2\ (a=2, n=10000)$. The final loss obtained by applying the decay schedule $\eta_t=b\eta_0/(1+t/t_h)$. The theoretical minimum $(b,t_h)=(0.5,10)$ (marked as a star) agrees with numerical results. (d) Anisotropic loss $\ell=\sum_{i=1}^n\frac{1}{2}a_i\theta_i^2\ (a_i=10^{-2+4i/n}, n=10000)$. We set $b=0.5$ and try $t_h=10,100,1000$. Small sharpness is slower to converge than large sharpness.
• Figure 4: Test the existence of entropic forces in LLMs. Left: Various learning rate schedules with different stable $\eta_{\rm max}=0.0003,0.0006,0.0012$ and the same $\eta_{\rm min}=0.0006$. Right: Plot validation losses against learning rate sums. Curves for different $\eta$ roughly align, suggesting slightly negative entropic forces, corresponding to a slightly narrowing valley along the river.
• Figure 5: Dependence of $\sigma$ on gradient noise $\sigma_g$ and sharpness $a$.