A Theory of Generalization in Deep Learning

Abstract: We present a non-asymptotic theory of generalization in deep learning where the empirical neural tangent kernel partitions the output space. In directions corresponding to signal, error dissipates rapidly; in the vast orthogonal dimensions corresponding to noise, the kernel's near-zero eigenvalues trap residual error in a test-invisible reservoir. Within the signal channel, minibatch SGD ensures that coherent population signal accumulates via fast linear drift, while idiosyncratic memorization is suppressed into a slow, diffusive random walk. We prove generalization survives even when the kernel evolves $\mathcal{O}(1)$ in operator norm, the full feature-learning regime. This theory naturally explains disparate phenomena in deep learning theory, such as benign overfitting, double descent, implicit bias, and grokking. Lastly, we derive an exact population-risk objective from a single training run with no validation data, for any architecture, loss, or optimizer, and prove that it measures precisely the noise in the signal channel. This objective reduces in practice to an SNR preconditioner on top of Adam, adding one state vector at no extra cost; it accelerates grokking by $5 \times$, suppresses memorization in PINNs and implicit neural representations, and improves DPO fine-tuning under noisy preferences while staying $3 \times$ closer to the reference policy.

• The paper presents a novel operator-theoretic framework that decomposes the network’s output space into signal channels, where learning is effective, and reservoirs, where residual errors remain test-invisible.
• It formalizes minibatch SGD as a drift-diffusion process that accumulates signal faster than noise, providing a concrete mechanism for benign overfitting and double descent.
• The theory yields actionable insights for designing population-risk-aware optimizers, with demonstrated empirical improvements in grokking, PINNs, and noisy preference optimization.

A Theory of Generalization in Deep Learning: Technical Summary

Introduction and Context

The paper "A Theory of Generalization in Deep Learning" (2605.01172) introduces a comprehensive, non-asymptotic framework for understanding generalization in deep neural networks, with particular attention to the full feature-learning regime. The approach characterizes generalization by decomposing the output space via the empirical Neural Tangent Kernel (NTK), introducing the notions of signal channels (directions which SGD training can influence effectively and are visible at test time) and reservoirs (high-dimensional subspaces orthogonal to signal where residual error is trapped and test-invisible). The work systematically explains phenomena including benign overfitting, double descent, implicit bias, and grokking as natural consequences of the output-space dynamics and kernel evolution.

Output-Space Dynamics: Signal Channel versus Reservoir

The central construction operates in the output space of the network along the realized optimization trajectory. For a training set $S = (z_1, \dots, z_n)$, the per-parameter Jacobian $\bm J_S$ and the empirical tangent kernel $K = \bm J_S \bm J_S^\top$ define the directions in output space where training can effectively induce movement. The trajectory-integrated kernel yields a cumulative dissipation Gramian $\mathcal W_S$, whose range corresponds to the signal channel and whose kernel is the reservoir.

• Signal channel: Directions in which loss is dissipated during training and which SGD/Adam can align with signal in the data.
• Reservoir: High-codimension space (determined by near-zero Gramian eigenvalues) where residual training error cannot be moved—and which does not influence test outputs, even under full feature learning and large kernel drift.

Key implication: Residual error in the reservoir cannot impact test predictions, providing a mechanism for benign overfitting and test-time stability despite overparameterization.

Minibatch SGD: Drift-Diffusion Separation and Noise Suppression

Within the signal channel, the SGD dynamic is formalized as a superposition of drift (population gradient) and diffusion (minibatch noise):

• Drift: Accumulates linearly along population-gradient directions, supporting the "memorization via structure" effect.
• Diffusion: Idiosyncratic per-example memorization is diffused by minibatch noise, accumulating at a slower, sublinear rate $O(\sqrt{\eta T / b})$.

Consequently, label noise fitted on signal-channel directions is systematically suppressed by SGD's centered minibatch fluctuations, and signal accumulates much faster than noise.

Train-Test Coupling under Feature Learning

A central result is a deterministic, trajectory-dependent linear relationship (operator-valued) between training and test displacement in output space, even as the tangent kernel evolves by $O(1)$ in operator norm (i.e., beyond the lazy regime):

• There exists an optimal train-to-test linear predictor $\bm A_\circ$ (constructed from the trajectory-specific Gramian and test transfer operator) such that

