Weak Correlations as the Underlying Principle for Linearization of Gradient-Based Learning Systems

Published: arXiv:2401.04013v3 [cs.LG], 12 Aug 2026 (ICLR 2026 conference paper)

The paper investigates the phenomenon of linearization in wide neural networks during gradient descent-based training. The authors establish that:

"We establish that for gradient descent-based learning, linearity is equivalent to weak correlations between the first and higher-order derivatives of the hypothesis function with respect to the parameters, at initialization."

The paper frames deep learning systems as non-linear dynamical physical systems, characterized by a multitude of interacting degrees of freedom, noting that dynamical physical systems, when expanded to an infinite number of degrees of freedom, tend to exhibit a simplified form of dynamics (Anderson, 1972).

The authors address the NTK inferiority paradox—the empirical observation that NTK-based learning is less effective than finite-width neural networks on real-world data (Lee et al., 2020; Fort et al., 2020)—and propose that weak derivative correlations may explain both linearization and its limitations.

The paper lists five primary contributions:

1. Equivalence of linearity and weak derivative correlations: "We establish that for gradient descent-based learning, linearity is equivalent to weak correlations between the first and subsequent derivatives of the hypothesis function with respect to its parameters at their initial values."

2. Direct proof for wide neural networks: "We prove directly that wide neural networks display this weak derivative correlations structure. By relying on, and extending the tensor programs formalism (Yang & Littwin, 2021), our approach uniformly addresses a broader spectrum of architectures at once than any other proof we are aware of."

3. Architectural modifications and linearization rate: We demonstrate how modifications in the architecture of linearizing learning systems, and more specifically, wide neural networks, affect the linearization rate.

4. Deviation bound for stochastic gradient descent: "We derive a bound on the deviation from linearization over time during learning when utilizing stochastic gradient descent. This is a generalization of the traditional result for deterministic gradient descent (Lee et al., 2019)."

5. Random tensor asymptotic behavior formalism: We introduce the notion of random tensor asymptotic behavior as an effective analytical tool to describe the asymptotic behavior of random tensors.

The authors introduce a novel mathematical framework for characterizing random tensors:

The primary norm used is the Subordinate Tensor Norm (Kreyszig, 1991):

∥M ∥ = sup M · v 1 ×... × v r v 1,..., v r ∈ S N1,..., S Nr

Definition 2.1 (Asymptotic Upper Bound of Random Tensors): A random tensor M is asymptotically upper bounded by f ∈ N if M = O(f), which holds if and only if:

∀g ∈ N s.t. f = o (g): lim n→∞ P (∥M n ∥ ≤ g (n)) = 1

Theorem 2.1 (Definite Asymptotic Bounds for Tensors): "Consider a random tensor M with limiting parameter n as described above. There exists f ∈ N serving as a tight asymptotic upper bound for M, satisfying: M = O (f) ∧ ∀g s.t. f ̸< g: M ̸= O (g). Furthermore, the asymptotic behavior of f is unique."

The proof uses Zorn's lemma to handle the fact that the order on N is not total (e.g., sin(πn) and cos(πn) are incomparable).

The paper considers a hypothesis function F: R N → f: X → Y with parameters θ ∈ R N, a convex cost function C, and single-input-batch gradient descent:

∆x s θ (s) = θ (s + 1) − θ (s) = −η∇C (F (θ) (x s), ŷ (x s))θ=θ(s)

The linearized (NTK) dynamics are defined as:

F lin (0) = F (θ 0), ∀s ∈ N 0: F lin (s + 1) = F lin (s) − Θ 0 (·, x s) C ′ (F lin (s) (x s), ŷ (x s))

with the kernel Θ defined as: Θ(θ)(x, x′) = η∇F(θ)(x) T ∇F(θ)(x′).

Definition 3.1 (Derivatives Correlations): For any positive integer d ∈ N and non-negative integer D ∈ N 0:

C D,d (θ) = (η D/2 + d / (D!d!)) × ∇×D+d F (θ) (∇F (θ))×d

Explicitly:

C D,d (θ) i 0,i 1...i d α 1...α D+d (x 0, x 1... x d) = (η D/2 + d / (D!d!)) × Σ α 1...α D+d=1 N ∇ α 1...α D+d F i 0 (θ) (x 0) · (∇ α 1 F i 1 (θ) (x 1) · · · ∇ α d F i d (θ) (x d))

