diff options
Diffstat (limited to 'sdil/probes.py')
| -rw-r--r-- | sdil/probes.py | 18 |
1 files changed, 16 insertions, 2 deletions
diff --git a/sdil/probes.py b/sdil/probes.py index 2d10aad..1698d75 100644 --- a/sdil/probes.py +++ b/sdil/probes.py @@ -61,17 +61,22 @@ def alignment_report(net, x, y, y_onehot, cfg): h = fwd["h"] u = fwd["u"] logits = h[-1] + context = h[-2] e = net.output_error(logits, y_onehot) c = e out = {"cos_r_negg": [], "cos_innovation_negg": [], "cos_apical_negg": [], "cos_Ac_negg": [], - "r_norm": [], "g_norm": [], "loss": loss} + "r_norm": [], "g_norm": [], "traffic_norm": [], + "traffic_residual_norm": [], "traffic_r2": [], "loss": loss} for l in range(net.L - 1): g = grads[l] negg = -g - teaching, a, innovation = core.teaching_signal(net, l, c, h[l + 1], cfg) + teaching, a, innovation = core.teaching_signal( + net, l, c, h[l + 1], cfg, context) Ac = c @ net.A[l].t() # pure vectorizer output + traffic = net.apical_traffic(l, h[l + 1], context) + unexplained = traffic - net.baseline(l, h[l + 1]) # 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]) @@ -81,6 +86,14 @@ def alignment_report(net, x, y, y_onehot, cfg): out["cos_Ac_negg"].append(_row_cos(Ac * gain, negg * gain)) out["r_norm"].append(teaching.norm(dim=1).mean().item()) out["g_norm"].append(g.norm(dim=1).mean().item()) + out["traffic_norm"].append(traffic.norm(dim=1).mean().item()) + out["traffic_residual_norm"].append(unexplained.norm(dim=1).mean().item()) + traffic_power = traffic.pow(2).mean() + if traffic_power > 0: + out["traffic_r2"].append( + (1.0 - unexplained.pow(2).mean() / traffic_power).item()) + else: + out["traffic_r2"].append(float("nan")) return out @@ -136,6 +149,7 @@ def _clone(net): n.P_bias = ([p.clone() for p in net.P_bias] if getattr(net, "P_bias", None) is not None else None) n.Bnuis = [b.clone() for b in net.Bnuis] + n.Bcontext = [b.clone() for b in net.Bcontext] n.mW = [torch.zeros_like(w) for w in net.W] n.mb = [torch.zeros_like(b) for b in net.b] return n |