🤖 machine learning

Does Weight Decay Enhance Training Stability?

This paper investigates how weight decay enhances training stability at the Edge of Stability by demonstrating that it slows progressive sharpening and induces an architecture-dependent phase transition driven by the global alignment of parameters and sharpness gradients, ultimately revealing that traditional curvature-based stability diagnostics may be unreliable under regularization.

Original authors: Marius Saether, Amir Kolic, Tomaso Poggio, Pierfrancesco Beneventano

Published 2026-05-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Marius Saether, Amir Kolic, Tomaso Poggio, Pierfrancesco Beneventano

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 Question: Is Weight Decay a "Stabilizer"?

In the world of training AI (deep learning), there is a common trick called Weight Decay. Think of it as a gentle "nudge" that keeps the AI's internal numbers from getting too big. For a long time, people thought this was just a rule to prevent the AI from memorizing the training data too perfectly (a concept called overfitting).

However, modern AI practitioners have started saying, "Hey, Weight Decay also makes the training process stable." It stops the AI from going crazy and crashing during learning.

This paper asks a simple question: Is this true? And if so, how does it actually work?

The authors found that the answer is yes, but the mechanism is much more surprising and complex than anyone expected.

The Setting: The "Edge of Stability"

To understand the paper, you first need to understand where AI training happens. The authors describe a state called the Edge of Stability (EoS).

The Analogy: The Tightrope Walker
Imagine an AI trying to learn is like a tightrope walker.

If the walker is too careful (too stable), they move too slowly and never get anywhere.
If they are too reckless, they fall off the rope.
The Edge of Stability is the sweet spot where the walker is wobbling back and forth, barely staying on the rope, but actually making progress.

In this state, the "sharpness" of the terrain (how steep the hills are) oscillates. The AI wobbles, but a self-correcting mechanism keeps it from falling.

The Discovery: Weight Decay Changes the Rules

The authors investigated what happens when you add Weight Decay (the "nudge") to this tightrope walker.

1. The Old Theory (The Naive Expectation)

Before this paper, people thought Weight Decay would simply act like a brake.

The Expectation: If the walker is wobbling at a certain height, adding Weight Decay should just lower the wobble slightly, like adding a little bit of friction. It was expected to be a simple, linear adjustment.

2. The New Reality (The Phase Transition)

The authors discovered that Weight Decay doesn't just add friction; it changes the physics of the tightrope, but only for certain types of AI architectures (specifically, simple Multi-Layer Perceptrons or MLPs).

The Analogy: The Magic Switch
Imagine the AI is a car driving on a bumpy road.

Without Weight Decay: The car bounces up and down wildly.
With a little Weight Decay: The car bounces a bit less.
With enough Weight Decay (The Phase Transition): Suddenly, the car doesn't just bounce less; it hits a "magic switch." The road underneath the car suddenly becomes flat and smooth. The car stops bouncing wildly and settles into a very low, calm groove.

The paper shows that for simple AI models (MLPs), once you turn the Weight Decay knob past a certain point, the AI stops wobbling at the "Edge of Stability" and settles into a much calmer state than anyone thought possible. It stabilizes at a level far lower than the old math predicted.

Note: For more complex models like CNNs (used in image recognition), Weight Decay acts more like a standard shock absorber, just dampening the bounces without changing the fundamental "road" the car is driving on.

How Does It Work? (The Secret Mechanism)

Why does this "magic switch" happen? The authors found it's due to a Global Interaction.

The Analogy: The Tethered Ball
Imagine the AI's learning process is a ball rolling in a valley.

Local Friction: Usually, we think of Weight Decay as just slowing the ball down locally (like sand on the ground).
Global Interaction: The authors found that Weight Decay actually attaches a long tether to the ball, connecting it to the center of the entire valley.

When the ball rolls, it's not just reacting to the hill it's on; it's being pulled by this tether to the center of the whole system.

If the tether is loose (low Weight Decay), the ball rolls normally.
If the tether is tight (high Weight Decay), it forces the ball to stay in a very specific, low-energy spot. It creates a "phase transition" where the ball is forced to settle much lower than it would have on its own.

