Publish SAIR Modular Arithmetic Challenge submission

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.FastModularModel",
+  "output_base": 2,
+  "framework": "pytorch",
+  "model_description": "A 260,929-parameter modulus-conditioned recurrent model with a 64-dimensional input projection and a two-layer, 96-hidden-unit bidirectional GRU. After pairwise a mod p and b mod p normalization, the learned cell repeatedly predicts the radix-2 Horner transition over binary state, multiplicand, and modulus channels, then emits the final state as base-2 digits.",
+  "training_description": "Trained independently from random initialization using AdamW on synthetic exact transitions from 8 through 2,048 bits. The curriculum mixed random moduli, small primes, wrap boundaries, on-policy Horner states, sparse powers, all-one prefixes, and width-boundary trajectories, with held-out exact-match and signed-margin checkpoint selection. The final checkpoint is a low-rate continuation hardened on anchor-biased structured trajectories. No pretrained, external, hand-initialized, or competitor weights were used."
+}

model.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from modchallenge.interface.base_model import ModularMultiplicationModel
+
+def ints_to_bits(values: list[int], device: torch.device, width: int) -> torch.Tensor:
+    """Convert nonnegative Python integers to fixed-width, MSB-first bits."""
+    byte_width = (width + 7) // 8
+    packed_bytes = bytearray().join(
+        int(value).to_bytes(byte_width, "big") for value in values
+    )
+    packed = torch.frombuffer(packed_bytes, dtype=torch.uint8)
+    packed = packed.reshape(len(values), byte_width).to(device=device)
+    shifts = torch.arange(7, -1, -1, device=device)
+    bits = ((packed[:, :, None] >> shifts) & 1).reshape(len(values), byte_width * 8)
+    return bits[:, byte_width * 8 - width :]
+
+
+def ints_to_digits(
+    values: list[int],
+    radix: int,
+    width: int,
+    device: torch.device,
+) -> torch.Tensor:
+    bits_per_digit = radix.bit_length() - 1
+    mask = radix - 1
+    rows = [
+        [
+            (value >> (bits_per_digit * position)) & mask
+            for position in range(width - 1, -1, -1)
+        ]
+        for value in values
+    ]
+    return torch.tensor(rows, dtype=torch.long, device=device)
+
+
+class TransitionCell(nn.Module):
+    def __init__(
+        self,
+        radix: int = 2,
+        dmodel: int = 32,
+        hidden: int = 64,
+        layers: int = 2,
+        bidirectional: bool = True,
+    ) -> None:
+        super().__init__()
+        self.input_projection = nn.Linear(3, dmodel)
+        self.digit_embedding = nn.Embedding(radix, dmodel)
+        self.recurrent = nn.GRU(
+            dmodel,
+            hidden,
+            num_layers=layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+        directions = 2 if bidirectional else 1
+        self.output = nn.Linear(directions * hidden, 1)
+
+    def forward(self, features: torch.Tensor, digits: torch.Tensor) -> torch.Tensor:
+        embedded = self.input_projection(features)
+        embedded = embedded + self.digit_embedding(digits)[:, None, :]
+        hidden, _ = self.recurrent(embedded)
+        return self.output(hidden).squeeze(-1)
+
+
+class FastModularModel(ModularMultiplicationModel):
+    def __init__(self) -> None:
+        self.model: TransitionCell | None = None
+        self.device: torch.device | None = None
+        self.radix = 2
+        self.max_width = 2048
+        self.bits_per_digit = 1
+
+    def load(self, model_dir: str) -> None:
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        checkpoint = torch.load(
+            Path(model_dir) / "weights.pt",
+            map_location=self.device,
+            weights_only=True,
+        )
+        config = checkpoint["config"]
+        self.radix = int(config["radix"])
+        self.bits_per_digit = self.radix.bit_length() - 1
+        self.max_width = int(checkpoint["max_width"])
+        self.model = TransitionCell(**config)
+        self.model.load_state_dict(checkpoint["state_dict"])
+        self.model.to(self.device)
+        self.model.recurrent.flatten_parameters()
+        self.model.eval()
+        if self.device.type == "cuda":
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
+    def preprocess_a(self, a: str) -> int:
+        return int(a)
+
+    def preprocess_b(self, b: str) -> int:
+        return int(b)
+
+    def preprocess_p(self, p: str) -> int:
+        return int(p)
+
+    @torch.inference_mode()
+    def predict_digits(self, a_enc, b_enc, p_enc) -> list[int]:
+        return self.predict_digits_batch([(a_enc, b_enc, p_enc)])[0]
+
+    @torch.inference_mode()
+    def predict_digits_batch(self, inputs) -> list[list[int]]:
+        output: list[list[int]] = [[0] for _ in inputs]
+        indices: list[int] = []
+        a_values: list[int] = []
+        b_values: list[int] = []
+        moduli: list[int] = []
+
+        for index, (a_enc, b_enc, p_enc) in enumerate(inputs):
+            p = int(p_enc)
+            if p < 2 or p.bit_length() > self.max_width:
+                continue
+            indices.append(index)
+            a_values.append(int(a_enc) % p)
+            b_values.append(int(b_enc) % p)
+            moduli.append(p)
+
+        if not indices:
+            return output
+
+        assert self.device is not None
+        effective_width = max(p.bit_length() for p in moduli)
+        effective_width = min(self.max_width, max(8, ((effective_width + 7) // 8) * 8))
+        digit_width = max(
+            1,
+            (max(value.bit_length() for value in b_values) + self.bits_per_digit - 1)
+            // self.bits_per_digit,
+        )
+
+        p_bits = ints_to_bits(moduli, self.device, effective_width).float()
+        x_bits = ints_to_bits(a_values, self.device, effective_width).float()
+        control_digits = ints_to_digits(
+            b_values,
+            self.radix,
+            digit_width,
+            self.device,
+        )
+        state = torch.zeros(
+            (len(indices), effective_width),
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        for position in range(digit_width):
+            state = self._step(
+                state,
+                x_bits,
+                p_bits,
+                control_digits[:, position],
+            )
+
+        rows = state.to(dtype=torch.int64).tolist()
+        for row_index, output_index in enumerate(indices):
+            output[output_index] = [int(bit) for bit in rows[row_index]]
+        return output
+
+    def max_batch_size(self) -> int:
+        return 256
+
+    def _step(
+        self,
+        state: torch.Tensor,
+        multiplicand: torch.Tensor,
+        modulus: torch.Tensor,
+        digit: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.model is not None
+        features = torch.stack((state, multiplicand, modulus), dim=-1)
+        if self.device is not None and self.device.type == "cuda":
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                logits = self.model(features, digit)
+        else:
+            logits = self.model(features, digit)
+        return (logits.float() > 0).float()

weights.pt

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