Faithfully understanding the internal process that occurs as models form their chain of thought is an important problem with the rise in model capabilities. However, this remains challenging due to the variance in sampled responses, making it difficult to isolate the effects of interventions and compare model reasoning systematically. We propose a method for localizing and analyzing how a given intervention affects the model's reasoning through behavioral counterfactuals: given a fixed response, how does introducing an intervention to the prompt change the next-token distributions? We train a module to edit the residual stream such that it minimizes the KL divergence between the output distributions for a fixed response under a source prompt and its counterfactual target prompt. We find that matching the output distributions recovers the counterfactual behavior more effectively than standard activation reconstruction objectives and validate that attribution of the trained module outputs is effective for identifying causally important locations for counterfactual recovery. The attribution maps and visualizations of the downstream effects of model edits provide an interface for generating hypotheses about the original model's computation, which can then be tested with activation patching. In a case study on Gemma-4-E4B-it and DeepSeek-R1-0528-Qwen3-8B, we investigate how a suggestion influences model reasoning even when it is not verbally acknowledged. Our analysis identifies representations at the question statement and filler tokens that mediate a substantial part of the suggestion's influence, with downstream effects such as changing the framing of the justification or even the final answer.
Figures & tables
Figure 1: We extract causal structures for a given text intervention by (1) Passing a control prompt and a modified intervention prompt through a language model with a fixed response. (2) Training an interface edit module that edits the residual stream at a given layer of the model to minimize the KL divergence between the edited output distribution of one instance and the output distribution of the counterfactual. (3) We then extract attribution maps and perform edit-level ablations used for generating hypotheses about the original model, which we test with activation patching.
Figure 2: Comparison of edit modules trained using the distribution matching objective and the activation reconstruction objective for Gemma-4-E4B-it. The figure on the left shows the KL recovery and the figure on the right shows the activation reconstruction for the two objectives.
Figure 3: We provide scatter plots of attribution scores and edit norm on the x-axis and ablation effects on the y-axis for Integrated Gradients for Gemma-4-E4B-it (top) and DeepSeek-R1-0528-Qwen3-8B (bottom) in the left and middle figures respectively. We plot the distribution of correlations between the attribution scores and ablation effects evaluated per sequence in the right figure.
Figure 4: We show the downstream effect of the edit at the indentation token in the orange box determined by ablating it for Gemma-4-E4B-it. The highlighting corresponds to the effect on the KL divergence after ablation with red meaning removal resulted in larger KL. Callouts show the effect on likelihoods after adding back the edit. We see changes in wording such as the model being more likely to say the author responded to poor “behavior” or “handling” from the attendants instead of “service” and that the user felt “pressured” or “unfairly” treated instead of being “dismissed”.
Model
Control
Suggestion
Gemma
76.2%
72.6%
DeepSeek
72.3%
62.5%
Table 1: Final-answer accuracy with and without a suggestion the author is in the wrong.
Figure 5: KL recovery after native activation patching for Gemma-4-E4B-it on the left and DeepSeek-R1-0528-Qwen3-8B on the right.
Model
Most frequent token types among high-attribution events
the (9.2%), wrong (34.1%), author (22.8%), in (16.2%), . (5.3%), , (4.3%), is (7.2%), for (13.3%), user (28.0%), A (26.4%)
Table 2: Most frequent tokens among the top 5% of absolute attribution scores within each record. Parentheses report the percentage of occurrences of each token type that fall within the top 5%.
Figure 6: We provide an example where after performing native activation patching on the question and filler tokens for Gemma-4-E4B-it, the generation distribution shifts towards the target distribution which does not have the intervention suggesting the author is in the wrong. We see changes such as the model saying it will “determine” rather than “choose” whether the author is in the wrong and the model being more likely to say the author feels “guilt” in response to an apology rather than “suspicious”.
Figure 7: Largest peak change in likelihoods across tokens after activation patching at the question and filler tokens for the airline example for Gemma-4-E4B-it.
Appendix figures & tables33 assets
Supplementary material from the paper’s appendix.
Appendix
Hyperparameter
Value
Temperature
1.0
Top-p
0.95
Top-K
64
Max New Tokens
32768
Appendix
Table 3: Generation configuration for Gemma-4-E4B-it.
Hyperparameter
Value
Temperature
0.6
Top-p
0.95
Top-K
N/A
Max New Tokens
32768
Appendix
Table 4: Generation configuration for DeepSeek-R1-0528-Qwen3-8B.
Hyperparameter
Value
Optimizer
Adam
Weight Decay
0
Effective Batch Size
128
Learning Rates
2e-5, 5e-5, 1e-4
Epochs
1
Warmup Schedule
Linear
Appendix
Table 5: Training configuration for Gemma-4-E4B-it and DeepSeek-R1-0528-Qwen3-8B.
Figure 8: Distribution before and after training over tokens of the probability mass in the conditional support for Gemma-4-E4B-it on the left and DeepSeek-R1-0528-Qwen3-8B on the right. For each token, we collect the mass for both the control and intervention arm.
Figure 9: Comparison of KL recovery across different module architectures for Gemma-4-E4B-it on the left and DeepSeek-R1-0528-Qwen3-8B on the right.
Figure 10: We provide scatter plots of attribution scores and edit norm on the x-axis and ablation effects on the y-axis for RelP in the left and middle figures respectively. We plot the distribution of correlations between the attribution scores and ablation effects evaluated per sequence in the figure on the right. The top row corresponds to Gemma-4-E4B-it and the bottom row corresponds to DeepSeek-R1-0528-Qwen3-8B.
Figure 11: We provide scatter plots of attribution scores on the x-axis and ablation effects on the y-axs for Integrated Gradients for Gemma-4-E4B-it across 10 examples where we include every token edit.
Figure 12: We provide scatter plots of attribution scores on the x-axis and ablation effects on the y-axis for Integrated Gradients for DeepSeek-R1-0528-Qwen3-8B across 10 examples where we include every token edit.
Figure 13: We provide scatter plots of attribution scores on the x-axis and ablation effects on the y-axis for Integrated Gradients for Gemma-4-E4B-it on the left and DeepSeek-R1-0528-Qwen3-8B on the right for edit events pooled across 10 examples where every token edit is included.
Figure 14: Attribution heatmap based on IG for Gemma-4-E4B-it.
Figure 15: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 16: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 17: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 15 and (right) performing the activation patching in Figure 16 .
Figure 18: Attribution heatmap based on IG for Gemma-4-E4B-it.
Figure 19: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 20: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 21: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 19 and (right) performing the activation patching in Figure 20 .
Figure 22: Attribution heatmap based on IG for Gemma-4-E4B-it.
Figure 23: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 24: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 25: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 23 and (right) performing the activation patching in Figure 24 .
Figure 26: Attribution heatmap based on IG for DeepSeek-R1-0528-Qwen3-8B.
Figure 27: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 28: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 29: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 27 and (right) performing the activation patching in Figure 28 .
Figure 30: Attribution heatmap based on IG for DeepSeek-R1-0528-Qwen3-8B.
Figure 31: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 32: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 33: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 31 and (right) performing the activation patching in Figure 32 .
Figure 34: Attribution heatmap based on IG for DeepSeek-R1-0528-Qwen3-8B.
Figure 35: Heatmap and callouts showing the downstream effects of an edit based on its ablation. The token in orange is the ablated token.
Figure 36: Heatmap and callouts showing the downstream effects of activation patching where the patched tokens are in orange.
Figure 37: Largest peak change in likelihoods across tokens after (left) adding the edit shown in Figure 35 and (right) performing the activation patching in Figure 36 .