---
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.