Publish addition-transformer run s85nnxtf

Loadable model and complete training record.

README.md

@@ -0,0 +1,69 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- arithmetic
+- interpretability
+- arxiv:2405.14813
+---
+
+# Fixed-width addition transformer
+
+Run `s85nnxtf` is a 1-block, bias-free causal transformer trained for
+4-digit base-10 addition. Operands are zero-padded and answers use
+5 digits, retaining overflow.
+
+## Results
+
+| Metric | Value |
+| --- | ---: |
+| Validation loss | 0.003520 |
+| Validation generated-token accuracy | 99.85% |
+| Validation exact-answer accuracy | 99.32% |
+| No-carry exact-answer accuracy | 97.27% |
+| Single-carry exact-answer accuracy | 100.00% |
+| Multiple-carry exact-answer accuracy | 98.44% |
+| Carry-chain exact-answer accuracy | 97.27% |
+
+## Training configuration
+
+- Updates: 10000
+- Optimizer: muon
+- Muon peak learning rate: 0.02
+- AdamW peak learning rate: 0.0003
+- Weight decay: 0.01
+- Warmup updates: 100
+- Minimum learning-rate ratio: 0.1
+- Initialization: normal
+- Seed: 0
+- Source commit: `unavailable`
+
+The complete resolved configuration, environment, metrics, source snapshot, and checkpoints are
+available in [`training/`](./training/). Machine-readable hashes and metrics are in
+[`export_manifest.json`](./export_manifest.json).
+
+## Loading
+
+This repository contains custom Transformers code. For reproducible or security-sensitive use,
+pin the commit revision printed by the uploader.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+revision = "PINNED_COMMIT_HASH"
+tokenizer = AutoTokenizer.from_pretrained(
+    "OWNER/REPO", trust_remote_code=True, revision=revision
+)
+model = AutoModelForCausalLM.from_pretrained(
+    "OWNER/REPO", trust_remote_code=True, revision=revision
+)
+inputs = tokenizer("0000 + 0000 =", return_tensors="pt")
+output = model.generate(**inputs, max_new_tokens=model.config.answer_digits, do_sample=False)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+## Intended use and limitations
+
+This model is intended for mechanistic-interpretability research on its configured fixed-width
+addition task. It is not a general arithmetic system: inputs outside the configured grammar or
+width are unsupported, and generated answers must not be treated as reliable calculations.

addition_transformer.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from .initialization import initialize_module
+from .model_config import AdditionModelConfig
+from .transformer_block import TransformerBlock, TransformerBlockOutput
+
+
+@dataclass(frozen=True)
+class TransformerOutput:
+    logits: torch.Tensor
+    block_outputs: tuple[TransformerBlockOutput, ...] = ()
+
+
+class AdditionTransformer(nn.Module):
+    def __init__(self, config: AdditionModelConfig, vocab_size: int) -> None:
+        super().__init__()
+        self.config = config
+        self.vocab_size = vocab_size
+
+        self.token_embedding = nn.Embedding(vocab_size, config.d_model)
+        self.position_embedding = nn.Embedding(config.max_seq_len, config.d_model)
+        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
+        self.unembedding = nn.Linear(config.d_model, vocab_size, bias=False)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        initialize_module(self, self.config)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        return_activations: bool = False,
+    ) -> TransformerOutput:
+        batch, seq_len = input_ids.shape
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f"Sequence length {seq_len} exceeds max_seq_len={self.config.max_seq_len}.")
+
+        positions = torch.arange(seq_len, device=input_ids.device)
+        x = self.token_embedding(input_ids) + self.position_embedding(positions)[None, :, :]
+        # Modula-inspired attention scaling and residual interpolation; see arXiv:2405.14813.
+        block_outputs: list[TransformerBlockOutput] = []
+        for block in self.blocks:
+            block_output = block(x, return_pattern=return_activations)
+            x = block_output.residual_after_mlp
+            if return_activations:
+                block_outputs.append(block_output)
+        logits = self.unembedding(x)
+
+        return TransformerOutput(
+            logits=logits,
+            block_outputs=tuple(block_outputs),
+        )
+
+    def symmetrized_mlp_tensor(self, detach: bool = True, layer: int = 0) -> torch.Tensor:
+        if not 0 <= layer < len(self.blocks):
+            raise IndexError(f"Layer {layer} is outside the model's {len(self.blocks)} layers.")
+        return self.blocks[layer].mlp.symmetrized_bilinear_tensor(detach=detach)
+
+    def analysis_tensors(self, detach: bool = True) -> dict[str, torch.Tensor]:
+        tensors: dict[str, torch.Tensor] = {}
+        for layer, block in enumerate(self.blocks):
+            prefix = f"blocks.{layer}"
+            bilinear_tensor = block.mlp.bilinear_tensor(detach=False)
+            tensors.update(
+                {
+                    f"{prefix}.attention.W_Q": block.attention.W_Q.weight,
+                    f"{prefix}.attention.W_K": block.attention.W_K.weight,
+                    f"{prefix}.attention.W_V": block.attention.W_V.weight,
+                    f"{prefix}.attention.W_O": block.attention.W_O.weight,
+                    f"{prefix}.mlp.W_1": block.mlp.W_1.weight,
+                    f"{prefix}.mlp.W_2": block.mlp.W_2.weight,
+                    f"{prefix}.mlp.W_O": block.mlp.W_O.weight,
+                    f"{prefix}.mlp.bilinear_tensor": bilinear_tensor,
+                    f"{prefix}.mlp.symmetrized_bilinear_tensor": 0.5
+                    * (bilinear_tensor + bilinear_tensor.transpose(-1, -2)),
+                }
+            )
+        return {name: tensor.detach() for name, tensor in tensors.items()} if detach else tensors

