PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivalence of Context and Parameter Updates in Modern Transformer Blocks

Adrian Goldwaser, Michael Munn, Javier Gonzalvo, Benoit Dherin

arXiv:2511.17864v3cs.LG

TL;DR

How context translates into effective parameter updates in modern LLMs remains unclear beyond vanilla transformer blocks. This paper derives patches for a Gemma-style block, extends the result to multilayer models through a controllability framework, and reports near-perfect logit matching with identical token generation in experiments.

  • Problem

    How context maps to implicit weight patches in modern architectures remained unanalysed beyond vanilla blocks, including architectures with gating, normalization, or multiple layers.

    Source: §2 · p. 2

  • Method

    The paper derives exact patches for a Gemma-style block, extends them inductively with an algorithm to multilayer models, and unifies the results through input and output controllability.

    Source: §1 · p. 1

  • Results

    The patched model without context achieves near-perfect logit matching and identical token generation to the original model with context on Gemma 3 text and image contexts and Falcon text contexts.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The theory extends implicit weight updates from single-layer vanilla transformers to complex, multilayer architectures used in modern LLMs.

    Source: §6 · p. 9

    Exact updates replicate the next-token distribution for a specific history, while naive use in multi-token free generation is disrupted by attention values of newly generated tokens.

    Source: §6.1 · p. 9

Abstract

from arXiv · show

Recent research has established that the impact of context in a vanilla transformer can be represented implicitly by forming a token-dependent, rank-1 patch to its MLP weights. This work extends that foundational theory to the diverse architectures of modern Large Language Models. We first demonstrate a precise, analytical solution for a Gemma-style transformer block, proving that the entire effect of a context can be perfectly mapped to rank-1 patches on its MLP weight matrices and a patch to the RMSNorm scale. We then generalize this result, providing a constructive proof and algorithm for multi-layer models. To unify these findings, we introduce a general framework centered on two core properties: input controllability and output controllability. We prove that a perfect implicit weight patch is possible for any MLP block where the inner function is input-controllable and the outer function is output-controllable. This provides a simpler and more powerful lens for understanding how transformer models transmute prompts into effective weights. This setup generalizes to a wide range of modern LLM architectures including gating, pre-/post-norm, mixture of experts and sequential/parallel transformer blocks.

1. Introduction

The paper addresses the lack of formal analysis of implicit weight patches in modern transformer architectures. It derives exact patches for a Gemma-style block, extends them to multilayer models, and proposes a controllability framework for broader architectures.

1. Introduction

  • Modern architectures use components such as gated MLPs and RMSNorm, but their compatibility with implicit weight patches had not been formally analyzed.The earlier proof applied to a vanilla transformer block.Source: §1 · p. 1
  • The paper derives exact parameter patches to absorb context into the MLP and normalization layers of a Gemma-style architecture.The single-block result is presented as a constructive proof.Source: §1 · p. 1
  • The single-block result extends inductively to deep, multilayer models, with a practical algorithm for computing a perfect patch across the network.The paper identifies this as Theorem 2 and Algorithm 1.Source: §1 · p. 1
  • Input controllability and output controllability underpin a unified theorem that generalizes the patch framework to architectures including Gemma, Llama, Falcon, Mistral, and MoE models.The framework is introduced as Definitions 1 and 2 and used in Theorem 3.Source: §1 · pp. 1–2
  • Experiments on Gemma 3 for text and image contexts and Falcon for text contexts show near-perfect logit matching and identical token generation without context.The patched model is compared with the original model using context.Source: §1 · p. 2

2. Background

Earlier accounts of in-context learning include Bayesian inference, gradient-descent analogies, and induction heads; prior mechanistic work formalized context as an implicit MLP patch for vanilla blocks. This paper addresses the remaining gap for modern architectures and distinguishes its focus on patch existence and form from work on patch reuse.

2. Background

  • Prior accounts explain in-context learning through Bayesian inference, a gradient-descent analogy, or induction heads; Dherin et al. (2025) instead frame context as an implicit MLP weight patch.The Bayesian, optimization, and induction-head accounts are presented as complementary theories.Source: §2 · p. 2
  • In a vanilla block, the contextual difference is absorbed by a rank-1 patch to the input weight matrix W1 and a patch to the output bias.This equivalence depends on the presence of the output bias b2.Source: §2 · p. 2
  • The vanilla-block result leaves open extensions to bias-free modern models, multilayer networks, normalization layers, and gating.Gemma and Llama are cited as modern models that have eliminated biases.Source: §2 · p. 2
  • Innocenti and Achour (2025) address Pre-LayerNorm and arbitrary positions but remain restricted to standard residual blocks with biases; this paper targets bias-free, gated, RMS-normalized architectures.The prior work also acknowledges a lack of exact correspondence to practical model architectures.Source: §2 · p. 2
  • Unlike work aggregating transient patches into reusable, token-independent patches, this paper focuses on whether patches exist and what form they take in modern architectures.It presents the implicit patch as a general principle rather than an artifact of vanilla transformers.Source: §2 · p. 2

