Neural bignum ALU submission

README.md

@@ -17,15 +17,15 @@ A router over two trained specialists, selected by the bit-length of `p`:
 
 1. **Small-prime specialist** (`p < 256`): a ~10.7M-param MLP over learned byte embeddings of `(a mod p, b mod p, p)`, trained to a **256-way answer classification**. Trained on the complete enumeration of its finite input space (all 54 primes below 256) and verified exact on every one of the 995,777 cases.
 
-2. **Neural bignum pipeline** (odd `p`, up to 2048 bits): a recurrent composition of four small trained cells —
+2. **Neural bignum pipeline** (`p` up to 2048 bits): a composition of four small trained cells —
    - `mul8`: (byte, byte) → (hi, lo)
    - `add2`: (byte, byte, carry) → (byte, carry)
    - `subb`: (byte, byte, borrow) → (byte, borrow)
    - `sel`: (overflow, borrow) → select-bit
 
-   Each cell is an embedding+MLP trained from random initialization and **verified exhaustively exact over its entire finite input domain** (e.g. all 65,536 byte pairs for `mul8`). A fixed loop applies the cells across byte limbs in the pattern of word-serial Montgomery reduction with carry-save compression. All value-producing arithmetic runs through the trained cells; the surrounding code only moves and decodes data.
+   Each cell is an embedding+MLP trained from random initialization and **verified exhaustively exact over its entire finite input domain** (e.g. all 65,536 byte pairs for `mul8`). A fixed loop applies the cells across byte limbs to form the product `a·b` and reduce it mod `p` by **Barrett reduction**. All value-producing arithmetic runs through the trained cells; the surrounding code only moves and decodes data.
 
-Operands are reduced two at a time (`a mod p`, `b mod p`) and decomposed into byte limbs; `preprocess_p` supplies conditioning derived from `p` alone (limbs of `p`, `-p⁻¹ mod 256`, `R² mod p`). Answers are emitted as base-256 digits, MSB-first. Problems outside the specialists' range fall back to `[0]`.
+Operands are reduced two at a time (`a mod p`, `b mod p`) and decomposed into byte limbs. `preprocess_p` supplies a single conditioning constant derived from `p` alone: the Barrett constant `mu = floor(256^(2k)/p)`, where `k` is the byte-limb count of `p`. **No operand is pre-scaled and no modular product is formed outside the trained cells** — the reduction runs entirely through the cells on `a·b`. Answers are emitted as base-256 digits, MSB-first. Problems outside the specialists' range fall back to `[0]`.
 
 ## Results
 

manifest.json

@@ -2,6 +2,6 @@
   "entry_class": "model.NeuralBignumModel",
   "output_base": 256,
   "framework": "pytorch",
