Spectral Edge Dynamics Reveal Functional Modes of Learning

This paper demonstrates that training dynamics during grokking concentrate along a small number of dominant spectral edge directions that represent low-dimensional functional modes over the input domain, which are invisible to standard mechanistic interpretability tools and whose structure is determined by the algebraic symmetries of the specific task.

The Gist

The Big Picture: Finding the "Secret Sauce" of Learning

Imagine you are watching a student learn a difficult math problem. At first, they are just memorizing answers by rote. Then, suddenly, something clicks. They stop memorizing and start understanding the pattern. In the world of AI, this sudden moment of understanding is called "Grokking."

For a long time, scientists have tried to figure out how this happens inside a neural network (the AI's brain). They usually look at the "hardware"—the specific neurons and connections—to see what changed. They ask, "Which lightbulb turned on?" or "Which wire got stronger?"

This paper says: Stop looking at the wires. Look at the music.

The authors argue that the most important changes during learning aren't happening in specific parts of the hardware. Instead, they are happening in the function—the overall "song" the AI is learning to sing. They found a way to listen to the AI's learning process and hear a specific, dominant melody that appears right when the AI finally "gets it."

---

The Core Concept: The "Spectral Edge"

To understand the paper, we need a new way of looking at the AI's changes.

The Analogy: The Orchestra and the Soloist
Imagine the AI's learning process as a massive orchestra playing a chaotic, noisy symphony. Every time the AI learns something, it tweaks its internal settings (weights). Most of these tweaks are just background noise—random adjustments that don't really matter.

However, the authors discovered that when the AI "groks" (suddenly understands), a Spectral Edge appears.

• The Bulk: This is the noisy orchestra playing random notes.
• The Spectral Edge: This is a tiny, distinct group of soloists that suddenly separate from the noise and start playing a clear, powerful melody.

The paper proves that this "soloist" isn't a single neuron or a specific wire. It's a functional pattern—a specific way the AI responds to different inputs.

The Discovery: It's About the "Shape" of the Answer

The researchers tested this on math problems like modular addition (e.g., $3 + 4 \pmod{10}$) and multiplication.

1. The "Wrong" Way to Look (The Hardware View)
If you try to find the "soloist" by looking at which neurons are firing, you get confused. The signal is spread out everywhere, like trying to find a single drop of ink in a swimming pool. Standard tools (like checking which "head" of the AI is working) fail because the learning isn't localized to one spot.

2. The "Right" Way to Look (The Function View)
Instead of looking at the neurons, the authors looked at how the AI's answer changes when you change the input.

• The Analogy: Imagine the AI is a weather forecaster. Instead of asking "Which sensor broke?", they asked, "How does the forecast change if the wind speed changes?"
• They found that for simple tasks (like addition), the AI's learning pattern looks like a perfect sine wave (a smooth, repeating wave).
• For multiplication, it looks like a sine wave only if you look at it through a special lens (a mathematical trick called a "discrete log").

The Key Insight: The AI isn't just memorizing numbers; it's learning to ride a specific mathematical wave. The "Spectral Edge" is the AI locking onto that wave.

The Hierarchy of Complexity

The paper shows that the "shape" of this learning wave depends on how complex the math problem is:

• Simple Tasks (Addition): The AI learns a single, perfect wave. It's like a singer hitting one perfect note.
• Medium Tasks (Subtraction): The AI learns a small chord (a few notes playing together). It's not just one wave, but a small family of them.
• Complex Tasks (Squares and Sums): The AI learns a complex composition. It's not a single wave or a simple chord. It's a mix of different waves interacting (like a jazz improvisation). The AI combines the "addition wave" and the "multiplication wave" to solve the harder problem.

The "Reuse" Experiment: The Lego Block Theory

One of the coolest parts of the paper is what happens when you train the AI on multiple tasks at once.

The Analogy:
Imagine you are teaching a robot to build a house.

• Task A: Build a door.
• Task B: Build a window.
• Task C: Build a whole room (which needs both).

If you teach the robot to build the room after it already knows how to build doors and windows, does it invent a new way to make a door? Or does it reuse the door it already knows how to build?

The paper found that the AI reuses the patterns. When learning the complex task ($x^2 + y^2$), the AI's "Spectral Edge" (its learning pattern) started looking exactly like the pattern it used for simple addition. It didn't reinvent the wheel; it grabbed the "addition wave" it had already learned and used it as a building block for the harder math.

Why Does This Matter?

1. We've been looking in the wrong place.
For years, AI researchers have tried to explain AI by dissecting the "neurons" (the hardware). This paper says: "The magic isn't in the neurons; it's in the function." The AI learns by finding the right mathematical "shape" or "wave" to solve the problem.

2. It explains how AI learns, not just that it learns.
We know AI eventually gets good at math. This paper shows the moment it happens: when the chaotic noise of learning suddenly organizes itself into a clean, low-dimensional wave (the Spectral Edge).

3. It suggests AI is building with "Functional Primitives."
Just like a human learns to walk, then run, then dance, the AI seems to learn simple mathematical "moves" (like the addition wave) and then combines them to do complex things.

Summary in One Sentence

This paper discovered that when AI models suddenly "get" a math problem, they aren't just tweaking random wires; they are locking onto a specific, low-dimensional mathematical "wave" or pattern, and they can reuse these waves to build more complex skills later on.

Technical Summary

1. Problem Statement

Neural network training trajectories are highly structured, often concentrating along a small number of dominant directions during phase transitions like grokking (the sudden jump from memorization to generalization). However, the nature of these dominant directions remains unclear.

• The Gap: Standard mechanistic interpretability tools (e.g., head attribution, sparse autoencoders, activation space analysis) operate in representation space (neurons, heads, features). These tools have failed to capture the structure of the dominant learning directions, suggesting a "category mismatch."
• The Question: Are these dominant directions localized circuits, interpretable features, or something else? The paper posits that they may be functional modes—structured perturbations of the model's input-output behavior—rather than localized parameter structures.

2. Methodology

The authors analyze transformer models trained on modular arithmetic tasks (addition, subtraction, multiplication, and quadratic forms $x^2+y^2$) modulo a prime $p=97$.

A. Spectral Edge Detection

• Data: Weight updates ($\delta\theta_t$) restricted to attention parameters over a sliding window of training steps.
• Metric: A Gram matrix $G_{ij} = \langle \delta\theta_i, \delta\theta_j \rangle$ is constructed. The eigenvalues ($\sigma_1 \ge \sigma_2 \ge \dots$) are analyzed.
• Definition: The Spectral Edge is defined as a small block of leading directions (e.g., top 2–3) that separate from the "bulk" spectrum. This is identified by a sharp decline in the spectral gap ratio ($\sigma_k / \sigma_{k+1}$) or a specific gap dynamic ($g_{23} = \sigma_2 - \sigma_3$).
• Observation: This edge emerges consistently during grokking but is absent in non-grokking regimes.

B. Functional Perturbation Analysis

Instead of analyzing the weight vectors directly, the authors map them to input space:

1. Perturbation Response: For a spectral edge direction $v_k$, they compute the change in the model's residual stream: $\Delta h_k(x) = h(x; \theta + \epsilon v_k) - h(x; \theta)$.
2. Scalar Field: They define a scalar functional mode $f_k(x) = \|\Delta h_k(x)\|^2$, representing the sensitivity of the model's output to perturbations along $v_k$ for each input $x$.
3. Basis Analysis: They analyze $f_k(x)$ using Fourier analysis in group-theoretic bases appropriate to the task (e.g., additive characters for addition, discrete logarithm characters for multiplication).

C. Experimental Setup

• Models: 2-layer Transformers ($d_{model}=128$).
• Tasks: Modular addition, subtraction, multiplication, $x^2+y^2$, and non-grokking controls.
• Multitask: Shared-trunk models trained on combinations of tasks (e.g., Add + Mul + $x^2+y^2$) to test for functional reuse.

3. Key Contributions

1. Robust Spectral Edge Detection: Confirmed that a spectral edge (a small set of leading update directions) reliably distinguishes grokking from non-grokking regimes across tasks and seeds.
2. Negative Results for Representation-Level Interpretability: Demonstrated that standard tools (head attribution, SAEs, activation rank analysis) fail to capture the spectral edge. The edge is diffuse in parameter/activation space but structured in function space.
3. Functional Structure in Symmetry-Adapted Bases: Showed that spectral edge directions correspond to low-dimensional functional subspaces. When analyzed in the correct group-theoretic basis, these directions collapse to specific Fourier modes.
4. Compositional Structure: For complex tasks like $x^2+y^2$, the structure is not a single harmonic mode but a low-dimensional subspace formed by cross-terms of additive and multiplicative features.
5. Evidence of Functional Reuse: In multitask training, the spectral edge of a composite task ($x^2+y^2$) aligns with the functional modes of its constituent tasks (addition/multiplication), proving that functional modes are reusable primitives.

4. Key Results

A. The Spectral Edge as a Discriminator

• In all 12 grokking runs, the spectral gap $g_{23}$ declined by 15–110×.
• In non-grokking controls, the decline was negligible (<2×).
• Three independent spectral analyses (Gram matrix, displacement PCA, weight-matrix SVD) yielded consistent temporal signatures.

B. Functional Modes vs. Task Symmetry

• Modular Addition: The spectral edge collapses to a single dominant Fourier mode ($\omega \approx 25-26$) in the additive basis. The top 3 directions share this frequency.
• Modular Multiplication: No structure appears in the additive basis. However, in the discrete-logarithm basis (matching the multiplicative group structure), the edge collapses to a single mode ($\omega = 29$). Concentration improves by 5.9× compared to the additive basis.
• Subtraction: The edge spans a small family of modes ($\omega \in \{6, 16, 32\}$) rather than a single mode, suggesting a multi-mode subspace.
• Quadratic ($x^2+y^2$): No single harmonic basis captures the edge. The structure is explained by cross-terms between additive and multiplicative features. A multivariate probe using cross-terms increased the $R^2$ from 0.04 to 0.16 (a 4× boost), confirming compositional learning.

C. Multitask Reuse

• In a shared-trunk model, the spectral edge of the $x^2+y^2$ head becomes 2.3× more aligned with the additive mode ($\omega=26$) characteristic of the addition task compared to a single-task model.
• This provides direct evidence that the network reuses functional primitives (the "addition circuit") when solving composite tasks.

D. Interference vs. Effective Weights

• Directions above the spectral edge behave as "effective weights" (coherent functional modes activated by non-overlapping input subsets).
• Directions below the edge behave as "interference weights" (diffuse, high-correlation parameter mass).

5. Significance and Implications

• Paradigm Shift: The paper argues for a shift from representation-level interpretability (finding circuits in neurons) to function-level interpretability (identifying low-dimensional subspaces over the input domain).
• Dynamical Selection: It suggests that training dynamics (SGD) actively select directions aligned with the task's natural eigenmodes (group characters), rather than learning uniformly.
• Basis Dependence: The "simplicity" of learning is not intrinsic to the parameters but depends on expressing the perturbation in the correct symmetry-adapted basis.
• Compositional Learning: The findings support the hypothesis that neural networks build complex capabilities by composing reusable functional primitives, a mechanism made visible through spectral edge analysis.
• Future Directions: The authors propose that this framework could be extended to language models to discover "functional modes" (e.g., syntactic or semantic structures) without prior knowledge of the basis, moving beyond scaling laws to understand what is being learned at each stage.

In summary, the paper establishes that the "spectral edge" is a geometric signature of learning that reveals low-dimensional functional subspaces governing generalization, which are invisible to standard mechanistic tools but become clear when analyzed through the lens of the task's algebraic structure.