From 1212a6319e69c89c5b71d38e126955d8e401f07e Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 03:13:24 -0500 Subject: experiments: audit residual block usage with lesions --- experiments/protocol_smoke.py | 14 +++++++- experiments/run.py | 83 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 94 insertions(+), 3 deletions(-) (limited to 'experiments') diff --git a/experiments/protocol_smoke.py b/experiments/protocol_smoke.py index 339ac47..18de4c9 100644 --- a/experiments/protocol_smoke.py +++ b/experiments/protocol_smoke.py @@ -7,7 +7,8 @@ import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.data import get_dataset_splits from sdil.core import SDILConfig, SDILNet -from experiments.run import calibration_work_per_event, fixed_training_probe +from experiments.run import (calibration_work_per_event, fixed_training_probe, + residual_lesion_report) def loader_labels(loader): @@ -65,11 +66,22 @@ def main(): } assert layerwise["batch_loss_evaluations"] == 9 assert abs(layerwise["forward_equivalent_batches"] - 11.0 / 3.0) < 1e-12 + + lesion_net = SDILNet([1, 4, 4, 4, 4, 2], act="relu", residual=True, device="cpu") + lesion_x = torch.linspace(0, 1, 16).view(-1, 1) + lesion_y = torch.zeros(16, dtype=torch.long) + lesion = residual_lesion_report( + lesion_net, [(lesion_x, lesion_y)], lesion_x[:8], fraction=1.0 / 3.0) + assert lesion["interior_layers"] == [1, 2, 3] + assert lesion["lesioned_layers"] == [3] + assert len(lesion["branch_to_skip_rms"]) == 3 + assert 0.0 <= lesion["lesion_eval_acc"] <= 1.0 print("validation split hash:", metadata["validation_index_sha256"]) print("train/validation/test: 59000/1000/10000; stratification exact") print("FashionMNIST recovery split: 55000/5000/10000; stratification exact") print("training-prefix diagnostic probe: deterministic; shuffle state unchanged") print("simultaneous/layerwise calibration cost accounting: exact") + print("final-third residual-block lesion: exact block selection and finite report") print("ALL PROTOCOL SMOKE CHECKS PASSED") diff --git a/experiments/run.py b/experiments/run.py index 0ace433..93c1862 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -12,12 +12,14 @@ 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, sdil_step, neutral_p_update @@ -96,6 +98,64 @@ def fixed_training_probe(train_loader, probe_bs, device): 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): """Hardware-independent causal-calibration work for one minibatch. @@ -242,9 +302,10 @@ def train(args): # 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": + if args.diagnostics != "none" or args.residual_lesion_fraction > 0: px, py = fixed_training_probe(train_loader, args.probe_bs, device) - poh = onehot(py, n_out, device=device) + if args.diagnostics != "none": + poh = onehot(py, n_out, device=device) log = {"args": vars(args), "split": split, "provenance": code_provenance(), "diagnostic_protocol": { @@ -419,6 +480,22 @@ def train(args): 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"] = { @@ -461,6 +538,8 @@ def get_args(): 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, -- cgit v1.2.3