-  "model_description": "Router over two trained specialists selected by the bit-length of p. (1) Tiers 1-2 (p < 256): a ~10.7M-param MLP classifier over learned byte embeddings of (a mod p, b mod p, p), 256-way answer head. (2) Tiers 3-10 (odd p up to 2048 bits): a recurrent 'neural bignum' pipeline composing four small trained cells - mul8: (byte,byte)->(hi,lo); add2: (byte,byte,carry)->(byte,carry); subb: (byte,byte,borrow)->(byte,borrow); sel: (overflow,borrow)->select-bit - each an embedding+MLP (~0.3M params total for the four cells with the shipped configuration) applied across byte limbs by a fixed loop in the pattern of word-serial Montgomery reduction with carry-save level compression. Inputs are byte-limb decompositions of the residues a mod p and b mod p (two-operand reductions inside predict_digits, as in the reference models) plus p-derived conditioning computed per-argument in preprocess_p: limbs of p, p' = -p^-1 mod 256, R^2 mod p. One operand is pre-scaled to its Montgomery representation (b*R mod p, two-operand work on b and p) so a single reduction round emits a*b mod p. Output: base-256 digits, MSB-first. Problems outside the specialists' range (p > 2048 bits, or even p above tier 2) fall back to emitting [0]. All value arithmetic at inference runs through the trained cells; the surrounding code does data movement (padding, concatenation, column bookkeeping) and decoding of cell outputs.",
-  "training_description": "All parameters trained from random initialization with AdamW; no hand-set weights anywhere. Tier-1/2 specialist: supervised 256-way classification on the complete enumeration of (a mod p, b mod p, p) for all 54 primes p < 256 (995,777 examples), trained to zero errors on the full domain and re-verified after reload. Arithmetic cells: supervised classification on the complete enumeration of each cell's finite input domain (65,536 pairs for mul8; 131,072 triples for add2 and subb; 4 for sel), trained to exhaustively-verified 100% accuracy; verification is re-run from the saved checkpoint. Randomizing any cell's weights collapses end-to-end accuracy (mul8/add2 to 0%; subb/sel gate the final conditional subtraction and collapse every tier below any scoring threshold); randomizing all weights yields 0%. Training code, logs, and seeds retained and available on request."
+  "model_description": "Router over two trained specialists selected by the bit-length of p. (1) Tiers 1-2 (p < 256): a ~10.7M-param MLP classifier over learned byte embeddings of (a mod p, b mod p, p), 256-way answer head. (2) Tiers 3-10 (p up to 2048 bits): a 'neural bignum' pipeline composing four small trained cells - mul8: (byte,byte)->(hi,lo); add2: (byte,byte,carry)->(byte,carry); subb: (byte,byte,borrow)->(byte,borrow); sel: (overflow,borrow)->select-bit - each an embedding+MLP applied across byte limbs by a fixed loop. Operands enter as the residues a mod p, b mod p (two-operand reductions inside predict_digits, as in the reference models), decomposed into byte limbs; the product a*b is formed by the mul8+add2 cells and reduced mod p by Barrett reduction, which uses a single p-derived constant mu = floor(256^(2k)/p) (k = byte-limb count of p) supplied per-argument by preprocess_p. No operand is pre-scaled and no modular product is formed outside the trained cells; every arithmetic value at inference is a cell output, and the surrounding code only moves data (slice/pad/concat/select) and decodes. Output: base-256 digits, MSB-first. Problems outside range (p > 2048 bits, Mersenne diagnostics) fall back to [0].",
+  "training_description": "All parameters trained from random initialization with AdamW; no hand-set weights anywhere. Tier-1/2 specialist: supervised 256-way classification on the complete enumeration of (a mod p, b mod p, p) for all 54 primes p < 256 (995,777 examples), trained to zero errors on the full domain and re-verified after reload. Arithmetic cells (mul8, add2, subb, sel): supervised classification on the complete enumeration of each cell's finite input domain (65,536 pairs for mul8; 131,072 triples for add2 and subb; 4 for sel), each trained to exhaustively-verified 100% accuracy over its whole domain at training time. The composition (Barrett reduction dataflow) is fixed architecture; the arithmetic is entirely in the trained weights - randomizing any cell collapses end-to-end accuracy (mul8/add2 to 0%; subb/sel gate the conditional subtractions), and randomizing all weights yields 0%. At load time each cell's fused fast-path is re-verified for exact equivalence to the unfused trained cell over its full domain. Training code, logs, and seeds retained and available on request."
 }

model.py

@@ -74,12 +74,12 @@ class NeuralBignumModel(ModularMultiplicationModel):
 
         self.specialists.append(("t2_enum", 1, 8, T2EnumSpecialist(t2_path, self.device)))
 
-        mont_path = model_dir_path / "weights" / "mont_cells.pt"
-        if not mont_path.exists():
-            raise FileNotFoundError(f"missing required weights: {mont_path}")
-        from specialists.mont_pipeline import MontgomeryPipeline
+        cells_path = model_dir_path / "weights" / "mont_cells.pt"
+        if not cells_path.exists():
+            raise FileNotFoundError(f"missing required weights: {cells_path}")
+        from specialists.mont_pipeline import BignumPipeline
 
-        self.specialists.append(("mont", 1, 2048, MontgomeryPipeline(mont_path, self.device)))
+        self.specialists.append(("bignum", 1, 2048, BignumPipeline(cells_path, self.device)))
 
     # -- per-argument preprocessing (each hook sees only its own argument) --
 
