Neural Thermodynamics: Entropic Forces in Deep and Universal Representation Learning

TL;DR

The paper proposes an entropic-force perspective on neural network learning, framing SGD dynamics with an entropic loss $F_{\eta,\gamma}(\theta)$ and an entropic force $\nabla S(\theta)$ that together shape optimization as a gradient flow with $\dot{\theta} = -\eta(\nabla L + \gamma\theta + \nabla S)$. By showing that discretization-induced entropy breaks continuous parameter symmetries while preserving discrete ones, the authors derive gradient balance (an equipartition of gradient energy) across layers and neurons, and prove universal representation alignment via a Platonic Representation Hypothesis, including a perfect alignment theorem for deep linear networks. The framework also connects to the edge-of-stability phenomenon, explaining how data/noise balance and entropy influence sharpness and stability during training. Across ReLU nets, self-attention, and Vision Transformers, entropy-driven dynamics predict representation alignment and universal structure, providing a thermodynamics-like foundation for emergent phenomena in deep learning. The work offers predictive insights for training dynamics and a principled route to understanding universal representations and optimization behavior across architectures.

Abstract

With the rapid discovery of emergent phenomena in deep learning and large language models, understanding their cause has become an urgent need. Here, we propose a rigorous entropic-force theory for understanding the learning dynamics of neural networks trained with stochastic gradient descent (SGD) and its variants. Building on the theory of parameter symmetries and an entropic loss landscape, we show that representation learning is crucially governed by emergent entropic forces arising from stochasticity and discrete-time updates. These forces systematically break continuous parameter symmetries and preserve discrete ones, leading to a series of gradient balance phenomena that resemble the equipartition property of thermal systems. These phenomena, in turn, (a) explain the universal alignment of neural representations between AI models and lead to a proof of the Platonic Representation Hypothesis, and (b) reconcile the seemingly contradictory observations of sharpness- and flatness-seeking behavior of deep learning optimization. Our theory and experiments demonstrate that a combination of entropic forces and symmetry breaking is key to understanding emergent phenomena in deep learning.

Figures (12)

• Figure 1: Entropic forces due to discretization error and stochasticity. Left: The learning dynamics of SGD at a large learning rate (LLR) and a small learning rate (SLR) is different. One can view the difference between LLR and SLR training as coming from an entropy term, which is an order $\eta$ force. After entropic correction, the difference between SLR and LLR is reduced to $O(\eta^2)$ and it becomes possible to analyze LLR SGD with gradient flow. Right: An example of entropic effect in neural network training. ResNet18 trained on CIFAR-10 with learning rate decay at $100$ and $150$ epochs. At the first learning rate drop (black dashed lines), the gradient (entropy) increases. This is unexpected and can only be explained by the entropic loss, where a large learning rate penalizes the entropy, and thus decreasing the learning rate leads to an increase in entropy. The second drop does not create too much effect because the learning rate is too small after the first drop.
• Figure 2: Layer and neuron gradient balance during training of a two-layer ReLU network. Here, every dot is a fixed time during training, where bluer dots are closer to the initialization, and redder dots are closer to convergence. Left 1-2: The entropy is strongly correlated with the neuron balance. As entropy decreases, the neuron balance improves. In contrast, the loss is not correlated to entropic effects at all. Right: Similarly, the layer balance is also correlated with entropy and not with loss function value.
• Figure 3: The representations of two $6$-layer networks independently trained on randomly transformed MNIST become perfectly aligned for every pair of layers. The figure shows the average alignment between the same or different layers of two networks. This alignment does not weaken even if the input is arbitrarily transformed (Theorem \ref{['theo: universal representation']}). The black dashed line shows the average alignment to the input data, which is significantly weaker. Left: linear network. Right: tanh network.
• Figure 5: The entropic theory predicts the boundary for the edge of stability (EOS) phenomenon cohen2021gradient. The theory shows that the imbalance of features and the uncertainty of labels make the model converge to sharper solutions. We run a two-layer linear network trained on a regression task. The Left panel plots the quantity $\eta \lambda_{\rm max}$ at convergence. For stability, $\eta \lambda_{\rm max}$ must stay (approximately) below $2$, and the black dotted line plots the theoretical boundary for $\eta \lambda_{\rm max}=2$. Middle: The same figure that emphasizes the phase boundary. The blue-red boundary empirically defined by the condition $\eta \lambda_{\rm max} = 2 -\epsilon$ with $\epsilon=0.1$ -- due to random sampling, the actual edge of stability is slightly smaller than 2 liu2021noise. Right: We control the learning rate and balance of the label noise for this experiment. As predicted, as the data noise becomes more balanced, the sharpness metric $\eta \lambda_{\rm max}$ gets smaller, indicating better dynamical stability during training.
• Figure 7: 4-layer linear networks and no weight decay trained on a teacher-student setting. Left: Layer imbalance for each layer, which verifies Theorem \ref{['theo:layer_balance']}. Middle: Neuron balance for each layer, which verifies Theorem \ref{['theo: neuron balance']}. The curves are smoothed and averaged over $5$ runs for better visualization. Right: Loss and entropy.

Theorems & Definitions (28)

• Theorem 1
• Definition 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Theorem 6
• Theorem 7
• Theorem 8: Perfect Platonic Representation Hypothesis
• Theorem 9