attention.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@dataclass(frozen=True)
+class AttentionOutput:
+    values: torch.Tensor
+    pattern: torch.Tensor | None = None
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, bias: bool = False) -> None:
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads.")
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads
+
+        self.W_Q = nn.Linear(d_model, d_model, bias=bias)
+        self.W_K = nn.Linear(d_model, d_model, bias=bias)
+        self.W_V = nn.Linear(d_model, d_model, bias=bias)
+        self.W_O = nn.Linear(d_model, d_model, bias=bias)
+        self.register_buffer("_causal_mask", torch.empty(0, 0, dtype=torch.bool), persistent=False)
+
+    def forward(self, x: torch.Tensor, return_pattern: bool = False) -> AttentionOutput:
+        batch, seq_len, _ = x.shape
+        q = self._split_heads(self.W_Q(x))
+        k = self._split_heads(self.W_K(x))
+        v = self._split_heads(self.W_V(x))
+
+        scores = torch.matmul(q, k.transpose(-1, -2)) / self.d_head
+        if self._causal_mask.shape[0] < seq_len or self._causal_mask.device != x.device:
+            self._causal_mask = torch.triu(
+                torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device),
+                diagonal=1,
+            )
+        causal_mask = self._causal_mask[:seq_len, :seq_len]
+        scores = scores.masked_fill(causal_mask, float("-inf"))
+        pattern = F.softmax(scores, dim=-1)
+
+        attended = torch.matmul(pattern, v) / 3.0
+        values = self.W_O(self._merge_heads(attended, batch, seq_len))
+        return AttentionOutput(values=values, pattern=pattern if return_pattern else None)
+
+    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = x.shape
+        return x.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+
+    def _merge_heads(self, x: torch.Tensor, batch: int, seq_len: int) -> torch.Tensor:
+        return x.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)

