Equivalence of Context and Parameter Updates in Modern Transformer Blocks

Summary

This paper extends the Dherin et al. (2025) result — that context in a vanilla transformer block can be represented as a rank-1 patch to the MLP weights — from the single vanilla block to the modern-transformer setting (gated MLPs, RMSNorm, no biases, pre- and post-norm, MoE, parallel attention/MLP), and from a single block to a full $L$-layer network. The conceptual move is to factor the problem into two abstract properties — input controllability of the inner MLP core and output controllability of the outer function — and to show that whenever these hold at each block, the entire attention-induced difference can be absorbed exactly by rank-1 patches to the MLP input matrices, a patch to the RMSNorm scale, and an optional patch to the down projection. This yields a Gemma-specific constructive proof (Theorem 1), an inductive multi-layer extension with a self-correcting layer-by-layer algorithm (Theorem 2, Algorithm 1), and a unified framework (Theorem 3) that in principle covers Gemma, Llama, Falcon, Mistral/Mixtral, Qwen, GPT-2, GPT-J, and MoE variants. A companion impossibility result (Appendix D) shows that RMS post-norm without a trainable scale genuinely lies outside the framework, giving the theory a clean architectural boundary. The mechanistic contribution is worthwhile: it plugs the specific gap where the earlier proof depended on a bias term modern architectures have removed, and it turns a proof about one block into a general recipe. A full reimplementation of Algorithm 1 driven from the paper text alone reproduces the central float32 equivalence result exactly (100% argmax token match on Gemma-3-1B/4B and Falcon-7B, with logit differences at float32 rounding noise), so the theoretical contribution replicates cleanly. The Section 4.2 numerical-stability narrative, however, does not: on a $5\times 20$-token rerun the naive bfloat16 update already achieved 100% token match rather than the paper's 87.5%, so the RMSNorm-inversion stable variant had no failure to rescue and in fact scored marginally worse. The main conceptual limitation is that the equivalence is per-token — patches must be recomputed at every generation step — so the framework is best read as a descriptive lens rather than a competitive inference method, a caveat the paper acknowledges but sometimes downplays. The empirical scope is also narrower than the universal framing suggests, with only two model families exercised and no MoE or parallel-block validation.

Strengths

• Conceptual contribution. The two-property abstraction (input controllability + output controllability) is a genuinely clarifying reframing. It reduces a per-architecture analysis to checking two structural properties of the block, and Table 1 shows the framework absorbs a variety of gating/normalisation/expert patterns as instantiations rather than special cases.
• Constructive proofs, not existence-only. Theorem 1, Theorem 2, and the Appendix E gradient-descent recovery all give explicit formulas for the required updates, and Algorithm 1 is self-correcting layer-by-layer so numerical drift at one layer is compensated at the next. This is a stronger form of equivalence than a pure existence result.
• Clean impossibility boundary. Appendix D's proof that RMS post-norm without a trainable scale cannot be output-controllable defines the exact edge of the framework and rules out a natural alternative parameterisation, which strengthens rather than weakens the main claim.
• Numerically stable variant. The RMSNorm-inversion update in Appendix B, backed by an analytical Lagrangian derivation and a monotonicity proof for the bisection root, is a substantive engineering piece — even where its empirical rescue narrative is contested (see Weaknesses), the derivation itself is careful.
• Multi-modal check. The image-context experiment on Gemma 3 4B (Figure 5) is a nice sanity check that nothing about the derivation depends on the modality of the tokens the attention sub-layer consumes; the token-match at 40/40 in the replication confirms this.
• Paper is implementable from the text. A full reimplementation of Algorithm 1 and Appendix B built purely from the paper (no released code) reproduces the headline float32 result on both Gemma-3-1B/4B and Falcon-7B. That is a real signal about the methodology's specificity.

Weaknesses

