Gradient Atoms: Unsupervised Discovery, Attribution and Steering of Model Behaviors via Sparse Decomposition of Training Gradients

This paper introduces Gradient Atoms, an unsupervised method that decomposes training gradients into sparse, interpretable components ("atoms") to automatically discover broad model behaviors and steer them via weight-space perturbations, eliminating the need for query-specific scoring and enabling scalable behavior analysis.

The Gist

Imagine you have a very smart robot that has been trained by reading thousands of different books, articles, and instructions. Now, you want to know: "What exactly did this robot learn from all that reading?"

Traditionally, scientists tried to answer this by asking, "Which single book made the robot learn to do math?" or "Which single book made it refuse to answer rude questions?" They would check every single book one by one against every behavior. It's like trying to figure out how a river flows by blaming a single raindrop. It's slow, expensive, and misses the big picture.

The authors of this paper, Gradient Atoms, say: "That's the wrong way to look at it."

Here is the simple breakdown of their new idea, using some everyday analogies:

1. The Problem: Blaming the Wrong Raindrop

When a robot learns to do arithmetic, it doesn't learn it from just one math problem. It learns it because hundreds of math problems all push the robot's brain in the exact same direction.

• Old Way: Trying to find the "magic book" that taught the robot math. (This is like trying to find the one specific raindrop that caused a flood).
• New Way: Looking at the direction the robot's brain moved when it saw all those math problems combined.

2. The Solution: Finding "Gradient Atoms"

The authors developed a method to take all the tiny nudges (gradients) the robot received from its training data and break them down into building blocks, which they call "Atoms."

Think of the robot's knowledge like a Lego castle.

• The old method tried to find which specific brick was responsible for the whole castle.
• The Gradient Atoms method takes the castle apart and sorts the bricks into piles based on their shape and color.
  • One pile is all the "Math bricks."
  • One pile is all the "Polite Refusal bricks."
  • One pile is all the "Code-writing bricks."

They did this without telling the computer what to look for. The computer just looked at the data and said, "Hey, these 500 groups of bricks seem to do the same thing."

3. The Discovery: What Did They Find?

They found 500 of these "Atoms." The best ones were very clear and easy to understand, even though the computer wasn't told what they were.

• The "Math Atom": A group of instructions that taught the robot to do arithmetic.
• The "Refusal Atom": A group that taught the robot to say, "I can't do that" when the instructions were vague.
• The "List Atom": A group that taught the robot how to make bullet points.

It's like finding a drawer in a messy workshop labeled "Screws" or "Nails" without anyone ever writing a label on it. The computer figured it out on its own.

4. The Superpower: Steering the Robot

