In-Context Algebra — part 2/5
Cleaned-text code points [15856, 31693). Bibliography omitted here; source.md includes it.
Model Performance on Subdistributions. We evaluate the model on 400 sequences sampled from each distribution 𝒟 i \mathcal{D}_{i} , and report results at k = 50 k=50 and k = 100 k=100 facts (Figure 3 c), with more results in Appendix C . We find the model gets near perfect performance on four of the five data distributions that we test: verbatim copying (100.0%), commutative copying (99.0%), identity element recognition (100.0%), and closure-based cancellation (97%). However, model performance on sequences that test associative composition is worse (60.2%), suggesting only partial learning of this property.
## 5 Causal Verification of Learned Mechanisms
Based on the results in Section 4.1 , we perform causal interventions to understand how the model mechanistically implements the algorithms with stronger empirical evidence: (1) verbatim copying, (2) commutative copying, (3) identity element recognition, and (4) closure-based cancellation.
### 5.1 Causal Interventions
In order to understand the internal computations underlying the model’s capabilities, we use causal interventions ( Vig et al., 2020 ; Meng et al., 2022 ; Geiger et al., 2025 ) to verify how the model implements the targeted behavior. This is typically done by implicating model components such as attention heads or directions in a model’s activation space ( Wang et al., 2023 ; Geiger et al., 2024 ; Mueller et al., 2025 ) . Similar to prior work, we quantify the importance of a component via its indirect effect ( Pearl, 2001 , IE;) . We compute IE as the change in probability of the target variable token v target v_{\text{target}} under some intervention across a pair of algebra sequences that differ in a meaningful way ( s clean s_{\text{clean}} , s corrupt s_{\text{corrupt}} ). Equation 4 shows an example of computing IE for an attention head a ( l , h ) a^{(l,h)} at layer l l , head h h , by patching its activations from s clean s_{\text{clean}} into s corrupt s_{\text{corrupt}} : IE ( l , h ) = P ( v target | a s clean ( l , h ) → s corrupt ) − P ( v target | s corrupt ) \text{IE}(l,h)=P(v_{\text{target}}|a_{s_{\text{clean}}}^{(l,h)}\rightarrow s_{\text{corrupt}})-P(v_{\text{target}}|s_{\text{corrupt}}) (4) where a s clean ( l , h ) → s corrupt a_{s_{\text{clean}}}^{(l,h)}\rightarrow s_{\text{corrupt}} indicates activations a ( l , h ) a^{(l,h)} are being patched (or replaced) from s clean s_{\text{clean}} into the same location in s corrupt s_{\text{corrupt}} . The average indirect effect (AIE) can be computed over a dataset 𝒟 \mathcal{D} as: AIE ( 𝒟 , l , h ) = 1 | 𝒟 | ∑ 𝒟 ( IE ( l , h ) ) \text{AIE}(\mathcal{D},l,h)=\frac{1}{|\mathcal{D}|}\sum_{\mathcal{D}}\left(\text{IE}(l,h)\right) (5)
### 5.2 Copying and Commutative Copying
In this subsection, we investigate how the model implements verbatim and commutative copying (see also Appendix D ). As shown in Section 4.1 , a large percentage of our training data ( ∼ 80 % {\sim}80\% ) can either be solved by verbatim copying or commutative copying (Figure 3 a), and the model achieves high performance ( 97 97 - 100 % 100\% ) when either form of copying is possible (see Figure 3 c).
Verbatim Copying. We search for attention heads responsible for copying the correct answer from the context by computing the average indirect effect (AIE) of each attention head a ( l , h ) a^{(l,h)} (layer l l , head h h ). We patch the activations of each head a ( l , h ) a^{(l,h)} from the final predictive token in s clean s_{\text{clean}} , taken from 𝒟 copy \mathcal{D}_{\text{copy}} , into the same token position in s corrupt s_{\text{corrupt}} , a randomly sampled sequence where copying is not possible, and measure its IE. We compute AIE ( 𝒟 copy , l , h ) \text{AIE}(\mathcal{D}_{\text{copy}},l,h) over 200 samples from 𝒟 copy \mathcal{D}_{\text{copy}} .
We find a single attention head (layer 3, head 6) with high AIE ( 0.91 0.91 ) that is primarily responsible for copying, with no other head having an AIE higher than 0.08 0.08 (Figure 13 a). We visualize the attention patterns of this copying head in Figure 4 a, and find that it attends to the answer-slot of duplicated facts (shown in green), much like the n-gram heads observed in Akyürek et al. (2024) . On 𝒟 copy \mathcal{D_{\text{copy}}} , head 3.6 strongly promotes the logit of the attended-to token (Figure 4 e) which can be seen by applying the model’s unembedding matrix to the attention head output U ( a ( l , h ) ) U(a^{(l,h)}) . This allows us to understand its output contribution in terms of vocabulary tokens ( Nostalgebraist, 2020 ; Elhage et al., 2021 ; Dar et al., 2023 ) . The top logit consistently matches the attended-to token (Figure 14 ).
Commutative Copying. In many groups, knowing that a b = c ab=c will also imply that b a = c ba=c . To investigate how the model implements such commutative copying, we compute the indirect effects of patching attention head activations from sequences in s clean ∈ 𝒟 commute s_{\text{clean}}\in\mathcal{D}_{\text{commute}} , to non-copying sequences s corrupt s_{\text{corrupt}} , where neither verbatim nor commutative copying are possible. We find the same pattern: head 3.6 is again the only head with strong AIE (0.48) for commutative copying (Figure 13 b). In the absence of duplicate facts, head 3.6 attends to the predictive token and answer-slot of the commutative fact (Figure 4 b) and similarly promotes the attended-to token (Figure 4 f).
Non-Copying Sequences. When neither copy-inducing fact is present in the context, head 3.6 often self-attends (Figures 4 c), not strongly promoting any token (Figure 4 g). However, when the query contains an identity fact, we find head 3.6 has an interesting identity demotion behavior (§ 5.3 ).
Corrupt Copying. While copying the answer-slot of a commutative fact can solve facts for abelian groups, this doesn’t work for non-commutative facts. When analyzing the copying behavior of head 3.6 on cyclic and dihedral groups separately, we find that more than 97 % 97\% of the time it promotes the token it attends to, even if that token is the wrong answer (see Figure 14 b). We illustrate this behavior in Figure 4 d, where we inject a duplicate fact with an incorrect answer and show that head 3.6 attends to both duplicates (red, green) and promotes both of their logits (Figure 4 h).
### 5.3 Identity Recognition
Our coverage analysis has revealed that when verbatim and commutative copying are no longer allowed, the identity algorithm can solve close to 30 % 30\% of all hold-out problems. In this section, we use data from 𝒟 identity \mathcal{D}_{\text{identity}} to study how the model solves sequences where the query is an identity fact. Recall that an identity element e ∈ G e\in G satisfies e ⋅ x = x ⋅ e = x e\cdot x=x\cdot e=x for all elements x ∈ G x\in G , so that if one variable in the question is known to be the identity, the answer is equal to the other variable.
Our experiments suggest that identity recognition emerges from the interaction of two complementary mechanisms: query promotion, that elevates both variables in the question as potential answers, and identity demotion that suppresses the known identity element. When both mechanisms activate simultaneously, the non-identity token is correctly selected.
Structure from PCA. First, we note that our transformer’s representations reveal a strong signal correlated with the presence of an identity element in the question. To analyze this, we use PCA to plot final-layer attention outputs at the predictive token position (the “ = = ” symbol) just before the model predicts an answer. There is a clear separation between facts containing identity elements (blue) and non-identity facts (red) along the first PCA dimension (Figure 5 a). This separation is invariant to the specific variables in the fact or the underlying group. This suggests the model has learned to recognize and solve identity facts differently from those without an identity element.
Query Promotion and Identity Demotion. To analyze the role of the final layer attention at predicting identity facts in Figure 5 b we use the logit lens ( Nostalgebraist, 2020 ) and find two heads whose logits correlate strongly with identity variables. Head 3.1 promotes both variables in a given fact, serving as a “query promotion” mechanism. This strategy of predicting that the answer is equal to the question is appropriate for problems in which one of the factors is the identity, although on its own it would have the undesirable effect of promoting the identity element itself as the answer.
On identity fact sequences, head 3.6 acts as an “identity demotion” mechanism, attending to previous identity facts in the context and suppressing the identity token’s logit (Figure 5 b, pink). Combined with the previous strategy, this serves to leave only the non-identity factor as the promoted answer.
Causal Verification. Our experiments suggest that the dominant PCA direction in representation space controls the query promotion submechanism. To understand the causal effects, we perform representation steering experiments along this learned direction (see Appendix E for details). When we intervene on the final layer attention output of a non-identity fact and steer it toward the identity cluster, the model begins producing equal logits for both query tokens (Figure 5 c: i vs. ii).
In addition to query promotion, we can also manipulate the model’s identity recognition by introducing false identity facts to influence the identity demotion signal. When we inject a fact incorrectly suggesting one of the query tokens is an identity element, the identity demotion head (3.6) responds by suppressing that token and causes a cleaner identity prediction (Figure 5 c, iii, iv).
On the other hand, when the model is presented with a false identity fact in-context while the query is a non-identity fact, it typically confuses the prediction. However, if we steer in the negative PCA direction (away from the identity cluster), the model recovers and correctly predicts the non-identity answer. These steering experiments demonstrate that the learned PCA direction has causal influence over the model’s identity reasoning, enabling us to both induce and suppress identity predictions.
### 5.4 Closure-Based Cancellation
The closure-based cancellation algorithm is a combination of two key submechanisms: (i) tracking which variables belong to the same group (i.e., the closure ), and (ii) systematically eliminating invalid answers using the cancellation law , which implies that for elements x , y , z ∈ G x,y,z\in G , if y ≠ z y\neq z then x y ≠ x z xy\neq xz and y x ≠ z x yx\neq zx (see Appendix A ).
We hypothesize the algorithm can be understood at a high level as computing the difference of two sets: S closure S_{\text{closure}} - S cancel S_{\text{cancel}} . Consider the sequence sampled from 𝒟 cancel \mathcal{D}_{\text{cancel}} shown in Equation 6 . For the final query p e = pe= , the closure contains all elements that have previously appeared in facts involving p p or e e , i.e., S closure = { p , e , f , a , n } S_{\text{closure}}=\{p,e,f,a,n\} . The cancellation law then eliminates candidates from facts that share a variable in the left- or right-slot: p p (from p f = p pf=p ), n n (from e e = n ee=n ), f f (from a e = f ae=f ), and e e (from p p = e pp=e ), leaving a a as the only valid answer (i.e., S cancel = { p , n , f , e } S_{\text{cancel}}=\{p,n,f,e\} ). “ p f = p , e e = n , p f = p , p f = p , a e = f , p p = e , p f = p , p n = f , p p = e , p e = ” \text{``}pf=p,ee=n,pf=p,pf=p,ae=f,pp=e,pf=p,pn=f,pp=e,pe=\text{''} (6) We use causal interventions to determine how the model implements these two submechanisms. Our analysis reveals evidence of both a closure subspace, that promotes the logits of variables in the same group, and an elimination subspace that demotes answers based on facts present in the context.
Closure Submechanism. The closure submechanism emerges naturally from autoregressive training: when predicting the right-slot of a fact like x y = xy= , the model must identify which variables could plausibly follow x x . These are precisely the elements that belong to the same group (i.e., the closure). In fact, when we analyze the model’s predictions at left-slot positions, we find nearly uniform logits across all elements previously associated with that variable, confirming the model has learned how to compute group closure (see Figure 15 , Appendix F ).
Inspired by previous work showing subspaces can encode high-level causal variables ( Geiger et al., 2024 ; Prakash et al., 2025 ) , we aim to identify a subspace W W that captures the model’s representation of the closure set, S closure S_{\text{closure}} . We construct counterfactual pairs ( s , s ′ ) (s,s^{\prime}) from 𝒟 cancel \mathcal{D}_{\text{cancel}} that have different closure and elimination sets ( S i , S i ′ S_{i},S^{\prime}_{i} ) such that under intervention, the expected counterfactual answer corresponds to a modified set difference: v CF = S closure − S cancel ′ v_{\text{CF}}=S_{\text{closure}}-S^{\prime}_{\text{cancel}} , where the closure set comes from s s and the elimination set comes from s ′ s^{\prime} . We perform subspace-level patching from s s into s ′ s^{\prime} and train W W to maximize the likelihood of producing the expected counterfactual output v CF v_{\text{CF}} (Equation 7 ). If the intervention causes the model to predict v CF v_{\text{CF}} , we take this as evidence that the subspace W W correctly represents the hypothesized closure set. P ( v CF | ( W a s l + ( I − W ) a s ′ l ) → a s ′ l ) P(v_{\text{CF}}|(Wa_{s}^{l}+(I-W)a_{s^{\prime}}^{l})\rightarrow a_{s^{\prime}}^{l}) (7) We can measure its accuracy as how often the model’s predicted answer under intervention matches the expected counterfactual target (Figure 17 ). We train a 16-dimensional W W on the model’s final layer attention output a l a^{l} , and find that it can achieve good intervention accuracy ( 99.8 % 99.8\% ) after a few epochs of training on 2000 data pairs (Figure 18 ). More details are in Appendix F.2 .
For the closure subspace, we train probes ( Alain & Bengio, 2017 ; Belinkov, 2022 ) to understand what the subspace has learned, and how it represents variables. We train each probe on the subspace to detect whether a variable is in the group closure or not. We find probes are able to identify when a variable is in the closure subspace with high accuracy (Figure 20 ), and that these variable-level probes partially align with the model’s unembedding matrix (Figure 21 ), furthering evidence that the closure subspace promotes variables it has seen before in the context.
Cancellation Submechanism. To understand the cancellation submechanism, we train a subspace using a similar construction to the above, but vary the patching setup. If this new subspace W ′ W^{\prime} captures the elimination set, then it should generate the counterfactual answer arising from the opposite set difference, where the closure comes from the corrupt sequence s ′ s^{\prime} , and the elimination set comes from the clean sequence s s , giving: ( S closure ′ − S cancel ) = v CF ′ (S^{\prime}_{\text{closure}}-S_{\text{cancel}})={v^{\prime}_{\text{CF}}} .
We similarly train this subspace and find it also achieves high intervention accuracy, indicating it successfully represents the elimination set (Figure 19 ). Intervening in this subspace causes variables to be less likely to be predicted by the model. Both subspaces are trained on a single attention layer, and we find they capture partial contributions from several attention heads (Appendix F.1 ).
## 6 Phase Transitions Correspond to Learning of Discrete Skills