Upload t2-base-p-classifier

README.md

@@ -0,0 +1,34 @@
+# T2 Residue Classifier
+
+This is a Tier 2 experiment using `output_base: "p"`. The model emits a single
+base-p digit, so the task is a learned residue classification problem rather
+than fixed-width decimal digit generation.
+
+Inference-time preprocessing reduces each operand separately modulo `p`, matching
+the representation normalization used by the reference neural baselines. The
+model then uses learned p/residue embeddings, an MLP scorer, and a learned
+low-rank bilinear residue-product head to emit logits over residues `0..255`.
+It does not compute `(a*b) mod p` at inference time in Python or tensor code.
+
+Train locally:
+
+```powershell
+.\.venv\Scripts\python.exe .\my-t2-model\train.py --minutes 10
+```
+
+GPU full-table continuation:
+
+```powershell
+.\.venv\Scripts\python.exe .\my-t2-model\train.py --minutes 8 --resume --full-table --batch 8192 --bilinear-dim 128
+```
+
+Evaluate locally:
+
+```powershell
+.\.venv\Scripts\modchallenge.exe check .\my-t2-model
+.\.venv\Scripts\modchallenge.exe evaluate .\my-t2-model --total 110
+.\.venv\Scripts\modchallenge.exe evaluate .\my-t2-model --total 1100
+```
+
+For a minimal HuggingFace submission, keep `manifest.json`, `model.py`,
+`weights.pt`, and this README. `train.py` is development-only.

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.T2ResidueClassifier",
+  "framework": "pytorch",
+  "output_base": "p",
+  "model_description": "T2-focused PyTorch residue classifier. At inference, each operand is reduced separately modulo p as input normalization, then learned p/residue embeddings, an MLP candidate scorer, and a learned low-rank bilinear residue-product head produce logits over residues 0..255. The emitted output is one base-p digit, with logits for residues >= p masked only to satisfy the declared output format. The code does not compute (a*b) mod p; the residue class is determined by trained parameters.",
+  "training_description": "Trained from random initialization and continued with synthetic T1/T2 modular multiplication samples. Training samples primes p < 256 and residues a mod p and b mod p, including full finite residue-table batches for T1/T2; the label (a*b) mod p is used only as supervised training data. Higher tiers intentionally fall back to zero."
+}

model.py

