🤖 machine learning

Focus and Dilution: The Multi-stage Learning Process of Attention

This paper reveals a recurrent focus-dilution cycle in Transformer training dynamics, explaining how attention learning on Markovian data progresses through distinct stages of rank condensation, frequency-driven focus, mass redistribution, and symmetry breaking to initiate subsequent learning cycles.

Original authors: Zheng-An Chen, Pengxiao Lin, Zhi-Qin John Xu, Tao Luo

Published 2026-05-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zheng-An Chen, Pengxiao Lin, Zhi-Qin John Xu, Tao Luo

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 a Transformer model (the brain behind modern AI chatbots) not as a static machine, but as a student learning a new language. This paper, titled "Focus and Dilution: The Multi-stage Learning Process of Attention," reveals that this student doesn't learn everything at once. Instead, it goes through a rhythmic, repeating cycle of zooming in on a specific word and then zooming out to look at the whole picture, over and over again.

Here is the story of how this learning process works, broken down into simple stages using everyday analogies.

The Big Picture: A Cycle of "Zoom In" and "Zoom Out"

The authors discovered that attention (the mechanism that lets the model decide which words to focus on) doesn't just get better linearly. It follows a Focus-Dilution Cycle:

Focus: The model locks onto the most common word it sees.
Dilution: The model realizes that locking onto just that one word isn't enough, so it "dilutes" its focus to start noticing other, less common words.
Repeat: Once it masters the new words, it locks onto the next most common one, and the cycle starts again.

The Four Stages of the Cycle

The paper breaks this cycle into four distinct phases, which the authors prove mathematically and verify with experiments.

Stage 1: The "Squish" (Initial Condensation)

The Analogy: Imagine a box of 100 different colored marbles (representing all the words in the vocabulary) that are scattered randomly. At the very beginning of training, the model is confused. Suddenly, the model "squishes" all these marbles into a single, flat line.
What happens: The complex internal parts of the model (called embeddings and projections) collapse into a simple, one-dimensional structure. However, the "attention" mechanism (the part that decides what to look at) stays frozen and inactive during this squish. The model is essentially organizing its basic vocabulary before it starts paying attention to specific words.

Stage 2: The "Spotlight" (Focus)

