"""Learnable soft path through ALL layers (pre-registration item 14). The frozen model = 35 same-typed functions on one residual bus. Instead of a hand-fixed loop over L14-30, learn a gate matrix g[t, l] in [0,1]: on loop iteration t, layer l's residual delta is scaled by g[t, l] (prompt positions only; generated/answer positions always run ungated). Gates are initialized as a Gaussian bump over depth centered mid-band, so at init a loop iteration is approximately the hand band pass — then SGD may move the compute envelope anywhere in [0, n_layers). The anchor merge adapter is kept at each iteration boundary for stability (rho < 1). Reading the result: if the learned envelope concentrates on the lens band, gradient descent independently rediscovers the workspace; if it wins with mass elsewhere, the lens placement story needs revision. E2B caveat (stated in advance): KV sharing makes attention deltas of layers >= 15 loop-inert; gate mass there is interpretable for MLP deltas only. """ import math import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint from loop_common import BAND, BandLooper, _text_model class PathGates(nn.Module): """g[t, l] = sigmoid(logits[t, l]); row 0 = warm sweep, rows 1..k = loops.""" def __init__(self, n_layers, k_max, mu=None, sigma=6.0, band=BAND): super().__init__() mu = (band[0] + band[1]) / 2 if mu is None else mu init = torch.empty(k_max + 1, n_layers) for l in range(n_layers): p = math.exp(-((l - mu) ** 2) / (2 * sigma ** 2)) p = min(max(p, 1e-3), 1 - 1e-3) init[:, l] = math.log(p / (1 - p)) self.logits = nn.Parameter(init) def g(self, t): return torch.sigmoid( self.logits[min(t, self.logits.shape[0] - 1)].float()) def envelope(self): with torch.no_grad(): return torch.sigmoid(self.logits.float()).tolist() class GatedLooper(BandLooper): """BandLooper over the FULL depth with per-iteration per-layer gates.""" def __init__(self, model, gates): tm = _text_model(model) super().__init__(model, band=(0, len(tm.layers) - 1)) self.gates = gates def _gated(self, h, calls, t, loop_mask=None): g = self.gates.g(t).to(h.device) x = h for i in range(self.l0, self.l1 + 1): args, kwargs = calls[i] out = self.tm.layers[i](x, *args, **kwargs) if isinstance(out, tuple): out = out[0] gi = g[i].to(x.dtype) if loop_mask is not None: # ungated (g=1) off the prompt span gi = torch.where(loop_mask[..., None], gi, torch.ones_like(loop_mask[..., None], dtype=x.dtype)) x = x + gi * (out - x) return x def loop_logits(self, adapter, input_ids, k, attention_mask=None, use_checkpoint=False, return_states=False, last_only=False, loop_mask=None, feedforward=False, bptt=None): calls, base_logits = self.capture(input_ids, attention_mask, logits_to_keep=1 if last_only else 0) if k == 0: return (base_logits, None) if return_states else base_logits del base_logits e = self._hin[self.l0].detach() def sweep(x, t): if use_checkpoint: return checkpoint( lambda x_: self._gated(x_, calls, t, loop_mask), x, use_reentrant=False) return self._gated(x, calls, t, loop_mask) s = sweep(e, 0) # warm sweep, t=0 (gates trainable here too) states = [s] n_nograd = max(0, k - bptt) if bptt else 0 for i in range(k): if i < n_nograd: with torch.no_grad(): x = adapter(e, s) if loop_mask is not None: x = torch.where(loop_mask[..., None], x, e) s = self._gated(x, calls, i + 1, loop_mask) s = s.detach() states.append(s) continue x = adapter(e, s) if loop_mask is not None: x = torch.where(loop_mask[..., None], x, e) s = sweep(x, i + 1) states.append(s) logits = self.suffix_logits(s, calls, last_only=last_only) return (logits, states) if return_states else logits @torch.no_grad() def generate_frozen_prompt(self, adapter, tok, input_ids, k, max_new_tokens=220, attention_mask=None, stop_strs=(), feedforward=False, conv_out=None): """Per-layer KV write-in: run the final gated sweep recording every layer's INPUT, then one native prefill with pre-forward hooks swapping each layer's hidden_states to the recorded stream — the cache then holds exactly the gated states; decode is native.""" if k == 0: return super().generate_frozen_prompt( adapter, tok, input_ids, 0, max_new_tokens=max_new_tokens, attention_mask=attention_mask, stop_strs=stop_strs) calls, _ = self.capture(input_ids, attention_mask, logits_to_keep=1) e = self._hin[self.l0] s = self._gated(e, calls, 0) for i in range(k): x = adapter(e, s) if i < k - 1: s = self._gated(x, calls, i + 1) # final sweep: record per-layer inputs of the gated stream xs = {} g = self.gates.g(k).to(x.device) h = x for i in range(self.l0, self.l1 + 1): xs[i] = h args, kwargs = calls[i] out = self.tm.layers[i](h, *args, **kwargs) if isinstance(out, tuple): out = out[0] h = h + g[i].to(h.dtype) * (out - h) del calls P = input_ids.shape[1] handles = [] for i in range(self.l0, self.l1 + 1): def pre(mod, args, kwargs, i=i): hh = kwargs.get("hidden_states", args[0] if args else None) if hh is not None and hh.shape[1] == P: # prefill only if "hidden_states" in kwargs: kwargs["hidden_states"] = xs[i].to(hh.dtype) return args, kwargs return (xs[i].to(hh.dtype),) + args[1:], kwargs return None handles.append(self.tm.layers[i].register_forward_pre_hook( pre, with_kwargs=True)) try: gen = self.model.generate( input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.pad_token_id or 0) finally: for hd in handles: hd.remove() return gen