Publish addition-transformer run rmb41c76

Loadable model and complete training record.

README.md

@@ -0,0 +1,70 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- arithmetic
+- interpretability
+- arxiv:2405.14813
+---
+
+# Fixed-width addition transformer
+
+Run `rmb41c76` 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.014791 |
+| Validation generated-token accuracy | 99.38% |
+| Validation exact-answer accuracy | 97.46% |
+| No-carry exact-answer accuracy | 95.70% |
+| Single-carry exact-answer accuracy | 99.22% |
+| Multiple-carry exact-answer accuracy | 97.27% |
+| Carry-chain exact-answer accuracy | 89.45% |
+
+## 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
+- Position encoding: full-head RoPE on queries and keys (theta 10000)
+- 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,82 @@
+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.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:
+        _, 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}.")
+
+        x = self.token_embedding(input_ids)
+        # Modula-inspired RoPE, 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,68 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .rotary_embedding import RotaryEmbedding
+
+
+@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,
+        max_seq_len: int,
+        rope_theta: float = 10_000.0,
+        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.rotary = RotaryEmbedding(self.d_head, max_seq_len, rope_theta)
+        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))
+        q, k = self.rotary(q, k)
+
+        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,46 @@
+{
+  "answer_digits": 5,
+  "architecture_version": 3,
+  "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": 3,
+  "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,
+  "position_encoding": "rope",
+  "residual_alpha": 0.5,
+  "rope_theta": 10000.0,
+  "rotary_dim": 16,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.15.0",
+  "use_cache": false,
+  "vocab_size": 13
+}

