arithmetic relu-2layer-seed2: checkpoint, config, training history

README.md

@@ -0,0 +1,90 @@
+---
+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.

config.json

@@ -0,0 +1,45 @@
+{
+  "d_model": 32,
+  "d_mlp": 64,
+  "n_heads": 4,
+  "d_head": 8,
+  "n_layers": 2,
+  "mlp_type": "relu",
+  "d_vocab": 13,
+  "n_ctx": 16,
+  "bias": false,
+  "norm": false,
+  "tie_unembed": false,
+  "model_seed": 2,
+  "data_config": {
+    "n_train": 5000000,
+    "n_val": 10000,
+    "n_test": 10000,
+    "seed": 1234,
+    "balanced_ops": true
+  },
+  "vocab": {
+    "plus": 10,
+    "minus": 11,
+    "equals": 12,
+    "n_operand_digits": 4,
+    "n_answer_digits": 5
+  },
+  "train_config": {
+    "lr": 0.02,
+    "betas": [
+      0.9,
+      0.98
+    ],
+    "weight_decay": 0.01,
+    "batch_size": 1024,
+    "warmup_steps": 200,
+    "lr_schedule": "cosine",
+    "lr_min_frac": 0.1,
+    "eval_every_steps": 200,
+    "seed": 2,
+    "max_steps": null,
+    "grad_clip": 1.0
+  },
+  "planned_repo_id": "itzPotato/arithmetic-relu-2layer-seed2"
+}

history.json

@@ -0,0 +1,213 @@
+{
+  "steps": [
+    200,
+    400,
+    600,
+    800,
+    1000,
+    1200,
+    1400,
+    1600,
+    1800,
+    2000,
+    2200,
+    2400,
+    2600,
+    2800,
+    3000,
+    3200,
+    3400,
+    3600,
+    3800,
+    4000,
+    4200,
+    4400,
+    4600,
+    4800,
+    4883
+  ],
+  "train_loss": [
+    1.5722284317016602,
+    0.628692090511322,
+    0.32719022035598755,
+    0.12961021065711975,
+    0.20075900852680206,
+    0.06205069273710251,
+    0.04518857225775719,
+    0.02749330922961235,
+    0.04564381390810013,
+    0.04888881370425224,
+    0.03712354227900505,
+    0.029449349269270897,
+    0.02453625202178955,
+    0.015443632379174232,
+    0.01518481969833374,
+    0.014887431636452675,
+    0.015574837103486061,
+    0.008070074021816254,
+    0.01065557636320591,
+    0.0158293005079031,
+    0.0032624073792248964,
+    0.0007429839461110532,
+    0.0011565109016373754,
+    0.006390308029949665,
+    0.0003882101736962795
+  ],
+  "val_loss": [
+    1.60806677734375,
+    0.5938065234375,
+    0.273192080078125,
+    0.11908873046875,
+    0.1072494384765625,
+    0.062252421875,
+    0.05844634765625,
+    0.03443075622558594,
+    0.0377742578125,
+    0.05226347839355469,
+    0.03548737426757813,
+    0.023215485534667967,
+    0.020486417846679688,
+    0.017016024780273437,
+    0.028689197387695313,
+    0.012433219299316406,
+    0.014155939636230469,
+    0.009152752685546875,
+    0.007206977081298828,
+    0.008440182037353515,
+    0.0041812696838378904,
+    0.0034816056442260744,
+    0.0021877392196655273,
+    0.0014594219970703124,
+    0.0009101939392089844
+  ],
+  "val_digit_acc": [
+    0.36168,
+    0.72716,
+    0.88346,
+    0.95826,
+    0.96028,
+    0.98076,
+    0.98198,
+    0.9901,
+    0.98874,
+    0.983,
+    0.99032,
+    0.99304,
+    0.9937,
+    0.9943,
+    0.99078,
+    0.99614,
+    0.99578,
+    0.997,
+    0.9979,
+    0.99746,
+    0.9985,
+    0.9991,
+    0.9993,
+    0.99956,
+    0.99966
+  ],
+  "val_seq_acc": [
+    0.0005,
+    0.224,
+    0.5383,
+    0.8081,
+    0.8187,
+    0.9079,
+    0.9178,
+    0.9529,
+    0.9462,
+    0.9183,
+    0.9541,
+    0.9661,
+    0.9695,
+    0.9721,
+    0.9549,
+    0.9812,
+    0.9799,
+    0.9854,
+    0.9903,
+    0.9876,
+    0.9929,
+    0.9956,
+    0.9965,
+    0.9978,
+    0.9983
+  ],
+  "val_sign_acc": [
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0,
+    0.0
+  ],
+  "lr": [
+    0.02,
+    0.019919114190346594,
+    0.01967791065343155,
+    0.019280724932227013,
+    0.018734696290875732,
+    0.018049639389099856,
+    0.017237867867379117,
+    0.016313973013888615,
+    0.015294561491574842,
+    0.014197956839627206,
+    0.013043870114744036,
+    0.011853045592292541,
+    0.010646887895751307,
+    0.009447077256643526,
+    0.00827517982051921,
+    0.007152260003589835,
+    0.006098501867759447,
+    0.0051328463196937115,
+    0.004272650655137428,
+    0.00353337656804348,
+    0.0029283122324319567,
+    0.002468333452453597,
+    0.002161708173895497,
+    0.0020139478709612964,
+    0.002
+  ],
+  "best_step": 4883,
+  "best_val_seq_acc": 0.9983,
+  "best_val_loss": 0.0009101939392089844,
+  "final": {
+    "loss": 0.0009101939392089844,
+    "digit_acc": 0.99966,
+    "seq_acc": 0.9983,
+    "sign_acc": 0.0,
+    "n": 10000
+  },
+  "total_steps": 4883,
+  "wall_clock_s": 258.4060242176056,
+  "metadata": {
+    "seed": 2,
+    "git_commit": "b757908e274621a108b0af21e5348ffba407a074-dirty",
+    "torch_version": "2.13.0+cu130",
+    "numpy_version": "2.4.6",
+    "python_version": "3.11.15",
+    "platform": "Linux-6.18.44-fc-v22-x86_64-with-glibc2.39",
+    "device": "cpu",
+    "wandb_mode": "disabled"
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:726a01d75e2dde1088d4cee0c8127c4c233071e104b274d001bd8d6670a28e34
+size 72096