diff options
Diffstat (limited to 'sdil')
| -rw-r--r-- | sdil/core.py | 35 | ||||
| -rw-r--r-- | sdil/probes.py | 21 |
2 files changed, 46 insertions, 10 deletions
diff --git a/sdil/core.py b/sdil/core.py index 3cd8a97..d7be8c3 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -219,6 +219,31 @@ class SDILNet: return a, a +def teaching_signal(net, l, c, h_l, cfg): + """Select the signal that drives a hidden layer's local plasticity. + + ``match_innovation_norm`` is a deliberately strong raw-apical control: it + preserves each sample's raw apical direction exactly while rescaling its + Euclidean norm to the norm of the corresponding innovation. Therefore a + raw-vs-residual difference under this control cannot be attributed merely + to the nuisance making the teaching vector larger. + + Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic + probes use one definition of every signal. + """ + a = net.apical(l, c, h_l) + innovation = a - net.baseline(l, h_l) + if cfg.use_residual: + teaching = innovation + elif cfg.raw_scale_control == "match_innovation_norm": + raw_norm = a.norm(dim=1, keepdim=True) + innovation_norm = innovation.norm(dim=1, keepdim=True) + teaching = a * (innovation_norm / raw_norm.clamp_min(1e-12)) + else: + teaching = a + return teaching, a, innovation + + # -------------------------------------------------------------------------- # node perturbation: causal estimate of -grad_{h_l} L used to TRAIN A_l # -------------------------------------------------------------------------- @@ -341,11 +366,15 @@ class SDILConfig: use_residual=True, learn_A=True, learn_P=True, pert_sigma=1e-2, pert_every=5, pert_ndirs=1, momentum=0.0, wd=0.0, settle_steps=0, kappa=0.0, feedback="error", p_update_on_neutral=True, - normalize_delta=False, pert_mode="layerwise"): + normalize_delta=False, pert_mode="layerwise", + raw_scale_control="none"): self.eta = eta self.eta_A = eta_A self.eta_P = eta_P self.use_residual = use_residual # ablation: residual vs raw apical + if raw_scale_control not in ("none", "match_innovation_norm"): + raise ValueError(f"unknown raw_scale_control: {raw_scale_control}") + self.raw_scale_control = raw_scale_control self.learn_A = learn_A # ablation: perturbation-trained vs fixed-random A self.learn_P = learn_P # ablation: predictor on/off self.pert_sigma = pert_sigma @@ -390,7 +419,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): # ---- compute innovations for hidden layers ---- r_list, a_list = [], [] for l in range(net.L - 1): - r, a = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) + r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg) r_list.append(r) a_list.append(a) @@ -496,7 +525,7 @@ def _settle(net, h, u, c, cfg): dt_over_tau = 1.0 / max(1, cfg.settle_steps) for _ in range(cfg.settle_steps): for l in range(net.L - 1): - r, _ = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) + r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg) target = net.act(u[l]) + cfg.kappa * r h[l + 1] = h[l + 1] + dt_over_tau * (-h[l + 1] + target) # recompute downstream pre-activations from settled hidden states diff --git a/sdil/probes.py b/sdil/probes.py index 16aca21..fad82ca 100644 --- a/sdil/probes.py +++ b/sdil/probes.py @@ -36,11 +36,20 @@ def true_hidden_grads(net, x, y): return [g.detach() for g in grads], loss.item() -def _row_cos(u, v, eps=1e-12): - """Mean per-sample cosine between rows of u and v.""" +def _row_cos(u, v): + """Mean per-sample cosine, excluding rows with an exact zero vector. + + Adding an absolute epsilon to the denominator makes cosine spuriously + depend on signal magnitude near convergence. In particular, a positive + per-sample rescaling can then appear to change direction. Exact zero rows + have no defined direction and are omitted instead. + """ num = (u * v).sum(1) - den = u.norm(dim=1) * v.norm(dim=1) + eps - return (num / den).mean().item() + den = u.norm(dim=1) * v.norm(dim=1) + valid = den > 0 + if not valid.any(): + return float("nan") + return (num[valid] / den[valid]).mean().item() @torch.no_grad() @@ -61,10 +70,8 @@ def alignment_report(net, x, y, y_onehot, cfg): for l in range(net.L - 1): g = grads[l] negg = -g - a = net.apical(l, c, h[l + 1]) # raw apical (with nuisance) + teaching, a, innovation = core.teaching_signal(net, l, c, h[l + 1], cfg) Ac = c @ net.A[l].t() # pure vectorizer output - innovation = a - net.baseline(l, h[l + 1]) - teaching = innovation if cfg.use_residual else a # include the phi' gain that the plasticity rule actually applies, # so the reported alignment is what the weight update "sees" gain = net.act_prime(u[l]) |