$$\bm U_Q(T) - \bm U_Q(s) = \bm A_\circ (\bm U_S(T) - \bm U_S(s))$$

for the range of cumulative dissipation.

• When the training loss is squared, test set motion is completely determined by training set motion projected onto the signal channel.
• The irreducible remainder (non-predictable part), quantified as $\bm R_\perp$, vanishes along the actual trajectory, extending the classical frozen-kernel result to the feature-learning regime.

This provides a concrete decomposition of test error into bias and signal-channel variance (the only surviving stochastic part), with non-interference from the reservoir.

Unified Explanation of Generalization Phenomena

This framework yields detailed mechanistic explanations for classical and modern phenomena:

• Benign overfitting: Label noise fit into the reservoir during training is unconditionally invisible at test, validating observed harmlessness of interpolation under overparameterization.
• Double descent: The transfer of noise and signal energy between reservoir and signal channel under varying model capacity and data size produces the characteristic risk curve.
• Implicit bias: The filling of the signal channel under gradient flow is governed by the spectral structure (largest eigenvalues first), autonomously selecting low-complexity solutions.
• Grokking: Delayed generalization is interpreted as the migration of signal from the reservoir to the signal channel as the kernel evolves through training.
• Ridge regression, minimum-norm solutions: The spectral filter controlling risk decompositions in the classical regime is a special case of the general trajectory-coupled analysis.

Population-Risk Objective and Training Algorithm

A significant practical contribution is the derivation of an exact first-order estimate of population risk using a single training trajectory, without validation data, for any architecture, loss, or optimizer.

• The population-risk decrement at each step is computed via an off-diagonal kernel block contraction (leave-one-out form), naturally realized as a per-parameter variance gate atop Adam.
• The update is as simple as adding an extra state vector: a parameter is updated only if the minibatch mean squared gradient exceeds its rescaled variance, generalizing signal-to-noise-style controls.
• This induces an SNR-like preconditioner that systematically suppresses memorization and enhances transfer.
• Empirical results:
  • Grokking in modular arithmetic tasks: grokking accelerated by $5\times$ compared to Adam.
  • PINNs and implicit neural representations: memorization is suppressed, and convergence steps reduced $\bm J_S$0 or more.
  • Noisy DPO preference optimization: maintains alignment closer to the reference policy and boosts reward accuracy while remaining up to $\bm J_S$1 closer to the reference.

Technical Implications

This operator-theoretic view unifies disparate findings under a single spectral framework tied to the actual path traversed by optimization, not relying on worst-case geometric, uniform-convergence, or capacity bounds, nor on frozen-kernel lazy approximations. The analysis is constructive, computable from observed gradients and Jacobians along any trajectory, and robust to strong kernel evolution and the full feature-learning dynamics characteristic of practical large-scale deep learning.

• Reservoir invariance and invisibility: Generalization results hold for all parameter preconditioners $\bm J_S$2, not just in vanilla SGD.
• Non-asymptotic, path-dependent guarantees: All bounds and decompositions hold for finite time horizons, arbitrary architectures, and realistic training runs.
• Deterministic, non-probabilistic control: Disentangles the sources of test error without distributional concentration or asymptotics.

Theoretical and Practical Outlook

This theory provides a rigorous basis for designing optimizers and generalization-aware training procedures. By exposing the exact conditions for test transfer, bias, and surviving variance, it supports avenues for:

• Automated capacity selection based on signal-channel occupation and visibility spectrum.
• Development of population-risk-aware optimizers with theoretical safety guarantees for memorization suppression.
• Deeper analyses of label-noise robustness, structure learning, and interpolation phenomena in non-lazy, realistic training regimes.
• Building bridges between continuous-time analyses (gradient flow ODEs) and practical algorithms (minibatch SGD/Adam).

Conclusion

"A Theory of Generalization in Deep Learning" (2605.01172) establishes a mathematically precise, operator-based foundation for generalization analysis in overparameterized neural networks, capturing signal-versus-noise dynamics beyond the simplistic assumptions of previous kernel-based models. It provides both an explanatory apparatus for multiple empirical phenomena and a practically validated population-risk training rule with strong empirical benefits. This work fundamentally advances our understanding of the structure of generalization in modern deep learning and enables both novel algorithmic development and future theoretical extensions.