3. Results: Implicit Updates in Modern Transformers

The paper establishes exact context-to-parameter equivalence for a Gemma-style transformer block and extends it to multi-layer models. It also gives a layer-by-layer algorithm for computing the updates.

3.1. The Gemma Block

  • A Gemma-style block can exactly reproduce its full-context output after updates to its MLP parameters and output scale, provided the stated activation components are nonzero.The gate and up-projection updates preserve the internal MLP activation, while the scale update absorbs the residual-path difference.Source: §3.1 · pp. 3–4

3.2. Extension to Multi-Layer Architectures

  • For an L-layer model whose blocks satisfy Theorem 1, updated parameters make the reduced-context model’s output identical to the full-context model’s output.The result is proved inductively across layers.Source: §3.2 · p. 4

3.3. Algorithm for Multi-Layer Updates

  • The multi-layer proof yields a practical algorithm: record full-context activations, then compute each layer’s update using the previous layer’s patched output.This sequential procedure is described as self-correcting in the presence of numerical errors.Source: §3.3 · p. 4
  • Applying all context updates only at the final layer is theoretically possible but presents practical issues in real-world scenarios.The paper points to Appendix C for experiments on this approach.Source: §3.3 · p. 4

4. Experimental Validation

The experiments test whether updated models without explicit context reproduce the token-by-token outputs of contextualized models. Float32 results are near-exact, while lower-precision behavior varies with the update and architecture.

4.1. Setup

  • For each generated token, the updated model recomputes weights to absorb the prompt and prior generated tokens, then compares next-token distributions on a shared generation history.When top-token choices diverge, the updated model follows the baseline's token so later distributions remain comparable.Source: §4.1 · p. 5

4. Experimental Validation

  • The evaluation compares contextualized generation with autoregressive generation from the updated model without explicit context, using token-level matching, logit distance, and total variation distance.The logit and probability-distribution metrics are compared at each generation step.Source: §4.2 · p. 6

4.2. Results

  • Falcon achieves near-identical output distributions and perfect token matching even in bfloat16, while its lack of post-norm reduces sensitivity to numerical issues.The Falcon evaluation uses text contexts.Source: §4.2 · p. 6
  • Float32 maintains perfect token matching with extremely small logit differences across text and image contexts; bfloat16 token agreement rises from 87.5% to 100% with the stable update.The stable bfloat16 result is reported as on par with float32.Source: §4.2 · p. 6
  • Gemma's standard update is numerically sensitive because the RMSNorm-scale update divides by potentially very small or zero values; a stable variant reduces, but cannot eliminate, low-precision issues.The authors attribute remaining bfloat16 issues to high dimensionality and low precision.Source: §4.2 · p. 6

5. A General Framework for Implicit Updates

The section frames exact context removal through input and output controllability: the inner and outer functions of a residual MLP block must permit corresponding parameter updates. It extends this result across common architectures, while identifying a post-norm constraint and connecting the updates to gradient descent.

5. A General Framework for Implicit Updates

  • A perfect implicit update exists for a residual MLP block when its inner function is input-controllable, its outer function is output-controllable, and the activations are non-zero.The Gemma-specific theorem is presented as a special case of this more general result.Source: §5 · p. 8
  • The constructive proof first updates inner-function parameters to match the intermediate vector, then updates outer-function parameters to absorb the residual shift.Input controllability handles the input change; output controllability handles the remaining output difference.Source: §5 · p. 8
  • The framework establishes controllability for matrix multiplications and norms, and output controllability for biases, matrix multiplication, elementwise multiplication, and mixture-of-experts components.It also provides an update for parallel transformer blocks and covers common architectures including Gemma, Llama, Falcon, Mistral/Mixtral, Qwen, GPT-2, and GPT-J.Source: §5 · p. 8

5.1. Limits of controllability

  • RMS post-norm without a trainable scaling vector m is outside the framework because its output is constrained to the L2 sphere and cannot absorb a shift with a different scale.The passage identifies this as a failure of output controllability.Source: §5.1 · p. 9

5.2. Connection to Implicit Gradient Updates

  • The context-equivalent updates can be formulated exactly as gradient-descent steps on a complex trace-loss objective.The derivation covers standard outer-weight architectures and a numerically stable Gemma variant involving RMSNorm inversion and scale vector m.Source: §5.2 · p. 9

6. Conclusion

