bilinear-attn-subtraction-borrow-1layer

A 1-layer transformer with a bilinear MLP and RoPE, trained on 3-digit integer subtraction with borrowing. 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

This task has no reference paper. It is built as the structural mirror of the addition-with-carry task in the same project, which follows arXiv 2401.07993 (Kruthoff, Carrying over algorithm in transformers) — same format, same vocabulary size, same split, same recipe, same architecture — so that a difference between the two models is attributable to the arithmetic and not to the setup.

• Digit-level tokenisation, vocabulary of 12: digits 0-9, then - and =.
• Sequence a a a - b b b = = = (10 tokens). Operands and the answer are left zero-padded to 3 digits; digits are in natural big-endian order, not reversed.
• Restricted to a >= b, so the difference is non-negative and needs no sign token, giving 500,500 pairs — the same count as the addition task.
• Loss is computed on the trailing 3 = positions only.

Split 30%/70% train/test over a seeded permutation, the same way the addition dataset is split.

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 architecture of the paper behind the companion addition task, which uses pre-LayerNorm, biases, dropout and a ReLU MLP.

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.9866
Final test accuracy (all digits correct)	0.9845
Final test accuracy (per digit)	0.9948

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 borrow structure

Strata mirror the addition task's five carry classes under the generate/propagate/kill translation: a column generates a borrow when a_i < b_i and propagates an incoming one when a_i == b_i, which is what a_i + b_i == 9 does on the addition side. A single aggregate number can hide the thing the task is about — consecutive borrowing is only ~4% of the data (identical to the addition task: the two stratum populations are the same sizes, class for class).

borrow class	all digits correct
NB	0.9808
B@1	0.9963
B@2	0.9980
B all	0.9959
B all con.	0.8254

Out of distribution

On the pairs the dataset excludes (a < b), the three emitted digits match the difference modulo 1000 for 0.8864 of them. The model has no sign token and cannot express a negative result at all — this says what the model extrapolates to, and is not a signed-arithmetic result.

Usage

from src.model.transformer import Transformer

model, config = Transformer.from_pretrained("itzPotato/bilinear-attn-subtraction-borrow-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.