In-Context Algebra

AI summary

Overview

Research area: Mechanistic interpretability of transformer language models, specifically how transformers implement in-context reasoning when token embeddings carry no fixed meaning. The paper sits at the intersection of mechanistic interpretability, in-context learning, and algebraic reasoning.

Technical level: Advanced. The paper assumes familiarity with transformer internals (attention heads, logit lens, activation patching), causal intervention methods (indirect effect / average indirect effect), basic group theory (cyclic and dihedral groups, identity elements, closure, cancellation laws), and concepts like PCA steering and probing.

One-sentence scope: The authors train small transformers to complete sequences of finite-group multiplication facts in which every token is a variable whose meaning is re-randomized per sequence, then use targeted data distributions and causal interventions to identify the specific algorithms the models learn.

What This Paper Is About

Most work on how language models do arithmetic assumes the answer can be read out of fixed information baked into token embeddings — for example, a token embedding that already encodes "divisible by two," or geometric representations of numbers. This paper removes that crutch entirely: it builds a task where vocabulary tokens are pure variables that mean nothing on their own and only acquire meaning from the other facts in the same sequence.

The goal is to ask what computational strategies a transformer develops when it cannot pre-encode solutions in its embeddings, and then to verify those strategies causally rather than just observing correlations.

Key Contributions

1. A new in-context algebra task. The authors define a task of simulating a mixture of finite groups in which each sequence samples a set of groups, assigns their elements to a non-overlapping set of vocabulary tokens through a random one-to-one latent mapping, and concatenates sampled facts into a sequence. Because the mapping is fixed within a sequence but varies between sequences, tokens act as variables without fixed global meaning.

2. Five targeted data distributions as causal tests. The authors construct D_copy, D_commute, D_identity, D_associate, and D_cancel, each built to isolate a hypothesized algorithm by controlling which kinds of supporting facts are present (duplicates, commutative pairs, identity-revealing facts, associativity triples, or membership-tracking facts). They also implement Python equivalents of each algorithm and measure the fraction of sequences each can theoretically solve ("coverage").

3. Identification and causal verification of three mechanisms beyond verbatim copying. Using activation patching and indirect-effect measurements, the authors isolate commutative copying, identity element recognition, and closure-based cancellation, and show that the first two of these (plus verbatim copying) are carried largely by specific attention heads.

4. A staged account of learning. The authors track seven behavioral metrics across training checkpoints and show that the loss curve breaks into five stages, each corresponding to the acquisition of a discrete skill.

Main Findings

• Near-perfect in-distribution accuracy. The main model (4 layers, 8 attention heads per layer, hidden size 1024, trained on sequences of k=200 facts, roughly 1000 tokens) achieves high accuracy on the task despite tokens having no stable meaning across sequences, and generalizes to unseen groups.

• Accuracy grows with context length, and group order matters. Held-out accuracy increases monotonically with the number of in-context facts. Smaller groups (e.g., C4, C6) reach high accuracy with only a few facts, while larger groups (e.g., C8, C10) need substantially more context to reach similar performance.

• Copying explains a lot of the easy data but not all of it. At k=200, about 90% of queries are trivially solvable by copying a previously seen fact, versus 45% at k=50. When the authors evaluate on held-out data where neither the final fact nor its commutative pair appears elsewhere, the model still reaches near-perfect accuracy, with an abrupt improvement during training (a phase transition) on these non-copyable queries.

• Generalization extends beyond the training groups. On the complete set of order-8 groups excluded from training, including C4×C2, Q8, and Z2³, the model achieves near-perfect performance. Hold-out performance is good for non-group structures such as semigroups, worse for quasigroups (though still non-trivial, particularly on cancellation data), and collapses on magmas. Generalization is consistently stronger for groups than for non-groups.

• Coverage of hypothesized algorithms on random training sequences. Applied sequentially, verbatim copying has an AUC of 67.9%, commutative copying adds 12.1%, identity recognition 4.2%, closure-based cancellation 2.7%, and associativity 3.6%, for a total coverage AUC of 90.4%. The model's own accuracy AUC is 92.4%, leaving roughly 2.0% AUC unexplained by these algorithms.

• Coverage when copying is impossible. On hold-out sequences where neither form of copying is possible, the model's AUC is 87.3%. Identity recognition solves 28.7% of cases, closure-based cancellation another 39.1%, and associativity 16.9%, for a total coverage AUC of 84.7%. The gap between model performance and algorithmic coverage here is 2.6%, occurring primarily on sequences with fewer facts.

• Performance on the targeted subdistributions. Evaluating on 400 sequences from each distribution at k=50 and k=100 facts, the model reaches 100.0% on verbatim copying, 99.0% on commutative copying, 100.0% on identity element recognition, and 97% on closure-based cancellation, but only 60.2% on associative composition.