The paper notes that C 0,1 = Θ(θ), the NTK itself. For computational tractability, the authors show these correlations can be computed via the chain rule with only O(N) elements rather than O(N d).

Theorem 3.1 (Fixed Weak Correlations and Linearization Equivalence): For sufficiently small learning rate η < η the, the following are equivalent:

1. m(n)-fixed weak derivatives correlation: C d = O(1/m(n)) and C D,d = O(1/√(m(n))) uniformly for every d, D ∈ N.

2. Simple linearity: For every fixed training step s ∈ N:

• F(θ(s)) − F lin(s) = O(1/m(n))

• ∇×D F(θ(s)) − ∇×D F(θ 0) = O(1/√(m(n)))

Theorem 3.2 (Exponential Weak Correlations and Linearization Equivalence): The following are equivalent:

1. m(n)-power weak derivatives correlation: For (D,d) ≠ (0,1): C D,d = O((1/√(m(n))) d)

2. Strong linearity: For every reparametrisation η → r(n)η and every fixed training step s:

• F(θ(s)) − F lin(s) = O(r(n)/m(n))

• ∇×D F(θ(s)) − ∇×D F(θ 0) = O(√(r(n))/m(n))

The proofs use Taylor series expansion around step s−1, yielding:

∆∇×D F (θ) = Σ d=1∞ (η D/2 / D!) C D,d (θ) (−C ′ (F (θ), ŷ))×d

The authors relate their findings to several prior results:

• Chizat et al. (2019): "A rescaling of η, such as η → r(n)η, can either promote or impede the process of linearization... This insight offers a deeper understanding of the findings presented by Chizat et al. (2019), specifically elucidating how an alteration of an external scale influences linearization by affecting the scales of higher-order correlations differently from those of lower-order correlations."

• Liu et al. (2020): "The derivative correlations can be interpreted as a spectral norm, but concerning solely the gradient when considered as a vector. This interpretation refines the results presented in Liu et al. (2020). Unlike their approach, which required this ratio to be small within a neighborhood (ball), our framework demands its minimization specifically at the initialization point."

• Huang & Yau (2020): "Their kernels can be expressed as linear combinations of our correlations. However, our result is more general, as it does not rely on the structural assumptions of wide neural networks, and also generalizes to finite learning rate GD."

• Dyer & Gur-Ari (2019): The main difference between our and their approach is that they measure the asymptotic behavior of correlations directly, rather than averaging their values.

Corollary 4.1 (Weakly Correlated PGDML Deviation Over Time): Given exponentially m(n)-weakly correlated PGDML with η < η cor, and assuming the linear solution converges exponentially fast with typical time T:

F (θ (s)) − F lin (s) = O (s 0 / m (n))

where the asymptotic bounds are uniform in s.

The proof uses induction and leverages the fact that in the NTK limit, during the initial phases of the learning process, the system converges towards the target function exponentially fast.

The paper proves that fully connected neural networks (FCNNs) are √n-weakly correlated and power-weakly correlated. Key results:

• Activation function condition: The linearization rate is governed by sup n∈N φ[n]/(n+1)!, requiring φ[n] ≤ O((n+1)!) over the relevant domain.

• Initialization schemes: The framework accommodates a broad spectrum of initialization schemes, extending beyond the Gaussian initialization predominantly employed in other studies.

• Appropriate initialization (Definition G.1): Requires independence of parameters, symmetry around 0 (vanishing odd moments), and uniform normalization of all moments.

• Learning rate normalization: η ∼ 1/n.

Lemma G.2 (Asymptotic Behavior of Layer Derivatives): For each natural number D and layer l = 1,..., L:

(η D/2 / √(N D)) O(∇×D F(l)) ≤ O(1) Uniformly

where N D = n l n l−1 D n D represents the asymptotic behavior of the number of elements in the derivatives.

The proof for FCNNs proceeds by induction, showing that for any layer l, the correlations can be expressed as combinations of previous layer correlations (Lemma G.4), with the key insight that:

for most indices, the relevant terms are offset by the irrelevant ones, keeping the overall expression small

The generalization to tensor programs formalism is stated: Our proof for FCNNs can be generalized for any wide network described by the tensor programs formalism... all such systems exhibit a wide semi-linear form by definition.

