From 0f107d1ee4e750421a0e95f8bd06bf0f629d2fa7 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 07:02:04 -0500 Subject: feat: multiplex causal probes across hidden layers --- sdil/core.py | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 52 insertions(+), 3 deletions(-) (limited to 'sdil/core.py') diff --git a/sdil/core.py b/sdil/core.py index 3115d46..66f2d0a 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -264,6 +264,50 @@ def node_perturbation_targets(net, x, y, sigma=1e-2, antithetic=True, n_dirs=1): return qs +def simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, + antithetic=True, n_dirs=1): + """Estimate every hidden-layer target with shared forward evaluations. + + Each direction injects an independent Rademacher perturbation at *every* + hidden activation during one forward pass. The scalar loss difference is + demultiplexed locally as q_l = -dL/dsigma * xi_l. Cross-layer terms add + variance but vanish in expectation because the xi_l are independent. This + changes causal-calibration cost from layerwise O(n_dirs * depth^2) tail + replays to O(n_dirs * depth) full-network work. + """ + with torch.no_grad(): + base = net.forward(x) + h = base["h"] + qs = [torch.zeros_like(hl) for hl in h[1:-1]] + L0 = loss_ce(h[-1], y) + for _ in range(n_dirs): + xis = [torch.empty_like(q).bernoulli_(0.5).mul_(2).sub_(1) for q in qs] + Lplus = _loss_with_hidden_noise(net, x, y, xis, sigma) + if antithetic: + Lminus = _loss_with_hidden_noise(net, x, y, xis, -sigma) + coeff = (Lplus - Lminus) / (2.0 * sigma) + else: + coeff = (Lplus - L0) / sigma + for l, xi in enumerate(xis): + qs[l].add_(-coeff.unsqueeze(1) * xi) + return [q / n_dirs for q in qs] + + +def _loss_with_hidden_noise(net, x, y, noises, scale): + """Forward loss with additive interventions after every hidden layer.""" + cur = x + for l in range(net.L): + pre = cur + u = pre @ net.W[l].t() + net.b[l] + if l < net.L - 1: + act = net.act(u) + cur = pre + net.res_alpha * act if net.residual and l >= 1 else act + cur = cur + scale * noises[l] + else: + cur = u + return loss_ce(cur, y) + + def _relayer_loss(net, h, start_idx, h_start_new, y): """Recompute loss when hidden activation at layer index `start_idx` (1-based into h) is replaced by h_start_new, propagating forward. @@ -288,7 +332,7 @@ 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): + normalize_delta=False, pert_mode="layerwise"): self.eta = eta self.eta_A = eta_A self.eta_P = eta_P @@ -298,6 +342,9 @@ class SDILConfig: self.pert_sigma = pert_sigma self.pert_every = pert_every # run node-perturbation every k steps self.pert_ndirs = pert_ndirs # directions averaged per perturbation + if pert_mode not in ("layerwise", "simultaneous"): + raise ValueError(f"unknown pert_mode: {pert_mode}") + self.pert_mode = pert_mode self.momentum = momentum self.wd = wd self.settle_steps = settle_steps # online apical control iterations @@ -345,8 +392,10 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): did_pert = cfg.learn_A and (step % cfg.pert_every == 0) qs = None if did_pert: - qs = node_perturbation_targets( - net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs) + estimator = (simultaneous_node_perturbation_targets + if cfg.pert_mode == "simultaneous" + else node_perturbation_targets) + qs = estimator(net, x, y, sigma=cfg.pert_sigma, n_dirs=cfg.pert_ndirs) # ================= forward weight updates ================= # hidden layers: three-factor local rule -- cgit v1.2.3