The paper extends implicit weight-update theory from vanilla blocks to multi-layer modern LLMs, with a constructive Gemma-style derivation and an algorithm for computing updates. Its framework describes per-token context effects, while the exact updates remain history-specific and do not provide a global patch for every query.

6. Conclusion

  • The framework generalizes implicit updates to multi-layer architectures such as Llama and Gemma, deriving rank-1 MLP updates and an algorithm for computing them.The paper also reports almost identical output distributions and the same textual generation on Gemma 3 1B and Falcon 7B.Source: §6 · p. 9
  • Input and output controllability provide a unified framework for analyzing how modern language-model architectures implicitly fine-tune themselves on context.The authors present it as a tool for understanding in-context learning and designing transformer architectures.Source: §6 · p. 9

6.1. Limitations

  • Exact updates replicate the next-token distribution for a specific history, but newly generated tokens’ attention values disrupt naive use of those updates for free generation.The updates must be recomputed at each step, so they do not form a global update that absorbs context for every query.Source: §6.1 · p. 9, §6 · p. 9

6.2. Future Work & Applications

  • A future context compiler could traverse an architecture’s computational graph and apply controllability updates optimized for constraints such as numerical drift.The paper also identifies reusable “thought patches” as a route toward extending single-token equivalence to unconstrained generation, and cites concurrent work exploring this aggregation (Mazzawi et al., 2025).Source: §6.2 · p. 9

Impact Statement

The paper aims to advance theoretical understanding and mechanistic interpretability of LLMs through foundational analytical work. It hopes the framework will support more transparent, efficient, and robust AI architectures.

Impact Statement

  • The paper’s goal is to advance machine learning through theoretical understanding and mechanistic interpretability of large language models.The authors characterize the work as foundational and analytical.Source: Impact Statement · p. 9
  • The authors hope their mathematical framework will facilitate more transparent, efficient, and robust AI architectures.They report no direct or immediate negative societal consequences.Source: Impact Statement · p. 10

A.1. Proofs of Controllability

The controllability proofs show how parameter changes can reproduce input or output changes across several transformer-block components. A numerically stable RMSNorm update instead prioritizes the down-projection and uses scale updates only for residual error.

A.1. Proofs of Controllability

  • Standard and gated MLPs are input-controllable when their input is nonzero, because rank-1 projection updates can reproduce the changed input arguments.The result also applies to pre-norm MLPs when the normalized input is nonzero.Source: §A.1 · p. 12
  • An outer linear weight matrix or bias can reproduce any desired output change, with a rank-1 matrix update or a direct bias update.The stated matrix condition is a nonzero input vector.Source: §A.1 · p. 12
  • Element-wise multiplication is output-controllable when every input element is nonzero, and an MoE layer is output-controllable when its experts are controllable and its gate values sum to a nonzero value.For the MoE case, the desired change is distributed across experts in proportion to the reciprocal of the gate sum.Source: §A.1 · p. 12
  • A parallel transformer block admits a perfect implicit update when its outer MLP function is output-controllable, allowing that branch to absorb the contextual difference from attention.The block’s shared input cancels when its outputs are equated.Source: §A.1 · p. 13

B.1. Numerically Stable Update via RMSNorm Inversion

  • The stable RMSNorm update avoids relying on a potentially unstable scale-vector division by applying a rank-1 down-projection update first and using the scale vector only for residual error.It finds a target pre-normalization vector at the original RMS that best approximates the desired output, then updates Wdown to reach that target.Source: §B.1 · p. 13

B.2. Derivation of Analytical RMSNorm Inversion

The appendix derives an RMS-constrained vector that best matches a target after scaled RMSNorm, then validates update fidelity across prompts and on Falcon. The multiplier solution is obtained through a root-finding procedure.

B.2.1. PROBLEM FORMULATION

  • The RMSNorm inversion problem seeks the vector closest in squared L2 error to a target after scaled RMS normalization, subject to a fixed RMS value.The appendix states that fixing the RMS constraint reduces the solution set to a single one.Source: §B.2 · p. 13, §B.2.1 · p. 14

B.2.2. SOLUTION VIA LAGRANGE MULTIPLIERS

  • The derivation uses Lagrange multipliers to obtain a component-wise form for the optimal normalized vector, with the multiplier chosen to satisfy the constraint.The resulting vector is rescaled by the target RMS C.Source: §B.2.2 · p. 14, §B.2.3 · p. 15

B.2.3. FINDING THE MULTIPLIER µ

  • The multiplier is found as a root of F; continuity and boundary behavior establish existence, while strict monotonicity establishes uniqueness.The appendix says the unique root can be found numerically, for example by bisection.Source: §B.2.3 · p. 15

