File size: 8,627 Bytes
9e551a1 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | # Tensor Roll — Recursive CUDA-Q Model Quantizer
# Copyright (C) 2026 SnapKitty Collective
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Recursive distillation with the multi-objective loss.
L = l1*L_logits + l2*L_hidden + l3*L_attention + l4*L_embedding
+ l5*L_task + l6*L_reconstruction + l7*L_roll
All seven terms are real computed quantities (no placeholders):
logits : tau^2 * KL(softmax(teacher/tau) || softmax(student/tau))
hidden : MSE(student_final @ Ph, teacher_final), Ph frozen random proj
attention : MSE(mean-head attention probs), layer-aligned
embedding : MSE(student_emb @ Pe, teacher_emb), Pe frozen random proj
task : cross-entropy on the synthetic grammar
reconstruction: MSE(student weights, projected-teacher init targets)
roll : MSE(achieved per-tensor energy, roll-plan energy) — monitored,
gradient detached (documented; proxy grad via reconstruction)
Recursive distillation metrics (spec section 8) are computed per recursion
level: model -> layer group -> tensor -> sub-block, each reporting teacher
tensor, student tensor, relative error, cosine similarity, KL divergence,
retained energy.
"""
from __future__ import annotations
import numpy as np
from . import core as C
from .model import backward, ce_loss, forward, kl_loss
DEFAULT_LAMBDAS = {
"logits": 1.0,
"hidden": 0.3,
"attention": 0.3,
"embedding": 0.1,
"task": 1.0,
"reconstruction": 0.1,
"roll": 0.1,
}
def _align(a: np.ndarray, shape: tuple) -> np.ndarray:
"""Shape alignment for comparison only: truncate or zero-pad. This is a
measurement convenience, documented here — not a claim about the model."""
a = np.asarray(a, dtype=np.float64)
out = np.zeros(shape, dtype=np.float64)
sl = tuple(slice(0, min(x, y)) for x, y in zip(a.shape, shape))
out[sl] = a[sl]
return out
def distill_step(s_params: dict, s_cfg, t_params: dict, t_cfg,
ids: np.ndarray, *, lambdas=None, tau: float = 2.0,
proj_h=None, proj_e=None, recon_targets=None,
roll_plan_energies=None, achieved_energies=None):
"""One recursive-distillation step. Returns (total_loss, terms, grads)."""
lambdas = {**DEFAULT_LAMBDAS, **(lambdas or {})}
B = ids.shape[0]
inputs, targets = ids[:, :-1], ids[:, 1:]
t_logits, t_cache = forward(t_params, t_cfg, inputs)
s_logits, s_cache = forward(s_params, s_cfg, inputs)
terms = {}
# L_task + L_logits -> dlogits
terms["task"], dl_task = ce_loss(s_logits, targets)
terms["logits"], dl_kl = kl_loss(s_logits, t_logits, tau)
dlogits = lambdas["task"] * dl_task + lambdas["logits"] * dl_kl
# L_attention: mean-over-heads, layer-aligned (first min(L) layers)
d_attn = {}
n_align = min(s_cfg.layers, t_cfg.layers)
attn_num = 0.0
for l in range(n_align):
ps = s_cache[f"L{l}"]["P"].mean(axis=1) # (B,T,T)
pt = t_cache[f"L{l}"]["P"].mean(axis=1)
diff = ps - pt
terms[f"attention_l{l}"] = float((diff ** 2).mean())
attn_num += terms[f"attention_l{l}"]
dmean = (2.0 * diff / diff.size) * lambdas["attention"] / max(n_align, 1)
d_attn[l] = np.broadcast_to(
dmean[:, None, :, :], s_cache[f"L{l}"]["P"].shape).copy()
terms["attention"] = attn_num / max(n_align, 1)
# L_hidden: final hidden, frozen random projection to teacher dim
d_hiddens = {}
if proj_h is not None:
fs = s_cache["f"] # (B,T,ds)
ft = t_cache["f"] # (B,T,dt)
pred = fs @ proj_h
diff = pred - ft
terms["hidden"] = float((diff ** 2).mean())
d_hiddens[s_cfg.layers - 1] = (
(2.0 * diff / diff.size) @ proj_h.T) * lambdas["hidden"]
else:
terms["hidden"] = 0.0
grads = backward(s_params, s_cfg, s_cache, dlogits,
d_hiddens=d_hiddens or None, d_attn=d_attn or None)
# L_embedding
if proj_e is not None:
es = s_params["emb"].astype(np.float64)
et = t_params["emb"].astype(np.float64)
diff = es @ proj_e - et
terms["embedding"] = float((diff ** 2).mean())
grads["emb"] = grads["emb"] + (
(2.0 * diff / diff.size) @ proj_e.T * lambdas["embedding"]
).astype(np.float32)
else:
terms["embedding"] = 0.0
# L_reconstruction: stay near the projected-teacher init targets
if recon_targets:
rn, rv = 0.0, 0
for k, tgt in recon_targets.items():
if k in grads and grads[k].shape == tgt.shape:
w = s_params[k].astype(np.float64)
diff = w - tgt.astype(np.float64)
rn += float((diff ** 2).mean())
rv += 1
grads[k] = grads[k] + (
2.0 * diff / diff.size * lambdas["reconstruction"]).astype(np.float32)
terms["reconstruction"] = rn / max(rv, 1)
else:
terms["reconstruction"] = 0.0
# L_roll: achieved vs planned per-tensor retained energy (monitored)
if roll_plan_energies and achieved_energies:
vals = [(achieved_energies[k] - roll_plan_energies[k]) ** 2
for k in roll_plan_energies if k in achieved_energies]
terms["roll"] = float(np.mean(vals)) if vals else 0.0
else:
terms["roll"] = 0.0
total = sum(lambdas[k] * terms[k] for k in
("logits", "hidden", "attention", "embedding", "task",
"reconstruction", "roll"))
terms["total"] = float(total)
return float(total), terms, grads
def _kl_vec(p_logits: np.ndarray, q_logits: np.ndarray) -> float:
p = p_logits.astype(np.float64).ravel()
q = q_logits.astype(np.float64).ravel()
p = p - p.max(); q = q - q.max()
lp = p - np.log(np.exp(p).sum()); lq = q - np.log(np.exp(q).sum())
ep = np.exp(lp)
return float((ep * (lp - lq)).sum())
def recursive_metrics(t_params: dict, s_params: dict) -> dict:
"""Per-recursion-level teacher/student comparison.
Levels: model -> layer_group -> tensor -> subblock. Student tensors are
shape-aligned to the teacher by truncate/zero-pad (comparison only).
Each level reports: rel_err, cosine_sim, kl_div, retained_energy,
teacher_params, student_params.
"""
out = {}
def pack(params, keys):
return np.concatenate([np.asarray(params[k]).ravel() for k in keys])
def level(name, t_keys, s_keys):
tv = pack(t_params, t_keys).astype(np.float64)
sv_raw = pack(s_params, s_keys).astype(np.float64)
sv = _align(sv_raw, tv.shape)
out[name] = {
"teacher_params": int(tv.size),
"student_params": int(sv_raw.size),
"rel_err": C.relative_error(tv, sv),
"cosine_sim": C.cosine_sim(tv, sv),
"kl_div": _kl_vec(tv / (np.abs(tv).max() + 1e-9),
sv / (np.abs(sv).max() + 1e-9)),
"retained_energy": float((sv ** 2).sum() / ((tv ** 2).sum() + 1e-300)),
}
t_keys = sorted(t_params.keys())
s_keys = sorted(s_params.keys())
level("model", t_keys, s_keys)
# layer groups
layers = sorted({k.split(".")[0] for k in t_keys if k.startswith("L")})
for i, lg in enumerate(layers):
tk = [k for k in t_keys if k.startswith(lg + ".")]
sk = [k for k in s_keys if k.startswith(lg + ".")]
if tk and sk:
level(f"layer_group_{i}", tk, sk)
# tensors (sample up to 12 for readability; full data in JSON)
for k in t_keys[:12]:
sk = [x for x in s_keys if x == k]
if sk:
level(f"tensor:{k}", [k], sk)
# sub-blocks of the largest teacher tensor
big = max(t_keys, key=lambda k: t_params[k].size)
W = np.asarray(t_params[big], dtype=np.float64).ravel()
sm = [x for x in s_keys if x == big]
if sm:
S = np.asarray(s_params[sm[0]], dtype=np.float64).ravel()
for b, (ws, ss) in enumerate(zip(np.array_split(W, 4),
np.array_split(S, 4))):
sa = _align(ss, ws.shape)
out[f"subblock:{big}[{b}]"] = {
"teacher_params": int(ws.size),
"student_params": int(ss.size),
"rel_err": C.relative_error(ws, sa),
"cosine_sim": C.cosine_sim(ws, sa),
"kl_div": _kl_vec(ws / (np.abs(ws).max() + 1e-9),
sa / (np.abs(sa).max() + 1e-9)),
"retained_energy": float((sa ** 2).sum() / ((ws ** 2).sum() + 1e-300)),
}
return out
|