bilinear_mlp.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class BilinearMLP(nn.Module):
+    """Bilinear MLP: W_O((W_1 x) * (W_2 x))."""
+
+    def __init__(self, d_model: int, d_hidden: int, bias: bool = False) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.d_hidden = d_hidden
+        self.W_1 = nn.Linear(d_model, d_hidden, bias=bias)
+        self.W_2 = nn.Linear(d_model, d_hidden, bias=bias)
+        self.W_O = nn.Linear(d_hidden, d_model, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.W_O(self.W_1(x) * self.W_2(x))
+
+    def bilinear_tensor(self, detach: bool = False) -> torch.Tensor:
+        tensor = torch.einsum(
+            "oh,hi,hj->oij",
+            self.W_O.weight,
+            self.W_1.weight,
+            self.W_2.weight,
+        )
+        return tensor.detach() if detach else tensor
+
+    def symmetrized_bilinear_tensor(self, detach: bool = False) -> torch.Tensor:
+        tensor = self.bilinear_tensor(detach=False)
+        symmetrized = 0.5 * (tensor + tensor.transpose(-1, -2))
+        return symmetrized.detach() if detach else symmetrized

config.json

@@ -0,0 +1,43 @@
+{
+  "answer_digits": 5,
+  "architecture_version": 2,
+  "architectures": [
+    "AdditionForCausalLM"
+  ],
+  "attention_logit_divisor": "d_head",
+  "attention_output_scale": 0.3333333333333333,
+  "auto_map": {
+    "AutoConfig": "configuration_addition.AdditionConfig",
+    "AutoModelForCausalLM": "modeling_addition.AdditionForCausalLM"
+  },
+  "base": 10,
+  "bias": false,
+  "bos_token_id": 0,
+  "d_mlp": 128,
+  "d_model": 64,
+  "digit_token_offset": 3,
+  "dtype": "float32",
+  "eos_token_id": null,
+  "equals_token_id": 2,
+  "export_format_version": 2,
+  "hidden_size": 64,
+  "init_mode": "normal",
+  "intermediate_size": 128,
+  "is_decoder": true,
+  "max_position_embeddings": 16,
+  "max_seq_len": 16,
+  "model_type": "fixed-width-addition",
+  "n_heads": 4,
+  "n_layers": 1,
+  "normalization": "none",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 1,
+  "operand_digits": 4,
+  "pad_token_id": null,
+  "plus_token_id": 1,
+  "residual_alpha": 0.5,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "vocab_size": 13
+}

configuration_addition.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from transformers import PreTrainedConfig
+
+
+class AdditionConfig(PreTrainedConfig):
+    model_type = "fixed-width-addition"
+
+    def __init__(
+        self,
+        operand_digits: int = 5,
+        d_model: int = 64,
+        n_heads: int = 4,
+        d_mlp: int = 128,
+        n_layers: int = 1,
+        max_seq_len: int = 19,
+        init_mode: str = "normal",
+        residual_alpha: float | None = None,
+        **kwargs,
+    ) -> None:
+        expected_residual_alpha = 1.0 / (2.0 * n_layers) if n_layers > 0 else 0.0
+        if residual_alpha is not None and residual_alpha != expected_residual_alpha:
+            raise ValueError(
+                f"residual_alpha must be 1 / (2 * n_layers) = {expected_residual_alpha}."
+            )
+        invariants = {
+            "answer_digits": operand_digits + 1,
+            "base": 10,
+            "vocab_size": 13,
+            "hidden_size": d_model,
+            "num_attention_heads": n_heads,
+            "intermediate_size": d_mlp,
+            "num_hidden_layers": n_layers,
+            "max_position_embeddings": max_seq_len,
+            "bias": False,
+            "normalization": "none",
+            "attention_logit_divisor": "d_head",
+            "attention_output_scale": 1.0 / 3.0,
+            "bos_token_id": 0,
+            "plus_token_id": 1,
+            "equals_token_id": 2,
+            "digit_token_offset": 3,
+            "architecture_version": 2,
+            "export_format_version": 2,
+            "use_cache": False,
+            "tie_word_embeddings": False,
+            "is_decoder": True,
+            "is_encoder_decoder": False,
+            "eos_token_id": None,
+            "pad_token_id": None,
+        }
+        for name, expected in invariants.items():
+            if name in kwargs and kwargs.pop(name) != expected:
+                raise ValueError(f"{name} is fixed at {expected!r} for this architecture.")
+        self.operand_digits = operand_digits
+        self.answer_digits = operand_digits + 1
+        self.base = 10
+        self.vocab_size = 13
+        self.d_model = d_model
+        self.hidden_size = d_model
+        self.n_heads = n_heads
+        self.num_attention_heads = n_heads
+        self.d_mlp = d_mlp
+        self.intermediate_size = d_mlp
+        self.n_layers = n_layers
+        self.num_hidden_layers = n_layers
+        self.max_seq_len = max_seq_len
+        self.max_position_embeddings = max_seq_len
+        self.init_mode = str(init_mode)
+        self.residual_alpha = expected_residual_alpha
+        self.bias = False
+        self.normalization = "none"
+        self.attention_logit_divisor = "d_head"
+        self.attention_output_scale = 1.0 / 3.0
+        self.bos_token_id = 0
+        self.plus_token_id = 1
+        self.equals_token_id = 2
+        self.digit_token_offset = 3
+        self.architecture_version = 2
+        self.export_format_version = 2
+        self.use_cache = False
+        self.tie_word_embeddings = False
+        self.is_decoder = True
+        self.is_encoder_decoder = False
+        self._validate_architecture()
+        super().__init__(
+            bos_token_id=self.bos_token_id,
+            eos_token_id=None,
+            pad_token_id=None,
+            tie_word_embeddings=False,
+            is_decoder=True,
+            **kwargs,
+        )
+
+    @property
+    def d_head(self) -> int:
+        return self.d_model // self.n_heads
+
+    def _validate_architecture(self) -> None:
+        if self.operand_digits < 1:
+            raise ValueError("operand_digits must be positive.")
+        if self.d_model < 1 or self.n_heads < 1 or self.d_mlp < 1 or self.n_layers < 1:
+            raise ValueError("Model dimensions must be positive.")
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads.")
+        required_length = 3 * self.operand_digits + 4
+        if self.max_seq_len < required_length:
+            raise ValueError(
+                f"max_seq_len={self.max_seq_len} is too small for {required_length} full-sequence tokens."
+            )
+        if self.init_mode not in {"normal", "orthogonal"}:
+            raise ValueError(f"Unsupported init_mode: {self.init_mode}")
+        expected_residual_alpha = 1.0 / (2.0 * self.n_layers)
+        if self.residual_alpha != expected_residual_alpha:
+            raise ValueError("residual_alpha must equal 1 / (2 * n_layers).")
+
+
+AdditionConfig.register_for_auto_class("AutoConfig")

export_manifest.json

@@ -0,0 +1,147 @@
+{
+  "environment": {
+    "cuda": "12.6",
+    "device": "cuda",
+    "device_name": "NVIDIA GeForce RTX 4060 Laptop GPU",
+    "git_commit": null,
+    "git_dirty": true,
+    "packages": {
+      "huggingface-hub": "1.27.0",
+      "numpy": "2.5.2",
+      "pytest": "9.1.1",
+      "safetensors": "0.8.0",
+      "sympy": "1.14.0",
+      "transformers": "5.15.0",
+      "wandb": "0.28.1"
+    },
+    "platform": "Windows-11-10.0.26200-SP0",
+    "python": "3.12.4",
+    "torch": "2.13.0+cu126"
+  },
+  "export_files_sha256": {
+    "README.md": "b64a05029567e76bcd032dc4ea57e18c1c805845a196ffab30af6b28a0d06165",
+    "addition_transformer.py": "af1e1e1641aeaf2f974252e7177730f6356278c4fff0f47686d88aeed5673089",
+    "attention.py": "d1c738ba96f937c39cd66ad470d786acfad70ac368ecb5c1c25290b20edc2106",
+    "bilinear_mlp.py": "95d8b248058bdbce7ea7c13ba9ebf49ee69b20b95714a43040f19c9c528225c7",
+    "config.json": "481d105a5848390efe0e98edae059cded1b22611902ea15d2ab8b079cea86517",
+    "configuration_addition.py": "7f152f6c5917adedfb211e9332a9db08780363671932190b7df1d782c4fa3de5",
+    "generation_config.json": "13cc88dd90bbb2a181ade617b6a41b35998e933cc1eb356f7906c320a7d04b5d",
+    "initialization.py": "01de663e2494370c6bb2a9d4f770411104dbc44cbb97917fa18ba89ed9959d32",
+    "model.safetensors": "68be5fc852029537ac74bc0f13af4533a3efda7a4aac85396bd3a94b62f47098",
+    "model_config.py": "211bc4cc3e7379dfe469d5acb3bcfbe4c65adf710af5f2dd21ea49212ccb08fb",
+    "modeling_addition.py": "b5349de597570045eb3c21803f6a4b077605f5468f99e26ddb1678964d01fcea",
+    "tokenization_addition.py": "bc63473854f8f8fecfc3375757d257dfb34a33a0c30ee5a1ddffb3b984b8602f",
+    "tokenizer_config.json": "630ab54705ab78ecddef9c0c6543f477036b5a71cb2736a3db59462f464cbd06",
+    "transformer_block.py": "a66ce47f5b504abcc83bcbdfecbc71e01a9968d10bde534ad375e0c7158f04dc",
+    "vocab.json": "d6fdf48e6e2087546f1bfb7285a71b1072f5452934e8a33e831d069f7faf56a8"
+  },
+  "final_metrics": {
+    "examples_processed": 2560000,
+    "grad/clip_fraction": 0.0,
+    "grad/global_norm": 0.17957071997225285,
+    "lr/adamw": 2.9999999999999997e-05,
+    "lr/muon": 0.002,
+    "test/carry_chain/exact_accuracy": 0.97265625,
+    "test/carry_chain/generated_token_accuracy": 0.99375,
+    "test/carry_chain/loss": 0.012557083368301391,
+    "test/carry_chain/position_0/accuracy": 0.99609375,
+    "test/carry_chain/position_1/accuracy": 0.9765625,
+    "test/carry_chain/position_2/accuracy": 0.99609375,
+    "test/carry_chain/position_3/accuracy": 1.0,
+    "test/carry_chain/position_4/accuracy": 1.0,
+    "test/carry_chain/teacher_forced_token_accuracy": 0.99453125,
+    "test/multiple_carry/exact_accuracy": 0.984375,
+    "test/multiple_carry/generated_token_accuracy": 0.996875,
+    "test/multiple_carry/loss": 0.005255512893199921,
+    "test/multiple_carry/position_0/accuracy": 1.0,
+    "test/multiple_carry/position_1/accuracy": 0.984375,
+    "test/multiple_carry/position_2/accuracy": 1.0,
+    "test/multiple_carry/position_3/accuracy": 1.0,
+    "test/multiple_carry/position_4/accuracy": 1.0,
+    "test/multiple_carry/teacher_forced_token_accuracy": 0.996875,
+    "test/no_carry/exact_accuracy": 0.97265625,
+    "test/no_carry/generated_token_accuracy": 0.99453125,
+    "test/no_carry/loss": 0.014795759320259094,
+    "test/no_carry/position_0/accuracy": 1.0,
+    "test/no_carry/position_1/accuracy": 0.9765625,
+    "test/no_carry/position_2/accuracy": 0.99609375,
+    "test/no_carry/position_3/accuracy": 1.0,
+    "test/no_carry/position_4/accuracy": 1.0,
+    "test/no_carry/teacher_forced_token_accuracy": 0.99453125,
+    "test/single_carry/exact_accuracy": 1.0,
+    "test/single_carry/generated_token_accuracy": 1.0,
+    "test/single_carry/loss": 0.0010194769129157066,
+    "test/single_carry/position_0/accuracy": 1.0,
+    "test/single_carry/position_1/accuracy": 1.0,
+    "test/single_carry/position_2/accuracy": 1.0,
+    "test/single_carry/position_3/accuracy": 1.0,
+    "test/single_carry/position_4/accuracy": 1.0,
+    "test/single_carry/teacher_forced_token_accuracy": 1.0,
+    "throughput/examples_per_second": 12624.497642389888,
+    "train/loss": 0.002294534471002407,
+    "train/token_accuracy": 0.9989843726158142,
+    "val/exact_accuracy": 0.9931640625,
+    "val/generated_token_accuracy": 0.99853515625,
+    "val/loss": 0.003520155092701316,
+    "val/position_0/accuracy": 0.99951171875,
+    "val/position_1/accuracy": 0.994140625,
+    "val/position_2/accuracy": 0.9990234375,
+    "val/position_3/accuracy": 1.0,
+    "val/position_4/accuracy": 1.0,
+    "val/teacher_forced_token_accuracy": 0.9986328125
+  },
+  "resolved_config": {
+    "data": {
+      "base": 10,
+      "digits": 4
+    },
+    "model": {
+      "d_mlp": 128,
+      "d_model": 64,
+      "init_mode": "normal",
+      "max_seq_len": 16,
+      "n_heads": 4,
+      "n_layers": 1
+    },
+    "train": {
+      "adamw_lr": 0.0003,
+      "batch_size": 256,
+      "checkpoint_every": 500,
+      "device": "auto",
+      "eval_batch_size": 512,
+      "eval_every": 100,
+      "eval_seed": 1,
+      "grad_clip_norm": 1.0,
+      "log_every": 10,
+      "min_lr_ratio": 0.1,
+      "muon_lr": 0.02,
+      "num_targeted_examples": 256,
+      "num_val_examples": 2048,
+      "optimizer": "muon",
+      "progress_bar": true,
+      "run_root": ".",
+      "seed": 0,
+      "steps": 10000,
+      "weight_decay": 0.01
+    },
+    "wandb": {
+      "entity": null,
+      "group": null,
+      "log_model": false,
+      "mode": "auto",
+      "name": null,
+      "project": "circuits-addition",
+      "tags": []
+    }
+  },
+  "run_id": "s85nnxtf",
+  "schema_version": 1,
+  "training_files_sha256": {
+    "checkpoints/final.pt": "8ab3e4443c586c75fde95d36fdfed473f4e6fbf374bd318e6165e875727c5285",
+    "environment.json": "6c7ded30b327853d18e87320549e118d92e48cb606ce96daa673c94eb07a965a",
+    "metrics.jsonl": "da503ac47bc1604fc84d07353c9eab7acc1b82be07dc846960e1979049e47d4f",
+    "resolved_config.toml": "c87577c6f48541b9243ba37952ba82eb11cc008c62bc075779423f8ba35e7c68",
+    "source_snapshot.zip": "a4aca21623a11159fef2e4d3dce007e03737b6909872cda92210c3903c5ea357"
+  },
+  "training_step": 10000
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 0,
+  "do_sample": false,
+  "max_new_tokens": 5,
+  "transformers_version": "5.15.0",
+  "use_cache": false
+}

initialization.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import math
+
+from torch import nn
+
+from .model_config import AdditionModelConfig
+
+
+def initialize_module(module: nn.Module, config: AdditionModelConfig) -> None:
+    init_mode = str(config.init_mode)
+    if init_mode == "normal":
+        _initialize_normal(module)
+    elif init_mode == "orthogonal":
+        _initialize_orthogonal(module)
+    else:
+        raise ValueError(f"Unsupported init mode: {init_mode}")
+
+
+def _initialize_normal(module: nn.Module) -> None:
+    for child in module.modules():
+        if isinstance(child, nn.Embedding):
+            nn.init.normal_(child.weight, mean=0.0, std=1.0 / math.sqrt(2.0))
+        elif isinstance(child, nn.Linear):
+            nn.init.normal_(child.weight, mean=0.0, std=1.0 / math.sqrt(child.in_features))
+            _zero_bias(child)
+
+
+def _initialize_orthogonal(module: nn.Module) -> None:
+    for child in module.modules():
+        if isinstance(child, nn.Embedding):
+            nn.init.normal_(child.weight, mean=0.0, std=1.0 / math.sqrt(2.0))
+        elif isinstance(child, nn.Linear):
+            gain = math.sqrt(child.out_features / child.in_features) if child.out_features > child.in_features else 1.0
+            nn.init.orthogonal_(child.weight, gain=gain)
+            _zero_bias(child)
+
+
+def _zero_bias(module: nn.Linear) -> None:
+    if module.bias is not None:
+        nn.init.zeros_(module.bias)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68be5fc852029537ac74bc0f13af4533a3efda7a4aac85396bd3a94b62f47098
+size 175592

model_config.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class AdditionModelConfig(Protocol):
+    d_model: int
+    n_heads: int
+    d_mlp: int
+    n_layers: int
+    max_seq_len: int
+    init_mode: str
+    residual_alpha: float

modeling_addition.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .addition_transformer import AdditionTransformer
+from .configuration_addition import AdditionConfig
+
+
+class AdditionForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = AdditionConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+
+    def __init__(self, config: AdditionConfig) -> None:
+        super().__init__(config)
+        self.model = AdditionTransformer(config, vocab_size=config.vocab_size)
+        self.post_init()
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # AdditionTransformer owns initialization so research and exported models remain identical.
+        return None
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.token_embedding
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.token_embedding = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.model.unembedding
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.model.unembedding = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        past_key_values: Any | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
+        if kwargs:
+            names = ", ".join(sorted(kwargs))
+            raise TypeError(f"Unsupported model inputs: {names}")
+        if past_key_values is not None or use_cache:
+            raise ValueError("AdditionForCausalLM does not implement a key-value cache.")
+        if attention_mask is not None:
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError("attention_mask must have the same shape as input_ids.")
+            if not bool(torch.all(attention_mask != 0)):
+                raise ValueError("Padding is unsupported; attention_mask must contain only ones.")
+
+        return_dict = self.config.return_dict if return_dict is None else return_dict
+        output_attentions = bool(output_attentions)
+        output_hidden_states = bool(output_hidden_states)
+        core_output = self.model(
+            input_ids,
+            return_activations=output_attentions or output_hidden_states,
+        )
+        loss = None
+        if labels is not None:
+            if labels.shape != input_ids.shape:
+                raise ValueError("labels must have the same shape as input_ids.")
+            loss = F.cross_entropy(
+                core_output.logits[:, :-1, :].contiguous().view(-1, self.config.vocab_size),
+                labels[:, 1:].contiguous().view(-1),
+                ignore_index=-100,
+            )
+
+        hidden_states = None
+        if output_hidden_states:
+            hidden_states = (
+                core_output.block_outputs[0].residual_pre_attention,
+                *(block.residual_after_mlp for block in core_output.block_outputs),
+            )
+        attentions = (
+            tuple(block.attention_pattern for block in core_output.block_outputs)
+            if output_attentions
+            else None
+        )
+
+        if not return_dict:
+            values = (core_output.logits, None, hidden_states, attentions)
+            return ((loss,) + values) if loss is not None else values
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=core_output.logits,
+            past_key_values=None,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> dict[str, torch.Tensor | None | bool]:
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": False,
+        }
+
+    def analysis_tensors(self, detach: bool = True) -> dict[str, torch.Tensor]:
+        return self.model.analysis_tensors(detach=detach)
+
+    def symmetrized_mlp_tensor(self, detach: bool = True, layer: int = 0) -> torch.Tensor:
+        return self.model.symmetrized_mlp_tensor(detach=detach, layer=layer)
+
+
+AdditionForCausalLM.register_for_auto_class("AutoModelForCausalLM")