• Precondition on f in Theorem 1 appears inverted. The theorem states the update holds and $f(W_{\text{gate}} z_C, W_{\text{up}} z_C)_i = 0$ for all $i$. Under that condition the MLP contributes nothing and the block collapses to the residual, defeating the theorem's purpose. Section 4's divide by 1 to avoid NaN fallback confirms the intended condition is $f(\ldots)_i \ne 0$. This sits at the entry point of the paper's central result and needs to be fixed.
• Definition 2 does not cover its own instantiations. As stated, output controllability requires the property to hold for any desired difference vector. The elementwise-multiply case used in Theorem 1 (Lemma 5) only satisfies this when no component of $\mathbf{h}$ vanishes; Theorem 3 patches this over with a global non-zero activations proviso, but the definition itself is over-strong and the proof of Theorem 3 never actually invokes the proviso. This should be made explicit at the level of Definition 2 or per-lemma.
• Table 1's shorthand hides an ordering. The elementwise-multiply row writes $\Delta \mathbf{m} = \Delta A_x(Y) \oslash \mathbf{h}$ with $\mathbf{h} = f(\mathbf{v})$, but the actual construction requires $\mathbf{h}$ to be the post-input-patch MLP output — i.e. $f$ evaluated at the full-context internal state after the Lemma 1/2 updates. A reader composing the table entries as-is could apply the $\mathbf{m}$ update against the un-patched MLP output, which would be wrong.
• Numerical-stability narrative in Section 4.2 does not reproduce. The headline 87.5% naive bfloat16, 100% under the stable variant is the paper's motivating evidence for Appendix B, but under an independent rerun of the same protocol the naive bfloat16 update already achieved 100% token match, and the stable variant scored 99% (marginally worse). The paper's 87.5% is exactly 7/8, which suggests the aggregate pools over very few tokens; Section 4.2 does not state the token count, so the claim is effectively uncheckable. Because the whole Section 4.2 arc (naive degrades, Appendix B rescues) rests on this single scalar, this warrants a rewrite rather than a footnote.
• Much smaller update norm from Appendix B only holds in bfloat16. In a matched rerun the stable variant's total update norm is $\sim 49\times$ smaller than naive in bfloat16 (as claimed) but $\sim 4.6\times$ larger than naive in float32, driven by the vector (RMSNorm-scale) component — the exact quantity Appendix B is designed to shrink. Figure 8b renders the same reversal for the scaling variant, so this is a consistent property of the run rather than noise.
• Scaled-variant minimizes extreme values in Δm is contradicted. The scaled variant's Δm maximum ($6.07\times 10^6$) exceeds naive's ($2.5\times 10^6$); the stable variant, not the scaled one, is what minimises both median and max of Δm. Token-match is saturated at 100% in every cell, so the token-level evidence does not discriminate between variants.
• Empirical scope is narrower than universal. Only two model families (Gemma, Falcon) are exercised, spanning gated-post-norm and pre-norm variants respectively. Mixtral (MoE, Lemma 6), GPT-J (parallel blocks, Lemma 7), Llama, Mistral, and Qwen — all listed as covered — are never validated. The interesting stress cases (near-zero gating gates in MoE, aggregation across experts, the parallel-block full-attention-absorption) are precisely where the framework's cleanliness could be tested, and are absent.
• No head-to-head with the vanilla-transformer patch. The empirical validation compares the new modern-architecture update to the contextualised baseline, but never to the natural strawman: applying the Dherin et al. (2025) vanilla-transformer patch directly to Gemma/Falcon (which should fail because of the missing bias, RMSNorm scale, and gating). Without that baseline, the reader has to trust rather than see that the new framework is required.
• Small-N and no variance reporting. The headline 87.5% agreement in bfloat16 naive is aggregated over ~5 short prompts with ~10–20 tokens each. There are no confidence intervals, no per-prompt breakdown, and no seed variance. Since precision behaviour is the axis of interest, this is thin evidence for a scalar quoted so prominently.
• Descriptive-vs-prescriptive framing is under-signalled. The Introduction and Section 3.3 present Algorithm 1 as if it were an inference procedure, but recomputing per-layer rank-1 patches at every generation step is strictly more expensive than running the model with context. The Limitations note this, but the main text does not — and the tool for understanding framing in the Conclusion is asserted without a mechanistic-interpretability case study to back it.
• RMSNorm treated as L2 normalisation inside the impossibility proof. Appendix D says the normalised vector $\mathbf{z}/\|\mathbf{z}\|_2$ lies on the unit hypersphere and bounds $\|g(z)\|_2$ by $\|\mathbf{m}\|_\infty$. Under the RMSNorm normalisation actually used (dividing by $\|\mathbf{z}\|_2/\sqrt{d}$), the correct bound is $\sqrt{d}\,\|\mathbf{m}\|_\infty$, which is what the next-page diameter bound $2\sqrt{d}\,\|\mathbf{m}\|_\infty$ tacitly relies on. The impossibility conclusion still holds, but the intermediate inequality is not internally consistent.
• Appendix E fixed preconditioner may quietly redefine the algorithm. The telescoping recovery of $\Delta \mathbf{m}$ requires holding $\mathbf{h}'_{\text{norm}}$ fixed at $\text{Norm}(\mathbf{h}_{\text{target},n})$ across all context steps. The paper asserts this is consistent with the algorithm's construction but does not prove that $\text{Norm}(\mathbf{h}_{\text{target},i})$ is genuinely $i$-independent; if it is not, the derivation is exact only for this specific preconditioner choice.
• Metric naming issues. Rank distance is defined as a Pearson correlation, which is a similarity — a value of $1$ means identical orderings, and it can be negative. This inverts the axis interpretation in Figure 8. The Top-1 difference formula as rendered does not match its prose description; the intended expression is $p_{\arg\max p} - q_{\arg\max p}$.
• Does not appear in any real-world architecture overreaches. The impossibility result's practical relevance rests on the claim that RMS post-norm without a trainable scale never occurs in practice. No citation or enumeration is given; RMSNorm/LayerNorm variants stripped of the affine parameter do exist in some designs. Softening to none of the major architectures we survey would be defensible.

