From ee57b4ca69ace14ffadfcfde947112278001feef Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 01:36:37 -0500 Subject: feat: model endogenous mixed apical traffic --- sdil/core.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++----------- sdil/probes.py | 18 +++++++++++-- 2 files changed, 80 insertions(+), 17 deletions(-) (limited to 'sdil') diff --git a/sdil/core.py b/sdil/core.py index d7be8c3..37c7165 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -100,7 +100,7 @@ class SDILNet: def __init__(self, sizes, act="tanh", device="cpu", seed=0, w_scale=1.0, a_scale=1.0, feedback="error", dtype=torch.float32, nuis_rho=0.0, nuis_seed=1234, residual=False, - predictor_mode="diagonal"): + predictor_mode="diagonal", traffic_mode="soma"): # sizes = [n_in, n_h1, ..., n_hk, n_out] self.sizes = list(sizes) self.L = len(sizes) - 1 # number of weight layers @@ -109,6 +109,9 @@ class SDILNet: self.dtype = dtype self.act, self.act_prime = ACTS[act] self.feedback = feedback + if traffic_mode not in ("none", "soma", "topdown", "mixed"): + raise ValueError(f"unknown traffic_mode: {traffic_mode}") + self.traffic_mode = traffic_mode # residual (skip) connections on the constant-width hidden blocks: the # modern way to train deep nets with NO normalization. # h_{l+1}=h_l+alpha*phi(u_l) for the interior blocks (input projection and @@ -140,10 +143,18 @@ class SDILNet: # population map. "full" preserves the original experimental ablation. self.nuis_rho = nuis_rho gn = torch.Generator(device="cpu").manual_seed(nuis_seed) + gc = torch.Generator(device="cpu").manual_seed(nuis_seed + 1) if predictor_mode == "diagonal": # Positive log-normal slopes model ordinary soma-dendrite coupling. self.Bnuis = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gn)).to( device=device, dtype=dtype) for i in range(self.L - 1)] + # A second, independently generated traffic family carries the + # network's own high-level state back along corresponding feature + # channels. In constant-width residual networks this models normal + # contextual/top-down traffic rather than a target constructed from + # the local predictor. P still sees only the same cell's soma. + self.Bcontext = [torch.exp(0.25 * torch.randn(sizes[i + 1], generator=gc)).to( + device=device, dtype=dtype) for i in range(self.L - 1)] self.P = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L - 1)] self.P_bias = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) @@ -151,6 +162,9 @@ class SDILNet: else: self.Bnuis = [he((sizes[i + 1], sizes[i + 1]), 1.0, gen=gn) for i in range(self.L - 1)] + context_dim = sizes[-2] + self.Bcontext = [he((sizes[i + 1], context_dim), 1.0, gen=gc) + for i in range(self.L - 1)] self.P = [torch.zeros((sizes[i + 1], sizes[i + 1]), device=device, dtype=dtype) for i in range(self.L - 1)] self.P_bias = None @@ -195,14 +209,45 @@ class SDILNet: p = torch.softmax(logits, dim=1) return p - y_onehot - def apical(self, l, c, h_l=None): - """a_l = A_l c (+ soma-predictable nuisance rho * Bnuis h_l).""" - a = c @ self.A[l].t() - if self.nuis_rho and h_l is not None: + def apical_traffic(self, l, h_l, context=None): + """Return ordinary non-teaching traffic in the apical compartment. + + ``soma`` models back-propagating somatic activity / normal + soma-dendrite coupling. ``topdown`` uses the network's own final hidden + state, generated independently of P, as contextual feedback. ``mixed`` + combines both at matched total scale. The legacy ``nuis_rho`` argument + is retained as the traffic amplitude so existing result protocols are + exactly backward compatible. + """ + if not self.nuis_rho or self.traffic_mode == "none": + return torch.zeros_like(h_l) + + if self.predictor_mode == "diagonal": + soma = self.Bnuis[l] * h_l + else: + soma = h_l @ self.Bnuis[l].t() + + if self.traffic_mode == "soma": + traffic = soma + else: + if context is None: + raise ValueError("topdown apical traffic requires a context state") if self.predictor_mode == "diagonal": - a = a + self.nuis_rho * self.Bnuis[l] * h_l + if context.shape[1] != h_l.shape[1]: + raise ValueError( + "diagonal topdown traffic requires constant hidden width; " + f"got context={context.shape[1]} and layer={h_l.shape[1]}") + topdown = self.Bcontext[l] * context else: - a = a + self.nuis_rho * (h_l @ self.Bnuis[l].t()) + topdown = context @ self.Bcontext[l].t() + traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2 + return self.nuis_rho * traffic + + def apical(self, l, c, h_l=None, context=None): + """a_l = A_l c + ordinary non-teaching apical traffic.""" + a = c @ self.A[l].t() + if h_l is not None: + a = a + self.apical_traffic(l, h_l, context) return a def baseline(self, l, h_l): @@ -211,15 +256,15 @@ class SDILNet: return self.P[l] * h_l + self.P_bias[l] return h_l @ self.P[l].t() - def innovation(self, l, c, h_l, use_residual=True): + def innovation(self, l, c, h_l, use_residual=True, context=None): """r_l = a_l - ahat_l (or raw a_l if use_residual=False -> ablation).""" - a = self.apical(l, c, h_l) + a = self.apical(l, c, h_l, context) if use_residual: return a - self.baseline(l, h_l), a return a, a -def teaching_signal(net, l, c, h_l, cfg): +def teaching_signal(net, l, c, h_l, cfg, context=None): """Select the signal that drives a hidden layer's local plasticity. ``match_innovation_norm`` is a deliberately strong raw-apical control: it @@ -231,7 +276,7 @@ def teaching_signal(net, l, c, h_l, cfg): Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic probes use one definition of every signal. """ - a = net.apical(l, c, h_l) + a = net.apical(l, c, h_l, context) innovation = a - net.baseline(l, h_l) if cfg.use_residual: teaching = innovation @@ -416,10 +461,12 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): if cfg.settle_steps > 0: h = _settle(net, h, u, c, cfg) + context = h[-2] + # ---- compute innovations for hidden layers ---- r_list, a_list = [], [] for l in range(net.L - 1): - r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg) + r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg, context) r_list.append(r) a_list.append(a) @@ -476,7 +523,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): zero_c = torch.zeros_like(c) for l in range(net.L - 1): if cfg.p_update_on_neutral: - target = net.apical(l, zero_c, h[l + 1]) # nuisance only + target = net.apical(l, zero_c, h[l + 1], context) # traffic only else: target = a_list[l] # full apical (eats signal) resid = target - net.baseline(l, h[l + 1]) @@ -494,11 +541,12 @@ def neutral_p_update(net, x, eta): with torch.no_grad(): fwd = net.forward(x) h = fwd["h"] + context = h[-2] B = x.shape[0] zc = torch.zeros(B, net.n_classes, device=x.device, dtype=x.dtype) for l in range(net.L - 1): hl = h[l + 1] - target = net.apical(l, zc, hl) # nuisance only (c=0) + target = net.apical(l, zc, hl, context) # ordinary traffic only (c=0) resid = target - net.baseline(l, hl) _update_predictor(net, l, resid, hl, eta) @@ -524,8 +572,9 @@ def _settle(net, h, u, c, cfg): h = list(h) dt_over_tau = 1.0 / max(1, cfg.settle_steps) for _ in range(cfg.settle_steps): + context = h[-2] for l in range(net.L - 1): - r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg) + r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg, context) 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 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 -- cgit v1.2.3