• A single copying head carries verbatim copying. Patching attention head activations from clean to corrupt sequences over 200 samples from D_copy, one head (layer 3, head 6) has an average indirect effect of 0.91, and no other head exceeds 0.08. It attends to the answer-slot of duplicated facts and strongly promotes the logit of the attended-to token.

• The same head handles commutative copying, but is not always right. Head 3.6 is again the only head with strong average indirect effect (0.48) on D_commute, attending to the predictive token and answer-slot of the commutative fact when no duplicate exists. When analyzed separately on cyclic and dihedral groups, it promotes the token it attends to more than 97% of the time even when that token is the wrong answer, and when a corrupt duplicate fact is injected it attends to both duplicates and promotes both of their logits.

• Copying takes precedence over commutativity. When both an exact fact and its commutative counterpart are present, the head attends to the exact fact's answer-slot rather than the commutative one.

• Identity recognition is a two-part mechanism. Final-layer attention outputs at the predictive token separate identity facts from non-identity facts along the first PCA dimension, invariant to the specific variables or group. Head 3.1 acts as "query promotion," elevating both variables in the question; head 3.6 acts as "identity demotion," attending to prior identity facts and suppressing the identity token's logit. Only the combination leaves the non-identity factor as the answer.

• Identity behavior can be induced and suppressed. Steering the final-layer attention output along the learned PCA direction toward the identity cluster makes the model produce near-equal logits for both query tokens. Injecting a false identity fact causes the demotion head to suppress that token and produce a cleaner identity prediction; when a false identity fact is injected for a non-identity query, steering in the negative PCA direction lets the model recover the correct answer.

• Closure-based cancellation decomposes into two subspaces. The hypothesized algorithm computes the difference of a closure set (S_closure) and a cancellation set (S_cancel). A 16-dimensional subspace trained on the final-layer attention output reaches 99.8% intervention accuracy after a few epochs on 2000 data pairs. A separate subspace captures the elimination set with high intervention accuracy. Both subspaces are trained on a single attention layer and draw partial contributions from several attention heads.

• The closure submechanism is measurable directly. At left-slot positions, the model's predictions are nearly uniform across all elements previously associated with that variable, confirming it has learned group closure. Probes trained on the closure subspace identify whether a variable is in the closure with high accuracy, and variable-level probes partially align with the model's unembedding matrix.

• Training proceeds in five distinct stages. Tracking seven metrics (evaluated with 128 randomly sampled data points from each targeted dataset) across checkpoints shows: (a) the first sharp loss drop corresponds to predicting structural tokens '=' and ','; (b) next the model learns group closure and the identity query-promotion submechanism, reaching around 50% on identity sequences; (c) the next sharp drop corresponds to verbatim copying, then commutative copying; (d) after copying, cancellation and identity sequences improve in parallel, which the authors hypothesize reflects shared "demotion" submechanisms (identity demotion and cancellation) complementing the earlier promotion and closure submechanisms; (e) accuracy on associative sequences increases last, after all other mechanisms are learned. Only the copying and structural-token stages show sharp transitions.

• The learned strategies differ from previously reported ones. Rather than geometric Fourier-basis representations, the models here rely on sparse relational patterns and symbolic, context-based strategies, which the authors attribute to the task structure.

Methodology in Plain English

The authors build a synthetic data generator. For each training sequence, they sample a small collection of finite groups from a fixed pool — cyclic groups C3 through C10 and dihedral groups D3, D4, and D5 — and randomly map every element of those groups onto a distinct token from a vocabulary of 16 variable tokens (plus '=' and ','). Because each group gets a non-overlapping set of tokens, the meaning of any token is determined entirely by the group structure visible within that one sequence. They then sample facts (products of group elements) from those groups and write them out as xy=z statements, concatenating many such facts into a sequence. This is a standard next-token prediction setup: the model sees the sequence and must predict the answer token of the final fact.

The authors train a small autoregressive transformer on this data and evaluate it in three regimes: ordinary in-distribution sequences, held-out sequences where the final fact and its commutative counterpart never appear elsewhere, and entirely unseen groups (including the order-8 groups and non-group structures like semigroups, quasigroups, and magmas).

To figure out how the model solves the task, they first reason about which algorithms could solve a given sequence. They write Python implementations of five candidate algorithms — verbatim copying, commutative copying, identity recognition, closure-based cancellation, and associative composition — and measure what percentage of sequences each can solve, applied in sequence so each only gets credit for cases earlier algorithms missed. This "coverage" number is compared against the model's actual accuracy to see how much is unexplained.

