111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
"""E1 (PLAN_SELFPACED): learned per-prompt halting over the workspace loop.
|
|
|
|
HaltingMergeAdapter = frozen-recipe MergeAdapter + a halting head read off
|
|
the workspace state at the last prompt position after each iteration:
|
|
|
|
p_i = sigmoid(w · [e_last ; s_hat_i,last] + b) (halt after iter i)
|
|
q_i = p_i * prod_{j<i}(1 - p_j), q_kmax += remainder (ACT-style)
|
|
|
|
Training (soft, no RL): s_mix = sum_i q_i * s_i -> suffix -> CE, plus a
|
|
compute penalty lambda * E[iters] = lambda * sum_i q_i * i.
|
|
The halting head's weight is zero-init and its bias starts at -6
|
|
(p ~ 0.0025), so at init >=98% of the mass sits on k_max: the model is
|
|
approximately the fixed-k merge (documented tolerance, not bit-exact).
|
|
|
|
Deploy: iterate until cumulative halt mass crosses 0.5; per-item k*.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
from loop_common import ALPHA, MergeAdapter
|
|
|
|
|
|
class HaltingMergeAdapter(MergeAdapter):
|
|
|
|
def __init__(self, d=1536, hidden=512, alpha=ALPHA, bias0=-6.0):
|
|
super().__init__(d=d, hidden=hidden, alpha=alpha)
|
|
self.halt = nn.Linear(2 * d, 1)
|
|
nn.init.zeros_(self.halt.weight)
|
|
nn.init.constant_(self.halt.bias, bias0)
|
|
|
|
def halt_prob(self, e_last, s_last):
|
|
"""p(halt | state) from the last prompt position. (B,) in (0,1)."""
|
|
e32, s32 = e_last.float(), s_last.float()
|
|
s_hat = s32 * (e32.norm(dim=-1, keepdim=True)
|
|
/ (s32.norm(dim=-1, keepdim=True) + 1e-6))
|
|
return torch.sigmoid(self.halt(torch.cat([e32, s_hat], -1))).squeeze(-1)
|
|
|
|
|
|
def halting_loop_logits(looper, adapter, input_ids, k_max, attention_mask,
|
|
prompt_last_idx, loop_mask=None, use_checkpoint=True):
|
|
"""Soft-halting teacher-forced logits.
|
|
|
|
prompt_last_idx: (B,) index of each item's last prompt position.
|
|
Returns (logits_of_mixed_state, q, exp_iters):
|
|
q (B, k_max) halting distribution over iterations 1..k_max.
|
|
"""
|
|
calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
|
|
e = looper._hin[looper.l0].detach()
|
|
B = e.shape[0]
|
|
bidx = torch.arange(B, device=e.device)
|
|
with torch.no_grad():
|
|
s = looper.band(e, calls)
|
|
s_mix = torch.zeros_like(e)
|
|
keep = torch.ones(B, device=e.device) # prob of not-yet-halted
|
|
qs = []
|
|
for i in range(k_max):
|
|
x = adapter(e, s)
|
|
if loop_mask is not None:
|
|
x = torch.where(loop_mask[..., None], x, e)
|
|
if use_checkpoint:
|
|
s = checkpoint(lambda x_: looper.band(x_, calls), x,
|
|
use_reentrant=False)
|
|
else:
|
|
s = looper.band(x, calls)
|
|
p = adapter.halt_prob(e[bidx, prompt_last_idx],
|
|
s[bidx, prompt_last_idx])
|
|
q_i = keep * p if i < k_max - 1 else keep # remainder to k_max
|
|
keep = keep * (1 - p)
|
|
s_mix = s_mix + q_i[:, None, None] * s.float()
|
|
qs.append(q_i)
|
|
q = torch.stack(qs, -1) # (B, k_max)
|
|
exp_iters = (q * torch.arange(1, k_max + 1, device=q.device)).sum(-1)
|
|
logits = looper.suffix_logits(s_mix.to(e.dtype), calls)
|
|
return logits, q, exp_iters
|
|
|
|
|
|
@torch.no_grad()
|
|
def halted_k_per_item(looper, adapter, input_ids, k_max, attention_mask,
|
|
prompt_last_idx, thresh=0.5, allow_k0=True):
|
|
"""Deploy-time halting: smallest k where cumulative halt mass >= thresh
|
|
(k_max if never). With allow_k0, the head is also consulted on the
|
|
pre-loop state s_0 — items it flags there get k*=0 (no looping at
|
|
all; the E0 lesson: easy items are safest untouched).
|
|
Returns (B,) ints in 0..k_max."""
|
|
calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1)
|
|
e = looper._hin[looper.l0].detach()
|
|
B = e.shape[0]
|
|
bidx = torch.arange(B, device=e.device)
|
|
s = looper.band(e, calls)
|
|
kstar = torch.full((B,), k_max, dtype=torch.long, device=e.device)
|
|
cum = torch.zeros(B, device=e.device)
|
|
keep = torch.ones(B, device=e.device)
|
|
if allow_k0:
|
|
p0 = adapter.halt_prob(e[bidx, prompt_last_idx],
|
|
s[bidx, prompt_last_idx])
|
|
kstar[p0 >= thresh] = 0
|
|
cum = cum + p0
|
|
keep = keep * (1 - p0)
|
|
for i in range(k_max):
|
|
x = adapter(e, s)
|
|
s = looper.band(x, calls)
|
|
p = adapter.halt_prob(e[bidx, prompt_last_idx],
|
|
s[bidx, prompt_last_idx])
|
|
cum = cum + keep * p
|
|
keep = keep * (1 - p)
|
|
newly = (cum >= thresh) & (kstar == k_max) & (i < k_max - 1)
|
|
kstar[newly] = i + 1
|
|
return kstar
|