arithmetic-relu-2layer-seed0

arithmetic-relu-2layer-seed0

A 2-layer, bias-free, norm-free decoder-only transformer with a relu MLP, trained on 4-digit signed addition and subtraction. One of twelve models ({relu, bilinear} x {1, 2} layers x seeds {0, 1, 2}) trained under an identical recipe so that MLP type is the only structural difference between them.

Task format

One token per digit, operator and =; operands zero-padded to 4 digits, answer to 5, preceded by a sign token:

1 2 3 4 + 0 5 6 7 = + 0 1 8 0 1     (16 tokens)
                      ^^^^^^^^^     loss is on these 5 digits only

Token ids: digits 0-9 are themselves, +=10, -=11, ==12, d_vocab=13.

These ids are this project's own numbering. They were intended to match a reference Hub model, but that model (melephant/1-layer-addition-v2) turns out to be addition-only with a 13-token vocabulary containing no - and no answer-sign token, so it cannot express this task and its tokenisation could not have been adopted. Compare token ids before using activations from this model alongside anyone else's.

Architecture

relu is W_out ReLU(W_in x); bilinear is W_out[(W_L x) * (W_R x)].

Training

Identical recipe for every one of the twelve models: AdamW, lr 0.02 (cosine, 200 warmup steps), batch 1024, weight decay 0.01, grad clip 1.0, one pass over 5,000,000 examples. The learning rate was chosen by a six-point probe run on both MLP variants and set to the highest rate at which both remain stable, so the shared recipe favours neither.

Model seed 0; data seed 1234, identical across all twelve. Best step 4600 of 4883.

Results

Accuracy by operator

seq acc requires all five digits correct. sign acc is ~0 by construction and is not a bug: the loss covers the five answer digits only, so the position predicting the sign token receives no gradient. The sign is teacher-forced as input, so digit accuracy is unaffected.

The 1-layer models have solved addition and not subtraction; the 2-layer models solve both. The second layer is buying borrow propagation specifically.

Loading

from src.pretraining.model import PretrainTransformer
model, config = PretrainTransformer.from_directory("path/to/checkpoint")

history.json carries the full training curve and the checkpoint-selection trace.