Context-Aware Mean Ablation in LLMs

• The paper introduces CAMA, which replaces a token’s hidden state with its task-conditioned average to preserve context while suppressing input-specific details.
• CAMA reveals a phase-based computation in LLMs where early layers handle general processing before a brief cross-token transfer enables precise arithmetic predictions.
• Empirical results on models like Llama-3-8B demonstrate that delaying cross-token interaction enhances arithmetic fidelity and isolates critical computation stages.

Context-Aware Mean Ablation (CAMA) is a representation-level intervention introduced in "All for One: LLMs Solve Mental Math at the Last Token With Information Transferred From Other Tokens" to determine when LLMs begin to use cross-token information in mental-math tasks such as A+B+C and A-B (Mamidanna et al., 11 Sep 2025). It replaces a token’s hidden state after a chosen number of early layers with the expected representation that token would have under the task distribution, conditioned on the token itself. In this formulation, early layers can retain task-relevant background computation while input-specific information from other tokens is suppressed. Within the paper’s broader All-for-One (AF1) analysis, CAMA is the mechanism that operationalizes a “waiting” phase before a brief information-transfer stage and a final phase in which computation is concentrated at the last token.

1. Motivation and problem setting

CAMA is motivated by a mechanistic question about causal transformers: although causal self-attention and multilayer perceptron layers allow every token, in principle, to access and compute from all preceding tokens from the first layer onward, practical computation may be organized differently (Mamidanna et al., 11 Sep 2025). The mental-math setting is especially suitable for testing this distinction because the prompt is short, highly structured, and the target is a one-token answer produced by next-token prediction without explicit reasoning.

The motivating hypothesis is that early layers may not yet perform input-specific arithmetic. Instead, they may carry out task-general processing such as recognizing that the prompt is a math expression, identifying operands and operators, or encoding numerical structure. The relevant empirical question is therefore not merely whether the model can solve arithmetic, but how much of the network’s depth is devoted to token-local background computation before genuine input-specific information transfer begins.

Standard interventions are presented as inadequate for this purpose. Copying the raw embedding forward is described as too shallow and out-of-distribution. Restricting self-attention to self may remove useful task-general computation and create unnatural states. Random-token replacement and random mean baselines are said to destroy contextual structure too aggressively. CAMA is designed to avoid these failure modes by suppressing instance-specific cross-token information while preserving the context-conditioned average computation that a token would normally undergo.

2. Formal definition and representational meaning

CAMA replaces the true hidden state of token position $t$ after a chosen waiting period with a conditional expectation over prompts that keep that token fixed while varying the rest of the prompt according to the task distribution (Mamidanna et al., 11 Sep 2025). Its defining equation is

$$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$

Here, $m(\vec{x}, t, l)$ denotes the model’s representation at token position $t$ after layer $l$, $\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)$ means that prompts are sampled from the task distribution while conditioning on token $x_t$, and $\tilde x_t^{(wait)}$ is the substituted “waiting” representation.

This construction has two explicit consequences. First, it preserves task-general contextual effects, because the expectation is taken over in-distribution prompts with the same token at the same position. Second, it erases input-specific information from other tokens, because the varying content of the remaining prompt is averaged out. The term context-aware refers precisely to this conditionality: the replacement is not an unconditional mean over arbitrary activations, but a mean over the conditional context distribution associated with the same token in the same task format.

Because the architecture is causal, the required marginalization for token $x_t$ only extends over the prefix up to $t-1$. In the math templates considered, most non-operand tokens are fixed, so the variability is concentrated in operand values. The appendix further notes that tokens can be grouped when earlier tokens are deterministically tied to later ones under the template; in the two- and three-operand templates, fixed template tokens and the following operand can be treated as a group to compute the conditional expectation efficiently. This suggests that the method exploits template regularity rather than attempting unrestricted averaging over open-ended language.

3. Placement in the network and relation to AF1

CAMA is applied as a layerwise waiting mechanism in the residual stream for the first wait layers (Mamidanna et al., 11 Sep 2025). It is therefore not an attention-head-only edit, but a token-level intervention on the residual representation. In the AF1 procedure, this waiting stage is followed by Attention-Based Peeking (ABP), which controls where information transfer across positions may occur.

ABP modifies the attention mask so that each query token may attend only to a specified subset $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$0. Attention matrix entries $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$1 are set to $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$2 when $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$3, so those edges disappear after softmax. Two regimes are central. In full-peeking, $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$4, which recovers standard causal attention. In self-peeking, $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$5, so a token attends only to itself. Attention to the BOS token is always preserved because attention sink behavior makes its removal destructive.

