In-Context Algebra

Project Website | Arxiv Preprint | Model Weights | OpenReview

Setup

We use conda as a package manager. The environment used for this project can be found in env/algebra.yml. To install, you can run:

conda env create -f algebra.yml
conda activate algebra

Code

Our code is split into the following directories:

• src/ — main codebase: training, data generation, coverage mechanisms, and analysis utilities
• experiments/ — notebooks for reproducing paper results and figures
• colab/ — self-contained Google Colab notebooks for running experiments on GPU
• scripts/ — training and evaluation scripts (training.py, train_mix_10.sh, eval_50.py)
• assets/ — figures and saved data from experiment runs
• docs/ — paper PDF and supplementary materials

Coverage Experiment

colab/colab_experiment.ipynb is a self-contained Google Colab notebook that runs a full algorithmic coverage analysis on the trained model. It measures how much of the model's correct predictions can be explained by 7 mechanisms:

Mechanism	Description
Copy	Query answer already appears as-is in a prior fact
Commutative Copy	Answer derivable by swapping operands in a prior fact
Identity	One operand is the identity element
Closure/Cancellation	Answer derivable via cancellation laws
Assoc. (1-step)	Answer derivable by single-step associativity chaining
Assoc. (multi-step)	Answer derivable by multi-step associativity chaining
Result Frequency	Correct answer token has appeared as a result in prior facts

To run: upload the repo to MyDrive/algebra2/ in Google Drive, then open colab/colab_experiment.ipynb in Colab. Results and figures are saved to assets/.

Coverage figures from various runs are in assets/ — see coverage_two_panel_7mech.png for the most recent result with all 7 mechanisms.

Data

The in-context algebra task involves simulating a mixture of finite algebraic structures. The models we analyze are trained on groups (mostly Cyclic + Dihedral groups), but we test our models on non-group structures as well. Our data-generating code + more details can be found in the src/tasks directory.

Models

Download Pre-Trained Algebra Models

You can download the model analyzed in our paper from algebra.baulab.info/weights/.

The mixrosette-facts-10-16-8heads/ directory contains the trained model weights (algebra_gpt_best.pt), training checkpoints, and a metadata.json file with training configuration details. See src/load_utils.py or the experiment notebooks for examples of how to load these models.

At a minimum, you'll need to download the metadata.json, and at least one set of model weights (e.g. algebra_gpt_best.pt). See below for a minimal file structure for running the experiment notebooks.

outputs/
└── mixrosette-facts-10-16-8heads/
    ├── models/
    │   └── algebra_gpt_best.pt    # Best checkpoint
    └── metadata.json              # Training configuration

Note: I'm also looking into hosting these models on HuggingFace and will update later with more info.

Train Your Own Algebra Model

To train your own algebra model with checkpoints, you can use the scripts/training.py script (see the companion script scripts/train_mix_10.sh for an example). We use wandb to track training metrics and details.

Citing our work

This work appeared at ICLR 2026. The paper can be cited as follows:

@inproceedings{todd2026incontext,
    title={In-Context Algebra},
    author={Eric Todd and Jannik Brinkmann and Rohit Gandikota and David Bau},
    booktitle={The Fourteenth International Conference on Learning Representations},
    url={https://openreview.net/forum?id=J2peqXPQbB},
    year={2026},
    note={arXiv:2512.16902}
}