diff options
| -rw-r--r-- | experiments/run.py | 11 | ||||
| -rw-r--r-- | experiments/smoke.py | 32 | ||||
| -rw-r--r-- | sdil/core.py | 79 | ||||
| -rw-r--r-- | sdil/probes.py | 18 |
4 files changed, 122 insertions, 18 deletions
diff --git a/experiments/run.py b/experiments/run.py index b412b04..cc3e475 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -65,7 +65,8 @@ def build(args, device): net = SDILNet(sizes, act=args.act, device=device, seed=args.seed, w_scale=args.w_scale, a_scale=args.a_scale, nuis_rho=args.nuis_rho, feedback=args.feedback, - residual=bool(args.residual), predictor_mode=args.predictor_mode) + residual=bool(args.residual), predictor_mode=args.predictor_mode, + traffic_mode=args.traffic_mode) if args.mode == "dfa": cfg = dfa_config(eta=args.eta, momentum=args.momentum) elif args.mode == "sdil": @@ -169,6 +170,9 @@ def train(args): rec["cos_apical_negg"] = al["cos_apical_negg"] rec["cos_Ac_negg"] = al["cos_Ac_negg"] rec["r_norm"] = al["r_norm"] + rec["traffic_norm"] = al["traffic_norm"] + rec["traffic_residual_norm"] = al["traffic_residual_norm"] + rec["traffic_r2"] = al["traffic_r2"] if args.mode == "sdil" and step % (args.log_every * 5) == 0: rec["ldr"] = probes.loss_decrease_ratio(net, px, py, poh, cfg, step) log["steps"].append(rec) @@ -201,6 +205,9 @@ def train(args): log["final"]["cos_innovation_negg"] = al["cos_innovation_negg"] log["final"]["cos_apical_negg"] = al["cos_apical_negg"] log["final"]["cos_Ac_negg"] = al["cos_Ac_negg"] + log["final"]["traffic_norm"] = al["traffic_norm"] + log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"] + log["final"]["traffic_r2"] = al["traffic_r2"] os.makedirs(args.outdir, exist_ok=True) outpath = os.path.join(args.outdir, f"{args.tag}.json") with open(outpath, "w") as f: @@ -252,6 +259,8 @@ def get_args(): p.add_argument("--p_warmup_steps", type=int, default=200) # pre-task neutral P warmup p.add_argument("--p_warmup_eta", type=float, default=0.05) p.add_argument("--nuis_rho", type=float, default=0.0) + p.add_argument("--traffic_mode", default="soma", + choices=["none", "soma", "topdown", "mixed"]) p.add_argument("--predictor_mode", default="diagonal", choices=["diagonal", "full"]) p.add_argument("--normalize_delta", type=int, default=0) p.add_argument("--settle_steps", type=int, default=0) diff --git a/experiments/smoke.py b/experiments/smoke.py index 0bdb0ca..83ee896 100644 --- a/experiments/smoke.py +++ b/experiments/smoke.py @@ -82,11 +82,43 @@ def check_neutral_predictor(): assert after < 0.03 +def check_topdown_predictor(): + """A local soma predictor should remove the predictable part of feedback + generated by the network's own high-level contextual state.""" + torch.manual_seed(13) + net = SDILNet([20, 32, 32, 32, 5], device="cpu", seed=6, nuis_rho=1.0, + predictor_mode="diagonal", residual=True, + traffic_mode="topdown") + x = torch.randn(256, 20) + + def unexplained_fraction(): + with torch.no_grad(): + h = net.forward(x)["h"] + context = h[-2] + fractions = [] + for l in range(net.L - 1): + traffic = net.apical_traffic(l, h[l + 1], context) + residual = traffic - net.baseline(l, h[l + 1]) + fractions.append((residual.pow(2).mean() / traffic.pow(2).mean()).item()) + return fractions + + initial = unexplained_fraction() + for _ in range(500): + neutral_p_update(net, x, 0.02) + final = unexplained_fraction() + print("CHECK0c top-down traffic unexplained power:") + for l, (before, after) in enumerate(zip(initial, final)): + print(f" layer {l}: {before:.4f} -> {after:.4f}") + assert after < before * 0.8 + assert final[-1] < 0.03 + + def main(): torch.manual_seed(0) dev = "cpu" check_residual_local_jacobian() check_neutral_predictor() + check_topdown_predictor() print("loading MNIST subset...") tr, te, n_in, n_out = get_dataset("mnist", batch_size=128, device=dev) xb, yb = next(iter(tr)) 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 |