bilinear-attn-addition-carry-1layer

A 1-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	1 × (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.9934
Final test accuracy (all digits correct)	0.9925
Final test accuracy (per digit)	0.9975

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.9904
C@1	0.9954
C@2	0.9989
C all	0.9999
C all con.	0.9241

Out of distribution

On the pairs the dataset excludes (a + b >= 1000), the three emitted digits match the sum modulo 1000 for 0.9333 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

from src.model.transformer import Transformer

model, config = Transformer.from_pretrained("itzPotato/bilinear-attn-addition-carry-1layer")

config.json carries train_frac and split_seed alongside the architecture, so the exact held-out split can be reproduced from the checkpoint alone.