🤖 machine learning

Gradient-Direction Sensitivity Reveals Linear-Centroid Coupling Hidden by Optimizer Trajectories

This paper demonstrates that analyzing loss gradients rather than optimizer updates reveals a strong, previously hidden coupling between gradient-direction sensitivity and linear-centroid hypothesis features, showing that constraining attention updates to low-rank subspaces accelerates grokking while the natural full-rank updates are highly rank-redundant.

Original authors: Yongzhong Xu

Published 2026-04-29

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Yongzhong Xu

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: Watching a Car Drive vs. Looking at the Road

Imagine you are trying to understand how a self-driving car learns to navigate a complex city. You want to know: Which specific turns and lanes are actually helping the car learn the route?

For a long time, researchers have been watching the car's actual movement (the optimizer's trajectory) to figure this out. They look at the path the car took over the last few minutes, draw a line through the middle of that path, and say, "Aha! The car is learning by moving in this specific direction."

This paper argues that looking at the car's movement is misleading.

The authors say that the car's movement is a messy mix of:

Momentum: The car is still coasting from turns it made 10 seconds ago.
Adaptive Scaling: The car is adjusting its speed based on how slippery the road feels.
Weight Decay: The car is trying to stay in the center of the lane even when it doesn't need to.
Conflicting Goals: In a multi-task setting, the car is trying to go to the grocery store, the gym, and the office all at once. The path it takes is a compromise, not a clear route to any single destination.

The paper claims that to truly understand what the car is learning, you shouldn't watch where the car went. Instead, you should look at the road signs and the map (the gradients) that the car was looking at right now.

The Core Discovery: The "Messy Path" vs. The "Clear Signal"

The researchers tested this on a computer model learning math problems (like addition and multiplication). They used a "test" to see if the model was forming specific features (like a mental rule for adding numbers).

1. The Old Way (Watching the Car):
When they analyzed the path the model took (the "Update"), the test results were weak and confusing.

The Result: The model seemed to be learning, but the "direction" of learning looked random. In complex tasks where the model had to do four things at once, the test failed completely. It looked like the model wasn't learning anything specific at all.
The Analogy: It's like trying to figure out a hiker's destination by looking at their muddy footprints. The footprints are a jumbled mess of sliding, stopping, and changing direction, so you can't tell which way they were actually trying to go.

2. The New Way (Looking at the Map):
The researchers switched to analyzing the gradients (the raw mathematical signal telling the model which way to move before the messy optimizer steps in).

The Result: Suddenly, the signal became incredibly clear. The "direction" of learning was 30 to 100 times stronger than before.
The Analogy: Instead of looking at the muddy footprints, they looked at the GPS coordinates the hiker was aiming for at that exact second. The direction was sharp, clear, and perfectly aligned with the destination.

Key Takeaway: The "optimizer" (the algorithm that updates the model) adds so much noise and history that it hides the true learning signal. You have to strip away the momentum and the history to see what the model is actually learning.

The Multi-Task Problem: The "Committee Meeting"

The paper also looked at a model trying to learn four different math problems at the same time (Add, Subtract, Multiply, and a complex square formula).

The Old View: When they looked at the combined path of the model, it looked like a disaster. The model was trying to satisfy four different "committees," and the resulting path was a compromise that satisfied none of them well. The diagnostic tool said, "This model isn't learning any specific feature."
The New View: When they looked at the "map" for each committee separately (calculating the gradient for addition, then for subtraction, etc.), they found that the model was actually learning very well for all four tasks.
The Analogy: Imagine a committee of four people trying to decide on a lunch order.
- The Old Way: You look at the final, messy compromise (e.g., "Let's get a salad with pizza toppings"). It looks like a bad decision that satisfies no one.
- The New Way: You listen to what each person individually wanted to order. You realize that Person A really wanted a salad, and Person B really wanted pizza. The "messy compromise" was just the result of them arguing, but their individual desires were clear and strong.

Conclusion: In multi-task training, the "compromise path" hides the fact that the model is successfully learning distinct features for each task. You must separate the tasks to see the learning.

The "Rank-3" Surprise: It's About the Size, Not the Direction

The researchers did one more experiment. They asked: "Do we need the model to learn in these specific directions we found, or just in any low-dimensional direction?"

