bilinear-attn-addition-carry-2layer

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.

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

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.

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

from src.model.transformer import Transformer

model, 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.