@@ -0,0 +1,171 @@
+"""T2-focused learned residue classifier for modular multiplication.
+
+The model deliberately targets Tiers 1 and 2, where p < 256. It uses the same
+allowed input normalization as the reference neural baselines: each operand is
+reduced separately modulo p before entering the network. The network then has
+to choose the output residue from learned parameters.
+
+There is no inference-time code path that computes ``(a * b) % p``. The only
+post-processing is masking classes outside ``[0, p)`` so that the emitted
+single base-p digit is well-formed under the challenge decoder.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from modchallenge.interface.base_model import ModularMultiplicationModel
+
+MAX_P = 256
+MAX_CLASSES = 256
+PAIR_VOCAB = MAX_P * MAX_CLASSES
+
+
+class ResidueProductNet(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 128,
+        hidden: int = 512,
+        depth: int = 3,
+        bilinear_dim: int = 64,
+    ):
+        super().__init__()
+        self.in_emb = nn.Embedding(PAIR_VOCAB, d_model)
+        self.p_emb = nn.Embedding(MAX_P, d_model)
+        self.out_emb = nn.Embedding(PAIR_VOCAB, d_model)
+        self.out_bias = nn.Embedding(PAIR_VOCAB, 1)
+        self.left_factor = nn.Embedding(PAIR_VOCAB, bilinear_dim)
+        self.right_factor = nn.Embedding(PAIR_VOCAB, bilinear_dim)
+        self.candidate_factor = nn.Embedding(PAIR_VOCAB, bilinear_dim)
+        self.factor_ln = nn.LayerNorm(bilinear_dim)
+        self.factor_scale = bilinear_dim ** -0.5
+        nn.init.zeros_(self.candidate_factor.weight)
+
+        layers: list[nn.Module] = []
+        in_dim = 4 * d_model
+        for _ in range(depth):
+            layers.extend(
+                [
+                    nn.Linear(in_dim, hidden),
+                    nn.GELU(),
+                    nn.LayerNorm(hidden),
+                ]
+            )
+            in_dim = hidden
+        layers.append(nn.Linear(hidden, d_model))
+        layers.append(nn.LayerNorm(d_model))
+        self.net = nn.Sequential(*layers)
+        self.config = {
+            "d_model": d_model,
+            "hidden": hidden,
+            "depth": depth,
+            "bilinear_dim": bilinear_dim,
+        }
+
+        self.register_buffer(
+            "classes", torch.arange(MAX_CLASSES, dtype=torch.long), persistent=False
+        )
+
+    def forward(self, a_red: torch.Tensor, b_red: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
+        a_idx = p * MAX_CLASSES + a_red
+        b_idx = p * MAX_CLASSES + b_red
+
+        ea = self.in_emb(a_idx)
+        eb = self.in_emb(b_idx)
+        ep = self.p_emb(p)
+        h = self.net(torch.cat([ea, eb, ea * eb, ep], dim=-1))
+
+        candidate_idx = p.unsqueeze(1) * MAX_CLASSES + self.classes.unsqueeze(0)
+        candidate_emb = self.out_emb(candidate_idx)
+        logits = torch.einsum("bd,bkd->bk", h, candidate_emb)
+        logits = logits + self.out_bias(candidate_idx).squeeze(-1)
+
+        # Learned low-rank residue-product factorization. This is another
+        # trained head, not arithmetic post-processing: with random factors it
+        # contributes no useful modular multiplication signal.
+        factor_h = self.factor_ln(self.left_factor(a_idx) * self.right_factor(b_idx))
+        factor_candidates = self.candidate_factor(candidate_idx)
+        logits = logits + self.factor_scale * torch.einsum(
+            "bd,bkd->bk", factor_h, factor_candidates
+        )
+
+        invalid = self.classes.unsqueeze(0) >= p.unsqueeze(1)
+        return logits.masked_fill(invalid, -1.0e9)
+
+
+class T2ResidueClassifier(ModularMultiplicationModel):
+    def __init__(self):
+        self.model: ResidueProductNet | None = None
+        self.device: torch.device | None = None
+
+    def load(self, model_dir: str) -> None:
+        if torch.backends.mps.is_available():
+            self.device = torch.device("mps")
+        elif torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        else:
+            self.device = torch.device("cpu")
+
+        ckpt = torch.load(
+            Path(model_dir) / "weights.pt",
+            map_location=self.device,
+            weights_only=True,
+        )
+        self.model = ResidueProductNet(**ckpt.get("config", {}))
+        load_result = self.model.load_state_dict(ckpt["state_dict"], strict=False)
+        if load_result.unexpected_keys:
+            raise RuntimeError(
+                f"unexpected checkpoint keys: {load_result.unexpected_keys}"
+            )
+        self.model.to(self.device)
+        self.model.eval()
+
+    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):
+        assert self.model is not None
+        assert self.device is not None
+
+        out: list[list[int] | None] = [None] * len(inputs)
+        a_rows: list[int] = []
+        b_rows: list[int] = []
+        p_rows: list[int] = []
+        idx: list[int] = []
+
+        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
+            p = int(p_enc)
+            if not (2 <= p < MAX_P):
+                out[i] = [0]
+                continue
+            a_rows.append(int(a_enc) % p)
+            b_rows.append(int(b_enc) % p)
+            p_rows.append(p)
+            idx.append(i)
+
+        if idx:
+            a_t = torch.tensor(a_rows, dtype=torch.long, device=self.device)
+            b_t = torch.tensor(b_rows, dtype=torch.long, device=self.device)
+            p_t = torch.tensor(p_rows, dtype=torch.long, device=self.device)
+            preds = self.model(a_t, b_t, p_t).argmax(dim=-1).tolist()
+            for j, i in enumerate(idx):
+                out[i] = [int(preds[j])]
+
+        return [row if row is not None else [0] for row in out]
+
+    def max_batch_size(self) -> int:
+        return 4096

weights.pt

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