diff options
| -rwxr-xr-x | experiments/nuisance_control_sweep.sh | 66 | ||||
| -rw-r--r-- | experiments/run.py | 5 | ||||
| -rw-r--r-- | experiments/smoke.py | 22 | ||||
| -rw-r--r-- | sdil/core.py | 35 | ||||
| -rw-r--r-- | sdil/probes.py | 21 |
5 files changed, 137 insertions, 12 deletions
diff --git a/experiments/nuisance_control_sweep.sh b/experiments/nuisance_control_sweep.sh new file mode 100755 index 0000000..ef715f4 --- /dev/null +++ b/experiments/nuisance_control_sweep.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +# Three-way residualization control with per-neuron soma-dendrite coupling: +# raw = unmodified apical teaching signal +# matched = raw direction, per-sample norm matched to the innovation +# residual = soma-dendritic innovation +# +# Usage: +# bash experiments/nuisance_control_sweep.sh \ +# <gpu> "<rhos>" "<seeds>" <epochs> <prefix> [pert_mode] [n_dirs] +# +# Example: +# bash experiments/nuisance_control_sweep.sh \ +# 5 "0 0.05 0.2 0.5" "0 1 2" 15 nuis_ctrl simultaneous 16 +set -eu + +cd "$(dirname "$0")/.." +PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3 +GPU="${1:?GPU index required}" +RHOS="${2:-0 0.05 0.2 0.5}" +SEEDS="${3:-0}" +EPOCHS="${4:-15}" +PREFIX="${5:-nuis_ctrl}" +PERT_MODE="${6:-simultaneous}" +N_DIRS="${7:-16}" + +export CUDA_VISIBLE_DEVICES="$GPU" +export OMP_NUM_THREADS=2 +mkdir -p results logs/nuisance + +run_one() { + local rho="$1" + local arm="$2" + local seed="$3" + local residual=0 + local scale_control=none + local rho_tag="${rho//./p}" + local tag="${PREFIX}_rho${rho_tag}_${arm}_s${seed}" + local result="results/${tag}.json" + local log="logs/nuisance/${tag}.log" + if [[ "$arm" == residual ]]; then + residual=1 + elif [[ "$arm" == matched ]]; then + scale_control=match_innovation_norm + fi + if [[ -f "$result" ]]; then + echo "skip $tag (result exists)" + return + fi + echo ">>> $tag $(date --iso-8601=seconds) gpu=$GPU" + "$PY" experiments/run.py \ + --mode sdil --dataset mnist --depth 3 --width 256 --epochs "$EPOCHS" \ + --seed "$seed" --nuis_rho "$rho" --use_residual "$residual" \ + --raw_scale_control "$scale_control" --predictor_mode diagonal \ + --p_warmup_steps 200 --p_warmup_eta 0.05 \ + --pert_mode "$PERT_MODE" --pert_ndirs "$N_DIRS" --log_every 100 \ + --tag "$tag" --outdir results > "$log" 2>&1 + grep -h DONE "$log" +} + +for seed in $SEEDS; do + for rho in $RHOS; do + for arm in raw matched residual; do + run_one "$rho" "$arm" "$seed" + done + done +done diff --git a/experiments/run.py b/experiments/run.py index 3cb2e00..eea3f91 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -65,7 +65,8 @@ def build(args, device): momentum=args.momentum, settle_steps=args.settle_steps, kappa=args.kappa, feedback=args.feedback, p_update_on_neutral=bool(args.p_neutral), - normalize_delta=bool(args.normalize_delta)) + normalize_delta=bool(args.normalize_delta), + raw_scale_control=args.raw_scale_control) else: raise ValueError(args.mode) return net, cfg @@ -174,6 +175,8 @@ def get_args(): 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("--raw_scale_control", default="none", + choices=["none", "match_innovation_norm"]) p.add_argument("--learn_A", type=int, default=1) p.add_argument("--learn_P", type=int, default=1) p.add_argument("--p_neutral", type=int, default=1) # P update on neutral (c=0) drive diff --git a/experiments/smoke.py b/experiments/smoke.py index 575f0bd..0bdb0ca 100644 --- a/experiments/smoke.py +++ b/experiments/smoke.py @@ -8,7 +8,7 @@ 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, simultaneous_node_perturbation_targets, - neutral_p_update) + neutral_p_update, teaching_signal) from sdil.baselines import dfa_config from sdil import probes from sdil.data import get_dataset, onehot @@ -164,6 +164,26 @@ def main(): assert al_raw["cos_r_negg"] == al_raw["cos_apical_negg"], ( "reported teaching alignment must honor use_residual=False") + # A magnitude-matched raw signal is still pointed along raw apical activity, + # but has exactly the innovation's norm for every sample. This isolates the + # directional effect of subtracting the soma-predictable component. + cfg_matched = SDILConfig(use_residual=False, learn_A=False, learn_P=True, + raw_scale_control="match_innovation_norm") + with torch.no_grad(): + fwd_n = net_n.forward(x) + error_n = net_n.output_error(fwd_n["h"][-1], yoh) + for l in range(net_n.L - 1): + matched, raw, innovation = teaching_signal( + net_n, l, error_n, fwd_n["h"][l + 1], cfg_matched) + assert torch.allclose(matched.norm(dim=1), innovation.norm(dim=1), + atol=1e-6, rtol=1e-5) + assert row_cos(matched, raw) > 0.99999 + al_matched = probes.alignment_report(net_n, x, y, yoh, cfg_matched) + for matched_cos, raw_cos in zip(al_matched["cos_r_negg"], + al_matched["cos_apical_negg"]): + assert abs(matched_cos - raw_cos) < 1e-6 + print(" matched raw: direction preserved; per-sample innovation norm matched") + # ---- CHECK 5: single-step loss-decrease ratio vs exact GD ---- print("\nCHECK5 single-step loss-decrease ratio vs exact GD:") ldr = probes.loss_decrease_ratio(net, x, y, yoh, cfg, step=100) diff --git a/sdil/core.py b/sdil/core.py index 3cd8a97..d7be8c3 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -219,6 +219,31 @@ class SDILNet: return a, a +def teaching_signal(net, l, c, h_l, cfg): + """Select the signal that drives a hidden layer's local plasticity. + + ``match_innovation_norm`` is a deliberately strong raw-apical control: it + preserves each sample's raw apical direction exactly while rescaling its + Euclidean norm to the norm of the corresponding innovation. Therefore a + raw-vs-residual difference under this control cannot be attributed merely + to the nuisance making the teaching vector larger. + + Returns ``(teaching, raw_apical, innovation)`` so training and diagnostic + probes use one definition of every signal. + """ + a = net.apical(l, c, h_l) + innovation = a - net.baseline(l, h_l) + if cfg.use_residual: + teaching = innovation + elif cfg.raw_scale_control == "match_innovation_norm": + raw_norm = a.norm(dim=1, keepdim=True) + innovation_norm = innovation.norm(dim=1, keepdim=True) + teaching = a * (innovation_norm / raw_norm.clamp_min(1e-12)) + else: + teaching = a + return teaching, a, innovation + + # -------------------------------------------------------------------------- # node perturbation: causal estimate of -grad_{h_l} L used to TRAIN A_l # -------------------------------------------------------------------------- @@ -341,11 +366,15 @@ 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, pert_mode="layerwise"): + normalize_delta=False, pert_mode="layerwise", + raw_scale_control="none"): self.eta = eta self.eta_A = eta_A self.eta_P = eta_P self.use_residual = use_residual # ablation: residual vs raw apical + if raw_scale_control not in ("none", "match_innovation_norm"): + raise ValueError(f"unknown raw_scale_control: {raw_scale_control}") + self.raw_scale_control = raw_scale_control self.learn_A = learn_A # ablation: perturbation-trained vs fixed-random A self.learn_P = learn_P # ablation: predictor on/off self.pert_sigma = pert_sigma @@ -390,7 +419,7 @@ def sdil_step(net, x, y, y_onehot, cfg, step, prev_error=None): # ---- compute innovations for hidden layers ---- r_list, a_list = [], [] for l in range(net.L - 1): - r, a = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) + r, a, _ = teaching_signal(net, l, c, h[l + 1], cfg) r_list.append(r) a_list.append(a) @@ -496,7 +525,7 @@ def _settle(net, h, u, c, cfg): dt_over_tau = 1.0 / max(1, cfg.settle_steps) for _ in range(cfg.settle_steps): for l in range(net.L - 1): - r, _ = net.innovation(l, c, h[l + 1], use_residual=cfg.use_residual) + r, _, _ = teaching_signal(net, l, c, h[l + 1], cfg) 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 16aca21..fad82ca 100644 --- a/sdil/probes.py +++ b/sdil/probes.py @@ -36,11 +36,20 @@ def true_hidden_grads(net, x, y): return [g.detach() for g in grads], loss.item() -def _row_cos(u, v, eps=1e-12): - """Mean per-sample cosine between rows of u and v.""" +def _row_cos(u, v): + """Mean per-sample cosine, excluding rows with an exact zero vector. + + Adding an absolute epsilon to the denominator makes cosine spuriously + depend on signal magnitude near convergence. In particular, a positive + per-sample rescaling can then appear to change direction. Exact zero rows + have no defined direction and are omitted instead. + """ num = (u * v).sum(1) - den = u.norm(dim=1) * v.norm(dim=1) + eps - return (num / den).mean().item() + den = u.norm(dim=1) * v.norm(dim=1) + valid = den > 0 + if not valid.any(): + return float("nan") + return (num[valid] / den[valid]).mean().item() @torch.no_grad() @@ -61,10 +70,8 @@ def alignment_report(net, x, y, y_onehot, cfg): for l in range(net.L - 1): g = grads[l] negg = -g - a = net.apical(l, c, h[l + 1]) # raw apical (with nuisance) + teaching, a, innovation = core.teaching_signal(net, l, c, h[l + 1], cfg) Ac = c @ net.A[l].t() # pure vectorizer output - innovation = a - net.baseline(l, h[l + 1]) - teaching = innovation if cfg.use_residual else a # 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]) |