1 files changed, 98 insertions, 0 deletions
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}