From 898a0570dfe619ec2bcf330b07dce9b23cb63d54 Mon Sep 17 00:00:00 2001 From: Void Agent Date: Fri, 14 Aug 2026 13:11:02 +0100 Subject: implement Experiment 1: data pipeline, tied-RNN w/ ACT halting, transformer baseline, train/eval/plot, 16 tests --- src/models/rnn.py | 98 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 src/models/rnn.py (limited to 'src/models/rnn.py') diff --git a/src/models/rnn.py b/src/models/rnn.py new file mode 100644 index 0000000..9947c66 --- /dev/null +++ b/src/models/rnn.py @@ -0,0 +1,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} -- cgit v1.2.3