nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

item 20 scored; item 21 (GSM carry-CoT + control) pre-registered and queued

Browse Source
This commit is contained in:
Nils
2026-07-16 02:42:46 +02:00
parent 4b9e168838
commit 1618c206ae
9 changed files with 7431 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/path_gates.py
+170
View File
@@ -0,0 +1,170 @@
"""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