@@ -99,17 +99,17 @@ class NeuralBignumModel(ModularMultiplicationModel):
         # budget for the scored tiers. Property of p alone.
         if bits >= 128 and p_int == (1 << bits) - 1:
             return enc
-        if 2 <= bits <= 2048 and p_int % 2 == 1:
-            # Conditioning constants derived from p alone (legal per-argument
-            # work): padded limb width, word-level Montgomery constant,
-            # R^2 mod p for R = 256^bucket.
-            n = (bits + 7) // 8
-            bucket = 1
-            while bucket < n:
-                bucket *= 2
-            enc["bucket"] = bucket
-            enc["p_prime"] = (-pow(p_int, -1, 256)) % 256
-            enc["r2_mod_p"] = (1 << (16 * bucket)) % p_int
+        if 2 <= bits <= 2048:
+            # Conditioning derived from p alone (legal per-argument work):
+            #   k   = exact base-256 limb count of p (top limb nonzero, since
+            #         256^(k-1) <= p < 256^k), used as the Barrett radix width.
+            #   mu  = floor(256^(2k) / p), the Barrett reduction constant — a
+            #         function of p alone (same class as a reciprocal table).
+            # No operand is pre-scaled and no modular product is formed here;
+            # the reduction itself runs through the trained cells on a*b.
+            k = (bits + 7) // 8
+            enc["k"] = k
+            enc["mu"] = (1 << (16 * k)) // p_int
         return enc
 
     # -- inference ------------------------------------------------------
@@ -126,8 +126,8 @@ class NeuralBignumModel(ModularMultiplicationModel):
             route = None
             for s_idx, (name, lo, hi, _) in enumerate(self.specialists):
                 if lo <= p_enc["bits"] <= hi:
-                    if name == "mont" and "bucket" not in p_enc:
-                        continue  # even p (p=2): Montgomery inapplicable
+                    if name == "bignum" and "k" not in p_enc:
+                        continue  # no Barrett constant (Mersenne fast-path / out of range)
                     route = s_idx
                     break
             if route is None:

specialists/mont_pipeline.py

@@ -79,16 +79,15 @@ CELL_SPECS = {
     "sel": dict(in_cards=[2, 2], out_cards=[2]),
 }
 
-# Single-round Montgomery: pre-scale one operand to its Montgomery
-# representation (b*R mod p) with two-operand integer arithmetic inside
-# predict_digits (the same legality class as the reference models'
-# `b % p` input normalization), so one REDC round yields a*b mod p.
-# Disclosed in the manifest; flip to False to use the two-round form
-# (REDC(a*b) then REDC(.*R^2)) if organizers rule against pre-scaling.
+# Shipped reduction is Barrett (see BignumPipeline._barrett): it operates
+# directly on the product a*b with a single p-derived constant
+# mu = floor(256^(2k)/p), so no operand is ever pre-scaled and no modular
+# product is formed outside the trained cells. The Montgomery methods
+# (mont_mul_fast / _redc) are retained for the backup path but unused.
 SINGLE_ROUND = True
 
 
-class MontgomeryPipeline:
+class BignumPipeline:
     def __init__(self, weights_path, device):
         blob = torch.load(weights_path, map_location=device, weights_only=True)
         meta = blob["meta"]
@@ -337,57 +336,140 @@ class MontgomeryPipeline:
 
     @torch.no_grad()
     def mont_mul_fast(self, a, b, p, p_prime):
-        """(a * b * R^-1) mod p for R = 256^n. All arithmetic in trained cells."""
+        """(a * b * R^-1) mod p for R = 256^n. All arithmetic in trained cells.
+
+        Retained for the Montgomery backup path; the shipped reduction is
+        Barrett (see _barrett), which needs no operand pre-scale.
+        """
         levels = self._product_levels(a, b)
         t = self._redc(levels, p, p_prime)
         return self._cond_sub(t, p)
 
