diff options
| -rwxr-xr-x | experiments/narrow_depth_sweep.sh | 6 | ||||
| -rw-r--r-- | experiments/run.py | 2 | ||||
| -rw-r--r-- | experiments/smoke.py | 8 | ||||
| -rw-r--r-- | sdil/core.py | 55 |
4 files changed, 65 insertions, 6 deletions
diff --git a/experiments/narrow_depth_sweep.sh b/experiments/narrow_depth_sweep.sh index 8f62a68..3e4e6b2 100755 --- a/experiments/narrow_depth_sweep.sh +++ b/experiments/narrow_depth_sweep.sh @@ -3,7 +3,7 @@ # # Usage: # bash experiments/narrow_depth_sweep.sh \ -# <gpu> "<depths>" "<modes>" <width> <seed> <epochs> <prefix> +# <gpu> "<depths>" "<modes>" <width> <seed> <epochs> <prefix> [pert_mode] [n_dirs] # # Example: # bash experiments/narrow_depth_sweep.sh 5 "30 60" "bp dfa sdil" 64 0 5 pilot_jac @@ -18,6 +18,8 @@ WIDTH="${4:-64}" SEED="${5:-0}" EPOCHS="${6:-5}" PREFIX="${7:-narrow}" +PERT_MODE="${8:-layerwise}" +N_DIRS="${9:-8}" export CUDA_VISIBLE_DEVICES="$GPU" export OMP_NUM_THREADS=2 @@ -41,7 +43,7 @@ run_one() { "$PY" experiments/run.py \ --mode "$mode" --dataset cifar10 --depth "$depth" --width "$WIDTH" \ --act tanh --residual 1 --epochs "$EPOCHS" --seed "$SEED" \ - --eta "$eta" --eta_A 0.02 --pert_ndirs 8 \ + --eta "$eta" --eta_A 0.02 --pert_ndirs "$N_DIRS" --pert_mode "$PERT_MODE" \ --log_every 500 --probe_bs 256 --tag "$tag" --outdir results \ > "$log" 2>&1 grep -h DONE "$log" diff --git a/experiments/run.py b/experiments/run.py index bb2d151..eb9f356 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -61,6 +61,7 @@ def build(args, device): use_residual=bool(args.use_residual), learn_A=bool(args.learn_A), learn_P=bool(args.learn_P), pert_sigma=args.pert_sigma, pert_every=args.pert_every, pert_ndirs=args.pert_ndirs, + pert_mode=args.pert_mode, momentum=args.momentum, settle_steps=args.settle_steps, kappa=args.kappa, feedback=args.feedback, p_update_on_neutral=bool(args.p_neutral), @@ -169,6 +170,7 @@ def get_args(): p.add_argument("--pert_sigma", type=float, default=1e-2) p.add_argument("--pert_every", type=int, default=4) p.add_argument("--pert_ndirs", type=int, default=4) + p.add_argument("--pert_mode", default="layerwise", choices=["layerwise", "simultaneous"]) p.add_argument("--use_residual", type=int, default=1) p.add_argument("--learn_A", type=int, default=1) p.add_argument("--learn_P", type=int, default=1) diff --git a/experiments/smoke.py b/experiments/smoke.py index 8196d4f..f440f7f 100644 --- a/experiments/smoke.py +++ b/experiments/smoke.py @@ -7,7 +7,8 @@ import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.core import (SDILNet, SDILConfig, sdil_step, - node_perturbation_targets, neutral_p_update) + node_perturbation_targets, simultaneous_node_perturbation_targets, + neutral_p_update) from sdil.baselines import dfa_config from sdil import probes from sdil.data import get_dataset, onehot @@ -104,6 +105,11 @@ def main(): print(f"CHECK1 node-pert n_dirs={ndirs:2d} cos(q,-g) per layer = " + " ".join(f"{c:+.3f}" for c in cs)) assert all(c > 0 for c in cs), "node perturbation q should align with -grad" + qs = simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, n_dirs=32) + cs = [row_cos(qs[l], -grads[l]) for l in range(len(qs))] + print("CHECK1 simultaneous n_dirs=32 cos(q,-g) per layer = " + + " ".join(f"{c:+.3f}" for c in cs)) + assert all(c > 0.2 for c in cs), "simultaneous perturbation should align in expectation" # ---- CHECK 2: SDIL overfits a fixed batch and alignment climbs ---- net = SDILNet(sizes, device=dev, seed=2) 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 |