experiments: audit residual block usage with lesions

1 files changed, 81 insertions, 2 deletions

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,