Why Does This Matter? (Function Space Stability)

The paper concludes that this "lowering of the wobble" is actually a good thing for the AI's brain.

The Analogy: Tuning a Radio
When the AI settles into this new, lower, calmer state, it doesn't just stop shaking. It actually tunes its "radio" (called the Neural Tangent Kernel or NTK) to a clearer frequency.

The "noise" (instability) is reduced.
The signal (learning the right features) becomes much stronger.

The authors show that by forcing the AI to settle at this lower, calmer level, Weight Decay helps the AI align its internal logic with the correct answers much better. It's like the AI isn't just "stable" because it stopped shaking; it's stable because it found a better, clearer path to the solution.

Summary of Key Findings

It's not just a brake: Weight Decay does more than just slow things down; it can fundamentally change the landscape the AI is learning on.
The "Phase Transition": In simple AI models, increasing Weight Decay past a critical point causes a sudden shift where the AI stabilizes at a much lower, calmer level than expected.
Global vs. Local: This happens because Weight Decay connects the AI's local movements to its global position (the "tether" effect).
Better Learning: This lower, calmer state isn't just "safe"; it actually helps the AI learn better features and aligns its internal "brain" more effectively with the task it's trying to solve.

In short: Weight Decay is a dynamic, active participant in training. It doesn't just "regularize" the AI; it acts like a sophisticated control system that can suddenly switch the AI into a super-stable, high-performance mode.

Technical Summary: Does Weight Decay Enhance Training Stability?

Problem Statement

In modern deep learning, weight decay is frequently credited by practitioners with "stabilizing" training dynamics, a role that diverges from its classical interpretation as a static regularization penalty for ill-posed inverse problems. While recent literature has established weight decay's role in feature selection and solution picking, the specific mechanisms by which it influences training stability remain underexplored.

Training stability is generally understood through three distinct but related notions:

Parameter-space dynamics: Controlling loss curvature and the Edge of Stability (EoS).
Function-space dynamics: Bounding the Neural Tangent Kernel (NTK) spectrum.
Signal propagation: Controlling pre-activation scales (e.g., via $\mu$ P).

The central question addressed is whether weight decay genuinely stabilizes training, and if so, through which mechanism. Specifically, the paper investigates how weight decay affects the EoS, where the maximum eigenvalue of the loss Hessian ( $\lambda_{\max}(\nabla^2 L)$ ) typically hovers around the instability threshold $2/\eta$ (where $\eta$ is the learning rate). Naive heuristics suggest that adding weight decay $\gamma$ should simply shift this threshold to $2/\eta - \gamma$ . However, empirical observations indicate more complex, non-linear behaviors that challenge this simple additive model.

Methodology

The authors employ a combination of empirical analysis and mathematical modeling to investigate weight decay's effects on EoS dynamics.

Empirical Setup

Architectures: Experiments are conducted on Multi-Layer Perceptrons (MLPs) with ReLU activations and Convolutional Neural Networks (CNNs).
Datasets: A subset of CIFAR-10 (5,000 samples).
Optimization: Full-batch Gradient Descent (GD) with Mean Squared Error (MSE) loss.
Metrics: The study tracks the evolution of loss sharpness ( $\lambda_{\max}$ ), the progressive sharpening coefficient ( $\alpha = -\langle \nabla L, \nabla S \rangle$ ), and global interaction terms ( $c_y = \langle \nabla S, \theta \rangle$ ).
Limitations: The analysis is restricted to full-batch GD and small-scale architectures due to the computational cost of Hessian-vector products required for sharpness estimation.

Theoretical Framework

Building on the self-stabilization framework of Damian et al. [23], the authors model gradient descent with weight decay as an underdamped harmonic oscillator.

Local Dynamics: They analyze the system in a local basis spanned by the top eigenvector of the Hessian ( $u$ ) and the gradient of sharpness ( $\nabla S$ ).
Global Interaction: Crucially, the model incorporates global interaction terms ( $c_x = \langle u, \theta^* \rangle$ and $c_y = \langle \nabla S, \theta^* \rangle$ ) arising from the weight decay penalty.
Derivation: The discrete update rules for the sharpness deviation ( $y_t$ ) and the unstable direction ( $x_t$ ) are derived, showing how weight decay introduces both local friction and global offsets.

