Label-NTK Alignments and A Tighter Convergence Bound in the NTK Regime

Technical Summary: Label-NTK Alignments and A Tighter Convergence Bound in the NTK Regime

Problem Statement
The Neural Tangent Kernel (NTK) framework has successfully explained the optimization of over-parameterized neural networks by characterizing training dynamics as approximately linear. While existing theories guarantee exponential convergence to a global minimum, these guarantees are often overly pessimistic compared to empirical observations. Specifically, classical convergence rates depend critically on the smallest eigenvalue ($\lambda_{\min}$) of the NTK matrix. In realistic settings, $\lambda_{\min}$ is typically extremely small, leading to theoretical convergence bounds that are orders of magnitude slower than actual training curves. Furthermore, these worst-case analyses implicitly assume that the data label vector aligns with the eigenvector corresponding to $\lambda_{\min}$, a scenario that rarely occurs in practice. This paper addresses the gap between theoretical bounds and practical training dynamics by investigating whether a tighter convergence bound can be derived that accounts for the specific relationship between data labels and the NTK eigen-spectrum.

Methodology
The authors propose a novel analysis based on the empirical and theoretical observation of "alignments" between data vectors (labels and residuals) and the NTK eigen-system.

1. Empirical Observation: Through experiments on Multi-Layer Perceptrons (MLPs) and Convolutional Neural Networks (CNNs) across datasets like CIFAR-10, the authors observe two distinct scaling laws:

   • Label-NTK Alignment: The squared projection of the label vector $y$ onto the $i$-th eigenvector $v_i$ scales with the square of the corresponding eigenvalue: $(v_i^\top y)^2 \sim \lambda_i^2$.
   • Residual-NTK Alignment: The squared projection of the initial residual vector $r_0 = f(w_0) - y$ onto the eigenvectors scales linearly with the eigenvalues: $(v_i^\top r_0)^2 \sim \lambda_i$.
     These relationships hold particularly in the small-eigenvalue regime, indicating that labels and residuals have negligible components in the directions of small eigenvalues.

2. Theoretical Justification: The authors provide a theoretical proof for the Label-NTK alignment under mild data assumptions, including unit input length, non-degeneracy, and the Lipschitz continuity of the ground-truth data-label function. By analyzing a two-point dataset and generalizing to $n$ points, they demonstrate that the Lipschitz property of the true function forces the label vector to align more strongly with larger eigenvalues, precisely predicting the observed scaling exponent of 2.

3. Derivation of New Bounds: Leveraging the Residual-NTK alignment as a condition (Condition 4.2), the authors derive a refined convergence bound for gradient descent. Instead of relying on $\lambda_{\min}$, the new bound depends on the trace of a function of the full NTK matrix: $L(w_t) \sim \text{tr}[(I - \eta K)^{2t} K]$.

Key Contributions

• Identification of Alignment Phenomena: The paper identifies and quantifies Label-NTK and Residual-NTK alignments, showing that data labels and residuals are naturally aligned with the top eigenvectors of the NTK, effectively avoiding the worst-case scenarios assumed in prior literature.
• Tighter Convergence Theory: The authors establish a new convergence theorem (Theorem 4.5) where the loss decay is governed by the full spectrum of the NTK rather than just the smallest eigenvalue. This results in a bound that closely matches empirical training curves, as demonstrated in Figure 1.
• Improved Generalization Bounds: By utilizing the Label-NTK alignment property, the paper derives improved generalization bounds for over-parameterized networks, showing that the test loss is upper bounded by $O(\text{tr}[K]/n)$, which is tighter than previous results dependent on $\lambda_{\min}$.
• Time Complexity Analysis: The paper provides a time complexity analysis under the assumption of power-law eigenvalue decay, showing that the number of iterations required to reach a target loss $\epsilon$ is significantly lower than the standard $O(\frac{1}{\eta \lambda_{\min}} \log \frac{1}{\epsilon})$ bound.

Results

• Empirical Validation: Experiments on MLPs and CNNs across multiple datasets confirm that the proposed theoretical bound tracks the actual training loss curves much more accurately than classical bounds. The theoretical curve using the trace-based bound decreases rapidly, matching the empirical data, whereas the classical bound remains nearly flat due to the tiny $\lambda_{\min}$.
• Scaling Laws: The experiments consistently show a linear relationship between $\log((v_i^\top y)^2)$ and $\log(\lambda_i)$ with a slope of 2, and between $\log((v_i^\top r)^2)$ and $\log(\lambda_i)$ with a slope of 1, validating the proposed scaling laws.
• Generalization: The derived generalization bound suggests that the test loss scales with the trace of the NTK, offering a more favorable scaling than previous worst-case analyses.

Significance
The paper claims that its primary significance lies in resolving the discrepancy between NTK-based theoretical predictions and practical training performance. By demonstrating that the worst-case assumption (label alignment with $\lambda_{\min}$) does not hold in practice due to inherent data properties (Lipschitz continuity), the authors provide a mathematically rigorous explanation for the fast convergence observed in deep learning. The work offers a "sharper" theoretical lens that incorporates full-spectral information, thereby bridging the gap between idealized linearized dynamics and the observed efficiency of gradient descent in over-parameterized regimes. The authors note that while their theory is limited to the NTK regime (sufficiently wide networks, low learning rates) and does not capture nonlinear feature learning dynamics like the Edge of Stability, it provides a necessary foundation for understanding optimization in the most mathematically tractable setting available.