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

Summary

This paper proposes a mechanistic account of how attention is learned in a one-layer Transformer trained by cross-entropy on Markovian data. The central conceptual move is to describe attention learning as a repeated focus–dilution cycle: after an initial rank-one condensation of the embedding and output projection, the attention parameters $(W_Q, W_K)$ align with a single unstable direction and produce a frequency-driven bias toward the high-frequency token (focus); this bias then triggers next-order corrections in the embeddings that redistribute mass and progressively weaken the very focus that produced it (dilution); a small asymmetry among low-frequency tokens finally lifts a degenerate critical point and opens new embedding directions, seeding the next cycle. The analytical spine is stage-wise linearization around a sequence of critical points, with a Lyapunov–Schmidt reduction handling the degenerate transition in Sec. 4.4. Relative to prior single-stage or reparameterization-based analyses, the paper explains both attention amplification and its subsequent dissipation inside one coupled model, without artificially decoupling $(W_0, W_1)$ from $(W_Q, W_K)$. Re-executing the released synthetic pipeline confirms the empirical backbone: the four stages, and their recursive repetition across successive low-frequency tokens, reproduce cleanly across four independent diagnostics, which raises confidence that the core phenomenon is real and not a plotting artifact. The main conceptual limitations are that the sharp results live in a very stylized setting — population gradient flow, a one-layer softmax-only Transformer, a symmetric two-group stationary distribution with the asymmetry analysis restricted to $d=3$ — while the experiments use Adam on a fully exponential distribution, and the more realistic extensions (residual, multi-layer, multi-head, WikiText, TinyStories) verify mostly a secondary signature (embedding retraction), with the real-world and multi-head evidence only partially reproducing under execution. Several theoretical steps (undefined symbols such as $p_{\mathrm{sym}}$ and the growth constant $c$, the ordering $T_p\le T_1$, the approximate conservation law, and the second-order vanishing of $\eta$) are asserted rather than closed. Overall the contribution is genuine and the mechanism is worth having in the literature; presentation, scope of claims, and reproducibility artifacts need tightening.

Strengths

• Conceptual contribution. Framing attention learning as a cyclical focus–dilution process, and identifying one mechanism as the driver of successive stages, is a genuinely new angle: the paper explains both concentration and its later dissipation inside one coupled model rather than relying on reparameterizations or proxy dynamics.
• Coupled, native analysis. The stage-wise linearization tracks the interaction between $(W_0,W_1)$ and $(W_Q,W_K)$ rather than decoupling them. The Jacobian block structure at the second critical point (contracting outer subsystem, unstable attention subsystem) and the invariant rank-one manifold in Sec. 4.3 give an operational picture of why dilution follows focus — mass redistribution on the manifold — rather than merely observing it.
• Handling of the degenerate stage. The Lyapunov–Schmidt reduction at the degenerate critical point of Sec. 4.4, separating a $\Theta(\delta)$ transverse instability from an $O(\delta^2)$ tangential one, is the technical highlight and gives a concrete mechanism for the emergence of a new embedding direction.
• Multi-signal, reproducible synthetic picture. Fig. 2 aligns cosine-similarity heatmaps, PCA of embedding snapshots, attention maps, attention entropy, and per-token embedding norms into one coherent four-stage narrative — and, on re-execution, all of these diagnostics reproduce the stated stages and the recursive repetition, in order.
• Cross-domain qualitative signal. The WikiText and TinyStories descriptions show a focus-then-dilution signature and per-token embedding retraction, giving at least qualitative evidence that the mechanism is not purely an artifact of the synthetic setup.

Weaknesses