They then build five specialized data distributions designed so that a specific hypothesized algorithm is the only viable route to the answer, and check whether the model succeeds on each. Finally, they use causal interventions: they take pairs of sequences that differ in a meaningful way, patch the activations of an individual attention head from one sequence into the other, and measure how much the probability of the correct answer changes — a quantity called the indirect effect, averaged over a dataset to give the average indirect effect. This identifies which heads matter for which behavior. For behaviors that don't localize to a single head — identity recognition and closure-based cancellation — they instead learn low-dimensional subspaces of the model's activations and test whether patching through those subspaces reproduces the expected counterfactual answer, and they use PCA and linear probes to characterize what those subspaces represent.

Why This Matters

Impact on research. The paper argues that the strategies transformers learn are a function of task structure. In settings where tokens have fixed, informative embeddings, models have repeatedly been shown to learn geometric, Fourier-style solutions. This work shows that when that information is removed, models instead develop sparse, symbolic, context-dependent mechanisms. That is a direct challenge to the assumption that any one account of "how transformers do arithmetic" generalizes across settings, and it provides a concrete, verified case study of genuine in-context symbol manipulation rather than retrieval.

The paper's code and data are released openly at algebra.baulab.info.

Real-world applications. The paper itself does not report deployed applications; the following are implications rather than claims made by the authors.

• Auditing models that must reason about fresh, user-defined entities: agents operating on schemas, identifiers, or entities defined only within a session (database keys, newly introduced named entities in a conversation, API objects) face exactly the problem this task formalizes — symbols with no pretrained meaning.
• Designing interpretability tooling for variable-binding behavior: the copying head, identity demotion head, and closure/cancellation subspaces give a worked template for locating context-dependent variable tracking inside a network, which could inform probes or monitors for models that manipulate user-supplied data.
• Evaluating generalization beyond the training distribution: the finding that models transfer to unseen groups but degrade on quasigroups and collapse on magmas is a concrete demonstration of how generalization quality tracks how "well-behaved" the underlying structure is.
• Curriculum and training-dynamics insight: the staged learning curve, where structural tokens and closure come first, copying next, demotion-style mechanisms in parallel after, and associativity last, suggests an ordering that could inform how training data is sequenced for tasks built on compositional symbol manipulation.

Industry relevance. Mechanistic interpretability of in-context reasoning is directly relevant to teams that need to know whether a model is genuinely reasoning about the data in front of it or merely retrieving a memorized pattern — a distinction that matters for reliability, for evaluating whether a model can be trusted on novel inputs, and for building monitoring tools around models that manipulate structured, user-supplied information.

Future Directions

• Explain the residual gaps. The hypothesized algorithms leave roughly 2.0% AUC unexplained on random training sequences and 2.6% on non-copyable hold-out sequences, with the hold-out gap concentrated on sequences with fewer facts. What other mechanisms account for this?
• Understand why associativity is only partially learned. Associative composition is the worst-performing targeted distribution at 60.2%, and it is the last skill to emerge in training. Is chaining facts together fundamentally harder for this architecture, or does the data distribution underrepresent it?
• Characterize the boundary of generalization. The model transfers to unseen order-8 groups but performs worse on quasigroups and fails on magmas. Mapping exactly which structural properties of groups are required for transfer would clarify what the model has actually internalized.
• Test whether the findings scale or transfer to pretrained language models. The study uses one small transformer configuration (4 layers, 8 heads per layer, hidden size 1024), with the authors reporting qualitatively similar patterns across multiple training runs. Whether larger models, or pretrained LLMs prompted in-context rather than trained from scratch, adopt the same copying/identity/cancellation mechanisms is left open.

Target Audience

Mechanistic interpretability researchers and graduate students who study how transformers implement reasoning internally; researchers working on in-context learning, variable binding, and symbol manipulation in neural networks; and anyone specifically interested in how models handle novel entities whose meaning is defined only within a prompt. The paper is also a useful methodological template for practitioners who want to see causal intervention, subspace patching, and targeted data distribution design applied together in one study. It is not an introductory read — it assumes fluency with transformer internals and causal mediation analysis, though the core intuitions about variables and copying are explained clearly enough for readers with a machine learning background to follow.

Authors’ abstract

We investigate the mechanisms that arise when transformers are trained to solve arithmetic on sequences where tokens are variables whose meaning is determined only through their interactions in-context. While prior work has studied transformers in settings where the answer relies on fixed parametric or geometric information encoded in token embeddings, we devise a new in-context reasoning task where the assignment of tokens to specific algebraic elements varies from one sequence to another. Despite this challenging setup, transformers achieve near-perfect accuracy on the task and even generalize to unseen groups. We develop targeted data distributions to create causal tests of a set of hypothesized mechanisms, and we isolate three mechanisms models consistently learn: commutative copying where a dedicated head copies answers, identity element recognition that distinguishes identity-containing facts, and closure-based cancellation that tracks group membership to constrain valid answers. Our findings show that the kinds of reasoning strategies learned by transformers are dependent on the task structure and that models can develop symbolic reasoning mechanisms when trained to reason in-context about variables whose meanings are not fixed.