Reproducibility & code

• Headline float32 equivalence. The paper's central claim — a Gemma model with patched weights, no context, tracks the contextualised baseline exactly under float32 — reproduces cleanly on a wholly re-implemented codebase built from the paper text alone. 200/200 argmax tokens matched on Gemma-3-1B-IT across the five paper prompts, 40/40 on the Gemma-3-4B image experiment, and 100/100 on Falcon-7B, with median $L_\infty$ logit difference $\sim 4\times 10^{-5}$ and median TVD $\sim 1.6\times 10^{-6}$.
• Section 4.2's 87.5% / 100% stability figures. Did not reproduce: naive bfloat16 came out at 100.0% token match (12.5 points above the paper) and the stable variant at 99.0% (below, not above, naive). At 100 generated tokens, bfloat16's logit noise ($L_\infty \approx 0.125$) never flipped an argmax on this prompt set. The paper's 87.5% is exactly 7/8, which is consistent with pooling over roughly eight tokens; Section 4.2 does not state the count, so the claim as written is not checkable.
• Per-token recomputation necessity. A fixed-patch ablation (compute the patch once, re-use it across generation) collapses to 5% token match and TVD $\approx 1.0$ at the first newly generated token, exactly as the Section 6.1 caveat predicts. This is a genuine addition — it tests a design choice the paper asserts rather than measures — and strengthens the descriptive-lens framing.
• Stable-variant norm claim splits by precision. Total-update norm medians reproduce the paper's ~49× smaller claim in bfloat16 but reverse in float32 (~4.6× larger), driven by the RMSNorm-scale component. The two stable conditions also disagree with each other by roughly $240\times$ despite the RMSNorm inversion running in float64 regardless of model dtype, which is hard to explain for a routine whose whole point is precision insensitivity.
• Scaled-variant Δm-extremes claim contradicted. The scaled variant's Δm maximum exceeds naive's; the stable variant, not the scaled one, is what minimises the tails.
• Layer-ablation coverage gap. The only the final layer is numerically problematic half of the layer-depth ablation cannot be exercised from the shipped configuration — the sweep tops out at start-layer 25 on a 26-block model, so the deepest configuration still patches two blocks. The few layers suffice half replicates cleanly down to two layers.
• Underspecified numerical protocol. The divide by 1 NaN fallback (Section 4.1), the bisection tolerance/bracket for the RMSNorm inversion (Appendix B), the token count behind the aggregates, decoding strategy, seed, and library revision are all left implicit. These are precisely the details that determine whether bfloat16 reaches or misses 100%. The reimplementation had to guess several of them.
• No released artefact. Neither paper.md nor its appendices point to a repository, artefact archive, or data-availability statement. The multimodal experiment additionally requires the specific image, the Gemma-3-4B image-tokeniser configuration, and the exact chat-template bytes, none of which are provided. Falcon parameter mapping is not enumerated either.

Recommended Changes

Essential

