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
1from src.model.transformer import Transformer
2
3model, 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.