The Analogy: Now that the marbles are lined up, the model turns on a spotlight. Because one word (let's say "the") appears much more often in the training data than others, the spotlight shines blindingly bright on "the" and ignores everything else.
What happens: The attention parameters wake up and grow rapidly. They lock onto the most frequent token (word) in the data. The model becomes a "one-trick pony," predicting the next word based almost entirely on this high-frequency token. This is the Focus phase.

Stage 3: The "Fog" (Dilution)

The Analogy: The spotlight is so bright it's blinding. But as the model keeps training, the "marbles" (the word representations) start to shift and move apart. The spotlight gets blurry. The model realizes that if it only looks at "the," it misses the nuance of other words. The intense focus on the single word starts to fade, like a spotlight turning into a soft, diffuse fog.
What happens: As the attention amplitude grows, it creates a feedback loop that pushes the representations of the "common" word and the "rare" words in opposite directions. The model's obsession with the frequent word weakens. The attention becomes diluted, spreading out to cover more words again.

Stage 4: The "Crack" (Emergence of New Directions)

The Analogy: The model is now in a foggy state, but it's stuck. It can't tell the difference between the rare words because they all look the same in the fog. To break out, the model needs a tiny nudge—a tiny asymmetry. Imagine two identical twins standing in the fog; if one sneezes (a tiny difference), the model can finally tell them apart.
What happens: In real data, no two words are perfectly identical in frequency. These tiny, natural differences act as the "sneeze." They break the symmetry, allowing the model to escape the "fog" and create new, distinct directions in its memory for the next set of words. This unlocks the ability to learn the next most frequent word, starting the cycle all over again.

Why This Matters (According to the Paper)

The authors tested this theory on two types of data:

Synthetic Data: Made-up sentences generated by a simple mathematical rule (Markov chains).
Real Data: Actual text from Wikipedia (WikiText) and children's stories (TinyStories).

In both cases, they saw the exact same pattern: the model zooms in on a frequent word, gets stuck, dilutes its focus, and then uses tiny differences in the data to zoom in on the next frequent word.

The Takeaway

The paper argues that learning in AI isn't a smooth, straight line. It is a staircase. The model climbs one step by focusing intensely on one pattern, then has to "dilute" that focus to make room for the next pattern. This cycle of Focusing and Diluting is the engine that drives the model to learn complex language structures, one layer of frequency at a time.

In short: The AI learns by obsessing over one thing, getting bored of it, spreading its attention out, and then obsessing over the next thing, repeating this dance until it understands the whole language.

Technical Summary: Focus and Dilution in Transformer Attention Learning

Problem Statement

While Transformer-based models have achieved remarkable success, the mechanistic understanding of their training dynamics—specifically how attention mechanisms evolve during optimization—remains limited. Existing analyses often rely on technical simplifications (e.g., reparameterizations or proxy dynamics) that obscure the native coupling between embeddings, projections, and attention parameters. Furthermore, empirical observations suggest that training occurs in multiple stages and that attention patterns can shift from highly concentrated to diffuse, yet a rigorous dynamical theory explaining this "focus-to-dilution" transition is lacking.

Methodology

The authors investigate the training dynamics of a single-layer Transformer trained via gradient flow on Markovian data. The study employs a population gradient-flow limit (large sample size $N$ and long context $s$ ) to derive deterministic dynamics.

Key methodological components include:

Model Setup: A one-layer Transformer with cross-entropy loss, trained on data generated by a Markov chain with a stationary distribution $\pi$ . The model includes an embedding matrix $W_0$ , an output projection $W_1$ , and attention matrices $W_Q, W_K$ .
Stage-wise Linearization: The core theoretical approach involves analyzing the gradient flow trajectory by linearizing the system around successive critical points (saddles). This allows the authors to decompose the complex non-linear dynamics into distinct, tractable stages.
Symmetry Breaking: To explain the emergence of new embedding directions after a degenerate phase, the authors introduce a small symmetry-breaking perturbation ( $\delta$ ) among low-frequency tokens and utilize Lyapunov–Schmidt reduction to analyze the bifurcation of critical points.
Empirical Validation: Theoretical predictions are validated using synthetic Markovian data and real-world corpora (WikiText and TinyStories), tracking loss curves, parameter condensation, and attention entropy.

Key Contributions and Theoretical Findings

The paper identifies a recurring focus–dilution cycle in attention learning, decomposing it into four distinct stages:

1. Initial Condensation

Dynamics: Under small initialization, the trajectory is governed by the linearization near the origin.
Result: The embedding ( $W_0$ ) and projection ( $W_1$ ) matrices rapidly condense into a rank-one structure driven by the stationary distribution $\pi$ of the data. Specifically, $W_0$ aligns with $\pi$ , and $W_1$ aligns with $\pi - \frac{1}{d}\mathbf{1}$ .
Attention State: The attention parameters ( $W_Q, W_K$ ) remain effectively frozen (order $O(\epsilon)$ ) because the leading-order driving term vanishes at the origin.

2. Focus Phase (Attention Growth)

Dynamics: The trajectory reaches a second critical point $\theta_c^1$ , which is a saddle point. The $(W_0, W_1)$ subsystem is contracting (or neutrally stable), while the attention subsystem possesses a unique unstable mode.
Result: The attention parameters ( $W_Q, W_K$ ) grow exponentially along the unstable eigendirection. This induces a high-frequency bias, where the attention mechanism concentrates heavily on the most frequent token in the stationary distribution.

3. Dilution Phase (Mass Redistribution)

Dynamics: Once attention aligns with the unstable direction, the system enters a rank-one invariant manifold. The authors analyze next-order perturbations in the embeddings induced by the evolving attention amplitude.
Result: A mass-redistribution mechanism emerges. The embeddings of the high-frequency token and the remaining tokens begin to move in opposite directions. This growing separation in the embedding space progressively weakens the high-frequency focus, causing the attention weights to diffuse (dilute) across tokens.

4. Emergence of New Directions (Symmetry Breaking)

Dynamics: The system becomes trapped at a degenerate critical point on the rank-one manifold where driving forces vanish for perfectly symmetric low-frequency tokens.
Result: By introducing a small asymmetry ( $\delta$ ) among low-frequency tokens, the degeneracy is lifted. This creates a fast transverse unstable mode of order $\Theta(\delta)$ , driving the trajectory away from the rank-one manifold and unlocking new embedding directions. This initiates the next focus–dilution cycle, allowing the model to learn distinctions between lower-frequency tokens.

Experimental Results

Synthetic Data: Experiments on synthetic Markov data confirm the four predicted stages: initial condensation, attention growth (focus), attention dilution, and the emergence of new directions. Visualizations of embedding trajectories via PCA show the predicted retraction and expansion patterns.
Real-World Data: On WikiText and TinyStories, the model exhibits similar dynamics. High-frequency tokens (e.g., "the", whitespace) initially attract attention, followed by a dilution phase where attention spreads, and subsequent learning of medium-frequency tokens. The authors note that the Markov approximation holds well enough in early training phases to serve as a proxy for real-world dynamics.

Significance and Claims

The paper claims to provide a mechanistic theory for the multi-stage training dynamics of attention, specifically explaining the transition from concentrated to diffuse attention patterns.

Theoretical Rigor: It offers a stage-wise analysis based on linearization at critical points, avoiding ad-hoc assumptions about parameter coupling.
Cycle Explanation: It identifies the "focus–dilution" cycle as a fundamental mechanism where attention amplifies frequency-driven preferences and subsequently weakens them as the embedding structure adapts.
Role of Asymmetry: It demonstrates how small data asymmetries are crucial for breaking degeneracy and enabling the model to learn beyond the primary high-frequency token.

The authors maintain modesty regarding their scope, noting that the formal analysis is derived for a single-layer Transformer with small initialization. They acknowledge that extending these results to multi-layer, multi-head architectures and incorporating components like LayerNorm remains a direction for future work. However, they argue that the qualitative signatures observed in their minimal setting persist in realistic data, suggesting the framework is a useful lens for interpreting early-stage attention dynamics.