configuration_addition.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import math
+
+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,
+        rope_theta: float = 10_000.0,
+        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,
+            "position_encoding": "rope",
+            "rotary_dim": d_model // n_heads if n_heads > 0 else 0,
+            "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": 3,
+            "export_format_version": 3,
+            "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.rope_theta = float(rope_theta)
+        self.position_encoding = "rope"
+        self.rotary_dim = d_model // n_heads if n_heads > 0 else 0
+        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 = 3
+        self.export_format_version = 3
+        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.")
+        if self.d_head % 2 != 0:
+            raise ValueError("d_head must be even for RoPE.")
+        if not math.isfinite(self.rope_theta) or self.rope_theta <= 0:
+            raise ValueError("rope_theta must be positive.")
+        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,149 @@
+{
+  "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": "3cbb9ada58204b34e2afd699fa0ff7bbe5702751ff025c81272bd005d5cc73a9",
+    "addition_transformer.py": "6c527805f122b294ce3041ea592fb90e8b65a5012767eefc8bbd54ac582cacbf",
+    "attention.py": "e593866a411e4a7fcf089320b0732551170cd1d22500aa78f32e0e4a0a4f12b5",
+    "bilinear_mlp.py": "95d8b248058bdbce7ea7c13ba9ebf49ee69b20b95714a43040f19c9c528225c7",
+    "config.json": "e1b7a20a05146ee66d2d5442a7eb0dabfdfee043a266c5f2f9074f348440a3b3",
+    "configuration_addition.py": "56d24dfcef47b17e4e2a24f8f8f08d834ed9ec4bb86738ee8077ac75607b7c95",
+    "generation_config.json": "13cc88dd90bbb2a181ade617b6a41b35998e933cc1eb356f7906c320a7d04b5d",
+    "initialization.py": "01de663e2494370c6bb2a9d4f770411104dbc44cbb97917fa18ba89ed9959d32",
+    "model.safetensors": "af93c22839a0d30a2b81ef7e7444a0e4dbb1d726fe20561d654b15fd397af55c",
+    "model_config.py": "2c1c01f64064830d9071aacd0c53219eb2f26365838d113713dffaba54ed6cd3",
+    "modeling_addition.py": "b5349de597570045eb3c21803f6a4b077605f5468f99e26ddb1678964d01fcea",
+    "rotary_embedding.py": "e5565ee1c072cc42f53af0cc51098f03f470591bd83454f1a162768d98a58b1a",
+    "tokenization_addition.py": "bc63473854f8f8fecfc3375757d257dfb34a33a0c30ee5a1ddffb3b984b8602f",
+    "tokenizer_config.json": "630ab54705ab78ecddef9c0c6543f477036b5a71cb2736a3db59462f464cbd06",
+    "transformer_block.py": "6566c6432e8d4868664be63e783e9f96e6dc138e39055dcd9c5caa70370ba9ed",
+    "vocab.json": "d6fdf48e6e2087546f1bfb7285a71b1072f5452934e8a33e831d069f7faf56a8"
+  },
+  "final_metrics": {
+    "examples_processed": 2560000,
+    "grad/clip_fraction": 0.0,
+    "grad/global_norm": 0.4769897431135178,
+    "lr/adamw": 2.9999999999999997e-05,
+    "lr/muon": 0.002,
+    "test/carry_chain/exact_accuracy": 0.89453125,
+    "test/carry_chain/generated_token_accuracy": 0.97734375,
+    "test/carry_chain/loss": 0.05059999823570251,
+    "test/carry_chain/position_0/accuracy": 0.9921875,
+    "test/carry_chain/position_1/accuracy": 0.9375,
+    "test/carry_chain/position_2/accuracy": 0.95703125,
+    "test/carry_chain/position_3/accuracy": 1.0,
+    "test/carry_chain/position_4/accuracy": 1.0,
+    "test/carry_chain/teacher_forced_token_accuracy": 0.97890625,
+    "test/multiple_carry/exact_accuracy": 0.97265625,
+    "test/multiple_carry/generated_token_accuracy": 0.99453125,
+    "test/multiple_carry/loss": 0.018936741352081298,
+    "test/multiple_carry/position_0/accuracy": 0.99609375,
+    "test/multiple_carry/position_1/accuracy": 0.98828125,
+    "test/multiple_carry/position_2/accuracy": 0.98828125,
+    "test/multiple_carry/position_3/accuracy": 1.0,
+    "test/multiple_carry/position_4/accuracy": 1.0,
+    "test/multiple_carry/teacher_forced_token_accuracy": 0.99453125,
+    "test/no_carry/exact_accuracy": 0.95703125,
+    "test/no_carry/generated_token_accuracy": 0.9890625,
+    "test/no_carry/loss": 0.022115638852119444,
+    "test/no_carry/position_0/accuracy": 0.9765625,
+    "test/no_carry/position_1/accuracy": 0.984375,
+    "test/no_carry/position_2/accuracy": 0.984375,
+    "test/no_carry/position_3/accuracy": 1.0,
+    "test/no_carry/position_4/accuracy": 1.0,
+    "test/no_carry/teacher_forced_token_accuracy": 0.990625,
+    "test/single_carry/exact_accuracy": 0.9921875,
+    "test/single_carry/generated_token_accuracy": 0.9984375,
+    "test/single_carry/loss": 0.0053884580731391905,
+    "test/single_carry/position_0/accuracy": 1.0,
+    "test/single_carry/position_1/accuracy": 0.9921875,
+    "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": 0.9984375,
+    "throughput/examples_per_second": 24623.5772583565,
+    "train/loss": 0.010908610420301557,
+    "train/token_accuracy": 0.9959374904632569,
+    "val/exact_accuracy": 0.974609375,
+    "val/generated_token_accuracy": 0.99375,
+    "val/loss": 0.014790849387645721,
+    "val/position_0/accuracy": 0.99560546875,
+    "val/position_1/accuracy": 0.9853515625,
+    "val/position_2/accuracy": 0.9892578125,
+    "val/position_3/accuracy": 0.99853515625,
+    "val/position_4/accuracy": 1.0,
+    "val/teacher_forced_token_accuracy": 0.9943359375
+  },
+  "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,
+      "rope_theta": 10000.0
+    },
+    "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": "rmb41c76",
+  "schema_version": 1,
+  "training_files_sha256": {
+    "checkpoints/final.pt": "364b23da837d6da7dd72bdbded43f31c4b33fcd985e571b70920e932d9c29397",
+    "environment.json": "6c7ded30b327853d18e87320549e118d92e48cb606ce96daa673c94eb07a965a",
+    "metrics.jsonl": "d086c27ee9929dcb23441aeaf872f482004fde560a1447a79c8929266c0597fe",
+    "resolved_config.toml": "314099bfa658adccfe74c984160caf91f0153c91fbc07447181255ddd86c73d7",
+    "source_snapshot.zip": "e5f98551caef8fe2ba80f851fa33b481a2c0dd651ff55102e13ebd4754f5a3b4"
+  },
+  "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:af93c22839a0d30a2b81ef7e7444a0e4dbb1d726fe20561d654b15fd397af55c
+size 171400