The AF1 identification pipeline has three stages. First, CAMA suppresses input-specific cross-token computation in the early layers. Second, ABP opens a limited transfer window in which only the last token may attend to preceding tokens, while other tokens remain restricted to self-attention plus BOS. Third, the last token is also forced into self-peeking, so the remaining layers perform only last-token computation. The resulting picture is the paper’s “All-for-One” organization: early token-local background computation, a short middle segment that transfers the necessary information into the final position, and a late phase in which the answer is computed almost entirely from the last token’s residual stream.

4. Empirical findings in mental-math models

The main empirical role of CAMA is to reveal when meaningful input-specific cross-token exchange becomes necessary (Mamidanna et al., 11 Sep 2025). On strong arithmetic models such as Llama-3-8B and Llama-3.1-8B, the paper reports that the model can wait for a surprisingly long time before such exchange is required. For Llama-3-8B on A+B+C, faithfulness stays high until around wait = 15, then drops sharply when the waiting period extends past that point. The interpretation given is that meaningful cross-token interaction is not needed in the early layers, but is needed starting around layer 15.

The paper also reports that only a very short transfer phase is necessary. In Llama-3-8B, as few as 2 layers of transfer are enough for high performance. After that, last-token self-computation suffices. Mechanistically, this supports the claim that compositional arithmetic is not broadly distributed across the full sequence from the outset. Instead, many layers appear to be spent on local or general preprocessing before a narrow late subgraph transfers operand information to the final position.

The reported AF1 subgraph is said to be sufficient and necessary for high model performance across a wide variety of mental-math tasks, to transfer across different models, and to work on a variety of input styles. Head analysis further shows that important transfer layers contain a few heads that attend from the last token to specific operands, reinforcing the interpretation that these layers pass operand values into the final token. A plausible implication is that CAMA does not merely degrade early computation; it helps separate token-local background structure from the later, sparse information-flow pattern that directly supports arithmetic prediction.

5. Comparison with alternative waiting interventions

The paper explicitly contrasts CAMA with several alternative waiting baselines and reports that all of them fail to achieve non-zero faithfulness in the AF1 discovery setting (Mamidanna et al., 11 Sep 2025).

Alternative	Definition in the paper	Stated problem
DEC	Direct Embedding Copy: reuse the original embedding $$\tilde x_t^{(wait)}=\mathbb{E}_{\vec{x}'\sim \mathbb P(\vec{x}\mid x_t)} \!\bigl[\,m(\vec{x}',\,t,\,wait)\bigr].$$6	Too unnatural; not in-distribution
RTMA	Random Token Mean Ablation	Destroys contextual structure too aggressively
SPAW	Self-Peek as Waiting	Eliminates inter-token interaction too aggressively
IFP	Isolated Forward Pass	Removes needed background computation

The paper organizes these failures around two required properties. One is in-distribution representation preservation. DEC and RTMA do not reflect the conditional distribution of token states encountered by the model and thus break the residual stream’s expected geometry. The other is preservation of background or task-general computation. SPAW and IFP remove inter-token interaction so aggressively that they also remove the evolving contextual representation needed to encode operand structure and arithmetic context.

CAMA is presented as effective because it lies between these extremes. It removes instance-specific cross-token information while keeping the expected context-conditioned activation. In the paper’s framing, that balance is what makes it suitable for probing early arithmetic computation rather than merely damaging the model in an off-manifold way.

6. Scope, limitations, and terminological distinction

CAMA is described as having several important assumptions and limits (Mamidanna et al., 11 Sep 2025). It depends on a “cooperative” tokenizer that gives numbers single tokens; this works for Llama, Pythia, and GPT-J in the reported setup, but not for models such as Qwen or Gemma where numbers may be split across multiple tokens. The current formulation is tailored to highly structured prompt templates whose task distribution is easy to sample. It is best suited to arithmetic tasks with clean computational boundaries, and the AF1 pattern is reported not to generalize to semantically richer tasks such as word problems or Python-program prompts. The method also assumes that averaging over the conditional prompt distribution meaningfully approximates the background computation a token would undergo in context. For prompts in which only operand values vary and the remainder of the template is fixed, this assumption is operationally tractable; for more open-ended language tasks, the paper states that direct application would require additional care.

The paper suggests that CAMA may generalize to other cases in which information fusion does not start in the first layer, but leaves that possibility for future work. This suggests a broader interpretability use case: identifying computations that are delayed in depth rather than continuously accumulated from the start. At the same time, the acronym is not unique. "CAMA" is also used in "CAMA: Enhancing Multimodal In-Context Learning with Context-Aware Modulated Attention" to denote Context-Aware Modulated Attention, a training-free logit-level attention calibration method for multimodal in-context learning in large vision-LLMs (Li et al., 21 May 2025). That method is conceptually distinct from Context-Aware Mean Ablation: the former calibrates attention logits during inference, whereas the latter replaces residual-stream representations with conditional expectations to probe when input-specific cross-token information becomes necessary.