tokenization_addition.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from transformers import PreTrainedTokenizer
+
+
+CANONICAL_VOCAB = {"<BOS>": 0, "+": 1, "=": 2, **{str(digit): digit + 3 for digit in range(10)}}
+
+
+class AdditionTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, vocab_file: str | None = None, **kwargs) -> None:
+        if vocab_file is None:
+            vocab = dict(CANONICAL_VOCAB)
+        else:
+            with Path(vocab_file).open("r", encoding="utf-8") as handle:
+                vocab = json.load(handle)
+        if vocab != CANONICAL_VOCAB:
+            raise ValueError("AdditionTokenizer requires the canonical 13-token vocabulary.")
+        self._vocab = vocab
+        self._ids_to_tokens = {token_id: token for token, token_id in vocab.items()}
+        kwargs.pop("bos_token", None)
+        kwargs.pop("eos_token", None)
+        kwargs.pop("pad_token", None)
+        kwargs.pop("unk_token", None)
+        super().__init__(
+            bos_token="<BOS>",
+            eos_token=None,
+            pad_token=None,
+            unk_token=None,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self._vocab)
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+        compact = "".join(text.split())
+        invalid = sorted(set(compact) - set("0123456789+="))
+        if invalid:
+            raise ValueError(f"Unsupported characters for addition tokenizer: {''.join(invalid)}")
+        return list(compact)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        try:
+            return self._vocab[token]
+        except KeyError as exc:
+            raise ValueError(f"Unknown addition token: {token!r}") from exc
+
+    def _convert_id_to_token(self, index: int) -> str:
+        try:
+            return self._ids_to_tokens[index]
+        except KeyError as exc:
+            raise ValueError(f"Unknown addition token ID: {index}") from exc
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return "".join(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is not None:
+            raise ValueError("AdditionTokenizer does not support sequence pairs.")
+        return [self.bos_token_id, *token_ids_0]
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        if already_has_special_tokens:
+            return [int(token_id == self.bos_token_id) for token_id in token_ids_0]
+        if token_ids_1 is not None:
+            raise ValueError("AdditionTokenizer does not support sequence pairs.")
+        return [1, *([0] * len(token_ids_0))]
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is not None:
+            raise ValueError("AdditionTokenizer does not support sequence pairs.")
+        return [0] * (len(token_ids_0) + 1)
+
+    def save_vocabulary(
+        self,
+        save_directory: str,
+        filename_prefix: str | None = None,
+    ) -> tuple[str]:
+        directory = Path(save_directory)
+        directory.mkdir(parents=True, exist_ok=True)
+        filename = f"{filename_prefix + '-' if filename_prefix else ''}vocab.json"
+        path = directory / filename
+        path.write_text(json.dumps(self._vocab, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+        return (str(path),)
+
+
+AdditionTokenizer.register_for_auto_class("AutoTokenizer")
+

tokenizer_config.json

@@ -0,0 +1,25 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<BOS>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_addition.AdditionTokenizer",
+      null
+    ]
+  },
+  "backend": "custom",
+  "bos_token": "<BOS>",
+  "eos_token": null,
+  "model_max_length": 16,
+  "pad_token": null,
+  "tokenizer_class": "AdditionTokenizer",
+  "unk_token": null
+}