They forced the model to only learn using a very small, simple "subspace" (a tiny 3-dimensional slice of its massive brain).

The Result: The model learned faster (about 2.3 times faster) when forced to use this tiny slice.
The Twist: It didn't matter which slice they used. Whether they used the "smart" slice they found with their new method, or a completely random slice, the model learned just as fast.
The Analogy: Imagine you are trying to fit a large suitcase into a small car trunk.
- You might think, "I need to fold the clothes in this specific way to fit them."
- But the experiment showed that as long as you just compress the suitcase (reduce the rank), it fits and gets there faster. It doesn't matter if you fold the shirts or the pants first; the act of compressing the space is what matters.

Conclusion: The specific "directions" the model learns are just a symptom of the learning process, not the cause. The real magic is that the model doesn't need its full, messy brain to learn these math tricks; it only needs a tiny, compressed version of it.

Summary for the General Audience

Stop watching the footprints: If you want to understand how AI learns, don't look at the path it took (the optimizer updates). That path is too messy with history and compromises.
Look at the map: Look at the raw instructions (gradients) the AI is following right now. This reveals the true "directions" of learning, which are much stronger and clearer.
Separate the tasks: If an AI is doing many things at once, the combined path looks like a failure. You have to look at each task individually to see that it's actually succeeding.
Simplicity wins: The AI learns faster when forced to be simple. It doesn't need complex, specific directions to learn; it just needs to be constrained to a smaller, simpler space.

The paper essentially tells us that our current tools for "interpreting" AI are looking at the wrong thing. By switching to a cleaner view of the data, we can see that AI is learning much more effectively and simply than we thought.

1. Problem Statement

In mechanistic interpretability, a central question is identifying which directions in parameter space correspond to the formation of specific features (circuits). A common diagnostic approach involves analyzing the low-rank structure of optimizer trajectories (parameter updates, $\Delta \theta$ ) using Singular Value Decomposition (SVD) over rolling windows. This method, known as Update SED (Spectral Edge Diagnostics), assumes that the directions in which the optimizer moves faithfully reflect the directions relevant to feature formation.

The authors argue that this assumption is flawed in modern training pipelines (using AdamW with momentum, adaptive scaling, and weight decay) and especially in multitask settings. They posit that optimizer updates aggregate and distort the underlying gradient signal, potentially obscuring or destroying the alignment between trajectory-based diagnostics and actual feature formation.

2. Methodology

The study investigates this hypothesis using grokking on modular arithmetic tasks (addition, subtraction, multiplication, and $a^2+b^2$ ) within a 2-layer Transformer architecture. The authors compare two classes of spectral estimators against the Linear Centroids Hypothesis (LCH), which posits that features are formed along directions that maximize the sensitivity of task-specific centroids.

Key Estimators Compared:

Update SED: Performs SVD on a rolling window of actual AdamW parameter updates ( $\Delta \theta_t$ ). This captures the "path" taken by the optimizer.
Gradient SED: Performs SVD on a rolling window of raw gradients ( $g_t = \nabla_\theta L$ $g_{t} = \nabla_{θ} L$ ) evaluated at the current parameters. This captures the "steepest descent" direction of the loss landscape, stripped of optimizer artifacts (momentum, adaptive scaling, weight decay).
- Multitask Extension: In multitask settings, the authors compute per-op gradients ( $g_{op}$ ) by backpropagating each task's loss separately, rather than using the aggregated total loss gradient.

Diagnostic Metric:

The coupling between a spectral direction $v_k$ and feature formation is measured by the perturbative ratio $R_k$ :
$R_k(t) = \frac{A(v_k, t)}{\text{median}_j A(r_j, t)}$
Where $A(v, t)$ measures the change in the task centroid when the attention parameters are perturbed along direction $v$ . A high $R_k$ indicates that the direction $v$ is highly sensitive to feature formation.

Causal Interventions:

The authors perform causal interventions by constraining attention updates to low-rank subspaces (Rank-3) defined by either the SED basis or a random basis. They test two projection methods:

Update Projection: Projecting the final AdamW update.
Gradient Projection: Projecting the gradient before the AdamW optimizer computes its momentum and adaptive scaling moments.

