rnn: fully-tied cell (read-in+compute+decode), lambda ramp, ACT regression test; design: prior art + review

1 files changed, 8 insertions, 1 deletions

diff --git a/src/train.py b/src/train.py
index f0c63e9..975598d 100644
--- a/src/train.py
+++ b/src/train.py
@@ -133,7 +133,14 @@ def main() -> None:
             loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
                 logits.shape[0], -1) * batch["y_mask"].float()
             loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1)
-            lam = cfg.halt_penalty if (cfg.halting and step >= cfg.halt_warmup_steps) else 0.0
+            if cfg.halting and step >= cfg.halt_warmup_steps:
+                if step >= cfg.halt_ramp_end_steps:
+                    lam = cfg.halt_penalty
+                else:
+                    frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps)
+                    lam = cfg.halt_penalty * frac
+            else:
+                lam = 0.0
             halt = out["halt_steps"]
             penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
             loss = loss_tokens + penalty