""" SDIL main training / diagnostics driver. Trains one of {bp, dfa, sdil} (sdil with ablation flags) on MNIST/FashionMNIST, logging the quantities that actually test the hypothesis: - train loss, test accuracy - per-hidden-layer cos(innovation r_l, -grad h_l) <- the headline metric - cos(raw apical a_l, -grad) and cos(A_l c, -grad) <- residualization ablation - single-step loss-decrease ratio vs exact GD Everything is JSON-logged for later plotting. """ import argparse import json import math import os import subprocess import sys import time import torch import torch.nn.functional as F sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.core import (SDILNet, SDILConfig, apical_calibration_step, sdil_step, neutral_p_update) from sdil.baselines import BPNet, dfa_config, evaluate from sdil.local_baselines import FANet from sdil import probes from sdil.data import (get_dataset_splits, onehot, make_hierarchical, make_teacher_student, make_tentmap, split_training_loader) REAL_DATASETS = ("mnist", "fmnist", "cifar10") SYNTHETIC_DATASETS = ("teacher", "hierarchical", "tentmap") def code_provenance(): """Best-effort source revision metadata for reproducible result files.""" root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) try: commit = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=root, check=True, capture_output=True, text=True).stdout.strip() dirty = bool(subprocess.run( ["git", "status", "--porcelain", "--untracked-files=no"], cwd=root, check=True, capture_output=True, text=True).stdout.strip()) return {"git_commit": commit, "git_dirty": dirty} except (OSError, subprocess.CalledProcessError): return {"git_commit": None, "git_dirty": None} def device_sync(device): if str(device).startswith("cuda") and torch.cuda.is_available(): torch.cuda.synchronize() def reset_peak_memory(device): if str(device).startswith("cuda") and torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats(torch.device(device)) def hardware_report(device): report = {"device": str(device), "torch_version": torch.__version__} if str(device).startswith("cuda") and torch.cuda.is_available(): cuda_device = torch.device(device) props = torch.cuda.get_device_properties(cuda_device) report.update({ "cuda_device_name": props.name, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "device_total_memory_bytes": props.total_memory, "peak_memory_allocated_bytes": torch.cuda.max_memory_allocated(cuda_device), "peak_memory_reserved_bytes": torch.cuda.max_memory_reserved(cuda_device), }) else: report.update({ "cuda_device_name": None, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), "device_total_memory_bytes": None, "peak_memory_allocated_bytes": None, "peak_memory_reserved_bytes": None, }) return report def fixed_training_probe(train_loader, probe_bs, device): """Return a deterministic diagnostic batch without advancing data order. All project loaders expose their in-memory tensors. Slicing those tensors avoids consuming the loader's shuffle generator and, importantly, keeps a held-out validation/test split out of gradient-alignment diagnostics. """ if not hasattr(train_loader, "x") or not hasattr(train_loader, "y"): raise TypeError("diagnostic probes require an in-memory training loader") n = min(probe_bs, len(train_loader.x)) if n <= 0: raise ValueError("diagnostic probe requires at least one training example") return train_loader.x[:n].to(device), train_loader.y[:n].to(device) @torch.no_grad() def residual_lesion_report(net, evaluation_loader, probe_x, fraction): """Measure whether a trained residual network uses its late blocks. Interior hidden layers are the constant-width residual blocks after the input projection. The lesion bypasses the final ``ceil(fraction * n)`` such blocks at inference, while branch/skip RMS ratios are measured on a fixed training-prefix probe. Training is never rerun after selecting the lesion, and the output readout is unchanged. """ if not getattr(net, "residual", False): raise ValueError("residual lesion requires a residual network") if not 0.0 < fraction <= 1.0: raise ValueError(f"lesion fraction must be in (0, 1], got {fraction}") interior = list(range(1, net.L - 1)) if not interior: raise ValueError("residual lesion requires at least one interior block") n_lesion = max(1, math.ceil(fraction * len(interior))) lesioned = interior[-n_lesion:] lesioned_set = set(lesioned) def forward(x, lesion=False, collect_ratios=False): h = x ratios = [] for layer in range(net.L - 1): pre = h u = pre @ net.W[layer].t() + net.b[layer] act = net.act(u) if layer >= 1: branch = net.res_alpha * act if collect_ratios: branch_rms = branch.square().mean().sqrt() skip_rms = pre.square().mean().sqrt() ratios.append(float(branch_rms / (skip_rms + 1e-12))) h = pre if lesion and layer in lesioned_set else pre + branch else: h = act logits = h @ net.W[-1].t() + net.b[-1] return logits, ratios correct = 0 total = 0 total_loss = 0.0 for x, y in evaluation_loader: logits, _ = forward(x, lesion=True) total_loss += F.cross_entropy(logits, y, reduction="sum").item() correct += (logits.argmax(1) == y).sum().item() total += y.shape[0] _, ratios = forward(probe_x, collect_ratios=True) return { "interior_layers": interior, "lesioned_layers": lesioned, "branch_to_skip_rms": ratios, "lesion_eval_acc": correct / total, "lesion_eval_loss": total_loss / total, } def calibration_work_per_event(net, cfg, pert_mode=None, pert_ndirs=None): """Hardware-independent causal-calibration work for one minibatch. Forward equivalents use affine multiply-add work as the denominator. The returned loss evaluations count calls producing per-example scalar losses; the caller multiplies by the actual minibatch size for scalar observations. """ directions = cfg.pert_ndirs if pert_ndirs is None else pert_ndirs mode = cfg.pert_mode if pert_mode is None else pert_mode if mode == "simultaneous": return { "batch_loss_evaluations": 2 * directions, # The implementation performs one duplicate clean forward plus one # batched +/- forward for every direction. "forward_equivalent_batches": 1.0 + 2.0 * directions, "perturbation_batch_expansion": 2 * directions, } full_work = sum(weight.numel() for weight in net.W) tail_work = 0 for hidden_layer in range(net.L - 1): tail_work += sum(net.W[j].numel() for j in range(hidden_layer + 1, net.L)) return { # One clean baseline loss plus +/- tail loss for every hidden layer. "batch_loss_evaluations": 1 + 2 * directions * (net.L - 1), "forward_equivalent_batches": 1.0 + 2.0 * directions * tail_work / full_work, "perturbation_batch_expansion": 1, } def build(args, device): sizes = [args.n_in] + [args.width] * args.depth + [args.n_out] 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, vectorizer_mode=args.vectorizer_mode) cfg = SDILConfig(eta=args.eta, momentum=args.momentum) return net, cfg if args.mode == "fa": net = FANet(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, 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) elif args.mode == "sdil": cfg = SDILConfig( eta=args.eta, eta_A=args.eta_A, eta_P=args.eta_P, 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), normalize_delta=bool(args.normalize_delta), raw_scale_control=args.raw_scale_control, vectorizer_optimizer=args.vectorizer_optimizer, vectorizer_eps=args.vectorizer_eps) elif args.mode == "nodepert": if args.pert_every != 1: raise ValueError("direct node perturbation requires --pert_every 1") cfg = SDILConfig( eta=args.eta, use_residual=False, learn_A=False, learn_P=False, pert_sigma=args.pert_sigma, pert_every=args.pert_every, pert_ndirs=args.pert_ndirs, pert_mode=args.pert_mode, momentum=args.momentum, normalize_delta=bool(args.normalize_delta), direct_node_pert=True) else: raise ValueError(args.mode) return net, cfg def load_task(args, device): """Load a real dataset or construct one fixed synthetic task. ``task_seed`` controls examples and the target function; ``seed`` controls only the student initialization. A depth sweep therefore compares models on exactly the same compositional problem instead of silently changing the teacher between model seeds. """ if args.dataset in REAL_DATASETS: train, validation, test, n_in, n_out, split = get_dataset_splits( args.dataset, batch_size=args.batch_size, device=device, train_limit=args.train_examples or None, val_examples=args.val_examples, split_seed=args.split_seed) if args.eval_split == "validation": if validation is None: raise ValueError("--eval_split validation requires --val_examples > 0") evaluation = validation else: evaluation = test split["evaluation_split"] = args.eval_split return train, evaluation, n_in, n_out, split common = dict( n_train=args.task_train_examples, n_test=args.task_test_examples, seed=args.task_seed, batch_size=args.batch_size, device=device, ) if args.dataset == "teacher": task = make_teacher_student( n_in=args.task_n_in, n_classes=args.task_classes, t_depth=args.teacher_depth, t_width=args.teacher_width, residual=bool(args.teacher_residual), **common) elif args.dataset == "hierarchical": task = make_hierarchical( levels=args.task_levels, n_classes=args.task_classes, **common) elif args.dataset == "tentmap": task = make_tentmap( levels=args.task_levels, n_in=args.task_n_in, **common) else: raise ValueError(f"unknown dataset: {args.dataset}") train, test, n_in, n_out = task validation = None split_details = {} if args.val_examples: train, validation, split_details = split_training_loader( train, args.val_examples, args.split_seed, args.batch_size) if args.eval_split == "validation": if validation is None: raise ValueError("--eval_split validation requires --val_examples > 0") evaluation = validation else: evaluation = test split = { "dataset": args.dataset, "task_seed": args.task_seed, "train_examples": len(train.x), "test_examples": args.task_test_examples, "evaluation_split": args.eval_split, "synthetic_generator": True, } split.update(split_details) return train, evaluation, n_in, n_out, split def train(args): device = args.device torch.manual_seed(args.seed) train_loader, eval_loader, n_in, n_out, split = load_task(args, device) args.n_in, args.n_out = n_in, n_out net, cfg = build(args, device) # Reset after persistent data/model allocation so the peak includes the # resident training state plus transient perturbation-batch expansion. reset_peak_memory(device) # Diagnostics use a fixed training prefix without advancing the shuffled # training iterator. Held-out data are reserved for metric evaluation. px = py = poh = None if args.diagnostics != "none" or args.residual_lesion_fraction > 0: px, py = fixed_training_probe(train_loader, args.probe_bs, device) if args.diagnostics != "none": poh = onehot(py, n_out, device=device) log = {"args": vars(args), "split": split, "provenance": code_provenance(), "diagnostic_protocol": { "probe_source": "training_prefix" if px is not None else None, "probe_examples": len(px) if px is not None else 0, "schedule": args.diagnostics_schedule, }, "steps": [], "final": {}} step = 0 prev_error = None device_sync(device) t0 = time.time() train_wall_s = 0.0 eval_wall_s = 0.0 diagnostics_wall_s = 0.0 inline_diagnostics_wall_s = 0.0 warmup_examples = 0 ordinary_forward_examples = 0 perturbation_events = 0 calibration_batch_loss_evaluations = 0 calibration_example_loss_evaluations = 0 calibration_forward_equivalent_examples = 0.0 perturbation_batch_expansion = 0 feedback_warmup_examples = 0 feedback_warmup_events = 0 # predictor warmup on neutral-period (c=0) drive, so P cancels the apical # nuisance before task plasticity relies on the residual (no-op when rho=0). if args.mode == "sdil" and args.learn_P and args.p_warmup_steps > 0 and args.nuis_rho > 0: device_sync(device) warmup_t0 = time.time() it = iter(train_loader) for _ in range(args.p_warmup_steps): try: wx, _ = next(it) except StopIteration: it = iter(train_loader) wx, _ = next(it) wx = wx.to(device) warmup_examples += wx.shape[0] neutral_p_update(net, wx, args.p_warmup_eta) device_sync(device) warmup_wall_s = time.time() - warmup_t0 else: warmup_wall_s = 0.0 # Feedback-first timescale separation. Causal perturbations fit A on a # stationary forward network before noisy predictions can move W into a # bad basin. The output readout is frozen as well. This is supervised # calibration work, not free initialization, so every forward and scalar # loss observation is included in the same hardware-independent ledger. if args.a_warmup_steps > 0: if args.mode != "sdil" or not args.learn_A: raise ValueError("--a_warmup_steps requires SDIL with --learn_A 1") device_sync(device) feedback_warmup_t0 = time.time() loader_state = (train_loader.g.get_state().clone() if hasattr(train_loader, "g") else None) rng_devices = ([torch.cuda.current_device()] if str(device).startswith("cuda") and torch.cuda.is_available() else []) # Warmup length must not silently change the minibatch order or random # directions used by the subsequent joint phase. fork_rng and restoring # the loader generator isolate those nuisance differences while keeping # the learned A parameters. try: with torch.random.fork_rng(devices=rng_devices): it = iter(train_loader) for _ in range(args.a_warmup_steps): try: wx, wy = next(it) except StopIteration: it = iter(train_loader) wx, wy = next(it) wx, wy = wx.to(device), wy.to(device) wyoh = onehot(wy, n_out, device=device) apical_calibration_step( net, wx, wy, wyoh, cfg, pert_mode=args.a_warmup_mode, pert_ndirs=args.a_warmup_ndirs) feedback_warmup_examples += wx.shape[0] feedback_warmup_events += 1 event = calibration_work_per_event( net, cfg, pert_mode=args.a_warmup_mode, pert_ndirs=args.a_warmup_ndirs) perturbation_events += 1 calibration_batch_loss_evaluations += event["batch_loss_evaluations"] calibration_example_loss_evaluations += ( event["batch_loss_evaluations"] * wx.shape[0]) calibration_forward_equivalent_examples += ( event["forward_equivalent_batches"] * wx.shape[0]) perturbation_batch_expansion = max( perturbation_batch_expansion, event["perturbation_batch_expansion"]) finally: if loader_state is not None: train_loader.g.set_state(loader_state) device_sync(device) feedback_warmup_wall_s = time.time() - feedback_warmup_t0 else: feedback_warmup_wall_s = 0.0 for epoch in range(args.epochs): device_sync(device) train_t0 = time.time() for x, y in train_loader: x, y = x.to(device), y.to(device) ordinary_forward_examples += x.shape[0] yoh = onehot(y, n_out, device=device) if args.mode == "bp": loss = net.bp_step(x, y, cfg.eta, momentum=cfg.momentum) elif args.mode == "fa": loss = net.fa_step(x, y, yoh, cfg.eta, momentum=cfg.momentum) else: loss, aux = sdil_step(net, x, y, yoh, cfg, step, prev_error=prev_error) prev_error = aux["error"] if args.mode in ("sdil", "nodepert") and aux["did_pert"]: event = calibration_work_per_event(net, cfg) perturbation_events += 1 calibration_batch_loss_evaluations += event["batch_loss_evaluations"] calibration_example_loss_evaluations += ( event["batch_loss_evaluations"] * x.shape[0]) calibration_forward_equivalent_examples += ( event["forward_equivalent_batches"] * x.shape[0]) perturbation_batch_expansion = max( perturbation_batch_expansion, event["perturbation_batch_expansion"]) if step % args.log_every == 0: rec = {"step": step, "epoch": epoch, "train_loss": float(loss)} inline_diagnostics = (args.diagnostics != "none" and args.diagnostics_schedule == "inline") if inline_diagnostics: device_sync(device) diagnostics_t0 = time.time() if args.mode == "fa" and inline_diagnostics: rec.update(probes.fa_alignment_report(net, px, py, poh)) elif args.mode == "nodepert" and inline_diagnostics: rec.update(probes.nodepert_alignment_report(net, px, py, cfg)) elif args.mode != "bp" and inline_diagnostics: al = probes.alignment_report(net, px, py, poh, cfg) rec["cos_r_negg"] = al["cos_r_negg"] rec["cos_innovation_negg"] = al["cos_innovation_negg"] 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 args.diagnostics == "full" and step % (args.log_every * 5) == 0): rec["ldr"] = probes.loss_decrease_ratio(net, px, py, poh, cfg, step) if inline_diagnostics: device_sync(device) elapsed = time.time() - diagnostics_t0 diagnostics_wall_s += elapsed inline_diagnostics_wall_s += elapsed log["steps"].append(rec) step += 1 if args.max_steps and step >= args.max_steps: break if args.max_steps and step >= args.max_steps: device_sync(device) train_wall_s += time.time() - train_t0 break device_sync(device) train_wall_s += time.time() - train_t0 should_evaluate = args.eval_every > 0 and (epoch + 1) % args.eval_every == 0 record = {"epoch_end": epoch, "step": step} msg = f"[{args.tag}] epoch {epoch} step {step} loss {loss:.4f}" if should_evaluate: device_sync(device) eval_t0 = time.time() acc, tloss = evaluate(net, eval_loader) device_sync(device) eval_wall_s += time.time() - eval_t0 metric = "val_acc" if args.eval_split == "validation" else "test_acc" msg += f" {metric} {acc:.4f}" record.update({"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss}) if args.eval_split == "test": record.update({"test_acc": acc, "test_loss": tloss}) else: record.update({"val_acc": acc, "val_loss": tloss}) if (args.diagnostics != "none" and args.diagnostics_schedule == "inline"): device_sync(device) diagnostics_t0 = time.time() if args.mode == "fa": al = probes.fa_alignment_report(net, px, py, poh) meancos = sum(al["cos_fa_negg"]) / len(al["cos_fa_negg"]) msg += f" mean_cos(fa,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_fa_negg']]}" elif args.mode == "nodepert": al = probes.nodepert_alignment_report(net, px, py, cfg) meancos = sum(al["cos_q_negg"]) / len(al["cos_q_negg"]) msg += f" mean_cos(q,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_q_negg']]}" elif args.mode != "bp": al = probes.alignment_report(net, px, py, poh, cfg) meancos = sum(al["cos_r_negg"]) / len(al["cos_r_negg"]) msg += f" mean_cos(r,-g) {meancos:+.3f} per-layer {['%.2f'%v for v in al['cos_r_negg']]}" device_sync(device) diagnostics_wall_s += time.time() - diagnostics_t0 print(msg, flush=True) log["steps"].append(record) device_sync(device) eval_t0 = time.time() acc, tloss = evaluate(net, eval_loader) device_sync(device) eval_wall_s += time.time() - eval_t0 log["final"] = {"eval_split": args.eval_split, "eval_acc": acc, "eval_loss": tloss, "wall_s": None} if args.eval_split == "test": log["final"].update({"test_acc": acc, "test_loss": tloss}) else: log["final"].update({"val_acc": acc, "val_loss": tloss}) if args.mode == "fa" and args.diagnostics != "none": device_sync(device) diagnostics_t0 = time.time() log["final"].update(probes.fa_alignment_report(net, px, py, poh)) device_sync(device) diagnostics_wall_s += time.time() - diagnostics_t0 elif args.mode == "nodepert" and args.diagnostics != "none": device_sync(device) diagnostics_t0 = time.time() log["final"].update(probes.nodepert_alignment_report(net, px, py, cfg)) device_sync(device) diagnostics_wall_s += time.time() - diagnostics_t0 elif args.mode != "bp" and args.diagnostics != "none": device_sync(device) diagnostics_t0 = time.time() al = probes.alignment_report(net, px, py, poh, cfg) log["final"]["cos_r_negg"] = al["cos_r_negg"] 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"]["r_norm"] = al["r_norm"] log["final"]["g_norm"] = al["g_norm"] log["final"]["traffic_norm"] = al["traffic_norm"] log["final"]["traffic_residual_norm"] = al["traffic_residual_norm"] log["final"]["traffic_r2"] = al["traffic_r2"] device_sync(device) diagnostics_wall_s += time.time() - diagnostics_t0 if args.residual_lesion_fraction > 0: device_sync(device) lesion_t0 = time.time() report = residual_lesion_report( net, eval_loader, px, args.residual_lesion_fraction) device_sync(device) eval_wall_s += time.time() - lesion_t0 report["lesion_acc_drop"] = acc - report["lesion_eval_acc"] log["final"]["residual_lesion"] = report log["lesion_protocol"] = { "fraction_of_interior_blocks": args.residual_lesion_fraction, "selection": "final_contiguous_interior_blocks", "probe_source": "training_prefix", "probe_examples": len(px), "evaluation_split": args.eval_split, } device_sync(device) log["final"]["wall_s"] = time.time() - t0 log["timing"] = { "warmup_wall_s": warmup_wall_s, "feedback_warmup_wall_s": feedback_warmup_wall_s, "training_loop_wall_s": train_wall_s, "diagnostics_wall_s": diagnostics_wall_s, "inline_diagnostics_wall_s": inline_diagnostics_wall_s, "optimizer_wall_s_excluding_inline_diagnostics": max( 0.0, train_wall_s - inline_diagnostics_wall_s), "evaluation_wall_s": eval_wall_s, } log["cost"] = { "train_steps": step, "ordinary_training_forward_examples": ordinary_forward_examples, "predictor_warmup_forward_examples": warmup_examples, "feedback_warmup_forward_examples": feedback_warmup_examples, "feedback_warmup_perturbation_events": feedback_warmup_events, "perturbation_events": perturbation_events, "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations, "calibration_example_loss_evaluations": calibration_example_loss_evaluations, "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples, "training_forward_equivalent_examples": ( ordinary_forward_examples + warmup_examples + feedback_warmup_examples + calibration_forward_equivalent_examples), "max_perturbation_batch_expansion": perturbation_batch_expansion, } log["hardware"] = hardware_report(device) os.makedirs(args.outdir, exist_ok=True) outpath = os.path.join(args.outdir, f"{args.tag}.json") with open(outpath, "w") as f: json.dump(log, f) print(f"[{args.tag}] DONE {args.eval_split}_acc={acc:.4f} -> {outpath}", flush=True) return log def get_args(): p = argparse.ArgumentParser() p.add_argument("--mode", default="sdil", choices=["bp", "fa", "dfa", "sdil", "nodepert"]) p.add_argument("--dataset", default="mnist", choices=list(REAL_DATASETS + SYNTHETIC_DATASETS)) p.add_argument("--depth", type=int, default=3) # hidden layers p.add_argument("--width", type=int, default=256) p.add_argument("--act", default="tanh", choices=["tanh", "gelu", "silu", "relu"]) p.add_argument("--residual", type=int, default=0) # skip connections (deep no-BN) p.add_argument("--residual_lesion_fraction", type=float, default=0.0, help="after training, bypass this fraction of final interior residual blocks") p.add_argument("--epochs", type=int, default=15) p.add_argument("--batch_size", type=int, default=128) p.add_argument("--train_examples", type=int, default=0, help="0 uses the full training split") p.add_argument("--val_examples", type=int, default=0, help="stratified validation examples held out from training") p.add_argument("--split_seed", type=int, default=2027) p.add_argument("--eval_split", default="test", choices=["validation", "test"]) p.add_argument("--task_seed", type=int, default=0, help="fixed target/data seed, separate from student --seed") p.add_argument("--task_train_examples", type=int, default=50000) p.add_argument("--task_test_examples", type=int, default=10000) p.add_argument("--task_levels", type=int, default=8) p.add_argument("--task_n_in", type=int, default=128) p.add_argument("--task_classes", type=int, default=10) p.add_argument("--teacher_depth", type=int, default=8) p.add_argument("--teacher_width", type=int, default=64) p.add_argument("--teacher_residual", type=int, default=1) p.add_argument("--eta", type=float, default=0.05) p.add_argument("--eta_A", type=float, default=0.02) p.add_argument("--eta_P", type=float, default=0.002) p.add_argument("--momentum", type=float, default=0.9) p.add_argument("--w_scale", type=float, default=1.0) p.add_argument("--a_scale", type=float, default=1.0) p.add_argument("--feedback_scale", type=float, default=1.0) 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("--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 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("--a_warmup_steps", type=int, default=0, help="A-only causal-calibration prefix with all forward weights frozen") p.add_argument("--a_warmup_ndirs", type=int, default=4) p.add_argument("--a_warmup_mode", default="layerwise", choices=["layerwise", "simultaneous"]) p.add_argument("--nuis_rho", type=float, default=0.0) p.add_argument("--traffic_seed", type=int, default=1234, help="ordinary apical-traffic projection seed; independent of model/data seeds") 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("--vectorizer_optimizer", default="sgd", choices=["sgd", "nlms"]) p.add_argument("--vectorizer_eps", type=float, default=1e-6) 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) p.add_argument("--feedback", default="error", choices=["error", "error_deriv"]) p.add_argument("--seed", type=int, default=0) p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--log_every", type=int, default=50) p.add_argument("--diagnostics", default="full", choices=["none", "alignment", "full"]) p.add_argument("--diagnostics_schedule", default="inline", choices=["inline", "final"], help="run probes throughout training or only after final evaluation") p.add_argument("--eval_every", type=int, default=1, help="0 evaluates only once at the end") p.add_argument("--max_steps", type=int, default=0) # 0 = no cap (smoke only) p.add_argument("--probe_bs", type=int, default=512) p.add_argument("--outdir", default="results") p.add_argument("--tag", default="sdil_run") return p.parse_args() if __name__ == "__main__": train(get_args())