🤖 machine learning

Weibull Weight-Scale Parameter Evolution under AdamW Training Dynamics

This paper analyzes the evolution of the Weibull weight-scale parameter $\lambda$ during AdamW training by deriving a three-force decomposition of the squared weight norm, demonstrating that an alignment force dominates the initial growth phase while a balance between alignment and weight decay drives subsequent relaxation, with a proposed spline displacement method enabling accurate reconstruction of these dynamics from sparse checkpoints.

Original authors: Tiexin Ding

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tiexin Ding

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

Imagine training a giant AI model like AdamW is like trying to build a massive, complex sculpture out of millions of tiny, invisible Lego bricks. The goal is to shape these bricks so they work together perfectly. But how do the bricks actually move and settle into their final positions?

This paper acts like a detective story, investigating why the size of these Lego bricks (the "weights") grows, overshoots, and then settles down during the training process. The authors use a special measuring tool called the Weibull scale (let's call it the "Size Gauge") to track this.

Here is the simple breakdown of what they found:

1. The Mystery of the "Size Gauge"

In previous studies, researchers noticed a strange pattern: The "Size Gauge" ( $\lambda$ ) of the AI's weights doesn't just grow steadily. It shoots up, gets too big (overshoots), and then shrinks back down to a comfortable resting size.

The Question: Why does this happen? What forces are pushing the weights up and pulling them down?

2. The Three Invisible Hands

The authors discovered that the movement of these weights is controlled by three distinct forces, like three people pushing and pulling a heavy box:

The Alignment Force (The "Pusher"): This is the main engine. It pushes the weights in the direction they want to go based on the learning signal. Think of it as a strong wind blowing the sculpture into shape. The paper found that during the "growth" phase, this force is doing 88% to 94% of the work. It's the dominant driver.
The Injection Force (The "Jitter"): This is a tiny, constant nudge caused by the randomness of the data (noise). It's like a slight vibration in the table. It's always there, but it's very small (only about 4% of the effort).
The Decay Force (The "Puller"): This is a built-in brake. It constantly tries to shrink the weights back to zero to keep the model simple and prevent it from getting too wild. It's like a rubber band pulling the sculpture back toward the center.

3. The Story of the Curve (Rise, Peak, Relax)

The paper explains the weird "up-and-down" curve of the Size Gauge using these three forces:

The Rise: At the beginning, the Pusher (Alignment) is much stronger than the Puller (Decay). The weights grow rapidly because the push is winning.
The Peak: As the model gets closer to being "finished," the Pusher gets tired (it aligns less perfectly), and the Puller gets stronger as the weights get bigger. Eventually, they meet in the middle. The Push and the Pull balance out perfectly. Growth stops. This is the peak.
The Relaxation: If the Puller ever gets slightly stronger than the Pusher, the weights shrink back a little until they find a new, stable balance. This creates the "floor" where the model settles.

4. The "Magic Bridge" (Connecting the Dots)

There was a problem: The math for these forces works on the total weight size (RMS), but the "Size Gauge" (Weibull) measures the distribution of the weights.

The Solution: The authors found a "magic bridge." They discovered that the shape of the weight distribution stays almost perfectly locked in place (like a rigid mold) throughout training. Because the shape doesn't change, the total size and the "Size Gauge" move in perfect lockstep. This allowed them to translate the math of the forces directly into the behavior of the Size Gauge.

5. Solving the "Missing Data" Puzzle

Real-world AI models are often released without the "internal diary" (optimizer moments) that shows exactly how the forces were applied at every step. We only have snapshots of the final weights.

The Trick: The authors invented a "Spline Displacement Method." Imagine you have a movie of a car, but you only have photos of it every 10 seconds. You can't see the speed between photos. But, if you draw a smooth curve connecting the dots (a spline), you can guess the speed very accurately.
The Result: This method allowed them to guess the "Pusher" force from sparse snapshots with 92–94% accuracy, which is twice as good as just guessing between two points.

6. A Glimpse at the Data

The authors also noticed something interesting: The height of the peak (how big the weights get before shrinking) seems to depend on what the AI is reading.

If the AI reads a single type of text (like a specific wiki), the weights get very large (high peak).
If the AI reads a mix of many different topics, the peak is lower.
Note: The authors say this is just a "hint" for future research. They didn't prove exactly why yet, but they suspect that reading diverse data creates conflicting directions that stop the weights from growing as huge.

Summary

In short, this paper explains that the "breathing" pattern of AI weights (growing, overshooting, and settling) isn't random. It is a precise dance between a strong pushing force (learning) and a steady pulling force (decay). When the push wins, weights grow. When they balance, the model settles. The authors also figured out how to watch this dance even when we don't have the full video, just a few snapshots.

Technical Summary: Weibull Weight-Scale Parameter Evolution under AdamW

Problem Statement
Understanding the evolution of neural network weights during training is central to optimization and generalization. A prior framework (Ding, 2026) established a two-parameter Weibull distribution to diagnose transformer weight magnitudes, identifying a "Transmission-class" regularity where the shape parameter $k$ remains locked near $k_0 \approx 1.20$ , while the scale parameter $\lambda$ follows a characteristic trajectory: rise, overshoot, relaxation, and floor. However, the mechanism driving this trajectory remained unexplained. While AdamW dynamics are defined at the level of individual parameters and optimizer states, and Weibull parameters are distribution-level summaries, a bridge was missing to explain why $\lambda$ rises and then relaxes. Specifically, it was unclear which optimizer-level forces move the model along the $\lambda$ coordinate.

Methodology
The paper bridges the gap between individual parameter dynamics and distribution-level statistics through a three-step approach:

Three-Force Decomposition: Starting from the decoupled AdamW update rule, the authors derive a leading-order decomposition of the squared weight norm dynamics ( $\Delta \Sigma W^2$ ) into three distinct forces:
- Alignment Force ( $F_{align}$ ): $-2\eta \langle W, \hat{u} \rangle$ , measuring the correlation between current weights and the adaptive update direction.
- Injection Force ( $F_{inj}$ ): $+\eta^2 \|\hat{u}\|^2$ , representing the squared magnitude of the adaptive step (stochastic noise).
- Decay Force ( $F_{decay}$ ): $-2\eta \lambda_{wd} \Sigma W^2$ , the structural contribution from decoupled weight decay.
The $k$ -Lock Bridge: To connect these squared-norm forces to the Weibull scale parameter $\lambda$ , the paper utilizes the observation that $k$ remains approximately constant ( $k \approx 1.20$ ) for Transmission-class components. This allows the RMS ( $\sigma$ ) and $\lambda$ to be related by a nearly fixed multiplicative factor ( $\sigma = \lambda \sqrt{\Gamma(1+2/k)}$ ), enabling the translation of force dynamics on $\sigma^2$ to the trajectory of $\lambda(t)$ .
Sparse Checkpoint Recovery (Spline Displacement): Since public model releases lack optimizer moments ( $\hat{m}, \hat{v}$ ), the authors introduce a method to recover force information from sparse weight checkpoints. By fitting a cubic spline to saved weight trajectories, they estimate the unit-step displacement required to compute the alignment force. This method is validated against ground-truth moments from self-trained models.

Key Results

Rise Phase Dominance: On self-trained Pythia-70M models with ground-truth optimizer moments, the alignment force dominates the rise phase, contributing 88–94% of the absolute force budget. This dominance is robust across four random seeds and remains stable even after removing "super-weight" outliers. The injection force contributes only ~4%, and decay ~2–7%.
Transition to Relaxation: The transition from weight-scale growth to relaxation corresponds to a shift in force balance. As training progresses toward saturation, the alignment force decreases while the decay force increases. Near the peak and into the relaxation phase, alignment and decay approach a balance ( $F_{align} \approx |F_{decay}|$ ), causing the net force to vanish and the weight scale to stabilize or shrink.
Cross-Architecture Robustness: The alignment-dominance mechanism holds for both Pythia-style (GPT-NeoX) and Llama-style architectures, with alignment shares remaining within the 88–94% range during the rise phase.
Recovery Accuracy: The proposed spline displacement method recovers the alignment force from sparse checkpoints with 92–94% accuracy, approximately double the performance of a naive two-point baseline.
RMS-to-Weibull Offset: The closed-loop integration of forces to predict $\lambda(t)$ shows a systematic offset of approximately 5–6% in densely sampled regions. This is decomposed into a ~4.6% bridge error (RMS-to-Weibull reconstruction) and ~1.9% integration error, confirming that the force dynamics govern the dominant component of the trajectory.
Data-Dependent Modulation: An exploratory observation suggests the absolute peak value of $\lambda(t)$ varies with training-data coherence. Models trained on single-domain data (wikitext-103) reach higher peak $\lambda$ values (~0.069–0.076) compared to those trained on multi-domain data (The Pile, ~0.035). A continue-training intervention on a Pile-trained model switching to wikitext-103 resulted in a 2.6x increase in $\lambda$ , suggesting data coherence modulates the alignment force.

Significance and Claims
The paper claims to provide the first mechanistic explanation for the observed rise, overshoot, and relaxation trajectory of the Weibull weight-scale parameter $\lambda$ during AdamW training. Its primary contributions are:

Connecting Levels of Description: It successfully bridges the gap between low-level AdamW update rules and high-level distribution-level diagnostics, showing that $\lambda(t)$ tracks the weight-scale component of a leading-order force budget.
Identifying the Driver: It identifies the alignment force as the dominant driver of weight-scale growth during the training rise phase, challenging the notion that weight decay or noise are the primary factors in this phase.
Explaining the Trajectory: It explains the overshoot and relaxation as a consequence of the transition from alignment dominance to an alignment-decay balance near saturation.
Practical Applicability: It provides a validated method (spline displacement) to recover force dynamics from real-world models where optimizer states are unavailable, extending the analysis beyond self-trained environments.

The authors explicitly note that the data-coherence observation is exploratory and that a full characterization of data diversity effects requires controlled follow-up studies. The primary contribution remains the force-budget explanation of weight-scale dynamics, with the data result serving as motivation for future controlled experiments.