• Generalization claim outruns the evidence. Sec. 5.2 concludes broadly that "the proposed mechanism remains relevant beyond our simplified theoretical setting," but of the three extensions only the residual case shows the full cycle; multi-layer changes the "detailed manifestation," and multi-head verifies only embedding retraction. On re-execution the multi-head claim did not reproduce at its stated configuration and its stated trend appears reversed (see Reproducibility & code). This is stated contribution 3, so the mismatch matters.
• Cyclicity is proved for one cycle only. The title and abstract advertise a recurrent cycle, but the theory establishes Stages I–IV once; the self-similar re-entry that makes it cyclical is only hypothesized ("we hypothesize that after Stage IV, the model has effectively converged on Token 0"). The recursion is observed empirically but not certified theoretically.
• Theory–experiment regime gaps. Every theorem concerns continuous population gradient flow (Prop. 3.3), yet the experiments use Adam with a fixed learning rate on mini-batches, with no plain-GD control to isolate optimizer effects. Separately, Props. 4.3–4.7 assume a symmetric low-frequency group $\pi_1>\pi_2=\cdots=\pi_d$, but the experiments use a fully exponential distribution in which every token has a distinct frequency, and the asymmetry analysis (Thm. 4.9) is only carried out for $d=3$ while experiments use $d\ge 4$.
• Load-bearing steps stated without closure. Several pivotal steps are asserted: the ordering $T_p\le T_1$ (Prop. B.2), the monotonicity of $z$ and $\|a\|^2+\|b\|^2$ (Prop. B.5), the approximate conservation law used to solve the reduced $z$-dynamics, the vanishing of $\eta$'s first- and second-order derivatives at the degenerate point, and $\partial_\delta^2(\partial L/\partial\Phi)=0$ in Lemma C.14.
• Undefined / garbled symbols. $p_{\mathrm{sym}}$ carries the stage I–II error estimates but is never defined; the plateau scale appears as both $\varepsilon^{1/n}$ and "$\varepsilon\,n$"; the growth constant $c$ in Prop. 4.4 is rendered unreadable and uses an undefined $\lambda$; and $\gamma_{i=1}/\beta_{i=1}$ literally reads as evaluation at $i=1$ when $i\ne 1$ is meant.
• Statement–assertion mismatches. Prop. 4.4 calls $J_{\mathrm{att}}$ positive semi-definite while claiming a unique unstable direction — consistent only if $J_{\mathrm{att}}$ is rank-one PSD, which is not stated. In Lemma B.1 the identity $\mathrm{Var}(p_{\mathrm{sym}})u_\perp = p_{\mathrm{sym},i=1}u_\perp$ is called "directly verifiable" without an explicit $p_{\mathrm{sym}}$.
• Single-trajectory empirical evidence. Figs. 2–5 present one trajectory per configuration with no seeds or error bars, while stage boundaries and windows (e.g. steps 350–550) are read off visually; the reproduction saw those step numbers shift (focus window ~300–400), underscoring the need for multi-seed bands.
• Local bookkeeping issues. $c_6$ carries $\mathrm{Var}(P_1)$ where symmetry with $c_3$ predicts $\mathrm{Var}(P_{i\ne 1})$; "The matrix form is of the shape" is duplicated exactly where a bilinear block is being defined; and the softmax-to-sigmoid reduction is used before it is motivated.

Reproducibility & code

Code was executed (reconstructing missing pieces from the manuscript). The core synthetic results reproduce; the extensions and real-data results are the weak points.

