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
|
"""Experiment configuration. Defaults = Experiment 1 (design/preregistration.md)."""
import argparse
import json
from dataclasses import dataclass, fields
@dataclass
class Config:
# data
vocab_mode: str = "digits" # "digits" | "integers"
task_mode: str = "next_prime" # "next_prime" | "is_prime"
range_start: int = 2
range_end: int = 100 # inclusive
holdout_frac: float = 0.30
train_frac: float = 1.0 # 1.0 = use full train split; 0.4/0.5 for E5
seed: int = 0
# model
model: str = "rnn" # "rnn" | "transformer"
d_model: int = 128
max_steps: int = 20 # K tied iterations (RNN)
halting: bool = True # False -> fixed-K ablation
halt_penalty: float = 0.01 # lambda on mean steps
halt_warmup_steps: int = 1000 # penalty = 0 before this (anti-collapse)
halt_ramp_end_steps: int = 5000 # penalty ramps linearly 0 -> halt_penalty between warmup and here
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)
weight_decay: float = 1.0
max_train_steps: int = 200_000
eval_every: int = 200
eval_schedule: str = "fixed" # "fixed" | "adaptive" (more frequent early, coarser later for 4M runs)
early_stop_em: float = 1.0
early_stop_patience: int = 5
batch_size: int = 32
eval_batch_size: int = 512 # chunked evaluation batch size to avoid GPU OOM
max_out_len: int = 6
log_n_examples: int = 10
out_dir: str = "runs"
device: str = "auto" # "auto" | "cuda" | "cpu" | "mps"
use_amp: bool = True # FP16 mixed precision on CUDA (Turing Tensor Cores)
compile_model: bool = False # torch.compile (reduce-overhead / CUDA Graphs)
@property
def vocab(self) -> int:
"""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
is_prime = [True] * (limit + 1)
is_prime[0] = is_prime[1] = False
for p in range(2, int(limit ** 0.5) + 1):
if is_prime[p]:
for m in range(p * p, limit + 1, p):
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
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:
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
def to_json(self) -> dict:
d = {f.name: getattr(self, f.name) for f in fields(self)}
d["vocab"] = self.vocab
d["eos_id"] = self.eos_id
d["pad_id"] = self.pad_id
return d
@classmethod
def from_json(cls, d: dict) -> "Config":
cfg = cls()
for f in fields(cls):
if f.name in d:
setattr(cfg, f.name, d[f.name])
return cfg
def save(self, path: str) -> None:
with open(path, "w") as fh:
json.dump(self.to_json(), fh, indent=2)
@classmethod
def load(cls, path: str) -> "Config":
with open(path) as fh:
return cls.from_json(json.load(fh))
def _bool_arg(s: str) -> bool:
return s.lower() in ("1", "true", "yes", "on")
def parse_args(argv=None) -> Config:
cfg = Config()
p = argparse.ArgumentParser(description="prime-grokking train")
p.add_argument("model_pos", nargs="?", default=None, help="model: rnn | transformer")
p.add_argument("seed_pos", nargs="?", default=None, help="seed (int)")
for f in fields(Config):
if f.type is bool:
p.add_argument(f"--{f.name}", default=None, type=_bool_arg)
elif f.type in (int, float, str):
p.add_argument(f"--{f.name}", default=None, type=f.type)
a = p.parse_args(argv)
if a.model_pos is not None:
cfg.model = a.model_pos
if a.seed_pos is not None:
cfg.seed = int(a.seed_pos)
for f in fields(Config):
v = getattr(a, f.name, None)
if v is not None:
setattr(cfg, f.name, v)
return cfg
|