C.1. Additional Prompts

  • Across four additional prompts and textual domains, float32 updates show near-zero Logit Difference and Total Variation Distance.Figures 6 and 7 report generation metrics for Prompts 1–4.Source: §C.1 · p. 15, Figure 6 · p. 16, Figure 7 · p. 16

C.2. Extended Metrics Analysis

  • Figure 8’s distance metrics confirm aligned probability landscapes, while its update-norm plots indicate that patches are generally sparse and low-magnitude.The figure also reports low distance metrics and a much smaller update from the norm minimization described in Appendix B.Source: §C.2 · p. 16, Figure 8 · p. 17

C.3. Empirical Validation on Falcon

  • Falcon’s updated models achieve token-level fidelity and negligible logit drift, with perfect token matching even in base bfloat16 without stabilization.The paper attributes this stability to Falcon’s pre-norm architecture without post-normalization scaling, which avoids approximate RMSNorm inversion.Source: §C.3 · p. 17

C.4. Ablation: Starting Layer Depth

Updating only the final layer causes numerical issues, while using a few layers preserves accuracy with minimal impact. Scaling and stable updates perform similarly, but differ in simplicity and update norms.

C.4. Ablation: Starting Layer Depth

  • Updating only the final layer presents numerical issues, but a few layers are enough for good accuracy.The comparison covers the naive and numerically stable updates as the starting layer varies.Source: §C.4 · p. 18
  • With fewer updated layers, matrix update norms decrease, but stable-regime vector norms increase by multiple orders of magnitude.The smallest norm update results from starting at earlier layers.Source: §C.4 · p. 18

C.5. Analysis of Scaling Update Variant

  • The scaled update minimizes extreme ∆m values and better preserves token distributions in lower precision; scaling and stable updates perform similarly.Scaling is simpler, while the stable update produces a much smaller norm update.Source: §C.5 · p. 18

D. Impossibility of Controllability for RMS Post-Norm Without a Trainable Scaling Vector

A fixed-scale RMSNorm at the output boundary can prevent exact context-equivalent updates because its output is bounded while the attention-side difference is not. The paper also frames context-equivalent updates as gradient steps, deriving cumulative updates by telescoping intermediate shifts.

D. Impossibility of Controllability for RMS Post-Norm Without a Trainable Scaling Vector

  • A post-normalization residual block with fixed-scale RMSNorm cannot guarantee an exact context-equivalent update.The MLP branch’s corrective output is bounded, while the attention difference can exceed that bound; a trainable scale vector avoids this obstruction.Source: D. Impossibility of Controllability… · p. 20

E. Deriving the Implicit Update as a Gradient Step

  • Context-equivalent updates can be viewed as gradient descent steps on an implicit objective measuring the gap between uncontextualized and contextualized states.The construction accumulates context stepwise from no context to the full context.Source: E. Deriving the Implicit Update as a… · p. 20

E.1. Notation and Setup

  • The trace pseudo-loss L(W) = −Tr(∆⊤W) has gradient −∆, so a descent step produces the additive update +∆.The trace inner product is defined as ⟨∆, W⟩ = Tr(∆⊤W).Source: §E.1 · p. 20

E.2. Case 1: Standard Outer Weight Matrix (e.g., Llama-style)

  • For a standard outer weight matrix, summing the stepwise updates telescopes to the cumulative gate-matrix update derived in the main text.The update absorbs the difference between the full-context and no-context normalized inputs.Source: §E.2 · p. 21

E.3. Case 2: Numerically Stable Update (e.g., Gemma-style)

For the numerically stable Gemma-style update, the down-projection targets an inverted pre-norm vector while the RMSNorm scale absorbs the remaining error. A fixed elementwise preconditioner allows the remainder changes to telescope into the stable scale update.

E.3. Case 2: Numerically Stable Update (e.g., Gemma-style)

  • The stable update assigns Wdown an optimal pre-norm target and uses the RMSNorm scale vector to absorb the remaining error.The input-matrix shifts remain as in the standard case, while output shifts target the inverted goals.Source: §E.3 · p. 21
  • A fixed elementwise preconditioner for the scale-vector update is needed for the summation to telescope.Without it, elementwise division would not commute with summation.Source: §E.3 · p. 21
  • Summed remainder changes telescope to yield the stable scale update ∆m = rn ⊘h′norm.The initial remainder error is zero.Source: §E.3 · p. 22
  • Scaling and stable updates perform similarly, while the scaling update is simpler and the stable update has a much smaller norm.Starting at earlier layers gives the smallest norm update; updating only the final layer presents numerical issues.Source: Figure 11 · p. 23, Figure 12 · p. 24, Figure 13 · p. 25
Loading 2511.17864v3…