This is the coolest part. Because they found these "Atoms" (the specific directions in the robot's brain), they can use them as steering wheels.

Imagine the robot is a car driving down a road.

• Old way: To make the car turn left, you have to find the specific driver who taught it to turn left and ask them to drive again.
• Gradient Atoms way: You just grab the "Left Turn" steering wheel (the Atom) and twist it.

They tested this by taking the "Bullet Point Atom" and twisting it.

• Before: The robot made bullet points 33% of the time.
• After: They twisted the wheel, and the robot made bullet points 94% of the time.

They also found the "Refusal Atom." By twisting it the other way, they made the robot stop refusing to answer questions entirely (going from 50% refusal down to 0%).

Why This Matters

• It's Unsupervised: You don't need to tell the computer, "Look for math." The computer finds the math patterns itself.
• It's Fast: Instead of checking millions of documents one by one, they find all the behaviors at once.
• It's Controllable: You can now dial behaviors up or down like a volume knob, simply by adjusting these "Atoms."

The Big Picture

The paper argues that we shouldn't try to blame individual documents for what a robot learns. Instead, we should look at the shared patterns (the atoms) that groups of documents create. By finding these patterns, we can understand what the robot learned and even control its behavior in a very precise way, all without needing a human to label everything first.

In short: They turned a messy pile of training data into a neat set of "control knobs" that let us see and steer exactly what the AI has learned.

Technical Summary

1. Problem Statement

The paper addresses limitations in Training Data Attribution (TDA), a field aimed at understanding which training documents cause specific model behaviors.

• The Flaw in Current TDA: Existing methods are supervised and per-document. They require a user to specify a target behavior (query) and then score every training document against it ($O(Q \times N)$ complexity).
• The Mismatch: The authors argue that fine-tuning does not learn from isolated documents. Instead, models learn shared update directions induced by clusters of documents (e.g., hundreds of arithmetic examples pushing weights in the same direction). Attributing a behavior to a single document is like attributing a river's course to a single raindrop.
• Practical Issues: Current methods are expensive, require prior knowledge of what behaviors to look for, and cannot surface unexpected behaviors.

2. Methodology: Gradient Atoms

The authors propose Gradient Atoms, an unsupervised framework that decomposes the space of training gradients into sparse, interpretable components ("atoms"). The pipeline consists of five steps:

1. Per-Document Gradient Extraction:
   • For a dataset of $N$ documents, compute the cross-entropy loss gradient ($g_i$) for each document with respect to all trainable parameters.
2. EKFAC Projection and Preconditioning:
   • Raw gradient spaces are anisotropic (dominated by high-curvature directions).
   • The method uses EKFAC (an approximate Fisher Information Matrix eigendecomposition) to project gradients into a preconditioned eigenspace.
   • This normalizes the space, ensuring that unit steps in any direction correspond to roughly equal loss changes, allowing the decomposition to capture functional directions rather than curvature artifacts.
3. Sparse Dictionary Learning:
   • Projected gradients are normalized and decomposed into a dictionary of atoms ($D$) and sparse coefficients ($\alpha$).
   • Equation: $\hat{g}_i \approx \sum \alpha_{ij} d_j$.
   • A sparsity penalty ensures each document is explained by only a few atoms, forcing each atom to isolate a single, distinct computational pattern rather than blending unrelated behaviors.
4. Coherence Scoring:
   • For each atom, the authors identify the "activating documents" (those with non-zero coefficients).
   • They compute a coherence score based on the cosine similarity of the raw (unprojected) gradients of these activating documents. High coherence indicates the atom captures a genuine shared computational motif.
5. Unprojection to Steering Vectors:
   • Atoms are unprojected back into the full parameter space to create steering vectors ($v_j$).
   • These vectors can be applied as weight perturbations ($\theta_{new} = \theta \pm \alpha \cdot v_j$) to steer model behavior.

3. Key Contributions

1. Reframing Attribution: Shifts the unit of analysis from "which document caused this?" to "what are the shared update directions induced by clusters of documents?"
2. Unsupervised Discovery: Introduces a method that discovers candidate model behaviors solely from training gradients without behavioral labels, measurement functions, or per-query scoring.
3. Actionable Steering: Demonstrates that these unsupervised atoms function as effective steering vectors, enabling controllable shifts in model behavior.

4. Experimental Results

The method was validated on a Gemma-3 4B IT model fine-tuned on 5,000 instruction-response pairs using LoRA.

A. Atom Discovery (Interpretability)

• Decomposition: 500 atoms were discovered from 5,000 gradients.
• Interpretability: The highest-coherence atoms corresponded to clear, interpretable task types without any labels.
  • Top examples: Short factual QA, grammar editing, Yes/No classification, simple arithmetic, and multi-category classification.
  • Granularity: The method distinguished between sub-tasks (e.g., different types of grammar correction or code generation) and formatting styles (bulleted vs. numbered lists).
  • Refusal: Successfully identified atoms corresponding to "systematic refusal" (e.g., asking for input when instructions are missing).

B. Behavioral Steering (Controllability)
The authors applied the unprojected atoms as weight perturbations to test if they could steer behavior. Five atoms were tested with varying perturbation strengths ($\alpha$).

• Bulleted Lists: Steering increased bullet usage from 33% to 94% (+61 percentage points) and suppressed it to 0%.
• Systematic Refusal: Suppressed from 50% to 0% (the model stopped asking for clarification and answered directly).
• Code Generation: Increased from 42% to 58% or decreased to 28%.
• Yes/No Classification: Suppressed from 39% to 0%.
• Numbered Lists: Suppressed from 58% to 8%.

Key Findings on Steering:

• Suppression > Amplification: It was generally easier to suppress a behavior (disrupting a pathway) than to amplify it (strengthening it against competitors).
• Coherence vs. Steerability: High coherence did not strictly predict high steerability. An atom with lower coherence (bulleted lists, 0.103) produced a larger steering effect (+61pp) than a high-coherence atom (Yes/No, 0.647).

5. Significance and Implications

• Efficiency: The method scales independently of the number of query behaviors. It surfaces all candidate behaviors in a single decomposition, avoiding the $O(Q \times N)$ cost of standard TDA.
• New Paradigm for Model Editing: It connects unsupervised behavior discovery with controllable model editing. Unlike previous steering methods that require hand-crafted measurement functions or contrastive pairs, Gradient Atoms derive steering vectors directly from the training dynamics.
• Insight into Learning: The results suggest that models learn "procedural knowledge" (task types like arithmetic or classification) as distinct computational pathways in weight space, which can be isolated and manipulated.
• Limitations: The current approach relies on instruction-following data (recovering task types rather than semantic preferences) and uses regex-based evaluation for surface formatting. Future work aims to scale to larger dictionaries and cross-model comparisons.

In conclusion, Gradient Atoms provides a powerful, unsupervised lens for understanding and controlling fine-tuned language models by treating training gradients as a mixture of sparse, interpretable "atoms" that define the model's learned capabilities.