model_config.py

@@ -0,0 +1,14 @@
+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
+    rope_theta: float
+    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")

rotary_embedding.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, d_head: int, max_seq_len: int, theta: float = 10_000.0) -> None:
+        super().__init__()
+        if d_head < 2 or d_head % 2 != 0:
+            raise ValueError("RoPE requires a positive, even attention head dimension.")
+        if max_seq_len < 1:
+            raise ValueError("max_seq_len must be positive.")
+        theta = float(theta)
+        if not math.isfinite(theta) or theta <= 0:
+            raise ValueError("RoPE theta must be positive.")
+
+        self.d_head = d_head
+        self.max_seq_len = max_seq_len
+        self.theta = theta
+        self.rope_dim = d_head // 2
+        self.register_buffer("_cos", torch.empty(0), persistent=False)
+        self.register_buffer("_sin", torch.empty(0), persistent=False)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if query.shape != key.shape:
+            raise ValueError("RoPE query and key tensors must have the same shape.")
+        return self.rotate(query), self.rotate(key)
+
+    def rotate(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4 or x.shape[-1] != self.d_head:
+            raise ValueError(f"RoPE expects shape [batch, heads, sequence, {self.d_head}].")
+        seq_len = x.shape[-2]
+        if seq_len > self.max_seq_len:
+            raise ValueError(f"Sequence length {seq_len} exceeds RoPE limit {self.max_seq_len}.")
+
+        cos, sin = self._cos_sin(x.device)
+        cos = cos[:, :, :seq_len].to(dtype=x.dtype)
+        sin = sin[:, :, :seq_len].to(dtype=x.dtype)
+        first_half = x[..., : self.rope_dim]
+        second_half = x[..., self.rope_dim :]
+        return torch.cat(
+            (
+                cos * second_half + sin * first_half,
+                -sin * second_half + cos * first_half,
+            ),
+            dim=-1,
+        )
+
+    def _cos_sin(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._cos.numel() == 0 or self._cos.device != device:
+            inverse_frequencies = self.theta ** (
+                -torch.arange(self.rope_dim, dtype=torch.float32, device=device) / self.rope_dim
+            )
+            frequencies = torch.outer(
+                torch.arange(self.max_seq_len, dtype=torch.float32, device=device),
+                inverse_frequencies,
+            )
+            self._cos = frequencies.cos()[None, None, :, :]
+            self._sin = frequencies.sin()[None, None, :, :]
+        return self._cos, self._sin

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:364b23da837d6da7dd72bdbded43f31c4b33fcd985e571b70920e932d9c29397
+size 370007

training/checkpoints/step_000500.pt

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

training/checkpoints/step_001000.pt

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

training/checkpoints/step_001500.pt

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

training/checkpoints/step_002000.pt

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

training/checkpoints/step_002500.pt

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

training/checkpoints/step_003000.pt

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

training/checkpoints/step_003500.pt

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

training/checkpoints/step_004000.pt

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

training/checkpoints/step_004500.pt

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

training/checkpoints/step_005000.pt

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

training/checkpoints/step_005500.pt

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

training/checkpoints/step_006000.pt

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

training/checkpoints/step_006500.pt

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

training/checkpoints/step_007000.pt

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

training/checkpoints/step_007500.pt

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

training/checkpoints/step_008000.pt

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

training/checkpoints/step_008500.pt

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

training/checkpoints/step_009000.pt

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

training/checkpoints/step_009500.pt

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

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,39 @@
+[data]
+digits = 4
+base = 10
+
+[model]
+d_model = 64
+n_heads = 4
+d_mlp = 128
+n_layers = 1
+max_seq_len = 16
+rope_theta = 10000.0
+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:e5f98551caef8fe2ba80f851fa33b481a2c0dd651ff55102e13ebd4754f5a3b4
+size 162656

transformer_block.py

@@ -0,0 +1,48 @@
+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,
+            config.max_seq_len,
+            config.rope_theta,
+            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
+}