diff options
Diffstat (limited to 'src/models/rnn.py')
| -rw-r--r-- | src/models/rnn.py | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/src/models/rnn.py b/src/models/rnn.py index 19123b7..44f0a0d 100644 --- a/src/models/rnn.py +++ b/src/models/rnn.py @@ -45,7 +45,8 @@ class TiedRNN(PrimeModel): return pe def _encode(self, x: torch.Tensor) -> torch.Tensor: - """Read input digits through the tied cell (pad positions inject nothing).""" + """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) @@ -53,7 +54,8 @@ class TiedRNN(PrimeModel): 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]) + 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]: @@ -72,7 +74,7 @@ class TiedRNN(PrimeModel): 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: + if t < cfg.min_steps - 1: p = p * 0.0 h_list.append(h) p_list.append(p) |