🤖 machine learning

A Theory of Generalization in Deep Learning

This paper presents a non-asymptotic theory of deep learning generalization based on the empirical neural tangent kernel's partitioning of the output space, which explains phenomena like benign overfitting and grokking while introducing a practical SNR-based preconditioner that accelerates training and suppresses memorization without requiring validation data.

Original authors: Elon Litman, Gabe Guo

Published 2026-05-05

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Elon Litman, Gabe Guo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Why Do Super-Complex AI Models Work?

Imagine you are teaching a student (a neural network) for a final exam. You give them a textbook with 100 examples. But here's the twist: the student has a photographic memory and is so smart that they can memorize every single word of the textbook, including the typos and the random scribbles in the margins.

In the past, scientists thought: "If a student memorizes the typos, they will fail the exam because the exam won't have those typos." This is the problem of overfitting.

However, in modern AI, we see something strange: these "super-memorizers" often ace the exam, even when the training data is messy. This paper provides a new map to explain how and why this happens, and it even gives us a new way to train them to do it faster and better.

1. The Two Rooms: The "Signal Channel" and the "Reservoir"

The authors imagine the AI's learning process as happening in a building with two distinct rooms.

Room A: The Signal Channel (The Stage)
This is where the "real" learning happens. It's like a stage where the student learns the actual plot of the story. When the AI moves in this direction, it is learning patterns that apply to the real world (the test).
- What happens here: The AI learns quickly and steadily. It's like a runner sprinting down a track.
Room B: The Reservoir (The Soundproof Basement)
This is a huge, dark basement where the AI stores the "noise"—the typos, the random scribbles, and the pure garbage in the data.
- The Magic Trick: The authors prove that this basement is soundproof. Even if the AI memorizes every single typo in the basement, no sound escapes. The test (the exam) cannot hear what happens in the reservoir.
- Result: The AI can memorize the noise without hurting its exam score, because the noise is trapped in a place the test can't see.

2. The Traffic Cop: How SGD Keeps Things Organized

How does the AI know which direction is the "Stage" and which is the "Basement"? The paper explains that the standard training method (called SGD or Stochastic Gradient Descent) acts like a clever traffic cop.

The Drift vs. The Shuffle:
- Real Signals (The Stage): When the AI sees a real pattern, the traffic cop pushes it forward in a straight, fast line (a "drift"). This accumulates quickly.
- Noise (The Basement): When the AI sees random noise, the traffic cop tells it to just shuffle around in place (a "random walk"). It moves, but it doesn't get anywhere useful.
- The Result: Over time, the real patterns pile up high, while the noise stays small and gets lost in the shuffle. The AI naturally separates the wheat from the chaff.

3. The "Grokking" Mystery Solved

You might have heard of a phenomenon called "Grokking." This is when an AI seems to be failing (memorizing the training data) for a long time, and then suddenly, out of nowhere, it "gets it" and starts solving the problem perfectly.