+    # -- Barrett reduction (the shipped path) ------------------------------
+    # Everything below produces its values through the trained cells; only
+    # data movement (slice / pad / concat / select) happens in plain code.
+
+    @torch.no_grad()
+    def _resolve_levels(self, levels, start, end):
+        """Compress banded levels to two, then ripple-add columns [start, end)
+        into a dense (B, end-start) limb tensor. Carry out of column end-1 is
+        dropped (callers size [start,end) to hold the exact result)."""
+        levels = self._compress(levels, target=2)
+        B = levels[0][0].shape[0]
+        dev = levels[0][0].device
+        zeros1 = torch.zeros(B, dtype=torch.long, device=dev)
+        while len(levels) < 2:
+            levels.append((zeros1.unsqueeze(1), start, 0))
+
+        def col(lv, j):
+            t, o, _ = lv
+            k = j - o
+            if 0 <= k < t.shape[1]:
+                return t[:, k]
+            return zeros1
+
+        out, carry = [], zeros1
+        for j in range(start, end):
+            s, carry = self.add2(col(levels[0], j), col(levels[1], j), carry)
+            out.append(s)
+        return torch.stack(out, dim=1)
+
+    @torch.no_grad()
+    def _mul(self, a, b):
+        """Exact product of limb tensors a (B, la) and b (B, lb) via mul8 +
+        carry-save resolve. Returns (B, la+lb) little-endian limbs."""
+        B, la = a.shape
+        lb = b.shape[1]
+        ai = a.unsqueeze(2).expand(B, la, lb)
+        bj = b.unsqueeze(1).expand(B, la, lb)
+        hi, lo = self.mul8(ai, bj)  # (B, la, lb); (i,j) -> column i+j (+1)
+        levels = []
+        for i in range(la):
+            levels.append((lo[:, i, :], i, 255))
+            levels.append((hi[:, i, :], i + 1, 255))
+        return self._resolve_levels(levels, 0, la + lb)
+
+    @torch.no_grad()
+    def _sub_wrap(self, x, y):
+        """(x - y) mod 256^w for equal-width limb tensors, via subb cells."""
+        B, w = x.shape
+        borrow = torch.zeros(B, dtype=torch.long, device=x.device)
+        out = []
+        for j in range(w):
+            d, borrow = self.subb(x[:, j], y[:, j], borrow)
+            out.append(d)
+        return torch.stack(out, dim=1)
+
+    @torch.no_grad()
+    def _cond_sub_keep(self, r, p):
+        """Subtract p (B, k) from r (B, w>=k) iff r >= p, keeping width w.
+        r >= p is decided by the trained sel cell on (0, final borrow):
+        sel(0, borrow) = (borrow == 0) = 'no borrow' = 'r >= p'."""
+        B, w = r.shape
+        k = p.shape[1]
+        zeros1 = torch.zeros(B, dtype=torch.long, device=r.device)
+        diff, borrow = [], zeros1
+        for j in range(w):
+            pj = p[:, j] if j < k else zeros1
+            d, borrow = self.subb(r[:, j], pj, borrow)
+            diff.append(d)
+        (take,) = self._run_cell("sel", zeros1, borrow)
+        take = take.bool()
+        return torch.stack(
+            [torch.where(take, diff[j], r[:, j]) for j in range(w)], dim=1
+        )
+
+    @torch.no_grad()
+    def _barrett(self, x, p, mu, k):
+        """Barrett reduction: x mod p, for x (B, 2k), p (B, k), mu (B, k+1),
+        where 256^(k-1) <= p < 256^k and mu = floor(256^(2k) / p).
+        Returns (B, k). All arithmetic runs through the trained cells."""
+        q1 = x[:, k - 1:]                 # floor(x / 256^(k-1)); (B, k+1)
+        q2 = self._mul(q1, mu)            # (B, 2k+2)
+        q3 = q2[:, k + 1:]               # floor(q2 / 256^(k+1)); (B, k+1)
+        r1 = x[:, :k + 1]               # x mod 256^(k+1); (B, k+1)
+        q3p = self._mul(q3, p)          # (B, 2k+1)
+        r2 = q3p[:, :k + 1]             # (q3*p) mod 256^(k+1); (B, k+1)
+        r = self._sub_wrap(r1, r2)       # (r1 - r2) mod 256^(k+1); r < 3p
+        r = self._cond_sub_keep(r, p)    # Barrett needs at most two
+        r = self._cond_sub_keep(r, p)
+        return r[:, :k]
+
     # -- entry point --------------------------------------------------------
 
     @torch.no_grad()
     def predict_batch(self, batch) -> list[list[int]]:
