diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/config.py | 21 | ||||
| -rw-r--r-- | src/data.py | 134 | ||||
| -rw-r--r-- | src/model_api.py | 4 | ||||
| -rw-r--r-- | src/models/transformer.py | 23 | ||||
| -rw-r--r-- | src/train.py | 18 |
5 files changed, 164 insertions, 36 deletions
diff --git a/src/config.py b/src/config.py index 945c259..915cdae 100644 --- a/src/config.py +++ b/src/config.py @@ -25,6 +25,9 @@ class Config: min_steps: int = 2 # ACT: halt prob forced to 0 for the first min_steps-1 steps n_layers: int = 2 n_heads: int = 4 + # scratchpad / token-space recurrence (Addendum 8, E8) + scratch_mode: str = "none" # "none" | "structured" | "filler" | "random_learned" | "random_frozen" | "random_noise" + scratch_len: int = 16 # length of filler / random token sequence # training lr: float = 1e-3 lr_decay: bool = False # cosine 1e-3 -> 1e-4 over the run (E3) @@ -45,7 +48,11 @@ class Config: @property def vocab(self) -> int: - """Token count. Integers mode: value tokens 0..next_prime(range_end), plus EOS (+1) and pad (+1).""" + """Token count. + - Integers mode: value tokens 0..next_prime(range_end) + EOS + pad. + - Digits mode (none): 10 digits + EOS = 11. + - Digits mode (scratchpad): 10 digits + EOS + SEP + PAUSE + 16 random (a-p) + 4 structured symbols (c,=,d,:) + 1 noise slot = 34. + """ if self.vocab_mode == "integers": # inline mini-sieve: next prime above range_end bounds the target domain limit = self.range_end + 100 @@ -57,15 +64,21 @@ class Config: is_prime[m] = False np_ = next(i for i in range(self.range_end + 1, limit + 1) if is_prime[i]) return np_ + 2 # values 0..np_, EOS=np_+1, pad=np_+2 - return 11 # digits 0-9 + EOS + if self.scratch_mode != "none": + return 34 # expanded vocabulary for scratchpad/filler/random tokens + return 11 # digits 0-9 + EOS @property def eos_id(self) -> int: - return self.vocab - 1 + if self.vocab_mode == "integers": + return self.vocab - 1 + if self.scratch_mode != "none": + return 10 # fixed ID 10 for EOS in digits mode + return self.vocab - 1 # 10 for digits mode none @property def pad_id(self) -> int: - return self.vocab # one extra embedding row reserved for pad + return self.vocab # one extra embedding row reserved for pad def to_json(self) -> dict: d = {f.name: getattr(self, f.name) for f in fields(self)} diff --git a/src/data.py b/src/data.py index 26fb6b9..411ba7e 100644 --- a/src/data.py +++ b/src/data.py @@ -7,6 +7,18 @@ import torch from src.config import Config +# Special token constants for digits mode (Addendum 8, E8) +EOS_ID = 10 +SEP_ID = 11 # '#' separator between scratchpad and final target +PAUSE_ID = 12 # '<p>' pause/filler token (Arm C) +RAND_START_ID = 13 # 'a' (token IDs 13..28 represent 'a'..'p' for Arms D1/D2) +SYM_C = 29 # 'c' +SYM_EQ = 30 # '=' +SYM_D = 31 # 'd' +SYM_COLON = 32 # ':' +NOISE_SLOT_ID = 33 # '<noise>' placeholder for Arm D3 continuous dynamic noise + + def sieve_primes(limit: int) -> list[int]: """All primes <= limit (inclusive).""" if limit < 2: @@ -33,15 +45,51 @@ def encode_int(n: int, cfg: Config) -> list[int]: return [int(d) for d in str(n)] +def make_structured_trace(n: int, target_prime: int, primes_list: list[int]) -> list[int]: + """Generates trace tokens for candidate search and trial division: + For each candidate c in [n+1 .. target_prime]: + c, =, digits(c), d, digits(p), :, 0/1, ... + """ + tokens = [] + for c in range(n + 1, target_prime + 1): + tokens.append(SYM_C) + tokens.append(SYM_EQ) + tokens.extend([int(d) for d in str(c)]) + for p in primes_list: + if p * p > c: + break + tokens.append(SYM_D) + tokens.extend([int(d) for d in str(p)]) + tokens.append(SYM_COLON) + if c % p == 0: + tokens.append(0) # divisible -> composite found, halt checks for c + break + else: + tokens.append(1) # not divisible -> check next prime + return tokens + + def decode_tokens(ts, cfg: Config) -> int: - """Decode a token sequence, stopping at EOS. -1 if nothing decodable.""" + """Decode a token sequence, stopping at EOS. If SEP_ID is present, decodes digits AFTER SEP_ID. + -1 if nothing decodable.""" if cfg.vocab_mode == "integers": return int(ts[0]) if len(ts) else -1 + + if cfg.scratch_mode != "none": + if SEP_ID in ts: + idx = ts.index(SEP_ID) + ts = ts[idx + 1:] + else: + return -1 + digits = [] for t in ts: if t == cfg.eos_id: break - digits.append(int(t)) + if 0 <= t <= 9: + digits.append(int(t)) + else: + break return int("".join(map(str, digits))) if digits else -1 @@ -67,7 +115,6 @@ def get_splits(cfg: Config) -> tuple[list[int], list[int]]: train, val = sorted(inputs[n_val:]), sorted(inputs[:n_val]) if cfg.train_frac < 1.0: n_tr = max(1, round(len(train) * cfg.train_frac)) - # deterministic subsample: seeded shuffle, take first n_tr sub = random.Random(cfg.seed + 1000) # distinct stream from split shuffle sub.shuffle(train) train = sorted(train[:n_tr]) @@ -75,29 +122,79 @@ def get_splits(cfg: Config) -> tuple[list[int], list[int]]: def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]: - """[(input_tokens, target_tokens+EOS), ...]. task_mode selects the target function.""" + """[(input_tokens, target_tokens+EOS), ...]. task_mode and scratch_mode select format.""" margin = max(100, int(cfg.range_end * 0.05) + 50) - primes = sieve_primes(cfg.range_end + margin) if cfg.task_mode != "is_prime" else [] + primes = sieve_primes(cfg.range_end + margin) if (cfg.task_mode != "is_prime" or cfg.scratch_mode == "structured") else [] out = [] for n in inputs: + x_toks = encode_int(n, cfg) if cfg.task_mode == "is_prime": p = 1 if is_prime_n(n) else 0 + ans_toks = encode_int(p, cfg) + [cfg.eos_id] + if cfg.scratch_mode == "none": + out.append((x_toks, ans_toks)) + elif cfg.scratch_mode == "filler": + filler = [PAUSE_ID] * cfg.scratch_len + y_toks = filler + [SEP_ID] + ans_toks + out.append((x_toks, y_toks)) + else: + out.append((x_toks, ans_toks)) else: p = next_prime(n, primes) - out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id])) + ans_toks = encode_int(p, cfg) + [cfg.eos_id] + if cfg.scratch_mode == "none": + out.append((x_toks, ans_toks)) + elif cfg.scratch_mode == "structured": + trace = make_structured_trace(n, p, primes) + y_toks = trace + [SEP_ID] + ans_toks + out.append((x_toks, y_toks)) + elif cfg.scratch_mode == "filler": + filler = [PAUSE_ID] * cfg.scratch_len + y_toks = filler + [SEP_ID] + ans_toks + out.append((x_toks, y_toks)) + elif cfg.scratch_mode in ("random_learned", "random_frozen"): + rng = random.Random(cfg.seed + n * 37) + rand_toks = [RAND_START_ID + rng.randint(0, 15) for _ in range(cfg.scratch_len)] + y_toks = rand_toks + [SEP_ID] + ans_toks + out.append((x_toks, y_toks)) + elif cfg.scratch_mode == "random_noise": + noise_toks = [NOISE_SLOT_ID] * cfg.scratch_len + y_toks = noise_toks + [SEP_ID] + ans_toks + out.append((x_toks, y_toks)) + else: + out.append((x_toks, ans_toks)) return out def _global_lengths(cfg: Config) -> tuple[int, int]: - """(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix).""" + """(in_max, out_max): fixed global lengths so batch layout == singleton layout.""" if cfg.vocab_mode == "integers": - return 1, 2 # [value], [value, EOS] + return 1, 2 + in_max = len(str(cfg.range_end)) if cfg.task_mode == "is_prime": - return len(str(cfg.range_end)), 2 # digit token "1"/"0" + EOS - margin = max(100, int(cfg.range_end * 0.05) + 50) - primes = sieve_primes(cfg.range_end + margin) - max_target = next_prime(cfg.range_end, primes) - return len(str(cfg.range_end)), len(str(max_target)) + 1 # digits + EOS + base_out = 2 + else: + margin = max(100, int(cfg.range_end * 0.05) + 50) + primes = sieve_primes(cfg.range_end + margin) + max_target = next_prime(cfg.range_end, primes) + base_out = len(str(max_target)) + 1 + + if cfg.scratch_mode == "none": + return in_max, base_out + elif cfg.scratch_mode in ("filler", "random_learned", "random_frozen", "random_noise"): + return in_max, cfg.scratch_len + 1 + base_out + elif cfg.scratch_mode == "structured": + margin = max(100, int(cfg.range_end * 0.05) + 50) + primes = sieve_primes(cfg.range_end + margin) + # sample max trace length across the entire range + max_trace_len = 0 + for n in range(cfg.range_start, cfg.range_end + 1): + p = next_prime(n, primes) + tr = make_structured_trace(n, p, primes) + if len(tr) > max_trace_len: + max_trace_len = len(tr) + return in_max, max_trace_len + 1 + base_out + return in_max, base_out def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor: @@ -110,16 +207,19 @@ def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor: def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]: - """Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max.""" - xs, ys = zip(*examples) + """Fixed global layout: x left-padded to in_max, y right-padded to out_max.""" in_max, out_max = _global_lengths(cfg) B = len(examples) x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long) y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long) + loss_mask = torch.zeros((B, out_max), dtype=torch.bool) for i, (xi, yi) in enumerate(examples): x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long) y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long) - # teacher-forced decoder input: BOS(=EOS reuse) then shifted y + if cfg.scratch_mode in ("filler", "random_learned", "random_frozen", "random_noise"): + loss_mask[i, cfg.scratch_len : len(yi)] = True + else: + loss_mask[i, : len(yi)] = True y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1) y_mask = y != cfg.pad_id - return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask} + return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask, "loss_mask": loss_mask} diff --git a/src/model_api.py b/src/model_api.py index 85b5d55..07d5d8a 100644 --- a/src/model_api.py +++ b/src/model_api.py @@ -39,7 +39,9 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int if x.device != device: x = x.to(device) x = pad_inputs(x, cfg) - max_len = max_len or cfg.max_out_len + from src.data import _global_lengths + _, out_max = _global_lengths(cfg) + max_len = max_len or max(getattr(cfg, "max_out_len", 6), out_max) B = x.shape[0] y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device) for _ in range(max_len): diff --git a/src/models/transformer.py b/src/models/transformer.py index 32e158b..2b986ce 100644 --- a/src/models/transformer.py +++ b/src/models/transformer.py @@ -32,7 +32,7 @@ class TransformerBaseline(PrimeModel): self.cfg = cfg d = cfg.d_model self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad - self.pos = nn.Embedding(64, d) # 3 in + 6 out max, generous + self.pos = nn.Embedding(256, d) # supports sequences up to 256 self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)]) self.ln_f = nn.LayerNorm(d) self.head = nn.Linear(d, cfg.vocab) @@ -44,14 +44,25 @@ class TransformerBaseline(PrimeModel): seq = torch.cat([x, y_in], dim=1) # (B, T_in+T_out) T = seq.shape[1] device = seq.device - e = self.embed(seq) + self.pos(torch.arange(T, device=device)).unsqueeze(0) + + # Token embeddings + tok_embed = self.embed(seq) + if cfg.scratch_mode == "random_frozen": + # detach embeddings for tokens 'a'..'p' (IDs 13..28) + is_rand = (seq >= 13) & (seq <= 28) + tok_embed = torch.where(is_rand.unsqueeze(-1), tok_embed.detach(), tok_embed) + elif cfg.scratch_mode == "random_noise": + # inject fresh Gaussian noise at NOISE_SLOT_ID (ID 33) + is_noise = (seq == 33) + noise = torch.randn_like(tok_embed) + tok_embed = torch.where(is_noise.unsqueeze(-1), noise, tok_embed) + + e = tok_embed + self.pos(torch.arange(T, device=device)).unsqueeze(0) # causal mask: input positions (j < T_in) fully visible; output positions causal # (True = blocked, per torch.nn.MultiheadAttention bool convention) blocked = torch.zeros(T, T, dtype=torch.bool, device=device) - for i in range(T): - for j in range(T): - if j >= T_in and j > i: - blocked[i, j] = True + causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1) + blocked[:, T_in:] = causal[:, T_in:] key_pad = seq == cfg.pad_id # (B,T) True = ignore h = e for blk in self.blocks: diff --git a/src/train.py b/src/train.py index a0874c4..2523325 100644 --- a/src/train.py +++ b/src/train.py @@ -38,6 +38,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach "y": torch.empty((0, out_max), dtype=torch.long, device=device), "y_in": torch.empty((0, out_max), dtype=torch.long, device=device), "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device), + "loss_mask": torch.empty((0, out_max), dtype=torch.bool, device=device), "size": 0, "is_cached": True, } @@ -50,6 +51,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach "y": batch["y"].to(device, non_blocking=True), "y_in": batch["y_in"].to(device, non_blocking=True), "y_mask": batch["y_mask"].to(device, non_blocking=True), + "loss_mask": batch["loss_mask"].to(device, non_blocking=True), "size": len(examples), "is_cached": True, } @@ -77,13 +79,13 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: x = gpu_ds["x"][bi: bi + eval_bs] y = gpu_ds["y"][bi: bi + eval_bs] y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs] - mask = gpu_ds["y_mask"][bi: bi + eval_bs] + mask = gpu_ds["loss_mask"][bi: bi + eval_bs] else: batch = make_batch(chunk_examples, cfg) x = batch["x"].to(device) y = batch["y"].to(device) y_in_eval = batch["y_in"].to(device) - mask = batch["y_mask"].to(device) + mask = batch["loss_mask"].to(device) with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): out = model(x, y_in_eval) @@ -93,10 +95,9 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: token_total += int(mask.sum().item()) B = x.shape[0] - max_out_len = getattr(cfg, "max_out_len", 6) from src.data import _global_lengths _, out_max = _global_lengths(cfg) - max_len = max(max_out_len, out_max) + max_len = max(getattr(cfg, "max_out_len", 6), out_max) y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device) with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): @@ -106,10 +107,11 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: y_in = torch.cat([y_in, nxt[:, None]], dim=1) gen = y_in[:, 1:].cpu() - for i, (xrow, target) in enumerate(chunk_examples): + for i, ex in enumerate(chunk_examples): + xrow, target = ex[0], ex[1] pred = decode_tokens(gen[i].tolist(), cfg) target_n = decode_tokens(target, cfg) - ok = pred == target_n + ok = (pred == target_n) and (pred != -1) em_correct += int(ok) per_example.append((tuple(xrow), target_n, pred, ok)) @@ -268,7 +270,7 @@ def main() -> None: x = train_gpu["x"][idx] y = train_gpu["y"][idx] y_in = train_gpu["y_in"][idx] - mask = train_gpu["y_mask"][idx] + mask = train_gpu["loss_mask"][idx] with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): out = model(x, y_in) @@ -317,7 +319,7 @@ def main() -> None: x = batch["x"].to(device) y = batch["y"].to(device) y_in = batch["y_in"].to(device) - mask = batch["y_mask"].to(device) + mask = batch["loss_mask"].to(device) with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): out = model(x, y_in) |