Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 vmal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,193 @@
 ---
 license: mit
+language:
+- en
+tags:
+- arithmetic
+- process-supervision
+- scratchpad
+- chain-of-thought
+- from-scratch
+- interpretability
+- small-language-model
+- looped-transformer
+pipeline_tag: text-generation
+library_name: transformers
+datasets:
+- vmal/3-digit-arithmetic-scratchpad-traces
 ---
+
+# 3-digit-basic-calc
+
+A from-scratch, **1.6M-parameter** looped decoder-only transformer for addition,
+subtraction, multiplication, and division on integer operands with up to three
+digits. It generates explicit algorithmic scratchpads and is evaluated on
+prompt-disjoint held-out expressions. No tools or external memory: next-token
+prediction over a 34-token character-and-control vocabulary.
+
+| Operation | Validation | Fresh, never-seen problems |
+|-----------|-----------:|---------------------------:|
+| Addition `+`       | 99.9% | 99.6% |
+| Subtraction `−`    | 99.7% | 99.7% |
+| Multiplication `×` | 100%  | 100% |
+| Division `÷`       | 96.8% | 93.1% |
+| **Overall**        | **99.1%** | **98.1%** |
+
+
+---
+
+## Quick start
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("vmal/3-digit-basic-calc", trust_remote_code=True).eval()
+tok = AutoTokenizer.from_pretrained("vmal/3-digit-basic-calc", trust_remote_code=True)
+
+print(model.solve(tok, "842/37"))    # -> 22.757
+print(model.solve(tok, "213*145"))   # -> 30885
+print(model.solve(tok, "-500+500"))  # -> 0
+print(model.solve(tok, "12/0"))      # -> NAN
+
+# See the model's actual reasoning (the scratchpad it generates):
+answer, trace = model.solve(tok, "842/37", return_trace=True)
+print(trace)
+# 842/37=<div><pos><state>842037<step>00080008<qmul>00000<rem>008
+# <step>00840084<qmul>20074<rem>010 ... <ans>022.757
+```
+
+Operands are integers in **[−999, 999]**. Whitespace and a trailing `=` are
+accepted; malformed expressions and out-of-range operands raise `ValueError`.
+Division answers are rounded half-up to three decimals, and division by zero
+returns `NAN`.
+
+Mixed-length batching is supported with the tokenizer's default left padding:
+
+```python
+batch = tok(["1+2=", "999/7="], return_tensors="pt", padding=True)
+raw_traces = model.generate(**batch, do_sample=False)
+```
+
+## Research iterations
+
+**Goal:** Develop a small transformer that generalizes the algorithms for
+addition, subtraction, multiplication, and division, targeting at least 95%
+accuracy per operation on prompt-disjoint held-out splits.
+
+**Iterations:**
+
+**A plain ~4M-parameter transformer.** The obvious start. It fit the training data but didn't generalize the algorithm; accuracy fell apart on unseen numbers. Scaling wasn't obviously the fix; the failures looked structural.
+
+**A looped transformer.** Sharing a small block across many iterations gives cheap depth, the right bias for an iterative procedure. It helped a little. It didn't make the model generalize.
+
+**Abacus embeddings.** Following *Transformers Can Do Arithmetic with the Right Embeddings* (McLeish et al., NeurIPS 2024), I added its three pieces: digit-position (abacus) embeddings, input injection, and a randomized-offset scheme for length generalization. This was the first thing that clearly moved the *hardest* cases. But it still failed to generalize division and multiplication, and even addition and subtraction stayed shakier than they should have. I also tried **GRPO** on top; it bought a small, real improvement, not the step change the problem needed.
+
+Every one of those was a reasonable bet, and none fixed the core issue. The model was always being asked to carry too much state (a remainder, a running carry) implicitly, inside its activations. More depth, better position embeddings, and RL were all trying to make it better at holding hidden state. The move that worked was to stop asking it to.
+
+**Process supervision (the scratchpad).** Instead of `842/37 → 22.757`, teach the model to emit the whole procedure as tokens, writing down the state at every step. Each token now predicts a local step from explicit state written in the preceding tokens, reducing the hidden-state burden. This was the turning point: the model started to *generalize* across all four operations. But two operations still fell short.
+
+**Multiplication stalled near 65%.** Measured per step, one step did all the damage: a single token had to sum up to three digit-products and a carry at once. Meanwhile training loss had collapsed to near zero while validation sat at 65%: the model could already fit the traces perfectly. Capacity was not the primary bottleneck; one step packed too much into one decision. So I broke the column sum into **term-level micro-steps** (emit each partial product, then the digit and carry). **Multiplication → 99.6%.**
+
+**Then division was the laggard.** A first-error diagnostic showed the errors sat almost entirely in the **interior long-division steps**, never the setup, never the rounding. The step was still doing two hard things at once: find the quotient digit *and* compute the new remainder. Same fix, **more micro-steps**: each position becomes form-the-partial → emit the quotient digit and its product → subtract for the remainder. **Division cleared its plateau and reached ~97%**, while the other operations remained near 99–100%.
+
+Two things ran underneath all of it. **Sequence length kept growing**: more micro-steps mean longer traces, so `max_seq_len` climbed 80 → 128 → 176 to fit them. Whenever the model lagged on a specific skill (non-terminating quotients or high-carry columns), the fix was usually to *improve the data distribution* for that skill, not to change the model.
+
+> **The lesson:** in these experiments, capacity was not the primary bottleneck; decomposition was decisive. Each plateau broke when the hardest *local function per step* was split into simpler pieces and the data was kept well-distributed. The 256-wide, 2-layer, 8-loop architecture stayed fixed through the scratchpad experiments.
+
+## Evaluation
+
+All numbers are greedy decoding on prompt-disjoint held-out data. Detailed
+counts, seeds, conversion checks, and timing are available in
+[`evaluation_results.json`](evaluation_results.json).
+
+### Locked test split
+
+The reserved test split was evaluated once after model selection:
+
+| Operation | Accuracy |
+|-----------|---------:|
+| `+` | 99.9% (999/1,000) |
+| `−` | 99.9% (999/1,000) |
+| `×` | 100.0% (1,000/1,000) |
+| `÷` | 95.2% (952/1,000) |
+| **Overall** | **98.75% (3,950/4,000)** |
+
+### Generality, never-seen problems
+
+Beyond the reserved test split, I sampled and evaluated **3,000 new, uniformly
+generated problems with operands up to three digits**, each verified to be
+absent from all 108,000 dataset prompts:
+
+| Operation | Accuracy |
+|-----------|---------:|
+| `+` | 99.6% |
+| `−` | 99.7% |
+| `×` | **100.0%** (750/750) |
+| `÷` | 93.1% |
+| **Overall** | **98.1%** |
+
+The fresh benchmark differs from the locked test by −0.3 points on addition,
+−0.2 on subtraction, 0.0 on multiplication, and −2.1 on division. The larger
+division gap is consistent with uniform sampling stressing a different mixture
+than the scenario-balanced held-out splits. Prompt exclusion rules out exact
+training-prompt memorization, but these results should not be interpreted as
+length generalization: operands remain within the trained three-digit range.
+
+### Where it's hard, and hardest
+
+- **Easiest:** multiplication (perfect on fresh samples) and addition/subtraction (~99.7%).
+- **Hardest:** multi-step division, especially with two- and three-digit
+  divisors. On the locked test set, 47 of 48 division failures first diverged
+  at quotient-digit/product selection and one at the resulting remainder.
+  Errors occurred across interior positions rather than only at the rounding
+  digit, so a wrong quotient digit can cascade into later remainders.
+- **Rare +/− failure:** near-cancellation with mixed signs (e.g. `239 + (−271)`), where the true result is tiny and the model occasionally loses sign or magnitude, about 0.3% of cases.
+- **`÷0`** is a single atomic `<nan>` token.
+
+### Is there more headroom?
+
+For division problems the model gets wrong at greedy decoding, sampling 32
+times still reaches the correct answer in about 78% of cases, so correct
+alternatives often exist in the distribution. Two conservative GRPO trials
+did **not** convert that headroom into better greedy validation: one stayed at
+96.8% division accuracy and one regressed to 96.6%. The checkpoint-selection
+gates therefore retained the supervised baseline released here.
+
+
+## Technical specifications
+
+| | |
+|---|---|
+| Parameters | 1,596,160 |
+| Architecture | looped decoder-only (2 layers × 8 shared loops) |
+| Hidden size / heads | 256 / 8 |
+| Vocabulary | 34 tokens (character + scratchpad control tokens) |
+| Context length | 176 |
+| Scratchpad digit order | reversed for `+`, `−`, `×`; ordinary order for `÷` |
+| Positional encoding | RoPE |
+| Normalization | RMSNorm |
+| Precision | fp32 |
+
+## Limitations
+
+- Operands are restricted to integers in `[-999, 999]`.
+- The fixed-width scratchpad representation does not generalize to four-digit
+  or longer operands.
+- Division reaches 95.2% on the locked test split but 93.1% on the fresh
+  uniform benchmark.
+- This is a research model, not a reliable production calculator.
+
+## Intended use
+
+Education and research. This model is a small, fully inspectable case study in process supervision, step decomposition, and rigorous evaluation of algorithmic generalization. It is intended for studying how explicit intermediate computation can support arithmetic reasoning in small transformers.
+
+## References
+
+- Nye et al., *Show Your Work: Scratchpads for Intermediate Computation* (2021)
+- McLeish et al., *Transformers Can Do Arithmetic with the Right Embeddings* (NeurIPS 2024)
+- Dehghani et al., *Universal Transformers* (2019)
+- Shao et al., *DeepSeekMath / GRPO* (2024)
+
+## License
+
+MIT.

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "ThreeDigitBasicCalcForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.ThreeDigitBasicCalcConfig",
+    "AutoModelForCausalLM": "modeling.ThreeDigitBasicCalcForCausalLM"
+  },
+  "bos_token_id": 1,
+  "dropout": 0.05,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "hidden_size": 256,
+  "max_position_embeddings": 176,
+  "model_type": "3-digit-basic-calc",
+  "n_loops": 8,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 2,
+  "pad_token_id": 0,
+  "rms_eps": 1e-06,
+  "rope_base": 10000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.3.0",
+  "unk_token_id": 3,
+  "use_cache": false,
+  "vocab_size": 34
+}

config.py

@@ -0,0 +1,50 @@
+# Copyright 2026. Released under the MIT license.
+"""Configuration for the 3-digit-basic-calc looped arithmetic transformer."""
+
+from transformers import PretrainedConfig
+
+
+class ThreeDigitBasicCalcConfig(PretrainedConfig):
+    """Config for a looped ("universal") decoder-only calculator transformer.
+
+    A small stack of ``num_hidden_layers`` decoder blocks is applied ``n_loops``
+    times with shared weights, giving an effective depth of
+    ``num_hidden_layers * n_loops`` at the parameter cost of the stack alone.
+    """
+
+    model_type = "3-digit-basic-calc"
+
+    def __init__(
+        self,
+        vocab_size: int = 34,
+        max_position_embeddings: int = 176,
+        hidden_size: int = 256,
+        num_attention_heads: int = 8,
+        num_hidden_layers: int = 2,
+        n_loops: int = 8,
+        dropout: float = 0.05,
+        rms_eps: float = 1e-6,
+        rope_base: float = 10000.0,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        unk_token_id: int = 3,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.n_loops = n_loops
+        self.dropout = dropout
+        self.rms_eps = rms_eps
+        self.rope_base = rope_base
+        self.unk_token_id = unk_token_id
+        kwargs.setdefault("tie_word_embeddings", False)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )

evaluation_results.json

@@ -0,0 +1,105 @@
+{
+  "schema_version": 1,
+  "model": {
+    "name": "3-digit-basic-calc",
+    "parameters": 1596160,
+    "source_checkpoint": "ckp/scratchpad_v3_256x2x8_bs32.pt",
+    "representation": "scratchpad-v3",
+    "decoding": "greedy"
+  },
+  "conversion_validation": {
+    "date": "2026-07-24",
+    "dataset_rows_tokenized": 108000,
+    "tokenizer_id_mismatches": 0,
+    "learned_tensors_compared": 24,
+    "learned_tensor_mismatches": 0,
+    "export_tensor_count": 26,
+    "derived_export_tensors": [
+      "decoder_layers.0.attn.rope.inv_freq",
+      "decoder_layers.1.attn.rope.inv_freq"
+    ],
+    "prompt_logit_max_abs_difference": 0.0,
+    "complete_162_token_logit_max_abs_difference": 0.0,
+    "argmax_mismatches": 0
+  },
+  "validation": {
+    "split": "val",
+    "rows": 4000,
+    "implementation": "exported Hugging Face architecture",
+    "device": "cpu",
+    "elapsed_seconds": 1513.414,
+    "answer": {
+      "+": {"correct": 999, "total": 1000, "accuracy": 0.999},
+      "-": {"correct": 997, "total": 1000, "accuracy": 0.997},
+      "*": {"correct": 1000, "total": 1000, "accuracy": 1.0},
+      "/": {"correct": 968, "total": 1000, "accuracy": 0.968},
+      "overall": {"correct": 3964, "total": 4000, "accuracy": 0.991}
+    },
+    "trace_exact": {
+      "+": {"correct": 999, "total": 1000, "accuracy": 0.999},
+      "-": {"correct": 997, "total": 1000, "accuracy": 0.997},
+      "*": {"correct": 1000, "total": 1000, "accuracy": 1.0},
+      "/": {"correct": 949, "total": 1000, "accuracy": 0.949},
+      "overall": {"correct": 3945, "total": 4000, "accuracy": 0.98625}
+    },
+    "parse_success": {"correct": 4000, "total": 4000, "accuracy": 1.0}
+  },
+  "locked_test": {
+    "split": "test",
+    "rows": 4000,
+    "device": "mps",
+    "source_report": "sample/scratchpad_v3_test_eval.json",
+    "answer": {
+      "+": {"correct": 999, "total": 1000, "accuracy": 0.999},
+      "-": {"correct": 999, "total": 1000, "accuracy": 0.999},
+      "*": {"correct": 1000, "total": 1000, "accuracy": 1.0},
+      "/": {"correct": 952, "total": 1000, "accuracy": 0.952},
+      "overall": {"correct": 3950, "total": 4000, "accuracy": 0.9875}
+    },
+    "sequence_exact": {"correct": 3925, "total": 4000, "accuracy": 0.98125},
+    "parse_success": {"correct": 4000, "total": 4000, "accuracy": 1.0}
+  },
+  "fresh_prompt_disjoint": {
+    "seed": 20260723,
+    "rows": 3000,
+    "sampling": "uniform random operands and operation-balanced, excluding all 108000 dataset prompts",
+    "implementation": "native checkpoint after exact exported-logit parity was established",
+    "answer": {
+      "+": {"correct": 747, "total": 750, "accuracy": 0.996},
+      "-": {"correct": 748, "total": 750, "accuracy": 0.9973333333},
+      "*": {"correct": 750, "total": 750, "accuracy": 1.0},
+      "/": {"correct": 698, "total": 750, "accuracy": 0.9306666667},
+      "overall": {"correct": 2943, "total": 3000, "accuracy": 0.981}
+    },
+    "trace_exact": {
+      "+": {"correct": 747, "total": 750, "accuracy": 0.996},
+      "-": {"correct": 748, "total": 750, "accuracy": 0.9973333333},
+      "*": {"correct": 750, "total": 750, "accuracy": 1.0},
+      "/": {"correct": 669, "total": 750, "accuracy": 0.892}
+    },
+    "elapsed_seconds": 1227.1992
+  },
+  "division_failure_diagnosis": {
+    "source": "locked_test",
+    "answer_failures": 48,
+    "first_wrong_component": {
+      "quotient_digit_or_product": 47,
+      "remainder": 1
+    },
+    "first_wrong_position": {
+      "1": 4,
+      "2": 11,
+      "3": 17,
+      "4": 10,
+      "5": 4,
+      "6": 2
+    }
+  },
+  "grpo": {
+    "baseline_division_accuracy": 0.968,
+    "trial_1_division_accuracy": 0.968,
+    "trial_2_division_accuracy": 0.966,
+    "selected_checkpoint": "supervised baseline",
+    "reason": "neither trial improved gated greedy validation"
+  }
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "max_new_tokens": 176,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.3.0"
+}

model.safetensors

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

modeling.py

@@ -0,0 +1,413 @@
+# Copyright 2026. Released under the MIT license.
+"""Self-contained looped decoder-only transformer for character-level arithmetic.
+
+This module intentionally has no dependency on the training repository so it can
+run from the Hub with ``trust_remote_code=True``. The building blocks mirror the
+original ``model.py`` (RMSNorm, RoPE, causal attention, GELU MLP) and the
+parameter names match the released checkpoint so weights load directly.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+
+import torch
+import torch.nn as nn
+
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .config import ThreeDigitBasicCalcConfig
+
+
+class WordEmbeddings(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x):
+        return self.embedding(x) * math.sqrt(self.d_model)
+
+
+class RMSLayerNormalization(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        mean = x.pow(2).mean(dim=-1, keepdim=True)
+        return self.weight * x * torch.rsqrt(mean + self.eps)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, h_dim: int, base: float = 10000.0):
+        super().__init__()
+        assert h_dim % 2 == 0
+        self.max_seq_len = max_seq_len
+        # Persistent so from_pretrained (low-mem loading) always materializes it;
+        # cos/sin are derived per forward to avoid stale non-persistent buffers.
+        inv_freq = 1.0 / (base ** (torch.arange(0, h_dim, 2).float() / h_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=True)
+        # Derived runtime values are deliberately not registered buffers: the
+        # checkpoint contains only inv_freq, while repeated loop passes reuse
+        # the trigonometric table.
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cache_key = None
+
+    def _cos_sin(self, length: int, device):
+        key = (length, device.type, device.index, self.inv_freq.dtype)
+        if self._cache_key == key:
+            return self._cos_cached, self._sin_cached
+        t = torch.arange(length, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        self._cos_cached = freqs.cos()
+        self._sin_cached = freqs.sin()
+        self._cache_key = key
+        return self._cos_cached, self._sin_cached
+
+    def forward(self, x, position_ids=None):
+        seq = x.shape[-2]
+        if position_ids is None:
+            position_ids = torch.arange(seq, device=x.device).unsqueeze(0)
+        position_ids = position_ids.to(device=x.device, dtype=torch.long)
+        required = int(position_ids.max().item()) + 1
+        if required > self.max_seq_len:
+            raise ValueError(
+                f"position {required - 1} exceeds the configured context "
+                f"length {self.max_seq_len}"
+            )
+        cos_table, sin_table = self._cos_sin(required, x.device)
+        cos = cos_table[position_ids].unsqueeze(1)
+        sin = sin_table[position_ids].unsqueeze(1)
+        x1 = x[..., 0::2]
+        x2 = x[..., 1::2]
+        out1 = x1 * cos - x2 * sin
+        out2 = x1 * sin + x2 * cos
+        return torch.stack([out1, out2], dim=-1).flatten(-2).type_as(x)
+
+
+class MultiHeadAttentionBlock(nn.Module):
+    def __init__(self, max_seq_len, d_model, h, dropout=0.0, base=10000.0):
+        super().__init__()
+        assert d_model % h == 0
+        self.heads = h
+        self.d_head = d_model // h
+        self.w_q = nn.Linear(d_model, d_model, bias=False)
+        self.w_k = nn.Linear(d_model, d_model, bias=False)
+        self.w_v = nn.Linear(d_model, d_model, bias=False)
+        self.w_o = nn.Linear(d_model, d_model, bias=False)
+        self.rope = RotaryEmbedding(max_seq_len, self.d_head, base)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None, position_ids=None):
+        b, s, d = x.shape
+        q = self.w_q(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
+        k = self.w_k(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
+        v = self.w_v(x).view(b, s, self.heads, self.d_head).transpose(1, 2)
+        q = self.rope(q, position_ids)
+        k = self.rope(k, position_ids)
+        attn_score = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
+        if mask is not None:
+            attn_score = attn_score.masked_fill(~mask, float("-inf"))
+            # Left-padded query rows before the first real token have no legal
+            # key. Their output is discarded; zeroing only those score rows
+            # avoids NaNs while leaving every real-token score bit-identical to
+            # the unpadded/native implementation.
+            no_legal_key = ~mask.any(dim=-1, keepdim=True)
+            attn_score = attn_score.masked_fill(no_legal_key, 0.0)
+        attn = torch.softmax(attn_score, dim=-1)
+        attn = self.dropout(attn)
+        out = (attn @ v).transpose(1, 2).contiguous().reshape(b, s, d)
+        return self.w_o(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.0):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.fc1 = nn.Linear(d_model, d_ff)
+        self.fc2 = nn.Linear(d_ff, d_model)
+        self.act = nn.GELU()
+
+    def forward(self, x):
+        return self.fc2(self.dropout(self.act(self.fc1(x))))
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, max_seq_len, d_model, attention_heads, d_ff,
+                 eps=1e-6, dropout=0.0, base=10000.0):
+        super().__init__()
+        self.layer_norm_1 = RMSLayerNormalization(d_model, eps)
+        self.layer_norm_2 = RMSLayerNormalization(d_model, eps)
+        self.attn = MultiHeadAttentionBlock(
+            max_seq_len, d_model, attention_heads, dropout, base)
+        self.ffn = MLP(d_model, d_ff, dropout)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask, position_ids=None):
+        x = x + self.dropout(
+            self.attn(self.layer_norm_1(x), mask, position_ids)
+        )
+        x = x + self.dropout(self.ffn(self.layer_norm_2(x)))
+        return x
+
+
+class ThreeDigitBasicCalcForCausalLM(PreTrainedModel, GenerationMixin):
+    """Looped decoder-only causal LM. Parameter names match the checkpoint."""
+
+    config_class = ThreeDigitBasicCalcConfig
+    _supports_cache_class = False
+
+    def __init__(self, config: ThreeDigitBasicCalcConfig):
+        super().__init__(config)
+        d = config.hidden_size
+        self.embedding = WordEmbeddings(d, config.vocab_size)
+        self.decoder_layers = nn.ModuleList([
+            DecoderBlock(config.max_position_embeddings, d,
+                         config.num_attention_heads, 4 * d,
+                         config.rms_eps, config.dropout, config.rope_base)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.step_emb = nn.Embedding(config.n_loops, d)
+        self.layer_norm = RMSLayerNormalization(d, config.rms_eps)
+        self.projection = nn.Linear(d, config.vocab_size, bias=False)
+        self.config.use_cache = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embedding.embedding
+
+    def set_input_embeddings(self, value):
+        self.embedding.embedding = value
+
+    def get_output_embeddings(self):
+        return self.projection
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        position_ids=None,
+        labels=None,
+        **kwargs,
+    ):
+        b, s = input_ids.shape
+        if s > self.config.max_position_embeddings:
+            raise ValueError(
+                f"sequence length {s} exceeds the configured context length "
+                f"{self.config.max_position_embeddings}"
+            )
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+        if attention_mask.shape != input_ids.shape:
+            raise ValueError("attention_mask must have the same shape as input_ids")
+        attention_mask = attention_mask.to(device=input_ids.device)
+        if position_ids is None:
+            position_ids = attention_mask.long().cumsum(dim=-1) - 1
+            position_ids = position_ids.clamp_min(0)
+        if position_ids.shape[-1] != s:
+            raise ValueError("position_ids must have the same sequence length as input_ids")
+
+        causal = torch.tril(
+            torch.ones(s, s, dtype=torch.bool, device=input_ids.device)
+        ).view(1, 1, s, s)
+        key_mask = attention_mask.to(torch.bool).view(b, 1, 1, s)
+        combined_mask = causal & key_mask
+        x = self.embedding(input_ids)
+        for t in range(self.config.n_loops):
+            x = x + self.step_emb.weight[t]
+            for decoder in self.decoder_layers:
+                x = decoder(x, combined_mask, position_ids)
+        x = self.layer_norm(x)
+        logits = self.projection(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            shift_attention = attention_mask[:, 1:].to(torch.bool)
+            shift_labels = shift_labels.masked_fill(
+                ~shift_attention | (shift_labels == self.config.pad_token_id),
+                -100,
+            )
+            loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        # No KV cache: the model recomputes the full sequence each step.
+        attention_mask = kwargs.get("attention_mask")
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+        position_ids = attention_mask.long().cumsum(dim=-1) - 1
+        position_ids = position_ids.clamp_min(0)
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+        }
+
+    @torch.no_grad()
+    def solve(self, tokenizer, expression: str, max_new_tokens: int = 256,
+              return_trace: bool = False):
+        """Compute a 3-digit arithmetic expression, e.g. ``model.solve(tok, "842/37")``.
+
+        Returns the human-readable answer string (or a ``(answer, trace)`` tuple
+        when ``return_trace=True``). The model emits an explicit scratchpad; this
+        decodes its final ``<ans>`` field back to an ordinary number.
+        """
+        a, op, b, prompt = _parse_expression(expression)
+        encoded = tokenizer(prompt, return_tensors="pt")
+        ids = encoded.input_ids.to(self.device)
+        remaining = self.config.max_position_embeddings - ids.shape[1]
+        if remaining <= 0:
+            raise ValueError("the expression leaves no context for generation")
+        max_new_tokens = min(max_new_tokens, remaining)
+        out = self.generate(
+            ids,
+            attention_mask=encoded.attention_mask.to(self.device),
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        if out[0, -1].item() != tokenizer.eos_token_id:
+            raise ValueError(
+                "generation reached the context limit before emitting <eos>"
+            )
+        trace = tokenizer.decode(out[0], skip_special_tokens=False)
+        answer = _decode_answer(
+            prompt,
+            tokenizer.decode(out[0], skip_special_tokens=True),
+            parsed_expression=(a, op, b),
+        )
+        return (answer, trace) if return_trace else answer
+
+
+_EXPRESSION_RE = re.compile(
+    r"^\s*(-?\d{1,3})\s*([+\-*/])\s*(-?\d{1,3})\s*=?\s*$"
+)
+
+
+def _parse_expression(expression: str):
+    if not isinstance(expression, str):
+        raise TypeError("expression must be a string")
+    match = _EXPRESSION_RE.fullmatch(expression)
+    if match is None:
+        raise ValueError(
+            "expected '<integer><operator><integer>' with operands from "
+            "-999 through 999"
+        )
+    a, op, b = int(match.group(1)), match.group(2), int(match.group(3))
+    prompt = f"{a}{op}{b}="
+    return a, op, b, prompt
+
+
+def _magnitude_plan(a: int, op: str, b: int):
+    if op in ("*", "/"):
+        effective = "mul" if op == "*" else "div"
+        if op == "/" and b == 0:
+            return effective, "positive", abs(a), 0
+        sign = "negative" if (a < 0) != (b < 0) else "positive"
+        return effective, sign, abs(a), abs(b)
+
+    left_term = a
+    right_term = b if op == "+" else -b
+    if left_term == 0 and right_term == 0:
+        return "add", "positive", 0, 0
+    if left_term == 0:
+        sign = "negative" if right_term < 0 else "positive"
+        return "add", sign, 0, abs(right_term)
+    if right_term == 0:
+        sign = "negative" if left_term < 0 else "positive"
+        return "add", sign, abs(left_term), 0
+    if (left_term < 0) == (right_term < 0):
+        sign = "negative" if left_term < 0 else "positive"
+        return "add", sign, abs(left_term), abs(right_term)
+    if abs(left_term) == abs(right_term):
+        return "sub", "positive", abs(left_term), abs(right_term)
+    if abs(left_term) > abs(right_term):
+        sign = "negative" if left_term < 0 else "positive"
+        return "sub", sign, abs(left_term), abs(right_term)
+    sign = "negative" if right_term < 0 else "positive"
+    return "sub", sign, abs(right_term), abs(left_term)
+
+
+def _decode_answer(
+    prompt: str,
+    decoded: str,
+    *,
+    parsed_expression=None,
+) -> str:
+    """Strictly validate a v3 scratchpad and decode its fixed answer field."""
+    if parsed_expression is None:
+        a, op, b, canonical_prompt = _parse_expression(prompt)
+    else:
+        a, op, b = parsed_expression
+        canonical_prompt = f"{a}{op}{b}="
+    if prompt != canonical_prompt:
+        raise ValueError("prompt is not in canonical form")
+    if not decoded.startswith(canonical_prompt):
+        raise ValueError("generated scratchpad does not begin with the prompt")
+
+    effective, sign, left, right = _magnitude_plan(a, op, b)
+    operation_token = {
+        "add": "<add>",
+        "sub": "<sub>",
+        "mul": "<mul>",
+        "div": "<div>",
+    }[effective]
+    sign_token = "<neg>" if sign == "negative" else "<pos>"
+    if effective == "div":
+        state = f"{left:03d}{right:03d}"
+    else:
+        state = f"{left:03d}"[::-1] + f"{right:03d}"[::-1]
+    prefix = (
+        re.escape(canonical_prompt)
+        + re.escape(operation_token + sign_token + "<state>" + state)
+    )
+
+    if effective in ("add", "sub"):
+        body = r"(?:<step>\d{5}){3}<ans>(?P<answer>\d{4})"
+    elif effective == "mul":
+        groups = []
+        for term_count in (1, 2, 3, 2, 1):
+            groups.append(rf"(?:<step>\d{{7}}){{{term_count}}}<col>\d{{3}}")
+        body = "".join(groups) + r"<ans>(?P<answer>\d{6})"
+    elif b == 0:
+        body = r"<ans>(?P<answer><nan>)"
+    else:
+        record = r"<step>\d{8}<qmul>\d{5}<rem>\d{3}"
+        body = rf"(?:{record}){{7}}<ans>(?P<answer>\d{{3}}\.\d{{3}})"
+
+    match = re.fullmatch(prefix + body, decoded)
+    if match is None:
+        raise ValueError("generated scratchpad does not match the v3 grammar")
+    internal = match.group("answer")
+    if internal == "<nan>":
+        return "NAN"
+    if op == "/":
+        integer, fraction = internal.split(".", 1)
+        fraction = fraction.rstrip("0")
+        magnitude = str(int(integer)) + (f".{fraction}" if fraction else "")
+    else:
+        magnitude = str(int(internal[::-1]))
+    if sign == "negative" and magnitude != "0":
+        return "-" + magnitude
+    return magnitude

notebook.ipynb

@@ -0,0 +1,144 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# \ud83e\uddee 3-digit-basic-calc \u2014 3-digit arithmetic with a 1.6M-parameter transformer\n",
+    "\n",
+    "A from-scratch transformer that does `+ \u2212 \u00d7 \u00f7` on 3-digit numbers by writing out\n",
+    "the algorithm step by step. This notebook loads the model and runs it.\n",
+    "\n",
+    "Model: [`vmal/3-digit-basic-calc`](https://huggingface.co/vmal/3-digit-basic-calc)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%pip install -q transformers==5.3.0 torch safetensors\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the model\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "\n",
+    "REPO = \"vmal/3-digit-basic-calc\"\n",
+    "model = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True).eval()\n",
+    "tok   = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)\n",
+    "print(\"parameters:\", sum(p.numel() for p in model.parameters()))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Solve some problems\n",
+    "\n",
+    "`model.solve(tokenizer, expression)` returns the human-readable answer.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "for e in [\"842/37\", \"213*145\", \"999*999\", \"-500+500\", \"3/31\", \"12/0\"]:\n",
+    "    print(f\"{e:>10} = {model.solve(tok, e)}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## See the model's reasoning\n",
+    "\n",
+    "Pass `return_trace=True` to get the raw scratchpad the model generates \u2014\n",
+    "each `<step>`/`<qmul>`/`<rem>` is one local computation.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "answer, trace = model.solve(tok, \"842/37\", return_trace=True)\n",
+    "print(\"answer:\", answer)\n",
+    "print()\n",
+    "for c in ['<div>','<mul>','<add>','<sub>','<state>','<step>','<qmul>','<rem>','<col>','<ans>']:\n",
+    "    trace = trace.replace(c, '\\n'+c+' ')\n",
+    "print(trace.strip())\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Quick in-range behavior check on random problems\n",
+    "\n",
+    "This small smoke test samples the supported operand range. It does not load\n",
+    "the training prompts, so it should not be described as a leakage-controlled\n",
+    "unseen benchmark; the card reports that benchmark separately.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "random.seed(0)\n",
+    "def truth(a, b, op):\n",
+    "    if op == '+': return str(a + b)\n",
+    "    if op == '-': return str(a - b)\n",
+    "    if op == '*': return str(a * b)\n",
+    "    if b == 0: return 'NAN'\n",
+    "    # Exact integer round-half-up to three decimals (no float/banker's rounding).\n",
+    "    denominator = abs(b)\n",
+    "    scaled, remainder = divmod(abs(a) * 1000, denominator)\n",
+    "    scaled += int(2 * remainder >= denominator)\n",
+    "    integer, fraction = divmod(scaled, 1000)\n",
+    "    magnitude = str(integer)\n",
+    "    if fraction:\n",
+    "        magnitude += '.' + f'{fraction:03d}'.rstrip('0')\n",
+    "    negative = (a < 0) != (b < 0)\n",
+    "    return '-' + magnitude if negative and scaled else magnitude\n",
+    "good=n=0\n",
+    "for _ in range(30):\n",
+    "    a=random.randint(-999,999); b=random.randint(-999,999); op=random.choice('+-*/')\n",
+    "    if op=='/' and b==0: b=7\n",
+    "    got=model.solve(tok, f'{a}{op}{b}'); exp=truth(a,b,op)\n",
+    "    good+= (got==exp); n+=1\n",
+    "print(f'{good}/{n} correct on random in-range problems')\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenizer.py

@@ -0,0 +1,108 @@
+# Copyright 2026. Released under the MIT license.
+"""Character/control-token tokenizer for the 3-digit-basic-calc calculator model.
+
+Digits, operators and symbols are single characters; the scratchpad grammar adds
+multi-character control tokens (``<add>``, ``<step>``, ``<qmul>`` ...). A minus
+sign is emitted as the unary token ``~`` at the start of an expression, after
+``=``, or immediately after a binary operator, so ``12*-34`` and ``5--3`` stay
+unambiguous while ``-`` remains the binary subtraction operator.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+from transformers import PreTrainedTokenizer
+
+
+SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]
+DIGITS = list("0123456789")
+OPERATORS = ["+", "-", "/", "*"]
+SYMBOLS = ["=", ".", "~"]
+CONTROL_TOKENS = [
+    "<add>", "<sub>", "<mul>", "<div>", "<pos>", "<neg>",
+    "<state>", "<step>", "<ans>", "<nan>", "<col>", "<qmul>", "<rem>",
+]
+TOKENS = SPECIAL_TOKENS + DIGITS + OPERATORS + SYMBOLS + CONTROL_TOKENS
+_UNARY_PREDECESSORS = {"=", "+", "-", "/", "*"}
+
+
+class ThreeDigitBasicCalcTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, vocab_file=None, **kwargs):
+        if vocab_file and os.path.isfile(vocab_file):
+            with open(vocab_file, encoding="utf-8") as handle:
+                self._vocab = json.load(handle)
+        else:
+            self._vocab = {tok: i for i, tok in enumerate(TOKENS)}
+        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
+        self._ordered_controls = sorted(CONTROL_TOKENS, key=len, reverse=True)
+        kwargs.setdefault("pad_token", "<pad>")
+        kwargs.setdefault("bos_token", "<bos>")
+        kwargs.setdefault("eos_token", "<eos>")
+        kwargs.setdefault("unk_token", "<unk>")
+        # Generation takes the logits from the final sequence position.  Left
+        # padding therefore keeps the final position on a real prompt token for
+        # every member of a mixed-length batch.
+        kwargs.setdefault("padding_side", "left")
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self):
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def _tokenize(self, text):
+        tokens = []
+        index = 0
+        while index < len(text):
+            control = next(
+                (t for t in self._ordered_controls if text.startswith(t, index)),
+                None,
+            )
+            if control is not None:
+                tokens.append(control)
+                index += len(control)
+                continue
+            ch = text[index]
+            if ch == "-" and (
+                    index == 0 or text[index - 1] in _UNARY_PREDECESSORS):
+                tokens.append("~")
+            else:
+                tokens.append(ch)
+            index += 1
+        return tokens
+
+    def _convert_token_to_id(self, token):
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index):
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def convert_tokens_to_string(self, tokens):
+        out = []
+        for tok in tokens:
+            if tok in SPECIAL_TOKENS:
+                continue
+            out.append("-" if tok == "~" else tok)
+        return "".join(out)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos = [self._vocab["<bos>"]]
+        if token_ids_1 is None:
+            return bos + token_ids_0
+        return bos + token_ids_0 + token_ids_1
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        path = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + "vocab.json",
+        )
+        with open(path, "w", encoding="utf-8") as handle:
+            json.dump(self._vocab, handle, ensure_ascii=False, indent=2)
+        return (path,)

tokenizer_config.json

@@ -0,0 +1,50 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "backend": "custom",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "model_max_length": 176,
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "tokenizer_class": "ThreeDigitBasicCalcTokenizer",
+  "unk_token": "<unk>",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenizer.ThreeDigitBasicCalcTokenizer",
+      null
+    ]
+  }
+}

vocab.json

@@ -0,0 +1,36 @@
+{
+  "<pad>": 0,
+  "<bos>": 1,
+  "<eos>": 2,
+  "<unk>": 3,
+  "0": 4,
+  "1": 5,
+  "2": 6,
+  "3": 7,
+  "4": 8,
+  "5": 9,
+  "6": 10,
+  "7": 11,
+  "8": 12,
+  "9": 13,
+  "+": 14,
+  "-": 15,
+  "/": 16,
+  "*": 17,
+  "=": 18,
+  ".": 19,
+  "~": 20,
+  "<add>": 21,
+  "<sub>": 22,
+  "<mul>": 23,
+  "<div>": 24,
+  "<pos>": 25,
+  "<neg>": 26,
+  "<state>": 27,
+  "<step>": 28,
+  "<ans>": 29,
+  "<nan>": 30,
+  "<col>": 31,
+  "<qmul>": 32,
+  "<rem>": 33
+}