diff options
Diffstat (limited to 'src/models/rnn.py')
| -rw-r--r-- | src/models/rnn.py | 54 |
1 files changed, 30 insertions, 24 deletions
diff --git a/src/models/rnn.py b/src/models/rnn.py index 9947c66..19123b7 100644 --- a/src/models/rnn.py +++ b/src/models/rnn.py @@ -1,14 +1,15 @@ -"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder. +"""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 + 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. +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 @@ -24,15 +25,16 @@ class TiedRNN(PrimeModel): 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.ln1 = nn.LayerNorm(d) self.halt_head = nn.Linear(d, 1) - self.decoder = nn.GRUCell(d, d) 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) @@ -42,35 +44,39 @@ class TiedRNN(PrimeModel): pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d)) return pe - def _initial_state(self, x: torch.Tensor) -> torch.Tensor: + def _encode(self, x: torch.Tensor) -> torch.Tensor: + """Read input digits through the tied cell (pad positions inject nothing).""" 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) + 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 = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + return 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,)).""" + 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 = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h)))) + h = self._cell_step(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)))) + h = self._cell_step(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) + # 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) @@ -86,13 +92,13 @@ class TiedRNN(PrimeModel): 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) + 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.decoder(e[:, t], h) + h = self._cell_step(h + e_out[:, t]) outs.append(self.out_head(h)) - logits = torch.stack(outs, dim=1) # (B,T_out,vocab) + logits = torch.stack(outs, dim=1) # (B,T_out,vocab) return {"logits": logits, "halt_steps": steps} |