From 9c50f31c66e788ff08eeac83f84e90e9bc1a921e Mon Sep 17 00:00:00 2001 From: Void Agent Date: Fri, 14 Aug 2026 13:14:09 +0100 Subject: rnn: fully-tied cell (read-in+compute+decode), lambda ramp, ACT regression test; design: prior art + review --- src/train.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) (limited to 'src/train.py') 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 -- cgit v1.2.3