3. Key Results

A. Single-Task Results: Gradient vs. Update

Update SED Failure: On single-task models, Update SED yields modest coupling ratios ( $\bar{R}_k \approx 3\text{--}9\times$ ). Addition appears anomalously coupled, while multiplication is weak.
Gradient SED Success: Replacing updates with raw gradients increases coupling ratios by 30–100 $\times$ ( $\bar{R}_k \approx 100\text{--}330\times$ $\overset{ˉ}{R}_{k} \approx 100 - 330 \times$ ).
- Multiplication becomes the strongest coupled task, reversing the Update SED ranking.
- The "op-spread" (variance between tasks) shrinks significantly.
Role of Rolling Window: The high coupling at $W=1$ (instantaneous gradient) suggests the gradient direction itself is the primary driver. The rolling window ( $W=20$ ) acts primarily as a denoising mechanism, reducing the ratio monotonically but maintaining a signal orders of magnitude higher than random.

B. Multitask Results: The Critical Failure of Aggregation

Update SED Collapse: In multitask training, Update SED fails entirely, yielding $\bar{R}_k \le 1$ (worse than random). The aggregated optimizer update, which balances four competing objectives, becomes orthogonal to any single task's feature-relevant subspace.
Gradient SED Recovery: Using per-op gradient SED recovers strong coupling ( $\bar{R}_k \approx 20\text{--}45\times$ $\overset{ˉ}{R}_{k} \approx 20 - 45 \times$ ).
- Crucially, the multitask gradient SED signal is stronger than the single-task Update SED signal.
- This proves that feature formation is robust in multitask settings but is hidden by the aggregation of competing gradients in the optimizer's update step.

C. Causal Interventions: Rank vs. Direction

Acceleration: Constraining updates to a Rank-3 subspace accelerates grokking by approximately 2.3 $\times$ .
Direction Independence: This acceleration occurs regardless of whether the subspace is the SED basis (feature-aligned) or a random basis.
- Gradient Projection: When projecting the gradient before optimization, removing the SED subspace (keeping random) has no negative effect, and keeping only random rank-3 accelerates training just as much as keeping the SED basis.
- Update Projection: Initial experiments suggested removing SED helped, but this was an artifact of AdamW's internal momentum/adaptive scaling leaking information back into the trajectory. Gradient projection eliminates this artifact.
Conclusion: The rank constraint (reducing redundancy) is the causal factor for acceleration, not the specific direction. The natural AdamW update on attention is "rank-redundant," containing mostly non-functional motion for the grokking transition.

4. Key Contributions

Methodological Shift: Demonstrates that gradient structure, not optimizer trajectories, is the correct object for spectral analysis of feature formation. Optimizer updates introduce qualitative distortions that hide feature-relevant directions.
Multitask Interpretability: Identifies gradient aggregation across competing tasks as a general obstruction to circuit-level interpretability. Task-resolved gradient analysis is necessary to recover feature structure in multitask training.
Causal Clarification: Shows that while gradient-derived directions are strong diagnostics of feature formation, they are not uniquely causal. The key factor is the low-rank nature of the update, not the specific alignment with the gradient basis.
Artifact Identification: Reveals that previous "rank-1 manifold" observations in grokking were partially obscured by the noise and distortion inherent in AdamW updates.

5. Significance

For Interpretability: This work challenges the standard practice of analyzing parameter update trajectories (e.g., PCA on weight differences) to find circuits. It argues that such methods systematically underestimate feature formation, especially in complex multitask or large-scale settings (like Mixture-of-Experts).
For Training Dynamics: It suggests that the "natural" path taken by optimizers like AdamW is often inefficient and redundant. Constraining optimization to lower ranks can significantly speed up learning without sacrificing performance.
Future Directions: The findings imply that for large-scale models (LLMs, instruction tuning), analyzing per-task gradients (or approximations like Lanczos on the empirical Fisher) is essential for understanding how specific circuits form, rather than relying on aggregated trajectory analysis.

In summary, the paper argues that the gradient is the signal, and the optimizer is the noise. To understand how neural networks learn features, one must analyze the raw gradient landscape, particularly decomposed by task, rather than the distorted path the optimizer takes through parameter space.