| --- |
| license: mit |
| library_name: pytorch |
| tags: |
| - modular-arithmetic |
| - scratchpad |
| - grokking |
| - transformer |
| - sair-challenge |
| --- |
| |
| # modmul-tier3 — learned modular multiplication |
|
|
| A submission to the **SAIR Modular Arithmetic Challenge**: compute `(a · b) mod p` |
| with the arithmetic produced entirely by **trained parameters** — no `%`, `//`, |
| Barrett, Montgomery, or CRT on the product. |
|
|
| ## What it does |
|
|
| Two trained networks behind one interface, routed by the size of `p`: |
|
|
| | `p` range | Model | |
| |---|---| |
| | `p < 512` (tiers 1–2) | joint-attention Transformer with a classification head over the residue | |
| | `512 ≤ p < 65536` (tier 3) | **autoregressive "abacus" decoder** that emits an interleaved modular-multiply scratchpad | |
| | `p ≥ 65536` | out of trained range → returns `[0]` | |
|
|
| Operands are reduced per-argument (`a mod p`, `b mod p`) before the network runs; |
| the **product's** reduction is learned. |
|
|
| ## The tier-3 idea |
|
|
| A one-shot 5-digit × 5-digit multiply is a hard learnability wall. Instead the model |
| runs **Horner's method**, reducing mod `p` at every step so no intermediate ever |
| exceeds ~6 digits — the regime the decoder handles reliably. Every intermediate is |
| written to the scratchpad and supervised: |
|
|
| ``` |
| BOS x MUL y MOD p EQ d:q1:r1:pp:t:q2:r2 STEP … EOS |
| ``` |
|
|
| per `y`-digit: shift-reduce (`q1,r1`), single-digit partial product (`pp`), |
| the explicit add `t = r1 + pp`, then add-reduce (`q2,r2`). Emitting `t` explicitly |
| was the decisive fix — leaving that addition implicit capped end-to-end accuracy. |
|
|
| ## Results |
|
|
| Verified through the official `modchallenge evaluate-hf` pipeline: |
|
|
| | Tier | Accuracy | |
| |---|---| |
| | 1 | 1.00 | |
| | 2 | 0.99 | |
| | **3** | **1.00** | |
|
|
| **`highest_tier_above_90 = 3`**, deterministic, static-compliance check passed. |
| A 30-seed local sweep gave htop90 = 3 on all 30 (tier-3 min 0.95 / mean 0.986). |
| |
| ## Files |
| |
| - `model.py` — the `EBMModMul` entry class (loader + abacus decoder) |
| - `manifest.json` — entry class, `output_base=10`, descriptions |
| - `weights.pt` — trained weights for both routed models |
| |