The authors conclude:

"We propose that the weak derivatives correlations (3.1) are the underlying structure behind this phenomenon. We demonstrate that this formalism is natural for analyzing linearization: (i) it allows us to determine whether, and how fast, a general system undergoes linearization; and (ii) it aids us in analyzing the deviations from linearization during learning."

They raise a pivotal question:

"Is the emergence of the weak correlations structure simply a tool to ensure a linear limit for overparameterized systems? Or do weak correlations indicate an absence of inherent biases, leading to linearization? If the latter is true, it suggests that in systems with pre-existing knowledge, specific non-linear learning methodologies reflecting those biases might be beneficial. That could partially explain why the NTK limit falls short in comparison to finite neural networks."

The paper presents two sets of experiments:

Using fully connected architectures with mini-batch gradient descent, learning rate η 0 = 1, on CIFAR10, MNIST, and FMNIST datasets with ReLU, Sigmoid, and Erf activations across 1, 2, and 3 layers. The difference function between f and f lin was measured, and second/third-order correlations were estimated by randomly sampling weights.

Using a two-hidden-layer MLP (784 → n → n → 10) with GELU activations, standard PyTorch parameterization (without explicit 1/√n NTK output scaling), on MNIST with MSE loss. The linearized model was computed via Jacobian-vector products. Widths ranged from 2048 to 19484, with 100 random seeds per configuration, and learning-rate scaling η(n) = η 0 n α for α ∈ 0, −0.25, −0.5, −0.75, −1.0.

The empirical slopes from log-log power-law fits were: ≈ −1.14 (α = 0), ≈ −1.14 (α = −0.25), ≈ −1.3 (α = −0.5), ≈ −1.73 (α = −0.75), and ≈ −2.72 (α = −1). The authors note: the empirical trends do not perfectly match the theoretical expectation: the observed slopes are flatter than predicted, attributing this to finite-size effects and the fact that our theoretical results provide upper bounds.

The paper enumerates key assumptions:

1. F, C, and φ are analytical in their parameters

2. All derivatives of φ are bounded as in Equation 23

3. Analysis restricted to single-batch stochastic gradient descent

4. C is convex (C″ positive definite)

5. Theorems apply only to PGDML systems (properly normalized)

6. Theorem 3.1 and Corollary 4.1 valid only for sufficiently small η

7. Corollary 4.1 assumes exponential decay of C′ and bounded C″

8. Theorems require all derivatives to remain fixed (with a noted generalization possibility for orthogonal changes)

This work is supported in part by Israeli Science Foundation excellence center, the US-Israel Binational Science Foundation, and the Israel Ministry of Science.

Based on the paper, I can improve AI systems in the following specific ways:

Improvement: Implement a monitoring system that tracks derivative correlations (defined in Equation 10) during training to predict when a network will enter the lazy NTK regime.

What the improved system can do:

• Before training, estimate the width needed for a target linearization level using the correlation bounds from Theorem 3.2

• During training, compute the second-order correlation (C0,2) using the efficient chain-rule method (Equation 13) to detect if the network is deviating from linear behavior

• Adjust learning rate dynamically based on correlation strength, since Theorem 3.2 shows that rescaling η by r(n) directly controls linearization rate

These improvements are directly actionable and grounded in the paper's theoretical results, particularly Theorems 3.1, 3.2, Corollary 4.1, and the tensor asymptotic behavior framework.

Deep learning models, such as wide neural networks, can be conceptualized as nonlinear dynamical physical systems characterized by a multitude of interacting degrees of freedom. Such systems in the infinite limit, tend to exhibit simplified dynamics. This paper delves into gradient descent-based learning algorithms, that display a linear structure in their parameter dynamics, reminiscent of the neural tangent kernel. We establish this apparent linearity arises due to weak correlations between the first and higher-order derivatives of the hypothesis function, concerning the parameters, taken around their initial values. This insight suggests that these weak correlations could be the underlying reason for the observed linearization in such systems. As a case in point, we showcase this weak correlations structure within neural networks in the large width limit. Exploiting the relationship between linearity and weak correlations, we derive a bound on deviations from linearity observed during the training trajectory of stochastic gradient descent. To facilitate our proof, we introduce a novel method to characterise the asymptotic behavior of random tensors.