• Headline synthetic cycle reproduces. The four-stage focus–dilution cycle and its recursive repetition were reproduced at full scale on the shipped pure_attn model: a two-phase rank collapse ($W_0,W_1$ to near rank-1 while $W_Q,W_K$ stay high-rank, then $W_Q,W_K$ collapse), a synchronous rise-then-decay of attention on the high-frequency token, the flat-then-rising orthogonal projection, and staggered per-token norm retraction/regrowth all match Fig. 2. Exact step numbers drifted (focus window ~300–400 vs the paper's 350–550), consistent with seed/schedule sensitivity.
• Analysis pipeline not shipped runnably. The Fig. 2 diagnostics (cosine-similarity heatmaps, PCA, attention maps, entropy, orthogonal-projection norms) ship only as commented-out/removed code and had to be rewritten from scratch; the entry point also fails on a missing draw_condense_heatmap import. The model path itself is intact, but Fig. 2 is not regenerable from released artifacts as-is.
• Residual model masks the key token. The residual signature reproduces only after re-implementing Eq. (231) by hand: the shipped residual model applies a padding mask (seq_k.eq(0)) that treats token id 0 — the high-frequency token the effect concentrates on — as padding, silently suppressing the phenomenon if run verbatim.
• Multi-head extension did not reproduce as claimed. The paper's Fig. 5 configuration is $L=8$, $H\in\{1,2,8\}$; only $8$L$1$H and a single-layer $2$-head proxy ran. The $8$L$1$H retraction was muted (not "remarkably well preserved"), and in the available proxy retraction grew stronger with more heads — opposite to the stated trend. No $H=8$ or $8$L$2$H/$8$L$8$H model was executed.
• Real-data pipeline absent. The WikiText driver (main_wikitext.py), tokenizer, and preprocessing were never shipped and had to be reconstructed; TinyStories was not run. Consequently the WikiText four-phase attention transition reproduced only three phases (secondary dilution absent), the high-frequency token became <unk> under a vocab-2000 word tokenizer rather than whitespace, medium-token norm retraction on real data was not clearly documented, and the TinyStories pseudo-2-gram claim (Fig. 3C) could not be attempted.
• Setup mismatches. The executed config produces a five-token exponential distribution ($\approx (0.751,0.188,0.047,0.012,0.003)$) rather than the stated four-token $\pi=(0.75,0.19,0.05,0.01)$, and the released training model includes a value projection $W_V$ absent from Def. 3.2. Neither prevented the headline results from reproducing, but both should be reconciled. The code pointer is a bare text anchor with no URL or commit.

Recommended Changes

Essential.

• Narrow the generalization claim. In Sec. 5.2 either add direct focus/dilution measurements under multi-head/multi-layer settings at the reported $L=8$, $H\in\{1,2,8\}$ configuration, or restrict the conclusion to embedding retraction — the only signature that carries over — and correct the direction of the head-count trend to match what is actually observed.
• Prove or soften the cyclicity claim. Add the reduction argument showing the post-Stage-IV sub-dynamics inherit the original assumptions (so the cycle genuinely recurs), or downgrade the "recurrent/cyclical" language in the title and abstract to reflect that recursion is empirical/hypothesized.
• Close the load-bearing proof steps. Supply the growth-rate lower bound delivering $T_p\le T_1$ (Prop. B.2); a forward-invariance proof for the monotonicity in Prop. B.5 invoking $c>0$; a precise statement and error bound for the approximate conservation law; the explicit computation that $\eta$'s first/second derivatives vanish at the degenerate point; and the term-by-term argument for $\partial_\delta^2(\partial L/\partial\Phi)=0$ in Lemma C.14.
• Define $p_{\mathrm{sym}}$ and fix exponent bookkeeping. State $p_{\mathrm{sym}}$, derive the two "directly verifiable" identities in Lemma B.1, and choose one canonical plateau scale ($\varepsilon^{1/n}$ or "$\varepsilon\,n$") used consistently for $\|a\|$ and $\|a\|^2+\|b\|^2$.
• Repair $c$ and reconcile $J_{\mathrm{att}}$. Rewrite the constant $c$ in Prop. 4.4 with every symbol defined (disambiguating $\lambda$ from the Markov parameter), and state the rank-one PSD structure (single positive eigenvalue) of $J_{\mathrm{att}}$ so the alignment argument is unambiguous.
• Ship a runnable synthetic pipeline. Fix the draw_condense_heatmap import, re-enable/add the analysis code producing Fig. 2's heatmaps, PCA, attention maps, entropy, and projection curves, align the shipped config to the stated distribution (or document the exponential variant), and provide a one-command reproduction.
• Fix the residual masking trap. Remove or document the seq_k.eq(0) mask so token id 0 is not treated as padding, and ship a config that selects the Eq. (231) model actually analyzed.
• Provide a real-data reproducibility bundle. Release the WikiText/TinyStories driver, tokenizer, preprocessing, configs, and Fig. 3 plotting scripts, and name the exact tokens used (continue, whitespace, etc.); confirm whether the four-phase (vs three-phase) attention transition holds under a longer run.
• Reconcile the model with the theory. Either remove $W_V$ from the trained model (matching Def. 3.2) or add a statement covering the $W_V$ variant.

Suggested.

• Bridge population GF and Adam. Add a plain-GD control on the synthetic data to confirm the four stages survive the theoretical optimizer, and comment on why Adam should track the GF timescales.
• Justify the experimental distribution. State whether the leading-order theory (symmetric low-frequency group, $d=3$ asymmetry) is expected to govern the fully exponential $d\ge 4$ experimental setting, and how the extra tokens map onto the theory.
• Add multi-seed evidence. Rerun Figs. 2–5 over several seeds and report bands over the stage transitions and step windows; ideally add a quantitative stage-boundary detector.
• Positive test of the Markov approximation. For WikiText/TinyStories, measure attention on the immediately preceding token vs earlier positions, or a Markov-fit residual, before concluding synthetic Markov data is a "suitable proxy."
• Sensitivity to init scale. Report a sweep in $\varepsilon$ around the $1/m^2$ setting to show the four-stage picture is not tied to a single scale.
• Clean up local notation. Adopt an unambiguous replacement for $\gamma_{i=1}/\beta_{i=1}$; fix the duplicated "The matrix form is of the shape"; check $c_6$ for $\mathrm{Var}(P_{i\ne 1})$; and add a one-line justification for the softmax-to-sigmoid reduction at its first use. Replace the bare code anchor with a URL and commit hash.