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
|
"""Transformer baseline: GPT-style causal decoder over [input digits | output digits].
Fixed d_model (=128) matching the RNN arm. Parameter counts are logged per run but
NOT gated to parity (design/preregistration.md: weight sharing is the studied variable).
"""
import torch
import torch.nn as nn
from src.config import Config
from src.model_api import PrimeModel
class CausalBlock(nn.Module):
def __init__(self, d: int, heads: int):
super().__init__()
self.ln1 = nn.LayerNorm(d)
self.ln2 = nn.LayerNorm(d)
self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, key_pad: torch.Tensor) -> torch.Tensor:
a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
attn_mask=attn_mask, key_padding_mask=key_pad, need_weights=False)
x = x + a
x = x + self.mlp(self.ln2(x))
return x
class TransformerBaseline(PrimeModel):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
d = cfg.d_model
self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad
self.pos = nn.Embedding(1024, d) # supports full probe traces up to 1024
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)
self.register_buffer(
"causal_triu",
torch.triu(torch.ones(1024, 1024, dtype=torch.bool), diagonal=1),
persistent=False,
)
def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
cfg = self.cfg
B, T_in = x.shape
T_out = y_in.shape[1]
seq = torch.cat([x, y_in], dim=1) # (B, T_in+T_out)
T = seq.shape[1]
device = seq.device
# 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)
blocked[:, T_in:] = self.causal_triu[:T, T_in:T]
key_pad = seq == cfg.pad_id # (B,T) True = ignore
h = e
for blk in self.blocks:
h = blk(h, blocked, key_pad)
# logits at position p predict token p+1 -> positions [T_in-1, T_in+T_out-2] predict y
logits = self.head(self.ln_f(h))[:, T_in - 1: T_in - 1 + T_out]
return {"logits": logits, "halt_steps": None}
|