training/checkpoints/final.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ab3e4443c586c75fde95d36fdfed473f4e6fbf374bd318e6165e875727c5285
+size 383443

training/checkpoints/step_000500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06d05526d4b34992b4da3a982cbb8badad4df5f1a9550617dda1c6fb4afa4a70
+size 381663

training/checkpoints/step_001000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:200f8aef8138075175f2a153cbcd97c1269a83b69990b9f7a9294dfab95b1808
+size 381727

training/checkpoints/step_001500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:772ae2a4867745aeb9931c210603c6d334140ac279a29b1b3f6ca0733c29fcc5
+size 381663

training/checkpoints/step_002000.pt

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

training/checkpoints/step_002500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2adc6ebebdd9016fb311619613154e1db9c570385594b48de2605296074d39fd
+size 381663

training/checkpoints/step_003000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7136124680252d2a8aa94cbe6c865215c345c90b517ce741000934610fbdbb5e
+size 381663

training/checkpoints/step_003500.pt

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

training/checkpoints/step_004000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cb9659ba205822deb527d40942c998040aa8f885d59200f15cb6ebbdbc7d1ec
+size 381599

training/checkpoints/step_004500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:559b0aeb57198d423779e49c4e96faaf88ce0f01062cf26b4c6584b01ec8a813
+size 381663

