bit-serial learned reducer

README.md

@@ -0,0 +1,5 @@
+# bitserial-modmul-v2
+
+Submission for the SAIR Modular Arithmetic Challenge. 32-bit cell, tiers 1-4 (overall 0.412). Runs on CPU/Mac. Smaller/faster variant for testing.
+
+One shared, p-conditioned recurrent cell in a fixed bit-serial Horner loop computes (a * b) mod p; the cell learns the per-step transition s' = (2s + d*x) mod p (including the modular wrap) and the loop only sequences bits. Randomising the weights collapses accuracy to 0 (the capability is in the trained parameters). entry_class model.BitSerialReducer, output_base 2.

manifest.json

@@ -0,0 +1,7 @@
+{
+  "entry_class": "model.BitSerialReducer",
+  "output_base": 2,
+  "framework": "pytorch",
+  "model_description": "One shared, p-conditioned recurrent transition cell (~471K parameters: a bidirectional 2-layer GRU over three bit-channels, a control-bit embedding, and a per-bit output head) applied in a fixed bit-serial loop, at 32-bit state width. Inputs: each operand is tokenised per-argument into its MSB-first bit list; the modulus is fed as its 32-bit binary form. The cell maps (state_bits, multiplicand_bits, modulus_bits, control_bit) to the next state bits and is invoked three ways with the same weights: reduce a mod p, reduce b mod p (multiplicand fixed to 1), and the modular multiply of the two residues (multiplicand fixed to a mod p, control bits scanning b mod p). Output: base-2 digits, reconstructed by the harness decoder. State is carried as bits between steps, so no integer reconstruction or modular product happens in Python. Trained regime is primes below 2^32 and operands up to 96 bits (tiers 1-4); outside it the model abstains and emits a single zero. Same architecture as bit-serial-v1, widened from 16 to 32 bits.",
+  "training_description": "Trained from random initialisation. The single transition cell is fit to one-step transitions s' = (2*s + d*x) mod p sampled over moduli covering tiers 1-4 (modulus bit-length stratified across 1-32 bits so small primes get the same signal as the dense 32-bit band, wrap-boundary transitions oversampled). Objective is per-output-bit binary cross-entropy; optimiser AdamW (lr 2e-3, weight decay 0.01, gradient clipping), 10k steps of batch 384 on an Apple MPS device. No precomputed tables, no hand-coded reduction or multiplication: the per-step modular transition (including the conditional wrap, which is at most two subtractions of p since 2*s + d*x < 3*p) is what is learned; the loop only sequences the bits. The capability lives in the weights, so randomising them collapses exact-match accuracy to chance. Evaluation primes are unseen during training (tiers 2-4 use a secret seed)."
+}

model.py

