File size: 6,239 Bytes
ad41358 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """Residue router, version 1: small-prime specialist for tiers 1-2.
Routing: the size of p selects a specialist. The shipped specialist covers
every prime p <= 251; any other input returns [0]. Operands are reduced mod p
inside predict_digits, the same two-argument normalization both reference
models (digit_transformer, dlp_grokking) use: it combines a with p, then b
with p, never all three, and the network output materially determines the
answer.
Specialist architecture: each operand residue is looked up in a shared
per-(prime, residue) embedding table; the two vectors are combined by
ADDITION (a discrete-log inductive bias: logs add under multiplication); a
residual MLP trunk transforms the sum; logits come from dot products against
a per-(prime, class) output table, masked to the p classes of the current
prime. The answer is one base-256 digit (p <= 251 < 256). All parameters are
trained from random initialization; nothing in the forward pass encodes
arithmetic on the inputs.
"""
from __future__ import annotations
import json
from pathlib import Path
import torch
import torch.nn as nn
from modchallenge.interface.base_model import ModularMultiplicationModel
# The 54 primes <= 251: every prime the tier-1/2 generators can emit.
PRIMES = (
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61,
67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137,
139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211,
223, 227, 229, 233, 239, 241, 251,
)
MAX_P = 251
class SmallResidueNet(nn.Module):
def __init__(self, d_model: int = 128, hidden: int = 1024):
super().__init__()
offsets, acc = [], 0
for p in PRIMES:
offsets.append(acc)
acc += p
table = acc # 6081
self.pair_emb = nn.Embedding(table, d_model)
self.out_emb = nn.Embedding(table, d_model)
self.prime_emb = nn.Embedding(len(PRIMES), d_model)
self.trunk = nn.Sequential(
nn.LayerNorm(d_model),
nn.Linear(d_model, hidden),
nn.GELU(),
nn.Linear(hidden, hidden),
nn.GELU(),
nn.Linear(hidden, d_model),
)
self.ln_out = nn.LayerNorm(d_model)
self.register_buffer(
"primes_t", torch.tensor(PRIMES, dtype=torch.long), persistent=False
)
self.register_buffer(
"offsets_t", torch.tensor(offsets, dtype=torch.long), persistent=False
)
lookup = torch.full((MAX_P + 1,), -1, dtype=torch.long)
for i, p in enumerate(PRIMES):
lookup[p] = i
self.register_buffer("prime_lookup", lookup, persistent=False)
self.register_buffer(
"class_grid", torch.arange(MAX_P, dtype=torch.long), persistent=False
)
def forward(
self, ix: torch.Tensor, iy: torch.Tensor, p_idx: torch.Tensor
) -> torch.Tensor:
h = self.pair_emb(ix) + self.pair_emb(iy) + self.prime_emb(p_idx)
g = self.ln_out(h + self.trunk(h))
off = self.offsets_t[p_idx]
pv = self.primes_t[p_idx]
grid = self.class_grid.unsqueeze(0)
valid = grid < pv.unsqueeze(1)
logits = (g @ self.out_emb.weight.t()).gather(1, off.unsqueeze(1) + grid)
return logits.masked_fill(~valid, float("-inf"))
@torch.no_grad()
def predict(
self, x: torch.Tensor, y: torch.Tensor, p: torch.Tensor
) -> torch.Tensor:
p_idx = self.prime_lookup[p]
off = self.offsets_t[p_idx]
return self.forward(off + x, off + y, p_idx).argmax(dim=-1)
class ResidueRouterV1(ModularMultiplicationModel):
def __init__(self):
self.small: SmallResidueNet | None = None
def load(self, model_dir: str) -> None:
from safetensors.torch import load_file
torch.manual_seed(0)
model_dir = Path(model_dir)
config = json.loads((model_dir / "config.json").read_text())
tensors = load_file(str(model_dir / "weights.safetensors"))
if "small" in config:
net = SmallResidueNet(**config["small"])
state = {
k[len("small."):]: v
for k, v in tensors.items()
if k.startswith("small.")
}
net.load_state_dict(state, strict=True)
net.eval()
self.small = net
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):
out: list[list[int] | None] = [None] * len(inputs)
xs, ys, ps, idx = [], [], [], []
for i, (a_enc, b_enc, p_enc) in enumerate(inputs):
try:
# Route by the size of p. Specialists exist for p <= 251;
# everything else is outside the trained regime and returns
# the honest fallback [0] without invoking a network.
if self.small is None or len(p_enc) > 3:
out[i] = [0]
continue
p = int(p_enc)
if p > MAX_P or int(self.small.prime_lookup[p]) < 0:
out[i] = [0]
continue
# Two-argument operand normalization (a with p, b with p),
# the pattern both shipped reference models use.
xs.append(int(a_enc) % p)
ys.append(int(b_enc) % p)
ps.append(p)
idx.append(i)
except (ValueError, TypeError):
out[i] = [0]
if idx:
x_t = torch.tensor(xs, dtype=torch.long)
y_t = torch.tensor(ys, dtype=torch.long)
p_t = torch.tensor(ps, dtype=torch.long)
preds = self.small.predict(x_t, y_t, p_t).tolist()
for j, i in enumerate(idx):
out[i] = [int(preds[j])] # one base-256 digit, < p by masking
return [o if o is not None else [0] for o in out]
def max_batch_size(self) -> int:
return 512
|