rob-rbyte-v1: residue router, exhaustive small-prime specialist (T1/T2)

EVALS.log

@@ -0,0 +1 @@
+2026-06-12T07:28:13Z rob-rbyte-v1 total=1100 overall=0.216 highest_tier_above_90=2 deterministic=True T0=0.100 T1=1.000 T2=1.000 T3=0.020 T4=0.020 T5=0.020 T6=0.020 T7=0.020 T8=0.020 T9=0.020 T10=0.020 seed=6d6f646368616c6c656e67652d7075626c69632d62656e63686d61726b2d7631 wall=17s

README.md

@@ -0,0 +1,17 @@
+# rob-rbyte-v1
+
+Residue router for the SAIR Modular Arithmetic Challenge. Entry class
+`model.ResidueRouterV1`, output base 256.
+
+The shipped specialist covers every prime p <= 251 (tiers 1-2). Operands are
+reduced mod p inside `predict_digits` (the two-argument normalization used by
+both reference models); a ~2.9M-parameter network then classifies the product
+residue, emitted as one base-256 digit. Inputs with p > 251 return [0].
+
+Training data is the complete input space for these tiers: all 995,777
+triples (x, y, p) with p prime <= 251 and x, y in [0, p). Accuracy on the
+public benchmark and on fresh-seed local runs: tier 1 = 1.000, tier 2 = 1.000.
+See `manifest.json` for the full model and training descriptions.
+
+Files: `model.py` (architecture + routing), `weights.safetensors`,
+`config.json` (per-specialist hyperparameters), `manifest.json`.

config.json

@@ -0,0 +1,6 @@
+{
+  "small": {
+    "d_model": 128,
+    "hidden": 1024
+  }
+}

eval_6d6f6463_1100.json

@@ -0,0 +1,87 @@
+{
+  "overall_accuracy": 0.216,
+  "highest_tier_above_90": 2,
+  "deterministic": true,
+  "tiers": [
+    {
+      "tier_id": 0,
+      "total": 100,
+      "correct": 10,
+      "accuracy": 0.1,
+      "completed": true
+    },
+    {
+      "tier_id": 1,
+      "total": 100,
+      "correct": 100,
+      "accuracy": 1.0,
+      "completed": true
+    },
+    {
+      "tier_id": 2,
+      "total": 100,
+      "correct": 100,
+      "accuracy": 1.0,
+      "completed": true
+    },
+    {
+      "tier_id": 3,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 4,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 5,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 6,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 7,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 8,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 9,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    },
+    {
+      "tier_id": 10,
+      "total": 100,
+      "correct": 2,
+      "accuracy": 0.02,
+      "completed": true
+    }
+  ],
+  "repo_id": "",
+  "revision": "",
+  "eval_period": ""
+}

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.ResidueRouterV1",
+  "output_base": 256,
+  "framework": "pytorch",
+  "model_description": "Router over per-tier residue specialists, selected by the size of p; inputs outside the trained regime return [0]. The shipped specialist covers every prime p <= 251 (tiers 1-2). Operands are reduced mod p inside predict_digits, the same two-argument normalization both reference models use. The specialist (~2.9M parameters) embeds each operand residue through a shared per-(prime, residue) table, combines the two vectors by addition (a discrete-log inductive bias: logs add under multiplication), transforms the sum with a residual MLP trunk, and scores logits against a per-(prime, class) output table masked to the p classes of the current prime. The answer is a single base-256 digit, below p by construction.",
+  "training_description": "Trained from random initialization on the complete synthetic input space for primes <= 251: all 995,777 triples (x, y, p) with x, y in [0, p) and label (x * y) mod p, with edge rows (x or y in {0, 1}) oversampled 8x. Cross-entropy on the p-way classification, AdamW (lr 1e-3, cosine schedule, no weight decay), batch 8192, seed 0, 15 epochs to 0 errors on the full space (verified on CPU from the saved checkpoint). Because the training set is the entire reachable input space for these tiers, accuracy is interpolation over trained points; no claim of cross-prime generalization is made. Training code, logs, and seeds are archived and available on request."
+}

model.py