training/checkpoints/step_005000.pt

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

training/checkpoints/step_005500.pt

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

training/checkpoints/step_006000.pt

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

training/checkpoints/step_006500.pt

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

training/checkpoints/step_007000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47b3d5544981cb3b76ac807ca3f5edfd0f43034a598629d9c6966487f78f39df
+size 381727

training/checkpoints/step_007500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fc1f3fd25e18774689a0b39bd0d84c665f0ceb4d42a043a9ff2c73cfed0b8c7
+size 381663

training/checkpoints/step_008000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97767f8940dbfbfd4482f74585735f89e88cadae42720a4cabd9df4f3b99f877
+size 381663

training/checkpoints/step_008500.pt

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

training/checkpoints/step_009000.pt

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

training/checkpoints/step_009500.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:778c9a47d189067cbd932e93f36a067acd67cc775d75ade7efe30cce8024fbf6
+size 381663

training/environment.json

@@ -0,0 +1,19 @@
+{
+  "cuda": "12.6",
+  "device": "cuda",
+  "device_name": "NVIDIA GeForce RTX 4060 Laptop GPU",
+  "git_commit": null,
+  "git_dirty": true,
+  "packages": {
+    "huggingface-hub": "1.27.0",
+    "numpy": "2.5.2",
+    "pytest": "9.1.1",
+    "safetensors": "0.8.0",
+    "sympy": "1.14.0",
+    "transformers": "5.15.0",
+    "wandb": "0.28.1"
+  },
+  "platform": "Windows-11-10.0.26200-SP0",
+  "python": "3.12.4",
+  "torch": "2.13.0+cu126"
+}

