Correlation flow governs learning at criticality

Technical Summary: Correlation Flow Governs Learning at Criticality

Problem Statement
The initialization of deep neural networks is critical for determining whether information and gradients can propagate across depth. While empirical success has driven the scaling of network size, a unified theory connecting initialization properties to learning dynamics remains elusive. Specifically, there is a lack of understanding regarding how the interplay between correlation propagation, the spectrum of the end-to-end Jacobian, and the Neural Tangent Kernel (NTK) is conditioned by initialization choices (Gaussian vs. orthogonal) in the limit of infinitely wide and infinitely deep networks. Previous work has established that "criticality" (a specific point in the weight-bias variance plane) allows for algebraic rather than exponential decay of correlations, but the precise relationship between this correlation flow and the resulting learning dynamics (governed by the NTK) has not been fully characterized. Furthermore, the role of orthogonal initialization in controlling finite-size corrections in this limit requires rigorous theoretical grounding.

Methodology
The authors employ a combination of mean-field theory and random matrix theory to analyze fully connected neural networks in the sequential limit, where the width $N$ is taken to infinity before the depth $L$ ($N \to \infty$ followed by $L \to \infty$).

1. Theoretical Framework:

   • Mean-Field Regime: The authors analyze the deterministic recursion relations for the covariance and correlation of pre-activations as depth increases. They identify three regimes: ordered (correlation $\to 1$), chaotic (correlation $\to$ constant $<1$), and a newly identified vanishing variance phase (where variance $\to 0$).
   • Critical Point: They locate the critical point at the intersection of the vanishing variance phase boundary and the "edge of chaos" line, defined by $\sigma_w = 1/\phi_1$ and $\sigma_b = 0$ (where $\phi_1$ is the first Taylor coefficient of the activation function).
   • Jacobian Analysis: Using free probability theory, they analyze the singular value spectrum of the layer-wise and end-to-end Jacobians. They derive scaling laws for the mean and variance of the Jacobian spectrum under both Gaussian and orthogonal initialization.
   • NTK Derivation: They express the NTK as a sum of covariance kernels pulled back by partial Jacobians. By analyzing the asymptotic behavior of these terms at criticality, they derive the relationship between the NTK and the fixed-point correlation.

2. Finite-Size Analysis:

   • The authors derive the leading-order finite-size corrections ($O(L/N)$) to the empirical kernel. They decompose these fluctuations into "Wick defects" (arising from non-Gaussianity of pre-activations) and "weight non-orthogonality" effects.
   • They compare how these errors accumulate under Gaussian versus orthogonal initialization.

3. Numerical Validation:

   • Theoretical predictions are validated using fully connected networks with sine (and some tanh) activation functions. Experiments track manifold propagation, correlation flow, Jacobian spectra, and NTK conditioning across varying depths and widths.

Key Contributions and Results

1. Criticality is Necessary for Correlation Propagation:
   The authors prove that non-trivial correlation propagation to infinite depth is possible only at the critical point. In all other regimes, correlations either collapse to 1 (ordered) or a constant independent of input (chaotic), or the variance vanishes exponentially, causing numerical collapse. At criticality, the correlation converges algebraically, preserving input dependence.

2. Algebraic Vanishing of the End-to-End Jacobian:
   Contrary to the hope that criticality might enable "dynamical isometry" (where the Jacobian spectrum concentrates near unity), the authors prove that at criticality, the end-to-end Jacobian vanishes algebraically with depth ($\Xi \propto 1/L$). This occurs even with orthogonal initialization. Consequently, dynamical isometry is absent in the sequential limit at criticality.

3. Equivalence of NTK and Output Correlation:
   The paper's central theoretical result is the proof that, in the infinite-depth limit at criticality, the NTK becomes exactly proportional to the fixed-point correlation matrix (specifically, $\Theta_L \to \frac{1}{A\sigma_w^2} \rho_\infty$). This establishes a direct link between information propagation (correlation) and learning dynamics (NTK spectrum). This equivalence holds regardless of the vanishing Jacobian, as the decay is compensated in the rescaled overlap Jacobians used in the NTK definition.

4. Role of Orthogonal Initialization:
   While orthogonal initialization cannot restore dynamical isometry, it plays a crucial role in controlling finite-size effects.

   • Under Gaussian initialization, finite-size fluctuations in the correlation and NTK scale as $O(\sqrt{L/N})$, accumulating with depth and leading to ill-conditioned spectra.
   • Under orthogonal initialization, the leading-order fluctuations arising from weight non-orthogonality are suppressed. The remaining fluctuations (Wick defects) scale as $O(1/\sqrt{N})$ and are independent of depth. This allows the sequential limit predictions to hold for much deeper networks.

5. Activation Function Constraints:
   The authors show that for information to propagate without collapsing to a trivial fixed point, the activation function must satisfy $\phi_2 = 0$ (i.e., it must be an odd function at the lowest order). If $\phi_2 \neq 0$, correlations collapse regardless of initialization.

Significance and Claims
The paper claims to provide a unified framework connecting information propagation, gradient control, and learning dynamics in deep networks. Its primary significance lies in:

• Unifying Correlation and Learning: It demonstrates that the NTK spectrum, which dictates learning speed and direction, is asymptotically determined by the correlation structure of the inputs at criticality. This explains why preserving input similarity through the network is essential for training.
• Justifying Critical Initialization: It provides a rigorous theoretical justification for initializing at the critical point, showing it is the only choice that preserves input-dependent correlations while maintaining a well-defined (though vanishing) Jacobian structure.
• Clarifying Orthogonal Initialization: It resolves the tension between the "frozen" (linear) and "feature-learning" regimes by showing that orthogonal initialization is the necessary condition to suppress the uncontrolled growth of finite-size fluctuations that would otherwise destroy the predictability of the NTK in deep networks.
• Practical Implications: The results suggest that for very deep networks (e.g., SIRENs, Implicit Neural Representations), critical initialization combined with orthogonal weights is the principled approach to ensure stable training dynamics and controlled spectral bias, where the network learns along the principal modes of the input dataset's correlation structure.

The authors note that these results rely on the sequential limit and specific activation function assumptions (bounded, differentiable, $\phi_2=0$), and they acknowledge that extending these findings to the joint limit ($L/N \to c$) or unbounded activations (like ReLU) remains an open area for future work.