bilinear-attn-addition-carry-2layer
A 2-layer transformer with a bilinear MLP and RoPE,
trained on 3-digit integer addition with carrying. This
checkpoint exists to be taken apart: it is used for weight-based composition
analysis, where the MLP is rewritten exactly as a tensor of pairwise input
interactions and decomposed by SVD, so composition between components can be
scored from the weights alone with no forward pass over data.
Analysis code: (code repository not yet published)
Task
Task format follows arXiv 2401.07993 (Kruthoff, Carrying over algorithm in
transformers), so results are comparable to it:
- Digit-level tokenisation, vocabulary of 12: digits
0-9,
then + and =.
- Sequence
a a a + b b b = = = (10 tokens). Operands and the sum
are left zero-padded to 3 digits; digits are in natural
big-endian order, not reversed.
- Restricted to
a + b < 1000, giving 500,500 pairs.
- Loss is computed on the trailing 3
= positions only.
Split 30%/70% train/test over a
seeded permutation. (The paper's own code takes train and test as two disjoint
30% slices and discards the remaining 40%; this uses all of it.)
Architecture
No biases and no normalisation layers anywhere, and a bilinear MLP. This is
load-bearing, not stylistic: a bias or a LayerNorm affine term would add
structure the interaction-tensor decomposition does not represent, and the MLP
would stop being exactly P((Wx) ⊙ (Vx)).
That means this model differs from the reference paper's architecture, which
uses pre-LayerNorm, biases, dropout and a ReLU MLP. The task is theirs; the model
is not.
| |
|---|
| Layers | 2 × (attention + bilinear MLP) |
d_model | 128 |
| Heads | 4 × d_head=32 |
| MLP | bilinear, h = (Wx) ⊙ (Vx), d_hidden=512 |
| Context | 10 |
| Vocab | 12 |
| Positional encoding | RoPE on queries and keys only |
| Biases / norm | none / none |
RoPE is applied to q and k and not to v. Rotating v would make the OV
circuit W_O @ W_V position-dependent, and that circuit is exactly what the
composition analysis reads. (The reference implementation does rotate v.)
Training
| |
|---|
| Optimiser | AdamW, lr 0.001, betas (0.9, 0.98) |
| Weight decay | 0.2 |
| Batching | minibatch, size 1024 |
| Seed | 0 |
| Final train accuracy (all digits correct) | 0.9966 |
| Final test accuracy (all digits correct) | 0.9959 |
| Final test accuracy (per digit) | 0.9986 |
Learning rate is 1e-3, not the paper's 1.4e-4: that value is tuned for a
LayerNorm'd ReLU model and is far too slow for this architecture.
Held-out accuracy by carry structure
Using the paper's strata. A single aggregate number can hide the thing the task
is about — consecutive carrying is only ~4% of the data.
| carry class | all digits correct |
|---|
| NC | 0.9940 |
| C@1 | 0.9982 |
| C@2 | 0.9973 |
| C all | 0.9994 |
| C all con. | 0.9755 |
Out of distribution
On the pairs the dataset excludes (a + b >= 1000), the three emitted digits match the sum modulo 1000 for 0.2359 of them. The model has only 3 output slots and so cannot emit the leading digit those sums need — this is not a length-generalisation result.
Usage
1from src.model.transformer import Transformer
2
3model, config = Transformer.from_pretrained("itzPotato/bilinear-attn-addition-carry-2layer")
config.json carries train_frac and split_seed alongside the architecture, so
the exact held-out split can be reproduced from the checkpoint alone.
