Precise Dynamics of Diagonal Linear Networks: A Unifying Analysis by Dynamical Mean-Field Theory (2510.01930v1)

Abstract: Diagonal linear networks (DLNs) are a tractable model that captures several nontrivial behaviors in neural network training, such as initialization-dependent solutions and incremental learning. These phenomena are typically studied in isolation, leaving the overall dynamics insufficiently understood. In this work, we present a unified analysis of various phenomena in the gradient flow dynamics of DLNs. Using Dynamical Mean-Field Theory (DMFT), we derive a low-dimensional effective process that captures the asymptotic gradient flow dynamics in high dimensions. Analyzing this effective process yields new insights into DLN dynamics, including loss convergence rates and their trade-off with generalization, and systematically reproduces many of the previously observed phenomena. These findings deepen our understanding of DLNs and demonstrate the effectiveness of the DMFT approach in analyzing high-dimensional learning dynamics of neural networks.

• The paper presents a DMFT-based framework that derives a low-dimensional stochastic process to capture DLNs’ high-dimensional gradient flow.
• It identifies distinct dynamical regimes based on initialization scales, revealing a sharp transition from lazy to rich phases in optimization.
• The analysis quantifies trade-offs between fast optimization and improved generalization, linking convergence rates to implicit bias.

Precise Dynamics of Diagonal Linear Networks: A Unified Dynamical Mean-Field Theory Analysis

Introduction and Motivation

This work presents a comprehensive dynamical analysis of Diagonal Linear Networks (DLNs) under gradient flow, leveraging Dynamical Mean-Field Theory (DMFT) to derive a low-dimensional effective process that captures the high-dimensional training dynamics. DLNs, despite their linearity in input, exhibit nontrivial training behaviors due to their nonlinear parameterization, making them a valuable theoretical model for studying implicit bias, incremental learning, and initialization-dependent phenomena in neural network optimization.

The paper addresses two central questions: (1) the emergence of distinct dynamical regimes and timescales under varying initialization scales, and (2) the characterization and convergence rates of the solutions to which DLNs are driven by gradient flow. The analysis unifies previously disparate observations in the literature and provides new quantitative insights into the trade-offs between optimization speed and generalization.

Model and Theoretical Framework

The DLN considered is a two-layer diagonal linear network parameterized as $\bw = \frac{1}{2}(\bu^2 - \bv^2)$, with $\bu, \bv \in \mathbb{R}^d$. The training objective is a regularized quadratic loss:

$L(\bu, \bv) = \frac{1}{2n} \|\by - \bX \bw\|_2^2 + \frac{\lambda}{2d}(\|\bu\|_2^2 + \|\bv\|_2^2)$

where $\bX$ is a Gaussian random design matrix, $\by$ is generated from a linear teacher with additive noise, and $\lambda \geq 0$ is the regularization parameter. The analysis is performed in the proportional asymptotic regime ($n, d \to \infty$ with $n/d \to \delta$).

DMFT is employed to reduce the high-dimensional gradient flow to a scalar stochastic process, yielding a closed system of integro-differential equations for the macroscopic observables (correlation and response functions). This reduction enables precise characterization of both transient and asymptotic behaviors.

Dynamical Regimes and Timescale Structure

The analysis reveals that the training dynamics of DLNs are highly sensitive to the scale of initialization $\alpha$:

• Large Initialization ($\alpha \gg 1$): The network initially operates in a "lazy" regime, behaving as a linear model with rapid loss decay but poor generalization. At a critical time $t_c \sim 2\log(\alpha)/\lambda$, a sharp transition ("grokking") occurs to a "rich" regime, where the dynamics become nonlinear and the solution acquires a sparsity bias.

Figure 1: Large initialization ($\alpha \gg 1$) leads to a two-phase dynamic: an initial lazy regime followed by a sharp transition to a rich, generalizing regime.

• Small Initialization ($\alpha \ll 1$): The dynamics exhibit an early plateau (search phase) with negligible loss reduction, followed by a descent phase on a timescale $\Theta(\log(1/\alpha))$ characterized by incremental learning—coordinates are activated sequentially, leading to sparser solutions.

  Figure 2: Training and test error dynamics for small $\alpha$ ($d=200$), showing the early plateau and subsequent incremental descent.

The timescale separation and the sharpness of the transition in the large initialization regime are validated by rescaling time appropriately, leading to a collapse of learning curves across different $\alpha$.

