Fluctuation Dissipation framework for SGD

---

tags: - colorclass/a thermodynamic theory of statistical learning ---# Fluctuation-Dissipation Framework for Neural Network Training

System Definition

Let $\theta \in {\mathbb{R}}^n$ represent the network parameters, evolving according to the stochastic differential equation:

$$d{\theta }_t=-M(t)\nabla F({\theta }_t)dt+\sqrt{2k_{B}T(t)M(t)}d{W}_t$$

where: - $M(t)$ is the mobility tensor modulated by cosine annealing - $F(\theta )$ is the free energy landscape - $T(t)$ is the effective temperature - $W_t$ is a Wiener process

Cosine-Annealed Mobility

The mobility tensor incorporates learning rate annealing:

$$M(t)=\eta (t)I_n$$

where $\eta (t)$ is the cosine-annealed learning rate:

$$\eta (t)={\eta }_{min}+\frac{1}{2}({\eta }_{max}-{\eta }_{min})\left(1+\cos \left(\frac{2\pi t}{T_{cycle}}\right)\right)$$

Detailed Balance Condition

The system satisfies detailed balance with respect to the Gibbs distribution:

$$P_{eq}(\theta )\propto \exp \left(-\frac{F(\theta )}{k_{B}T}\right)$$

Fluctuation-Dissipation Relation

The noise correlation matrix $B(t)$ satisfies:

$$B(t)B(t{)}^T=2k_{B}T(t)M(t)$$

Training Dynamics

The Fokker-Planck equation describing the evolution of parameter distribution $P(\theta ,t)$:

$$\frac{\partial P}{\partial t}=\nabla \cdot \left(M(t)\nabla F(\theta )P+k_{B}T(t)M(t)\nabla P\right)$$

Information Capacity Saturation

The effective dimensionality $d_{eff}$ of learned features relates to the Fisher Information Matrix $\mathcal{I}$:

$$d_{eff}=\text{Tr}\left(\mathcal{I}(\theta )\right)={\mathbb{E}}_{p(x)}\left[{\nabla }_{\theta }{f}_{\theta }(x{)}^T{\nabla }_{\theta }{f}_{\theta }(x)\right]$$

where $f_{\theta }(x)$ is the network function.

Convergence Analysis

The expected convergence rate under temperature annealing:

$$\mathbb{E}\left[\Vert {\theta }_t-{\theta }^{\ast }{\Vert }^2\right]\le C\exp \left(-\frac{2}{T_{cycle}}{\int }_0^t\eta (s)ds\right)+\frac{k_{B}T(t)}{{\lambda }_{min}}$$

where ${\lambda }_{min}$ is the minimum eigenvalue of the Hessian at optimum ${\theta }^{\ast }$.