The Paper's Explanation:
Imagine the AI is slowly moving the "Signal" from the Soundproof Basement up onto the Stage.
- At first, the AI is stuck in the basement, memorizing noise.
- Slowly, the "kernel" (the AI's internal map) evolves.
- Eventually, the real signal finally migrates from the basement to the stage.
- Grokking is just the moment the signal arrives on the stage. It's not magic; it's just the signal finally catching up to the test.

4. The New Tool: "Population Risk" Training

The authors didn't just explain the theory; they built a practical tool based on it.

The Problem: Usually, to train an AI, we need a "validation set" (a practice exam) to check if it's learning the right things. If we don't have one, we might accidentally teach it the noise.
The Solution: They created a new training rule that acts like a Self-Correcting Filter.
- Instead of just looking at the whole batch of data, this new method looks at every single example and asks: "If I removed this one example, would the AI still learn the same thing?"
- If the answer is "No, it's just memorizing this specific noise," the filter blocks that update.
- If the answer is "Yes, this is a real pattern," the filter allows the update.

The Analogy:
Imagine a teacher grading a student.

Old Way (AdamW): The teacher looks at the whole test and says, "You got 90% right, good job!" (Even if the student cheated on 10 questions).
New Way (Population Risk): The teacher looks at each question and asks, "Did you learn this concept, or did you just memorize the answer key?" If it's just memorization, the teacher ignores that point. This forces the student to learn the concepts faster.

5. What Did They Achieve?

The paper tested this new method on three difficult tasks where AI usually fails or gets stuck:

Physics Simulations (PINNs): When training AI to solve physics equations with noisy data, the new method reached the correct answer 2.4 times faster than standard methods.
Math Puzzles (Grokking): On a modular division math problem, the AI reached 95% accuracy in 5,950 steps instead of the usual 29,450 steps. It "grokked" 5 times faster.
AI Chatbots (DPO): When fine-tuning a chatbot with messy human feedback (where people disagree on what's good), the new method learned better preferences while staying much closer to the original, safe behavior of the bot.

Summary

This paper tells us that deep learning works because the training process naturally sorts "real learning" from "memorized noise" into two separate rooms. The "noise" gets trapped in a soundproof basement where it can't hurt the AI's performance.

By understanding this, the authors built a new training tool that acts like a smart filter, automatically ignoring the noise and focusing only on the real signals. This makes AI learn faster, solve harder problems, and avoid the "memorization trap" without needing extra data to check its work.

Technical Summary: A Theory of Generalization in Deep Learning

Problem Statement

Deep learning models with parameter counts exceeding training examples can memorize arbitrary noise yet generalize well on real data. Classical capacity-based bounds (e.g., VC dimension, Rademacher complexity) are vacuous at practical scales, while Neural Tangent Kernel (NTK) theories describe only the "lazy" regime where kernels remain frozen. Modern architectures operate in the full feature-learning regime, where the kernel evolves significantly (by $O(1)$ in operator norm) during training. There is a lack of a non-asymptotic theory that explains generalization in this dynamic regime, unifies disparate phenomena like benign overfitting and grokking, and provides a practical method to train directly on population risk without validation data.

Methodology

The authors develop a theory based on output-space dynamics along the realized training trajectory. The core framework decomposes the output space using the cumulative dissipation Gramian ( $W_S$ ), derived from integrating the evolution of the empirical tangent kernel ( $K_{SS}$ ) and the output-gradient propagator ( $P_g$ ).

1. Signal Channel and Reservoir Decomposition

The theory partitions the output space into two orthogonal subspaces:

Signal Channel ( $\text{range}(W_S)$ ): Directions where training dissipates loss.
Reservoir ( $\ker(W_S)$ ): Directions where training dissipates no loss.

A key theoretical result is Reservoir Test-Invisibility (Proposition 3.2): The test transfer operator $G_Q$ vanishes on the reservoir. Consequently, any residual error or label noise trapped in the reservoir during training is invisible to the test set.

2. Drift-Diffusion Separation in SGD

Inside the signal channel, the paper analyzes Minibatch Stochastic Gradient Descent (SGD). It proves a separation of timescales:

Drift: The coherent population signal accumulates linearly ( $O(T)$ ) along signal directions.
Diffusion: The centered fluctuation (noise) behaves as a martingale, accumulating only as a random walk ( $O(\sqrt{T})$ ).
This ensures that genuine signal dominates idiosyncratic memorization over time, suppressing noise in the signal channel.

3. Train-Test Coupling under Feature Learning

Even when the kernel evolves by $O(1)$ , the paper proves that test motion is determined exactly by training motion within the signal channel. By normalizing the training and test displacement operators by the cumulative dissipation ( $W_S^{1/2}$ ), the test displacement $U_Q(T) - U_Q(s)$ can be expressed as a linear function of the training displacement $U_S(T) - U_S(s)$ plus an irreducible remainder $R_\perp$ . Under squared loss, this remainder vanishes, establishing an exact coupling:
$U_Q(T) - U_Q(s) = A^\circ (U_S(T) - U_S(s))$
where $A^\circ$ is the optimal linear predictor derived from the realized trajectory.

4. Population-Risk Training Algorithm

Leveraging the exchangeability of training data, the authors derive an unbiased estimator for the rate of population-risk decrease using a leave-one-out (LOO) perspective.

Objective: They define a per-parameter gate based on the ratio of the squared mean gradient ( $\mu_k^2$ ) to the batch variance ( $\sigma_k^2$ ).
Rule: A parameter $k$ is updated only if $\mu_k^2 > \sigma_k^2 / (b-1)$ .
Implementation: This results in a modification to the Adam optimizer requiring only one additional state vector (tracking streaming variance). This "population-risk gate" suppresses updates on parameters where the batch signal is dominated by noise, effectively preventing the fitting of incoherent noise.

Key Contributions

Non-Asymptotic Generalization Theory: A framework that handles full feature learning (evolving kernels) rather than relying on frozen-kernel approximations. It explains generalization via the geometric separation of signal and noise into test-visible and test-invisible subspaces.
Unified Explanation of Phenomena: The theory provides a single mechanistic explanation for:
- Benign Overfitting: Noise trapped in the reservoir (test-invisible).
- Double Descent: Noise moving between reservoir and signal channels as model capacity changes.
- Implicit Bias: The spectral schedule of $W_S(t)$ filling the signal channel from largest to smallest eigenvalues.
- Grokking: Signal migrating from the reservoir into the signal channel as the kernel evolves over time.
Exact Population-Risk Objective: A method to derive an exact population-risk objective from a single training run without validation data. This objective measures precisely the noise in the signal channel.
Practical Algorithm: A "Population-Risk Training" algorithm that acts as an SNR preconditioner on top of Adam. It adds minimal computational overhead (one state vector) and requires no hyperparameter tuning for the gate threshold.

Results

The proposed method was evaluated across three regimes where standard empirical-risk training is known to fail or overfit:

Physics-Informed Neural Networks (PINNs): On a noisy initial condition advection problem, the population-risk method reached a target test error ( $\ell_2 \le 0.40$ ) 2.4 $\times$ faster than the best learning-rate-tuned AdamW baseline. It suppressed noise fitting that caused standard training to diverge from the physical solution.
Grokking (Modular Arithmetic): On a modular division task ( $a \cdot b^{-1} \mod 97$ ) known for delayed generalization, the method achieved 95% held-out accuracy at step 5,950, compared to 29,450 for AdamW. This represents a 4.9 $\times$ reduction in training steps, effectively collapsing the grokking delay.
Noisy Preference Alignment (DPO): Fine-tuning Qwen2.5-0.5B-Instruct with 30% swapped preference labels, the method improved final reward accuracy from 0.566 to 0.641 while staying 3.05 $\times$ closer to the reference policy in terms of mean absolute reward drift. It maintained sustained high accuracy where AdamW failed to stabilize.

Significance and Claims

The paper claims to provide a unified theoretical lens that resolves the tension between overparameterization and generalization without relying on asymptotic limits or frozen kernels. By identifying the "reservoir" as a test-invisible sink for noise, it explains why overfitting training data does not necessarily imply overfitting test data.

The authors assert that their derived population-risk objective is not merely a heuristic but an exact, unbiased rate of population-risk decrease computable from a single run. The practical significance lies in the ability to train directly on population risk, accelerating convergence (e.g., 5 $\times$ faster grokking), suppressing memorization in implicit neural representations, and improving robustness in preference alignment, all without the need for validation sets or complex hyperparameter search. The theory bridges the gap between the "lazy" NTK regime and the "feature-learning" regime, showing that generalization mechanisms persist even when the kernel drifts significantly.