Key Contributions and Results

1. Weight Decay Slows Progressive Sharpening

Across both MLPs and CNNs, increasing the weight decay parameter $\gamma$ robustly slows the rate of progressive sharpening. This delays the onset of the Edge of Stability (EoS) and flattens the sharpness trajectory during early training.

2. Architecture-Dependent Phase Transitions

Once the network enters the EoS regime, weight decay induces divergent behaviors based on architecture:

CNNs: Weight decay acts primarily as a dampener. It reduces the amplitude of sharpness oscillations, causing the system to settle faster, but the stabilization threshold remains close to the naive expectation of $2/\eta - \gamma$ .
MLPs: A striking phase transition occurs. When $\gamma$ exceeds a critical threshold $\bar{\gamma}$ , the stabilization level of sharpness drops significantly below the theoretical $2/\eta - \gamma$ boundary. For example, with $\eta=0.02$ and $\gamma=0.02$ , the sharpness stabilizes around 80, whereas the naive threshold predicts $\approx 99.98$ . Despite this lower threshold, the training dynamics retain EoS signatures (non-monotonic loss, persistent oscillations).

3. Mechanism: Global Alignment and Limit Cycle Collapse

The paper identifies the global alignment between the parameter vector and the sharpness gradient as the driver of the phase transition.

The system is modeled as an oscillator where the amplitude of oscillation $\Delta$ depends on the term $(\alpha + \gamma^2 - \gamma c_y)$ .
A critical bifurcation occurs when the global interaction term $c_y$ crosses a threshold $c_y^{crit} = \frac{\alpha}{\gamma} + \gamma$ .
In MLPs: Sufficient weight decay causes $c_y$ to cross this critical threshold during progressive sharpening, leading to a "limit cycle collapse." The system detaches from the standard EoS boundary and stabilizes at a much lower curvature.
In CNNs: The global interaction term remains below the critical threshold, preventing the phase transition and resulting only in dampened oscillations.

4. Function-Space Stability (NTK)

The paper demonstrates that the depressed sharpness threshold in parameter space translates to stability in function space.

By lowering the EoS threshold, weight decay restricts the spectral radius of the empirical Neural Tangent Kernel (NTK).
Using a rank-1 NTK approximation, the authors show that a lower sharpness threshold (which reduces the scaling factor $c_t^2$ ) increases the ratio $\alpha_t = \|v_t\|^2 / c_t^2$ .
This increase in $\alpha_t$ shifts the alignment of the target vector towards the leading eigenvectors of the NTK, thereby improving feature learning and kernel target alignment (KTA).

5. Re-evaluating Curvature as a Stability Diagnostic

The authors conclude that curvature thresholds derived from convex or quadratic heuristics (e.g., expecting stability only when $\lambda_{\max} < 2/\eta - \gamma$ ) are unreliable diagnostics under regularization. Training can exhibit characteristic EoS behavior even when the sharpness is significantly lower than these naive thresholds.

Significance and Claims

The paper claims to provide the first mathematical framework that accurately models the "folklore" observation that weight decay enhances training stability. Its significance lies in:

Unifying Perspectives: It bridges parameter-space dynamics (EoS), function-space dynamics (NTK), and signal propagation, showing how weight decay modulates all three.
Mechanistic Insight: It moves beyond viewing weight decay as a static penalty, identifying it as a dynamic control signal that couples local geometry to global parameter alignment.
Refining Stability Diagnostics: It challenges the community to look beyond simple curvature thresholds to diagnose training stability, highlighting that "stabilization" can manifest as a phase transition to a lower curvature regime rather than a simple shift of the instability boundary.
Architecture Sensitivity: It highlights that the effects of weight decay are not uniform but depend critically on the architecture's interaction with global alignment terms.

The authors explicitly state that their findings support the practitioner's intuition that weight decay stabilizes training, but clarify that this occurs through a complex interplay of local friction and global interactions that fundamentally alters the optimization landscape at the Edge of Stability.