In-Context Algebra — part 4/5
Cleaned-text code points [47688, 63435). Bibliography omitted here; source.md includes it.
In this section, we provide more details about our training setup.
Model Training Details. We train autoregressive transformer models ( Vaswani et al., 2017 ) , on in-context algebra sequence data sampled as described in Section 2 , with a batch size of 128 128 , and sequences with 200 200 facts. Our vocabulary consists of 16 16 variable tokens, a predictive token, and a separator token ( N = 18 N=18 in total), and each fact is made up of 5 5 tokens. As an additional hold-out scheme, we select one variable token (e.g., a a ) that is never assigned as an identity element during training. This ensures that we can test both per-sequence hold-outs (where copying is not possible), and variable-assignment hold-outs where the exact assignment of elements to variables is unseen.
Tooling and Compute. Our transformer implementation is based on code adapted from nanoGPT ( Karpathy, 2022 ) ; we use Rotary Positional Embeddings (RoPE) ( Su et al., 2024 ) instead of learned positional embeddings. We use the AdamW optimizer ( Loshchilov & Hutter, 2019 ) with a learning rate of 10 − 5 10^{-5} , and 1000 1000 warmup steps. The primary model we study in the main paper has 4 4 layers, with 8 8 attention heads per layer, and a hidden state dimension of 1024 1024 . As shown in Figure 7 , we observe qualitatively similar patterns across different seeds, both in terms of phase transitions in the loss and hold-out evaluation accuracy. We usually see convergence (eval. accuracy ≥ 99 % \geq 99\% ) between 30,000 30{,}000 and 75,000 75{,}000 steps and save the checkpoint with the best evaluation accuracy.
We train our models with either NVIDIA A100 80GB GPUs or NVIDIA A6000 48GB GPUs. Training statistics are logged using Weights and Biases ( Biewald, 2020 ) . Experiments are implemented using NNsight ( Fiotto-Kaufman et al., 2025 ) and PyTorch ( Paszke et al., 2019 ) , and run on workstations with NVIDIA A6000 48GB GPUs. We use SymPy ( Meurer et al., 2017 ) for simulating various group structures for our in-context algebra setting and use a custom implementation for magmas, semigroups, and quasigroups, with quasigroup generation based on Jacobson & Matthews (1996) ’s method for Latin squares.
### C.1 Model Architecture Hyperparameters
In this subsection, we study the effect of three hyperparameters that govern model capacity: (1) number of layers, (2) number of attention heads per layer, and (3) hidden state dimension. The training loss and evaluation accuracy of each hyperparameter sweep is shown in Figure 8 , and the corresponding breakdown of model performance by metric and hyperparameter configuration is shown in Table 1 , where results are reported using the checkpoint with the highest evaluation accuracy.
For each hyperparameter setting, there are noticeable drops in loss throughout training which correspond to the learning of discrete skills relevant for the task (§ 6 ), and is consistent with similar findings of phase transitions and skill-learning in prior work ( Olsson et al., 2022 ; Nanda et al., 2023 ; Singh et al., 2024 ; Chen et al., 2024 ; Kangaslahti et al., 2025 ; Kim et al., 2025 ; Hoogland et al., 2025 ) . In general, models with more capacity (i.e., more layers, larger hidden dimension, or more heads) learn the task more quickly and have shorter loss plateaus. Having fewer layers corresponds to longer loss plateaus, and delayed generalization (Figure 8 , left). Models with smaller model dimensions ( ∼ d ≤ 256 {\sim}d\leq 256 ) fail to generalize well when trained on the same amount of data (Figure 8 , middle). Using hidden size d = 128 d=128 achieves accuracy near random guessing and d = 256 d=256 only achieves ∼ 60 % {\sim}60\% accuracy while the models with hidden size d ≥ 512 d\geq 512 achieve 95 % 95\% evaluation accuracy or above. Models trained with only two attention heads per layer exhibit delayed generalization compared to four or eight heads per layer (Figure 8 , right).
### C.2 Group Sampling Probability and Task Diversity
To generate an in-context algebra sequence, we first sample a set of groups from the training distribution 𝒢 \mathcal{G} (see Section 2 ). In this subsection, we provide more details about our sampling procedure, and investigate the effect of the group sampling probability p mix p_{\text{mix}} as a training hyperparameter.
When sampling a mixture of groups 𝒢 s \mathcal{G}_{s} , for a sequence s s , the first group is sampled uniformly from 𝒢 \mathcal{G} . After an initial group is chosen, additional groups are iteratively sampled with replacement with probability p mix p_{\text{mix}} , continuing while the total order is less than or equal to the number of variables | V | = N |V|=N or a random draw from the interval [ 0 , 1 ] [0,1] exceeds p mix p_{\text{mix}} . A new group is added to 𝒢 s \mathcal{G}_{s} only if the total order of 𝒢 s \mathcal{G}_{s} would remain less than or equal to N N . Thus, choosing p mix = 0 p_{\text{mix}}=0 results in sequences containing exactly one group, while p mix = 1 p_{\text{mix}}=1 maximizes the number of groups mixed within each sequence, up to total order | V | |V| . Thus, p mix p_{\text{mix}} can be thought of as a measure of in-context task diversity. The algebra model we study in the main paper uses p mix = 0.7 p_{\text{mix}}=0.7 .
In Figure 9 a, we show loss curves and evaluation accuracy for transformer models trained with the same seed but different values of p mix ∈ { 0.0 , 0.3 , 0.5 , 0.7 , 1.0 } p_{\text{mix}}\in\{0.0,0.3,0.5,0.7,1.0\} . Training loss curves follow a consistent pattern: models trained with lower values of p mix p_{\text{mix}} have steeper early drops, but longer loss plateaus. Higher values of p mix p_{\text{mix}} correspond to shorter loss plateaus, but higher training loss. However, lower train loss does not necessarily correspond to higher evaluation accuracy (Figure 9 b).
Recall that our evaluation data excludes copying and commutative copying sequences (Section 3 ). We find that models trained with higher values of p mix p_{\text{mix}} tend to achieve better held-out evaluation accuracy, even though they have higher training loss (Figure 9 b). One reason for this might be that sequences generated using higher values of p mix p_{\text{mix}} have more groups per sequence, and thus more in-context task diversity. This aligns with findings from previous work showing that higher task diversity leads to more robust generalization ( Raventos et al., 2023 ; Kirsch et al., 2024 ; Ye et al., 2024 ; Park et al., 2025b ; Wurgaft et al., 2025 ) . Similarly, since repetition is more likely to happen in sequences with fewer groups (i.e., lower values of p mix p_{\text{mix}} ), models trained with lower sampling probabilities have lower task diversity (e.g., copying is much more common as a possible solution).
An additional benefit of training with higher mixing probabilities (more groups per sequence) is that models tend to achieve high evaluation accuracy (generalize) faster than lower mixing probabilities. This was initially surprising, but is consistent with Kim et al. (2025) who show that increased task diversity actually shortens loss plateaus. While having more groups in a sequence is a more difficult problem, Figure 9 b shows this configuration learns more quickly. Using a mixing probability of p mix = 0.0 p_{\text{mix}}=0.0 , where only a single group is sampled per sequence, has the slowest time to held-out generalization, while higher values of p mix p_{\text{mix}} begin to generalize sooner.
### C.3 Performance on Groups and Non-Group Algebraic Structures
In this section, we compare the model’s performance in-distribution to the model’s performance on groups not seen during training, as well as non-group structures. Figure 10 shows the model’s performance on in-distribution cyclic and dihedral groups. Performance typically increases with the number of facts in the sequence, and groups with more elements take longer to achieve perfect accuracy. For C 3 C_{3} , the performance actually begins to decreases after 25 facts.
For unseen group structures of order 8, the model still performs very well (Figure 11 ). The holdout accuracy is similar to that of in-distribution groups. The model is also able to solve the semigroup with near-perfect accuracy with enough facts in the sequence, while hold-out performance (where copying is not possible) on quasigroups and magmas is significantly worse.
Finite quasigroups are naturally solvable via the cancellation law, thus we evaluate on a subset of quasigroup sequences where cancellation can solve the problem. We find that on this subset, the model does much better than the overall hold-out accuracy previously reported, providing evidence that some mechanisms (i.e., closure-based cancellation) learned by the model are generalizable symbolic mechanisms that do not depend on the specific algebraic structure the data is sampled from, as long as the data possesses that property.
### C.4 Additional Performance on Data Subsets
We extend Figure 3 c to show performance on data subsets for varying number of facts (Table 2 ).
## Appendix D Additional Results on Copying
In this section, we provide additional experimental details and results related to the copying and commutative copying mechanisms. Figure 13 a shows a heatmap of the average causal effect of patching from verbatim copying sequences into no-copy sequences for each attention head in the model, computed as AIE ( 𝒟 copy , l , h ) \text{AIE}(\mathcal{D}_{\text{copy}},l,h) for an attention head at layer l l and head index h h . Figure 13 b shows a similar heatmap of average causal effects for each attention head when patching from commutative copying sequences into no-copy sequences, similarly computed as AIE ( 𝒟 commute , l , h ) \text{AIE}(\mathcal{D}_{\text{commute}},l,h) for each layer l l and head index h h . In each case, we identify a single attention head (layer 3, head index 6), which has a much higher AIE than other heads and is primarily responsible for copying behavior.
To understand the behavior of these heads under various data settings, we characterize each head’s output using direct logit attribution, or the “logit lens” ( Nostalgebraist, 2020 ; Elhage et al., 2021 ) . We apply the unembedding matrix to each attention head output (i.e., U ( a ( l , h ) CLOSE U(a^{(l,h)} ) and compute how often the token with the highest decoded logit matches the target token.
Figure 14 a shows how often each attention head promotes the correct answer token when using only cyclic groups to generate copying sequences. This is computed using 200 prompts for each of 3 prompt distributions: sequences where verbatim copying is possible ( s ∈ 𝒟 copy s\in\mathcal{D}_{\text{copy}} ), sequences where commutative copying is possible ( s ∈ 𝒟 commute s\in\mathcal{D}_{\text{commute}} ), and sequences where neither form of copying is possible. The highest logit promoted by the copying head (layer 3, head 6) almost always matches the target answer for both verbatim (green) and commutative copying sequences (purple), but almost never on non-copying sequences (black).
However, performing this same analysis on sequences sampled using only dihedral groups yields a slightly different result (Figure 14 b). When verbatim copying is possible, we still see head 3.6’s top logit matches the correct answer token more than 99 % 99\% of the time (green), as expected. However, on sequences sampled from 𝒟 commute \mathcal{D}_{\text{commute}} , this value drops to 32.5 % 32.5\% (purple). If we instead measure whether the highest decoded logit matches the most attended-to token , this happens 97 % 97\% of the time for head 3.6 (shown in red). This is curious because for these sequences, head 3.6 seems to be “blindly” copying the symbol it attends to even though it is not the correct answer. While this strategy would solve any commutative pair of facts, it cannot solve non-commutative facts found in dihedral groups. In Figure 4 d, we show a related behavior where head 3.6 will attend to and promote the answers of injected, incorrect facts in addition to correct ones.
## Appendix E Identity Recognition Details
In this section, we provide additional details about the steering experiments mentioned in Section 5.3 . We use them study whether the separation of identity and non-identity facts seen via PCA in Figure 5 a is causally relevant to downstream predictions.
To extract the PCA steering direction, we sample 1000 1000 sequences, ensuring that each sequence has 200 200 facts, and gather final attention layer outputs at the predictive token position (the “=” symbol) for each sequence. We then perform PCA on the activations, and find that the variation along the first principle component separates identity facts from non-identity facts (Figure 5 a). We denote this PCA direction as 𝐯 1 \mathbf{v}_{1} , and evaluate its role through several experiments. In each experiment we use a sample of 5000 5000 non-identity sequences, each containing 200 200 facts.
PCA Steering Controls Query Promotion. We first test the effect of steering non-identity sequences towards the identity cluster along 𝐯 1 \mathbf{v}_{1} . Concretely, this is done for non-identity sequences by adding 5 𝐯 1 5\mathbf{v}_{1} to the final attention layer output at the final prediction token. While the default behavior of the model is to predict an answer that is not present in the query fact, we find that steering towards the identity cluster using 𝐯 1 \mathbf{v}_{1} causes the model to predict query variables 99.2 % 99.2\% of the time over the 5000 5000 non-identity sequences we evaluate on. Upon examining the model’s output logits, we find that both query variables are strongly promoted as likely answers (with near-equal logits). This suggests that 𝐯 1 \mathbf{v}_{1} encodes a signal for “query promotion”. Interestingly, query promotion by itself is a sufficient criteria to distinguish between identity and non-identity facts, since only identity facts will have one of the query symbols as an answer.
Adding False Identity Facts Controls Identity Predictions. While steering non-identity sequences with 𝐯 1 \mathbf{v}_{1} can promote both query variables, it does not reliably promote one over the other. Recall that the other submechanism identified as important for solving identity facts was identity demotion. Next, we study whether we can cause the model to treat arbitrary non-identity facts as identity facts through a combination of PCA steering and injecting false identity facts into the context.
Concretely, given a non-identity sequence with a query x y = xy{=} , we replace a random fact in the sequence with a false fact indicating that either the left-slot ( x x ) or right-slot ( y y ) should “act” as an identity element (e.g., x ? = ? x?{=}? , or ? y = ? ?y{=}? , where ‘ ? ? ’ could be any variable). We then steer the model towards the identity cluster using 𝐯 1 \mathbf{v}_{1} as before, and measure whether the model predicts the non-identity variable (e.g., y y or x x ) under this counterfactual setup.
We find that using PCA steering and injecting a false left-slot identity fact causes the model predict the right variable 96.1 % 96.1\% of the time. And similarly, PCA steering along with injecting a false right-slot identity fact causes the model to predict the left variable 79.6 % 79.6\% of the time. These results show that we can control whether a model treats a fact as an identity fact (i.e., uses its identity mechanisms) through a combination of two simple interventions: PCA steering along 𝐯 1 \mathbf{v}_{1} to induce query promotion, and false identity fact injection to suppress a specific variable.