File size: 2,307 Bytes
2246e67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Minimal loader for the fitted reference operators (A_l, b_l).
The reference predicts each layer's next state from its current one:
h_{l+1} ~= A_l @ h_l + b_l
and the *innovation* is the part it does not predict:
eps_l = h_{l+1} - (A_l @ h_l + b_l)
Extract your own hidden states (`output_hidden_states=True`), pick the matching
`condition` and `variant`, and subtract. No dependency on the fitting pipeline.
ops = load_operators("operators_gemma.pt")
eps = innovation(ops, h_l, h_next, layer=12) # h_*: [..., d]
Use `variant="raw"` for hidden_states as returned; `variant="normed"` if you feed
`input_layernorm_l(h_l)` (what block l actually reads — matters most for Gemma's
(1+gamma) RMSNorm). `condition` selects which fit: PT-COMP (base weights),
IT-COMP (instruct weights, no template), IT-CHAT (instruct + chat template).
"""
import torch
def load_operators(path):
"""Load an operators_<family>.pt bundle (dict: meta + per-condition A/b)."""
return torch.load(path, map_location="cpu", weights_only=False)
def _block(ops, condition, variant):
if condition not in ops:
raise KeyError(f"condition {condition!r} not in {[k for k in ops if k!='meta']}")
return ops[condition][variant]
def predict_next(ops, h, layer, condition="IT-COMP", variant="raw"):
"""A_l @ h + b_l for h of shape [..., d]. Returns the predicted h_{l+1}."""
blk = _block(ops, condition, variant)
A = blk["A"][layer].to(h.dtype).to(h.device)
b = blk["b"][layer].to(h.dtype).to(h.device)
return h @ A.T + b
def innovation(ops, h_l, h_next, layer, condition="IT-COMP", variant="raw"):
"""eps_l = h_{l+1} - (A_l @ h_l + b_l). h_l, h_next: [..., d]."""
return h_next - predict_next(ops, h_l, layer, condition, variant)
if __name__ == "__main__":
import sys
ops = load_operators(sys.argv[1] if len(sys.argv) > 1 else "operators_gemma.pt")
m = ops["meta"]
print(f"family={m['family']} d={m['d']} n_layers={m['n_layers']}")
print(f"conditions={list(m['conditions'])} variants={list(m['variants'])}")
d, nL = m["d"], m["n_layers"]
h = torch.randn(4, d) # 4 fake token states
eps = innovation(ops, h, torch.randn(4, d), layer=nL // 2)
print(f"demo innovation shape: {tuple(eps.shape)}")
|