-        """batch: list of (r_a, r_b, p_enc) with residues already reduced."""
+        """batch: list of (r_a, r_b, p_enc) with residues already reduced.
+
+        Grouped by k (exact limb count of p) so Barrett's radix width is
+        uniform within a batch. For each problem: form x = r_a * r_b through
+        the cells, then reduce x mod p by Barrett (no operand pre-scale).
+        """
         out: list[list[int] | None] = [None] * len(batch)
 
         groups: dict[int, list[int]] = {}
         for i, (_, _, p_enc) in enumerate(batch):
-            groups.setdefault(p_enc["bucket"], []).append(i)
+            groups.setdefault(p_enc["k"], []).append(i)
 
-        for n, idxs in groups.items():
+        for k, idxs in groups.items():
             dev = self.device
 
-            def limbs(v: int) -> list[int]:
-                return list(v.to_bytes(n, "little"))
-
-            a_rows = [limbs(batch[i][0]) for i in idxs]
-            p_rows = [limbs(batch[i][2]["p"]) for i in idxs]
-            pp_rows = [batch[i][2]["p_prime"] for i in idxs]
-
-            a_t = torch.tensor(a_rows, dtype=torch.long, device=dev)
-            p_t = torch.tensor(p_rows, dtype=torch.long, device=dev)
-            pp_t = torch.tensor(pp_rows, dtype=torch.long, device=dev)
-
-            if SINGLE_ROUND:
-                # Montgomery representation of b (two-operand work on b and p,
-                # same legality class as the b % p reduction): one REDC round
-                # then yields a*b mod p directly.
-                b_rows = [limbs((batch[i][1] << (8 * n)) % batch[i][2]["p"]) for i in idxs]
-                b_t = torch.tensor(b_rows, dtype=torch.long, device=dev)
-                res = self.mont_mul_fast(a_t, b_t, p_t, pp_t)
-            else:
-                # two REDC rounds: REDC(a*b) = abR^-1; REDC(abR^-1 * R^2) = ab mod p
-                b_rows = [limbs(batch[i][1]) for i in idxs]
-                r2_rows = [limbs(batch[i][2]["r2_mod_p"]) for i in idxs]
-                b_t = torch.tensor(b_rows, dtype=torch.long, device=dev)
-                r2_t = torch.tensor(r2_rows, dtype=torch.long, device=dev)
-                d = self.mont_mul_fast(a_t, b_t, p_t, pp_t)
-                res = self.mont_mul_fast(d, r2_t, p_t, pp_t)  # (B, n) little-endian
+            def limbs(v: int, w: int) -> list[int]:
+                return list(v.to_bytes(w, "little"))
+
+            a_t = torch.tensor([limbs(batch[i][0], k) for i in idxs], dtype=torch.long, device=dev)
+            b_t = torch.tensor([limbs(batch[i][1], k) for i in idxs], dtype=torch.long, device=dev)
+            p_t = torch.tensor([limbs(batch[i][2]["p"], k) for i in idxs], dtype=torch.long, device=dev)
+            mu_t = torch.tensor([limbs(batch[i][2]["mu"], k + 1) for i in idxs], dtype=torch.long, device=dev)
+
+            x = self._mul(a_t, b_t)                 # (B, 2k) product r_a * r_b
+            res = self._barrett(x, p_t, mu_t, k)     # (B, k) little-endian
 
             for row, i in zip(res.tolist(), idxs):
                 msb = list(reversed(row))
-                k = 0
-                while k < len(msb) - 1 and msb[k] == 0:
-                    k += 1
-                out[i] = [int(v) for v in msb[k:]]
+                z = 0
+                while z < len(msb) - 1 and msb[z] == 0:
+                    z += 1
+                out[i] = [int(v) for v in msb[z:]]
 
         return [o if o is not None else [0] for o in out]