In-Context Algebra
Eric Todd, Jannik Brinkmann, Rohit Gandikota, David Bau
TL;DR
This paper explores how transformers learn symbolic algebraic reasoning when variables' meanings vary across sequences, revealing learned mechanisms like copying, identity recognition, and group-based cancellation.
Contribution
It introduces a new in-context algebra task and demonstrates that transformers can develop symbolic reasoning strategies in this setting.
Findings
Transformers achieve near-perfect accuracy on the in-context algebra task.
Models learn specific mechanisms such as copying answers and recognizing identity elements.
Reasoning strategies depend on task structure and can include symbolic mechanisms.
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…
Peer Reviews
Decision·ICLR 2026 Poster
- There are a large number of experimental results. The models design a novel in-context algebra task, and perform lots of experiments. Unlike previous settings, this task isolates in-context arithmetic. It elicits symbolic mechanisms; the authors very convincingly show this. - Likewise, the authors bring a number of methods to bear to prove out the mechanisms described above, including computing the indirect effects at the level of attention heads.
- This task is novel, however it strikes me as contrived and very toy. I struggle to see (a) how these findings will generalize to more interesting settings, (b) what this tells us about models that we did not already know, or (c) how this work convincingly tests interpretability methods.
- The problem is interesting and well-motivated. The paper is generally well-written and clearly presented. - The methodology is sound, and the conclusions are moderate and well-supported by the empirical evidence. - The in-context algebra task provides a reasonable testbed for studying symbolic reasoning, which ablates the effect of learned geometric embeddings. - The analysis of the closure-based cancellation mechanism is technically solid and one of the paper’s strongest contributions.
- The task probes a narrow type of symbolic reasoning. The five hypothesized strategies are relatively simple and do not fully capture the richer “meaning-free” symbolic manipulation motivating the study. The contribution would be stronger if the task enabled the discovery of more complex or novel mechanisms. - The causal analysis fails to identify a concrete mechanism for associative composition, arguably the most conceptually interesting of the five. The mechanisms that are identified for the
The work is well-motivated and grounded in relevant literature on mechanistic interpretability and LLM arithmetic problem solving. The claims regarding learning learning mechanisms are largely validated using causal experiments. The experiment design is thorough and sound. There is sufficient reproducibility information. The paper is well-written and easy to follow for those in the field.
The scope of transformer model is limited to one architecture and size. It would be interesting to see how the results change as different numbers of layers (but also attention heads, representation dimension) are considered. One might predict that transformer layer depth allows for more complex sequential operations. For example, depth might lead to better performance on associative composition, which did not appear to benefit from more in-context examples. There is limited to no discussion o
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsTopic Modeling · Child and Animal Learning Development · Multimodal Machine Learning Applications