Figure 3: For large $\alpha$, rescaling time by $\alpha^2$ collapses the initial descent of training and test errors onto the predicted lazy solution.

Figure 4: For large $\alpha$, rescaling time by $\log(\alpha)$ aligns the transition times to the rich phase, confirming the predicted scaling.

Figure 5: After shifting time by the transition time $2\log(\alpha)/\lambda$, the post-transition dynamics collapse, indicating $O(1)$ convergence in the rich phase.

Figure 6: For small $\alpha$, rescaling time by $\log(1/\alpha)$ collapses the descent phase, confirming the predicted timescale for incremental learning.

Asymptotic Behavior: Fixed Points and Convergence Rates

The DMFT analysis yields a precise characterization of the fixed points of the dynamics:

• Regularized Case ($\lambda > 0$): The fixed point corresponds to the solution of an $\ell_1$-regularized regression problem, with the degree of sparsity controlled by the initialization scale.
• Unregularized, Overparameterized Case ($\lambda = 0, \delta < 1$): The solution is the minimum-norm interpolator with a norm $J_\alpha$ that interpolates between $\ell_2$ (large $\alpha$) and $\ell_1$ (small $\alpha$). As $\alpha \to 0$, the solution approaches the minimum $\ell_1$-norm interpolator, yielding improved generalization for sparse targets.

Figure 7: Test errors at fixed points for varying $\alpha$; smaller initialization yields better generalization in the overparameterized regime.

Figure 8: Fixed points for $\lambda > 0$; left: train/test errors vs. $\delta$, right: distribution of $\bw(\infty)$ showing sparsity.

Figure 9: Fixed points for $\lambda = 0, \delta < 1$; left: train/test errors vs. $\alpha$, right: distribution of $\bw(\infty)$.

The convergence rate of the loss is also derived:

• Regularized Case: The convergence is subexponential due to the presence of arbitrarily slow paths near the threshold.
• Unregularized Case: The loss converges exponentially, $L(t) \sim \exp(-2\gamma t)$, with the rate $\gamma$ increasing monotonically with $\alpha$. Thus, smaller initialization improves generalization but slows optimization.

Figure 10: Convergence rates of training errors for Gaussian data; smaller $\alpha$ leads to slower convergence, confirming the trade-off.

Figure 11: Theoretical convergence rates $\gamma$ for $\delta = 2$ and $\delta = 0.5$; $\gamma$ increases monotonically with $\alpha$.

Trade-off Between Generalization and Optimization Speed

A central result is the explicit trade-off between generalization and optimization speed: smaller initialization leads to better generalization (sparser solutions) but slower convergence. This is quantitatively established via the dependence of the fixed point and convergence rate on $\alpha$.

Figure 12: Training and test error dynamics for large $\alpha$ ($d=200$), illustrating the trade-off between rapid optimization and poor generalization in the lazy regime.

Figure 13: Convergence of the loss for $\delta = 2$; empirical curves match theoretical predictions for the convergence rate.

Rigorous Justification and Universality

The DMFT equations are rigorously justified for truncated DLNs, and the results are shown to be universal with respect to the input distribution. Experiments on real-world gene expression data and non-Gaussian synthetic data confirm the theoretical predictions, up to finite-sample effects.

Figure 14: Empirical distribution of whitened gene expression data, exhibiting heavier tails than the standard Gaussian.

Implications and Future Directions

The analysis demonstrates that DMFT provides a powerful and unifying framework for understanding the high-dimensional dynamics of neural network training, even in models with nontrivial parameterization-induced implicit bias. The explicit trade-off between generalization and optimization speed suggests that "better" solutions (in terms of generalization) are inherently harder to reach, a principle that may extend to broader classes of models and optimization algorithms.

Potential future directions include:

• Extending the DMFT analysis to deep linear and nonlinear networks, as well as to architectures such as transformers.
• Analyzing the impact of stochastic optimization (SGD) and other algorithmic choices on the implicit bias and dynamical regimes.
• Investigating the universality of the observed trade-offs and phase transitions in more complex, realistic settings.

Conclusion

This work provides a unified, quantitative, and rigorous analysis of the training dynamics of DLNs, elucidating the interplay between initialization, implicit bias, and optimization timescales. The DMFT approach not only reproduces and explains a range of previously observed phenomena but also yields new insights into the fundamental trade-offs governing high-dimensional learning dynamics. The results have broad implications for the theoretical understanding of neural network optimization and generalization, and the DMFT methodology is poised to play a central role in future analyses of complex learning systems.