bilinear-attn-modular-addition-p113

A 1-layer transformer with a bilinear MLP, trained on modular addition (p = 113). 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 head-to-MLP composition can be scored from the weights alone with no forward pass over data.

Analysis code: (code repository not yet published)

Task

Sequence [a, b, "="] → predict (a + b) mod 113 at the final position. All 113² = 12769 pairs are enumerated and split 30%/70% train/test, so the model must generalise rather than memorise.

Architecture

No biases and no normalisation layers anywhere. 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)).

Layers	1 attention + 1 bilinear MLP
d_model	128
Heads	4 × d_head=32
MLP	bilinear, h = (Wx) ⊙ (Vx), d_hidden=512
Context	3
Vocab	114 (tokens 0..112 plus "=")
Embeddings	learned token + positional, untied unembedding
Biases / norm	none / none

Training

Optimiser	AdamW, lr 0.001, betas (0.9, 0.98)
Weight decay	1.0
Batching	full batch
Seed	0
Final train accuracy	1.0000
Final test accuracy	0.9979

Was the grokking transition reached?

Yes. Train accuracy crossed 99% at epoch 400; test accuracy crossed 99% at epoch 5000, a gap of 4600 epochs. The model memorised first and generalised later, which is the transition.

Usage

from src.model.transformer import Transformer

model, config = Transformer.from_pretrained("itzPotato/bilinear-attn-modular-addition-p113")

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