diff options
| -rw-r--r-- | experiments/run.py | 7 | ||||
| -rw-r--r-- | experiments/smoke.py | 24 | ||||
| -rw-r--r-- | sdil/core.py | 53 | ||||
| -rw-r--r-- | sdil/probes.py | 4 |
4 files changed, 82 insertions, 6 deletions
diff --git a/experiments/run.py b/experiments/run.py index 93c1862..3b0f961 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -190,7 +190,8 @@ def build(args, device): if args.mode == "bp": net = BPNet(sizes, act=args.act, device=device, seed=args.seed, w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual), - predictor_mode=args.predictor_mode) + predictor_mode=args.predictor_mode, + vectorizer_mode=args.vectorizer_mode) cfg = SDILConfig(eta=args.eta, momentum=args.momentum) return net, cfg if args.mode == "fa": @@ -198,12 +199,14 @@ def build(args, device): w_scale=args.w_scale, nuis_rho=0.0, residual=bool(args.residual), predictor_mode=args.predictor_mode, + vectorizer_mode=args.vectorizer_mode, b_scale=args.feedback_scale) return net, SDILConfig(eta=args.eta, momentum=args.momentum) 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, + vectorizer_mode=args.vectorizer_mode, traffic_mode=args.traffic_mode, nuis_seed=args.traffic_seed) if args.mode == "dfa": cfg = dfa_config(eta=args.eta, momentum=args.momentum) @@ -583,6 +586,8 @@ def get_args(): 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("--vectorizer_mode", default="linear", + choices=["linear", "soma_gated", "context_gated"]) p.add_argument("--normalize_delta", type=int, default=0) p.add_argument("--settle_steps", type=int, default=0) p.add_argument("--kappa", type=float, default=0.0) diff --git a/experiments/smoke.py b/experiments/smoke.py index 025d3cc..fcc0e9c 100644 --- a/experiments/smoke.py +++ b/experiments/smoke.py @@ -132,6 +132,29 @@ def check_traffic_seed_isolation(): print("CHECK0d traffic seed changes feedback only: passed") +def check_state_conditioned_vectorizers(): + """Gated apical features must start as linear feedback and learn locally.""" + torch.manual_seed(17) + x = torch.randn(48, 6) + y = torch.randint(0, 2, (48,)) + yoh = onehot(y, 2) + for mode in ("soma_gated", "context_gated"): + net = SDILNet([6, 8, 8, 2], act="relu", device="cpu", seed=9, + residual=True, vectorizer_mode=mode) + fwd = net.forward(x) + c = net.output_error(fwd["h"][-1], yoh) + context = fwd["h"][-2] + before = net.vectorizer(0, c, fwd["h"][1], context) + assert torch.allclose(before, c @ net.A[0].t()) + gates_before = [gate.clone() for gate in net.A_gate] + cfg = SDILConfig(eta=0.01, eta_A=0.02, pert_every=1, + pert_ndirs=2, pert_mode="simultaneous") + sdil_step(net, x, y, yoh, cfg, step=0) + assert any(not torch.equal(old, new) + for old, new in zip(gates_before, net.A_gate)) + print("CHECK0e state-conditioned vectorizers: zero-init and local calibration passed") + + def main(): torch.manual_seed(0) dev = "cpu" @@ -139,6 +162,7 @@ def main(): check_neutral_predictor() check_topdown_predictor() check_traffic_seed_isolation() + check_state_conditioned_vectorizers() 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 8b3f923..e66dc39 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -102,7 +102,8 @@ 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", traffic_mode="soma"): + predictor_mode="diagonal", traffic_mode="soma", + vectorizer_mode="linear"): # sizes = [n_in, n_h1, ..., n_hk, n_out] self.sizes = list(sizes) self.L = len(sizes) - 1 # number of weight layers @@ -132,8 +133,24 @@ class SDILNet: self.W = [he((sizes[i + 1], sizes[i]), w_scale) for i in range(self.L)] self.b = [torch.zeros(sizes[i + 1], device=device, dtype=dtype) for i in range(self.L)] - # apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l), hidden layers only (l=1..L-1) + # Apical vectorizer A_l : c (n_classes) -> hidden velocity (n_l), + # hidden layers only. A fixed linear map cannot represent credit that + # changes across activation regions. Optional zero-initialized gated + # terms preserve the initial linear model while letting perturbation + # calibration learn state-dependent feedback locally. + if vectorizer_mode not in ("linear", "soma_gated", "context_gated"): + raise ValueError(f"unknown vectorizer_mode: {vectorizer_mode}") + self.vectorizer_mode = vectorizer_mode self.A = [he((sizes[i + 1], self.n_classes), a_scale) for i in range(self.L - 1)] + if vectorizer_mode == "linear": + self.A_gate = None + elif vectorizer_mode == "soma_gated": + self.A_gate = [torch.zeros_like(weight) for weight in self.A] + else: + context_dim = sizes[-2] + self.A_gate = [torch.zeros((sizes[i + 1], self.n_classes * context_dim), + device=device, dtype=dtype) + for i in range(self.L - 1)] if predictor_mode not in ("diagonal", "full"): raise ValueError(f"unknown predictor_mode: {predictor_mode}") self.predictor_mode = predictor_mode @@ -245,9 +262,29 @@ class SDILNet: traffic = topdown if self.traffic_mode == "topdown" else (soma + topdown) / _SQRT2 return self.nuis_rho * traffic + def vectorizer(self, l, c, h_l=None, context=None): + """Return the learned instructional component of apical activity. + + ``soma_gated`` adds ``h_i G_i c`` independently in every cell. + ``context_gated`` adds a linear readout of ``c outer context``. Both + vanish when c=0, so neutral-period predictor identification remains + unchanged. + """ + value = c @ self.A[l].t() + if self.vectorizer_mode == "linear": + return value + if self.vectorizer_mode == "soma_gated": + if h_l is None: + raise ValueError("soma-gated vectorizer requires somatic activity") + return value + h_l * (c @ self.A_gate[l].t()) + if context is None: + raise ValueError("context-gated vectorizer requires a context state") + features = (c.unsqueeze(2) * context.unsqueeze(1)).flatten(1) + return value + features @ self.A_gate[l].t() + 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() + a = self.vectorizer(l, c, h_l, context) if h_l is not None: a = a + self.apical_traffic(l, h_l, context) return a @@ -510,8 +547,16 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): # ================= apical vectorizer A via node perturbation ======= if did_pert: for l in range(net.L - 1): - dA = (qs[l] - r_list[l]).t() @ c / B # (n_l, n_classes) + calibration_error = qs[l] - r_list[l] + dA = calibration_error.t() @ c / B # (n_l, n_classes) net.A[l] += cfg.eta_A * dA + if net.vectorizer_mode == "soma_gated": + dgate = (calibration_error * h[l + 1]).t() @ c / B + net.A_gate[l] += cfg.eta_A * dgate + elif net.vectorizer_mode == "context_gated": + features = (c.unsqueeze(2) * context.unsqueeze(1)).flatten(1) + dgate = calibration_error.t() @ features / B + net.A_gate[l] += cfg.eta_A * dgate # ================= predictor P (neutral) ========================== # KEY identification condition. P must learn the soma->apical coupling diff --git a/sdil/probes.py b/sdil/probes.py index 1698d75..0075bca 100644 --- a/sdil/probes.py +++ b/sdil/probes.py @@ -74,7 +74,7 @@ def alignment_report(net, x, y, y_onehot, cfg): negg = -g teaching, a, innovation = core.teaching_signal( net, l, c, h[l + 1], cfg, context) - Ac = c @ net.A[l].t() # pure vectorizer output + Ac = net.vectorizer(l, c, h[l + 1], context) # 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, @@ -145,6 +145,8 @@ def _clone(net): n.W = [w.clone() for w in net.W] n.b = [b.clone() for b in net.b] n.A = [a.clone() for a in net.A] + n.A_gate = ([a.clone() for a in net.A_gate] + if getattr(net, "A_gate", None) is not None else None) n.P = [p.clone() for p in net.P] n.P_bias = ([p.clone() for p in net.P_bias] if getattr(net, "P_bias", None) is not None else None) |