A Theory of Generalization in Deep Learning
Elon Litman, Gabe Guo
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 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 , suppresses memorization in PINNs and implicit neural representations, and improves DPO fine-tuning under noisy preferences while staying closer to the reference policy.
AI Impact Assessments
(1 model)Scientific Impact Assessment
Core Contribution
This paper proposes a non-asymptotic theory of generalization in deep learning built around the empirical neural tangent kernel (eNTK) partitioning the output space into a "signal channel" (where training dissipates loss) and a "reservoir" (where residual error is trapped but test-invisible). The key claims are: (1) reservoir directions cannot affect test predictions (Proposition 3.2), (2) within the signal channel, SGD's drift-diffusion separation suppresses memorization (Theorem 4.1), (3) train-test coupling holds exactly under squared loss even with O(1) kernel drift (Theorem 5.1), and (4) this framework yields a practical population-risk optimizer derived from leave-one-out exchangeability arguments (Theorem 6.2, Corollary 6.3).
The paper attempts to unify several disparate phenomena—benign overfitting, double descent, implicit bias, and grokking—under one spectral decomposition framework. The practical contribution is an SNR-gated Adam variant that updates parameter k only when μ²_k > σ²_k/(b-1).
Methodological Rigor
Theoretical strengths: The output-space formulation is clean. The cumulative dissipation Gramian W_S(s,T) as an integral over the propagator-weighted kernel is a natural object that absorbs kernel evolution. Proposition 3.2 (reservoir test-invisibility) follows elegantly from the shared J_S^⊤ factor between K_SS and K_QS. Theorem 5.1's proof that ker(D) = ker(W) under quadratic loss is tight and uses the explicit propagator structure well.
Theoretical concerns: Several claims require scrutiny:
Experimental concerns: The experiments are limited in scale. Grokking on modular arithmetic (Figure 4) uses a 2-layer Transformer on a toy task. The PINN experiment is a simple linear advection equation. The DPO experiment fine-tunes a 0.5B model. While these demonstrate the algorithm works, they don't validate the theory at scales where generalization puzzles are most acute. The 5× grokking speedup is striking but the baseline comparison (AdamW only) is incomplete—other regularization approaches and learning rate schedules could narrow this gap.
Potential Impact
The *theoretical* framework, if the stability assumptions can be validated, provides an appealing unification lens. The four-cell decomposition (Figure 1) is pedagogically powerful and could influence how researchers think about generalization failure modes.
The *practical* contribution—the SNR gate on Adam—is remarkably simple (one extra state vector, one line change). If it proves robust across architectures and scales, it could see adoption. The PINN and INR denoising results address real pain points in those communities (sensitivity to noise in boundary/initial conditions). The DPO result is timely given the importance of preference learning under noisy annotations.
However, the claim of deriving "an exact population-risk objective from a single training run with no validation data, for any architecture, loss, or optimizer" overstates what is proven. The exactness holds under specific conditions (fresh batches, quadratic loss for coupling), and the first-order approximation discards O(η²) terms that may matter in practice.
Timeliness & Relevance
The paper addresses a genuine need: understanding when and why deep networks generalize despite overparameterization remains a central question. The feature-learning regime (beyond lazy/NTK) is where the field needs theoretical tools most. The practical connections to PINNs, INRs, and RLHF alignment are timely applications.
Strengths
1. Elegant mathematical framework: The output-space formulation with W_S, D, G operators is natural and the reservoir/signal channel partition is intuitive.
2. Unification: Connecting benign overfitting, double descent, implicit bias, and grokking through one spectral filter is intellectually satisfying.
3. Theory-to-practice pipeline: Going from pathwise operators to a one-line optimizer modification is impressive.
4. Comprehensive appendices: The proofs are detailed and well-organized.
Limitations
1. Stability assumptions are unverified: The replace-two stability and fresh-batch independence are strong assumptions that limit the theory's applicability.
2. Scale of experiments: No experiments on ImageNet-scale vision or billion-parameter language models where generalization is most puzzling.
3. Squared loss restriction: Theorem 5.1 (exact coupling) requires quadratic loss; cross-entropy, the workhorse of modern DL, gets only approximate treatment.
4. Comparison depth: The algorithm is compared only against AdamW; comparisons with SAM, weight averaging, or other regularization methods are absent.
5. Verifiability of framework: The pathwise operators are defined along the training trajectory and require ODE solves to compute—making empirical validation of intermediate theoretical claims difficult at scale.
Overall Assessment
This is an ambitious paper with a clean theoretical framework and a useful practical contribution. The mathematical development is substantial and largely rigorous within its assumptions. However, the gap between the theory (which assumes stability conditions and quadratic loss) and the practical claims (any architecture, loss, optimizer) is not fully bridged. The experiments, while positive, are too small-scale to be convincing about the theory's explanatory power for modern deep learning. The SNR-gated optimizer is the most immediately impactful contribution.
Generated May 5, 2026
Comparison History (94)
Paper 1 combines high novelty (a non-asymptotic generalization theory spanning feature learning, evolved NTK, and multiple deep learning phenomena) with unusually broad cross-field impact (theory + practical optimizer-style objective applicable to any architecture/loss/optimizer). Its claimed methodological rigor (proofs, derived objective) and wide applicability to modern training pathologies make it highly timely and likely to influence both theory and practice. Paper 2 is a strong, elegant, and useful methodological extension for equivariant ML force fields with clear applications, but its impact is narrower to computational chemistry/materials compared to Paper 1’s potential to reshape general deep learning understanding and tooling.
Paper 2 offers a unifying foundational theory addressing deep learning's biggest mysteries, such as grokking and double descent, and translates this into a universally applicable optimization improvement. While Paper 1 makes impressive strides in scaling neural quantum states, Paper 2 possesses vastly broader impact across the entire machine learning community. Its ability to derive a population-risk objective without validation data and improve diverse applications like DPO fine-tuning and PINNs guarantees widespread, transformative utility across multiple scientific domains.
Paper 1 presents a comprehensive theoretical framework for generalization in deep learning that unifies multiple previously disparate phenomena (benign overfitting, double descent, grokking, implicit bias) under a single theory. It extends to the full feature-learning regime beyond the lazy/NTK limit, derives a practical validation-free objective, and demonstrates broad applications across architectures. Its breadth of impact across deep learning theory and practice, combined with its foundational theoretical contributions, gives it higher potential impact than Paper 2, which, while practically valuable, addresses a more specific problem in LLM reward modeling.
While Paper 2 offers foundational theoretical insights into neural network expressivity, Paper 1 demonstrates higher potential impact by bridging a unifying theory of generalization with immediate practical applications. Paper 1 not only explains critical modern phenomena like grokking and double descent, but also translates its theory into a zero-cost optimizer modification that improves DPO fine-tuning and suppresses memorization. This rare combination of deep theoretical novelty, timely relevance to current LLM training pipelines, and immediate real-world algorithmic utility gives it a broader and more significant scientific impact across both theoretical and applied domains.
Paper 1 provides a foundational framework connecting neural network expressivity to computational/algorithmic complexity rather than just regularity, yielding a broad characterization of universal approximation with explicit complexity bounds. This reframes how we understand NN expressivity with wide-ranging theoretical implications. Paper 2 addresses generalization theory with practical applications, but its NTK-based framework with claimed explanations for many disparate phenomena (benign overfitting, double descent, grokking, etc.) in a single theory raises concerns about overreach. Paper 1's rigorous complexity-theoretic perspective offers more durable foundational contributions.
Paper 1 presents a comprehensive theoretical framework for generalization in deep learning that unifies multiple disparate phenomena (benign overfitting, double descent, grokking, implicit bias) under one coherent theory, extends to the full feature-learning regime, and derives a practical training objective from theory. Its breadth of impact spans deep learning theory and practice across multiple domains (PINNs, LLM fine-tuning, neural representations). Paper 2 offers a useful architectural insight (tapered width allocation) but is narrower in scope—an empirical design principle for language models with incremental improvements. Paper 1's theoretical depth and cross-domain applicability give it substantially higher potential impact.
Paper 2 has higher estimated impact due to its broader and deeper contribution: a non-asymptotic generalization theory spanning the feature-learning regime, with unifying explanations for multiple major phenomena (benign overfitting, double descent, implicit bias, grokking). It also proposes a practical, architecture- and optimizer-agnostic objective/preconditioner with demonstrated benefits across diverse domains (PINNs, implicit representations, DPO). If rigorous and correct, it can influence both theory and practice widely. Paper 1 is a useful, low-cost architectural heuristic, but its scope is narrower and more incremental.
Paper 1 presents a comprehensive theoretical framework that unifies multiple fundamental phenomena in deep learning (benign overfitting, double descent, grokking, implicit bias) under a single theory, extends to the full feature-learning regime, and derives a practical training objective with strong empirical results across diverse tasks. Its breadth of impact across deep learning theory and practice, combined with its novel unification of disparate phenomena and immediately applicable optimizer improvements, gives it substantially higher potential impact than Paper 2, which, while innovative in scalable geometric modeling, addresses a more specialized problem with narrower applicability.
Paper 1 presents a comprehensive theoretical framework for generalization in deep learning that unifies multiple disparate phenomena (benign overfitting, double descent, grokking, implicit bias) under a single theory, extends to the full feature-learning regime, and derives a practical training objective from theory. Its breadth of impact spans foundational theory and practical applications across multiple domains (PINNs, DPO fine-tuning, neural representations). Paper 2 provides valuable mechanistic insights into RL post-training for reasoning but is narrower in scope, focusing on one model and one domain (math reasoning), with more incremental contributions.
Paper 1 presents a comprehensive theoretical framework for generalization in deep learning that unifies multiple previously disparate phenomena (benign overfitting, double descent, grokking, implicit bias) under one coherent theory. It extends beyond the lazy/NTK regime to full feature learning, derives a practical validation-free risk estimator, and demonstrates broad empirical impact across diverse applications. Paper 2 is a solid contribution to RNN training methodology, but its scope is narrower—focused on one architectural family—and its impact may be limited as Transformers increasingly dominate. Paper 1's theoretical breadth and practical applicability give it higher potential impact.