training/resolved_config.toml

@@ -0,0 +1,38 @@
+[data]
+digits = 4
+base = 10
+
+[model]
+d_model = 64
+n_heads = 4
+d_mlp = 128
+n_layers = 1
+max_seq_len = 16
+init_mode = "normal"
+
+[train]
+steps = 10000
+batch_size = 256
+optimizer = "muon"
+muon_lr = 0.02
+adamw_lr = 0.0003
+weight_decay = 0.01
+min_lr_ratio = 0.1
+grad_clip_norm = 1.0
+log_every = 10
+eval_every = 100
+checkpoint_every = 500
+progress_bar = true
+run_root = "."
+seed = 0
+eval_seed = 1
+device = "auto"
+num_val_examples = 2048
+num_targeted_examples = 256
+eval_batch_size = 512
+
+[wandb]
+mode = "auto"
+project = "circuits-addition"
+tags = []
+log_model = false

training/source_snapshot.zip

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

transformer_block.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from .attention import CausalSelfAttention
+from .bilinear_mlp import BilinearMLP
+from .model_config import AdditionModelConfig
+
+
+@dataclass(frozen=True)
+class TransformerBlockOutput:
+    residual_pre_attention: torch.Tensor
+    attention_out: torch.Tensor
+    attention_pattern: torch.Tensor | None
+    residual_after_attention: torch.Tensor
+    mlp_out: torch.Tensor
+    residual_after_mlp: torch.Tensor
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: AdditionModelConfig) -> None:
+        super().__init__()
+        self.residual_alpha = config.residual_alpha
+        self.attention = CausalSelfAttention(config.d_model, config.n_heads, bias=False)
+        self.mlp = BilinearMLP(config.d_model, config.d_mlp, bias=False)
+
+    def forward(self, x: torch.Tensor, return_pattern: bool = False) -> TransformerBlockOutput:
+        attention_output = self.attention(x, return_pattern=return_pattern)
+        residual_after_attention = torch.lerp(x, attention_output.values, self.residual_alpha)
+        mlp_out = self.mlp(residual_after_attention)
+        residual_after_mlp = torch.lerp(residual_after_attention, mlp_out, self.residual_alpha)
+        return TransformerBlockOutput(
+            residual_pre_attention=x,
+            attention_out=attention_output.values,
+            attention_pattern=attention_output.pattern,
+            residual_after_attention=residual_after_attention,
+            mlp_out=mlp_out,
+            residual_after_mlp=residual_after_mlp,
+        )

vocab.json

@@ -0,0 +1,15 @@
+{
+  "+": 1,
+  "0": 3,
+  "1": 4,
+  "2": 5,
+  "3": 6,
+  "4": 7,
+  "5": 8,
+  "6": 9,
+  "7": 10,
+  "8": 11,
+  "9": 12,
+  "<BOS>": 0,
+  "=": 2
+}