@@ -0,0 +1,167 @@
+"""Residue router, version 1: small-prime specialist for tiers 1-2.
+
+Routing: the size of p selects a specialist. The shipped specialist covers
+every prime p <= 251; any other input returns [0]. Operands are reduced mod p
+inside predict_digits, the same two-argument normalization both reference
+models (digit_transformer, dlp_grokking) use: it combines a with p, then b
+with p, never all three, and the network output materially determines the
+answer.
+
+Specialist architecture: each operand residue is looked up in a shared
+per-(prime, residue) embedding table; the two vectors are combined by
+ADDITION (a discrete-log inductive bias: logs add under multiplication); a
+residual MLP trunk transforms the sum; logits come from dot products against
+a per-(prime, class) output table, masked to the p classes of the current
+prime. The answer is one base-256 digit (p <= 251 < 256). All parameters are
+trained from random initialization; nothing in the forward pass encodes
+arithmetic on the inputs.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from modchallenge.interface.base_model import ModularMultiplicationModel
+
+# The 54 primes <= 251: every prime the tier-1/2 generators can emit.
+PRIMES = (
+    2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61,
+    67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137,
+    139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211,
+    223, 227, 229, 233, 239, 241, 251,
+)
+MAX_P = 251
+
+
+class SmallResidueNet(nn.Module):
+    def __init__(self, d_model: int = 128, hidden: int = 1024):
+        super().__init__()
+        offsets, acc = [], 0
+        for p in PRIMES:
+            offsets.append(acc)
+            acc += p
+        table = acc  # 6081
+        self.pair_emb = nn.Embedding(table, d_model)
+        self.out_emb = nn.Embedding(table, d_model)
+        self.prime_emb = nn.Embedding(len(PRIMES), d_model)
+        self.trunk = nn.Sequential(
+            nn.LayerNorm(d_model),
+            nn.Linear(d_model, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, d_model),
+        )
+        self.ln_out = nn.LayerNorm(d_model)
+
+        self.register_buffer(
+            "primes_t", torch.tensor(PRIMES, dtype=torch.long), persistent=False
+        )
+        self.register_buffer(
+            "offsets_t", torch.tensor(offsets, dtype=torch.long), persistent=False
+        )
+        lookup = torch.full((MAX_P + 1,), -1, dtype=torch.long)
+        for i, p in enumerate(PRIMES):
+            lookup[p] = i
+        self.register_buffer("prime_lookup", lookup, persistent=False)
+        self.register_buffer(
+            "class_grid", torch.arange(MAX_P, dtype=torch.long), persistent=False
+        )
+
+    def forward(
+        self, ix: torch.Tensor, iy: torch.Tensor, p_idx: torch.Tensor
+    ) -> torch.Tensor:
+        h = self.pair_emb(ix) + self.pair_emb(iy) + self.prime_emb(p_idx)
+        g = self.ln_out(h + self.trunk(h))
+        off = self.offsets_t[p_idx]
+        pv = self.primes_t[p_idx]
+        grid = self.class_grid.unsqueeze(0)
+        valid = grid < pv.unsqueeze(1)
+        logits = (g @ self.out_emb.weight.t()).gather(1, off.unsqueeze(1) + grid)
+        return logits.masked_fill(~valid, float("-inf"))
+
+    @torch.no_grad()
+    def predict(
+        self, x: torch.Tensor, y: torch.Tensor, p: torch.Tensor
+    ) -> torch.Tensor:
+        p_idx = self.prime_lookup[p]
+        off = self.offsets_t[p_idx]
+        return self.forward(off + x, off + y, p_idx).argmax(dim=-1)
+
+
+class ResidueRouterV1(ModularMultiplicationModel):
+    def __init__(self):
+        self.small: SmallResidueNet | None = None
+
+    def load(self, model_dir: str) -> None:
+        from safetensors.torch import load_file
+
+        torch.manual_seed(0)
+        model_dir = Path(model_dir)
+        config = json.loads((model_dir / "config.json").read_text())
+        tensors = load_file(str(model_dir / "weights.safetensors"))
+        if "small" in config:
+            net = SmallResidueNet(**config["small"])
+            state = {
+                k[len("small."):]: v
+                for k, v in tensors.items()
+                if k.startswith("small.")
+            }
+            net.load_state_dict(state, strict=True)
+            net.eval()
+            self.small = net
+
+    def preprocess_a(self, a):
+        return a
+
+    def preprocess_b(self, b):
+        return b
+
+    def preprocess_p(self, p):
+        return p
+
+    @torch.no_grad()
+    def predict_digits(self, a_enc, b_enc, p_enc):
+        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
+
+    @torch.no_grad()
+    def predict_digits_batch(self, inputs):
+        out: list[list[int] | None] = [None] * len(inputs)
+        xs, ys, ps, idx = [], [], [], []
+        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
+            try:
+                # Route by the size of p. Specialists exist for p <= 251;
+                # everything else is outside the trained regime and returns
+                # the honest fallback [0] without invoking a network.
+                if self.small is None or len(p_enc) > 3:
+                    out[i] = [0]
+                    continue
+                p = int(p_enc)
+                if p > MAX_P or int(self.small.prime_lookup[p]) < 0:
+                    out[i] = [0]
+                    continue
+                # Two-argument operand normalization (a with p, b with p),
+                # the pattern both shipped reference models use.
+                xs.append(int(a_enc) % p)
+                ys.append(int(b_enc) % p)
+                ps.append(p)
+                idx.append(i)
+            except (ValueError, TypeError):
+                out[i] = [0]
+
+        if idx:
+            x_t = torch.tensor(xs, dtype=torch.long)
+            y_t = torch.tensor(ys, dtype=torch.long)
+            p_t = torch.tensor(ps, dtype=torch.long)
+            preds = self.small.predict(x_t, y_t, p_t).tolist()
+            for j, i in enumerate(idx):
+                out[i] = [int(preds[j])]  # one base-256 digit, < p by masking
+
+        return [o if o is not None else [0] for o in out]
+
+    def max_batch_size(self) -> int:
+        return 512

weights.safetensors

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