• Fix the Theorem 1 precondition. Change $f(W_{\text{gate}} z_C, W_{\text{up}} z_C)_i = 0$ for all $i$ to $\ne 0$ (or equivalent) so it states the condition needed for the elementwise-multiply update rather than trivialising the block — this is the single-block equivalence result on which the whole paper rests.
• Tighten Definition 2 (or Theorem 3's proof). Either add a non-degeneracy clause to the definition of output controllability so the elementwise-multiply operator satisfies it as stated, or explicitly invoke the non-zero-activation proviso inside the two-step proof of Theorem 3 and per-lemma in Table 1.
• Rewrite Section 4.2's stability narrative around a stated token count. Replace the single scalar 87.5% → 100% with a per-prompt / per-seed breakdown at a fixed, disclosed generation length, and state the number of tokens the aggregate pools over. As reported, the 87.5% could not be reproduced (naive bfloat16 sat at 100% on a $5\times 20$-token rerun, so the stable variant had no failure to rescue), and the recovery arc collapses without a stated protocol.
• Rescope the much smaller update claim (Figure 8 / Appendix B). State that the norm reduction is a bfloat16-only phenomenon or investigate why the RMSNorm-inversion path, run in float64 regardless of model dtype, produces a $\sim 4.6\times$ larger stable-variant norm than naive in float32 — the current claim is contradicted by the paper's own experimental design under a matched rerun.
• Correct the scaled-variant claim in Appendix C.5. Either specify the extremal quantity for which the scaled variant is minimal, or restate as reduces the bulk (median) of Δm without controlling its tails.
• Extend the layer-depth ablation to actually cover last-layer-only. The claim updating only the final layer is numerically problematic is not testable from the shipped sweep; add the start_layer = last-block configuration or restate the claim to reflect from two layers onward, accuracy is essentially unchanged.
• Add a vanilla-patch baseline. In Figures 3, 4, and 9, include a naive port of Dherin et al. (2025) baseline (no RMSNorm scale patch, no gating-aware split) so the reader can see how much error the modern-architecture treatment removes.
• Exercise at least one MoE and one parallel-block architecture. Even a single Mixtral prompt (Lemma 6, router-gate variation) and a single GPT-J prompt (Lemma 7) would substantiate the unified framing; without them the empirical evidence covers only two of the seven architectures Section 5 claims to cover.
• Release a reference implementation. Publish a minimal repository covering Algorithm 1 on Gemma-3-1B and Falcon-7B, the RMSNorm-inversion bisection, and the exact prompts, seeds, decoding settings, and library revisions used for Figures 3–13. This is what turns the qualitative reproducibility of the theory into numeric reproducibility of the empirics.

Suggested

• Annotate Table 1 with update ordering. Note explicitly that the elementwise-multiply row assumes the input matrices have already been patched (so $\mathbf{h}$ is the post-patch MLP output), preventing an obvious mis-composition.
• Frame Algorithm 1 as descriptive. Add one sentence to Section 3.3 or the Introduction stating that the algorithm is a proof-of-existence and analytical tool, not a competitive inference procedure — this is currently only in the Limitations.
• Reconcile the L2/RMS confusion in Appendix D. State the correct $\sqrt{d}$ factor in the intermediate bound so it agrees with the subsequent $2\sqrt{d}\,\|\mathbf{m}\|_\infty$ diameter bound, and note that the diameter is an upper bound rather than tight.
• Prove or restate the fixed-preconditioner step in Appendix E. Either show that $\text{Norm}(\mathbf{h}_{\text{target},i})$ is genuinely $i$-independent, or explicitly acknowledge that the gradient-step recovery is exact only under the chosen fixed preconditioner.
• Rename rank distance and fix the Top-1 formula. Call the correlation-based quantity rank correlation (or plot $1 - \rho$), and typeset Top-1 so it unambiguously reads $p_{\arg\max p} - q_{\arg\max p}$.
• Soften the no real-world architecture claim. Replace with none of the major architectures we survey plus a footnote surveying the RMSNorm-without-affine variants that do exist.
• Document the divide-by-1 fallback, the bisection tolerance, and the layer-index scheme. Give the threshold, the elementwise/matrix scope, an estimate of how often the fallback fires, the bisection bracket and tolerance, and the mapping from L1..L25 to Gemma-3-1B's actual block depth (plus the prompts used for the ablation).
• Add the fixed-patch collapse curve to the main text. The patch once, re-use across generation ablation collapses to 5% token match / TVD $\approx 1$ at the first newly generated token — a much sharper illustration of the per-token equivalence boundary than the prose statement in Limitations.
• Add a small mechanistic-interpretability case study. Even one example inspecting the direction or components of a rank-1 update for a semantically loaded prompt would demonstrate the tool for understanding framing that the Conclusion asserts.