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
|
"""Weight-tied RNN: ONE 2-layer cell reused for input read-in, K compute steps, and output decoding.
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)
Everything recurrent is the SAME cell (per design review: no un-tied GRU decoder,
no mean-pool blur — order reaches the cell via sinusoidal position added per digit).
ACT halting applies only to the K compute steps.
"""
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.cell_ln = nn.LayerNorm(d)
self.cell_w1 = nn.Linear(d, d)
self.cell_w2 = nn.Linear(d, d)
self.ln1 = nn.LayerNorm(d)
self.halt_head = nn.Linear(d, 1)
self.out_head = nn.Linear(d, cfg.vocab)
def _cell_step(self, h: torch.Tensor) -> torch.Tensor:
return h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h))))
@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 _encode(self, x: torch.Tensor) -> torch.Tensor:
"""Read input digits through the tied cell. Pad positions are exact no-ops
(state update masked) so batch composition cannot change an example's state."""
B, T = x.shape
d = self.cfg.d_model
pos = self._sinusoidal(T, d).to(x.device) # (T,d)
e = self.embed(x) # (B,T,d)
mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1)
h = torch.zeros(B, d, device=x.device)
for t in range(T):
h_new = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t])
h = mask[:, t] * h_new + (1 - mask[:, t]) * h # pad step = no-op
return h
def _run_compute(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""K tied compute steps with ACT learned halting. Returns (final_state, mean_steps)."""
cfg = self.cfg
B = h0.shape[0]
device = h0.device
if not cfg.halting:
h = h0
for _ in range(cfg.max_steps):
h = self._cell_step(h)
steps = h0.new_full((B,), float(cfg.max_steps))
return h, steps
h_list, p_list = [], []
h = h0
for t in range(cfg.max_steps):
h = self._cell_step(h)
p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,)
if t < cfg.min_steps - 1:
p = p * 0.0
h_list.append(h)
p_list.append(p)
# ACT aggregation (Graves 2016): w_t = p_t * prod_{s<t}(1-p_s); weights + remainder = 1
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._encode(x)
h, steps = self._run_compute(h)
# decode output digits through the SAME tied cell (teacher-forced during training)
e_out = self.embed(y_in) # (B,T_out,d)
outs = []
for t in range(y_in.shape[1]):
h = self._cell_step(h + e_out[:, t])
outs.append(self.out_head(h))
logits = torch.stack(outs, dim=1) # (B,T_out,vocab)
return {"logits": logits, "halt_steps": steps}
|
