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
|
"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder.
Spec pseudocode (design/experiment-spec.md):
state = embed(input_number)
for step in range(max_steps):
state = step_module(state) # same weights every iteration
if halt_condition(state): break
output = project(state)
Initial state: masked mean-pool of (digit embedding + sinusoidal positional encoding)
passed through a small MLP, so digit ORDER reaches the tied cell.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from src.config import Config
from src.model_api import PrimeModel
class TiedRNN(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.in_proj = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
self.ln1 = nn.LayerNorm(d)
self.cell_ln = nn.LayerNorm(d)
self.cell_w1 = nn.Linear(d, d)
self.cell_w2 = nn.Linear(d, d)
self.halt_head = nn.Linear(d, 1)
self.decoder = nn.GRUCell(d, d)
self.out_head = nn.Linear(d, cfg.vocab)
@staticmethod
def _sinusoidal(T: int, d: int) -> torch.Tensor:
pe = torch.zeros(T, d)
pos = torch.arange(T).float().unsqueeze(1)
i = torch.arange(d).float().unsqueeze(0)
pe[:, 0::2] = torch.sin(pos / 10000 ** (2 * i[:, 0::2] / d))
pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d))
return pe
def _initial_state(self, x: torch.Tensor) -> torch.Tensor:
B, T = x.shape
mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1)
pos = self._sinusoidal(T, self.cfg.d_model).to(x.device) # (T,d)
e = self.embed(x) + pos.unsqueeze(0) # (B,T,d)
h = (e * mask).sum(1) / mask.sum(1).clamp(min=1) # (B,d)
return self.in_proj(h)
def _run_cell(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Tied cell x K steps. Returns (final_state (B,d), mean_steps (B,))."""
cfg = self.cfg
B = h0.shape[0]
device = h0.device
if not cfg.halting:
h = h0
for _ in range(cfg.max_steps):
h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
steps = h0.new_full((B,), float(cfg.max_steps))
return h, steps
# ACT: run all K steps, accumulate weighted average (K=20 -> no early break needed)
h_list, p_list = [], []
h = h0
for t in range(cfg.max_steps):
h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,)
if t < cfg.min_steps:
p = p * 0.0
h_list.append(h)
p_list.append(p)
final = torch.zeros_like(h0)
steps = torch.zeros(B, device=device)
remaining = torch.ones(B, device=device)
for t in range(cfg.max_steps):
p = p_list[t]
w = remaining * p
final = final + w.unsqueeze(-1) * h_list[t]
steps = steps + (t + 1) * w
remaining = remaining * (1 - p)
final = final + remaining.unsqueeze(-1) * h_list[-1]
steps = steps + remaining * cfg.max_steps
return final, steps
def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
cfg = self.cfg
h = self._initial_state(x)
h, steps = self._run_cell(h)
# autoregressive digit decoder, teacher-forced during training
e = self.embed(y_in) # (B,T_out,d)
outs = []
for t in range(y_in.shape[1]):
h = self.decoder(e[:, t], h)
outs.append(self.out_head(h))
logits = torch.stack(outs, dim=1) # (B,T_out,vocab)
return {"logits": logits, "halt_steps": steps}
|