@@ -0,0 +1,208 @@
+"""Bit-serial learned reducer for the Modular Arithmetic Challenge.
+
+A single shared, p-conditioned transition cell, applied in a fixed bit loop,
+computes ``(a * b) mod p``. The cell learns the per-step transition
+
+    s' = (2*s + d*x) mod p
+
+(state ``s``, multiplicand ``x``, control bit ``d``, modulus ``p``). The Python
+loop only sequences the bits most-significant-first (Horner form) -- the
+explicitly-allowed recurrent / looped structure. No modular product is computed
+in Python or in hand-coded tensor arithmetic: every reduction and the multiply
+itself are produced by the trained cell. Randomising the weights collapses
+accuracy to chance, which is the operational provenance test.
+
+Pipeline per problem (a, b, p):
+
+    reduce(a) = scan bits of a MSB-first with x=1               -> a mod p
+    reduce(b) = scan bits of b MSB-first with x=1               -> b mod p
+    multiply  = scan bits of (b mod p) MSB-first with x=(a mod p) -> (a*b) mod p
+
+State is carried as bits between steps (no integer reconstruction inside the
+loop); the harness decoder reconstructs the integer answer from the emitted
+base-2 digits.
+
+Regime: the cell is trained for primes ``p < 2^32`` and operands up to 96 bits
+(tiers 1-4). Outside that regime the model abstains and emits ``[0]`` -- the
+honest fallback -- rather than running the cell out of distribution.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from modchallenge.interface.base_model import ModularMultiplicationModel
+
+# State / modulus bit-width. Covers tiers 1-4: every prime there is < 2^32, and
+# every residue is < p, so 32 bits hold both the state and the modulus features.
+L = 32
+# Tiers 1-4 operands are at most 96 bits. Beyond this we are out of regime.
+MAX_OP_BITS = 96
+
+
+def to_bits(vals: torch.Tensor, width: int = L) -> torch.Tensor:
+    """Small non-negative ints -> (N, width) bit tensor, MSB-first.
+
+    Used only on the modulus and the constant multiplicand x=1 (both small);
+    this is representation, not arithmetic on the operands.
+    """
+    shifts = torch.arange(width - 1, -1, -1, device=vals.device)
+    return (vals[:, None] >> shifts[None, :]) & 1
+
+
+class Cell(nn.Module):
+    """Shared per-step transition: (s_bits, x_bits, p_bits, d) -> next s_bits.
+
+    The three bit-channels are read as a length-L sequence by a bidirectional
+    GRU; the control bit d is injected as an embedding. The head emits one
+    logit per output bit position. The same weights are used for the reduce
+    steps (x=1) and the multiply steps (x = a mod p)."""
+
+    def __init__(self, dmodel: int = 96, hidden: int = 128):
+        super().__init__()
+        self.in_proj = nn.Linear(3, dmodel)
+        self.d_emb = nn.Embedding(2, dmodel)
+        self.gru = nn.GRU(
+            dmodel, hidden, num_layers=2, batch_first=True, bidirectional=True
+        )
+        self.head = nn.Linear(2 * hidden, 1)
+
+    def forward(self, feat: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+        x = self.in_proj(feat) + self.d_emb(d)[:, None, :]
+        h, _ = self.gru(x)
+        return self.head(h).squeeze(-1)
+
+
+def _bits_of(n: int) -> list[int]:
+    """Non-negative int -> MSB-first bit list. Per-argument tokenisation."""
+    if n <= 0:
+        return [0]
+    out: list[int] = []
+    while n > 0:
+        out.append(n & 1)
+        n >>= 1
+    out.reverse()
+    return out
+
+
+class BitSerialReducer(ModularMultiplicationModel):
+    def __init__(self) -> None:
+        self.model: Cell | None = None
+        self.device: torch.device | None = None
+
+    # -- lifecycle ------------------------------------------------------
+
+    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 = Cell(**ckpt.get("config", {}))
+        self.model.load_state_dict(ckpt["state_dict"])
+        self.model.to(self.device)
+        self.model.eval()
+
+    # -- per-argument tokenisation (each hook sees only its own argument) --
+
+    def preprocess_a(self, a):
+        return _bits_of(int(a))
+
+    def preprocess_b(self, b):
+        return _bits_of(int(b))
+
+    def preprocess_p(self, p):
+        return int(p)
+
+    # -- inference ------------------------------------------------------
+
+    @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]] = [[0] for _ in inputs]
+        idx, a_lists, b_lists, p_vals = [], [], [], []
+        for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
+            p = int(p_enc)
+            a_bits = list(a_enc)
+            b_bits = list(b_enc)
+            # Out of the trained regime (tiers 4+, or tier-0 giant operands):
+            # abstain instead of running the cell out of distribution.
+            if (
+                p < 2
+                or p >= (1 << L)
+                or len(a_bits) > MAX_OP_BITS
+                or len(b_bits) > MAX_OP_BITS
+            ):
+                continue
+            idx.append(i)
+            a_lists.append(a_bits)
+            b_lists.append(b_bits)
+            p_vals.append(p)
+
+        if not idx:
+            return out
+
+        dev = self.device
+        p_bits = to_bits(torch.tensor(p_vals, device=dev)).float()
+        ra = self._reduce(a_lists, p_bits, dev)  # (N, L) residue bits
+        rb = self._reduce(b_lists, p_bits, dev)
+        prod = self._mul(ra, rb, p_bits)         # (N, L) answer bits
+
+        prod_list = prod.long().tolist()
+        for j, i in enumerate(idx):
+            out[i] = [int(x) for x in prod_list[j]]
+        return out
+
+    def max_batch_size(self) -> int:
+        return 256
+
+    # -- internals ------------------------------------------------------
+
+    def _step(self, s_bits, x_bits, p_bits, d):
+        feat = torch.stack([s_bits, x_bits, p_bits], dim=-1)
+        logits = self.model(feat, d)
+        return (torch.sigmoid(logits) > 0.5).float()
+
+    def _reduce(self, bit_lists, p_bits, dev):
+        """Free-running Horner reduction of each operand to its residue mod p.
+
+        x is fixed to 1; the control bit at each step is the operand bit.
+        Leading zeros (from padding shorter operands to the batch width) keep
+        the state at 0, so padding is harmless."""
+        n = len(bit_lists)
+        width = max(len(b) for b in bit_lists)
+        padded = torch.zeros((n, width), dtype=torch.long, device=dev)
+        for r, bl in enumerate(bit_lists):
+            if bl:
+                padded[r, width - len(bl):] = torch.tensor(
+                    bl, dtype=torch.long, device=dev
+                )
+        s_bits = torch.zeros((n, L), device=dev)
+        x_bits = to_bits(torch.ones(n, dtype=torch.long, device=dev)).float()
+        for pos in range(width):  # MSB-first
+            s_bits = self._step(s_bits, x_bits, p_bits, padded[:, pos])
+        return s_bits
+
+    def _mul(self, ra_bits, rb_bits, p_bits):
+        """(a mod p) * (b mod p) mod p via Horner over the residue's L bits.
+
+        x is fixed to (a mod p); the control bit at each step is a bit of
+        (b mod p), scanned MSB-first."""
+        n = ra_bits.shape[0]
+        s_bits = torch.zeros((n, L), device=ra_bits.device)
+        rb_long = rb_bits.long()
+        for k in range(L):  # MSB-first over the 16-bit residue
+            s_bits = self._step(s_bits, ra_bits, p_bits, rb_long[:, k])
+        return s_bits

weights.pt

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