| --- |
| library_name: pytorch |
| tags: |
| - arithmetic |
| - interpretability |
| - bilinear |
| - transcoder |
| --- |
| |
| # arithmetic-relu-2layer-seed2 |
|
|
| 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 |
|
|
| | | | |
| |---|---| |
| | d_model | 32 | |
| | d_mlp | 64 | |
| | n_heads | 4 | |
| | d_head | 8 | |
| | n_layers | 2 | |
| | MLP | `relu` | |
| | parameters | 17,728 | |
| | bias / norm | none | |
| |
| `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 2; data seed 1234, identical across all twelve. |
| Best step 4883 of 4883. |
|
|
| ## Results |
|
|
| | split | loss | digit acc | seq acc | sign acc | |
| |---|---:|---:|---:|---:| |
| | validation | 0.0009 | 0.9997 | **0.9983** | 0.0000 | |
| | test | 0.0009 | 0.9997 | **0.9983** | 0.0000 | |
|
|
| ## Accuracy by operator |
|
|
| | operator | seq acc | digit acc | loss | |
| | --- | ---: | ---: | ---: | |
| | addition | 0.9988 | 0.9998 | 0.0005 | |
| | subtraction | 0.9978 | 0.9996 | 0.0013 | |
|
|
| `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 |
|
|
| ```python |
| 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. |
|
|
