experiment: freeze stagewise causal capture runner

2 files changed, 338 insertions, 0 deletions

diff --git a/experiments/analyze_oral_a_v6_calibration.py b/experiments/analyze_oral_a_v6_calibration.py
new file mode 100644
index 0000000..8972138
--- /dev/null
+++ b/experiments/analyze_oral_a_v6_calibration.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Validate and gate stagewise causally whitened feedback capture."""
+import argparse
+import json
+import math
+import os
+
+
+SPLIT_HASH = "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input", default="results/oral_a_v6_calibration/result.json")
+    parser.add_argument(
+        "--out", default="results/oral_a_v6_calibration_gate.json")
+    args = parser.parse_args()
+    with open(args.input) as handle:
+        record = json.load(handle)
+    if record.get("protocol") != "oral_a_v6_stagewise_whitened_causal_capture_v1":
+        raise ValueError("unexpected V6 protocol")
+    expected = {
+        "depth": 20, "width": 16, "seed": 0, "loader_seed": 0,
+        "batch_size": 128, "train_limit": 10000,
+        "val_examples": 5000, "split_seed": 2027,
+        "normalization": "batchnorm", "residual_scale": 1.0,
+        "feedback_scale": 1.0, "sigma": 0.01,
+        "perturb_seed": 5000, "events_per_stage": 20,
+        "readout_relative_ridge": 1e-6,
+        "conv_diagonal_relative_ridge": 1e-3,
+        "alignment_probe": 64, "calibration_augmentation": False,
+    }
+    if record.get("settings") != expected:
+        raise ValueError("V6 settings drift")
+    if record["provenance"]["git_tracked_dirty"]:
+        raise ValueError("V6 result came from a tracked-dirty tree")
+    if record["split"]["validation_index_sha256"] != SPLIT_HASH:
+        raise ValueError("V6 split drift")
+    if record["test_examples_touched"] or record["validation_endpoints_observed"]:
+        raise ValueError("V6 touched a held-out endpoint")
+    work = record["work"]
+    audit = record["method_audit"]
+    fixed = record["fixed_hfa"]
+    learned = record["learned_scib"]
+    finite_metrics = [
+        fixed["early_third_alignment"], fixed["all_layer_alignment"],
+        learned["early_third_alignment"], learned["all_layer_alignment"],
+        learned["min_feedback_forward_norm_ratio"],
+        learned["max_feedback_forward_norm_ratio"],
+    ]
+    checks = {
+        "finite": bool(record["finite"])
+                  and all(math.isfinite(value) for value in finite_metrics),
+        "exactly_19_stages": work["stages"] == 19,
+        "exactly_380_edge_events": work["edge_events"] == 380,
+        "exactly_760_batch_loss_queries": (
+            work["logical_batch_loss_queries"] == 760),
+        "exactly_48640_per_example_observations": (
+            work["per_example_causal_observations"] == 48640),
+        "stage_order_readout_then_18_to_1": (
+            audit["stage_order"] == ["readout"] + list(range(18, 0, -1))),
+        "forward_state_bitwise_fixed": (
+            audit["forward_state_max_absolute_difference"] == 0.0),
+        "zero_forward_weight_reads_in_fit": (
+            audit["forward_weight_reads_in_feedback_fit"] == 0),
+        "zero_reverse_mode_learning_operations": (
+            audit["reverse_mode_learning_operations"] == 0),
+        "early_third_at_least_0.10": (
+            learned["early_third_alignment"] >= 0.10),
+        "all_layer_at_least_0.20": (
+            learned["all_layer_alignment"] >= 0.20),
+        "early_gain_over_fixed_hfa_at_least_0.08": (
+            learned["early_third_alignment"]
+            - fixed["early_third_alignment"] >= 0.08),
+        "feedback_norm_ratios_in_0.1_to_3": (
+            learned["min_feedback_forward_norm_ratio"] >= 0.1
+            and learned["max_feedback_forward_norm_ratio"] <= 3.0),
+    }
+    output = {
+        "protocol": "oral_a_v6_stagewise_whitened_causal_capture_gate_v1",
+        "status": "passed" if all(checks.values()) else "failed",
+        "checks": checks, "fixed_hfa": fixed, "learned_scib": learned,
+        "work": work, "source_commit": record["provenance"]["git_commit"],
+        "source_result": args.input,
+        "conditional_short_task_gate_open": all(checks.values()),
+        "confirmation_test_seeds_touched": False,
+        "review_score_before": 5, "review_score_after": 5,
+        "score_change_rule": "causal capture alone cannot raise score",
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
+    with open(args.out, "w") as handle:
+        json.dump(output, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+    print(json.dumps({
+        "status": output["status"], "checks": checks,
+        "fixed_hfa": fixed, "learned_scib": learned,
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()
+
diff --git a/experiments/oral_a_v6_calibration_screen.py b/experiments/oral_a_v6_calibration_screen.py
new file mode 100644
index 0000000..fec8645
--- /dev/null
+++ b/experiments/oral_a_v6_calibration_screen.py
@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+"""Run the frozen stagewise causally whitened no-KP capture screen."""
+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.conv import (CIFARHierarchicalFAResNet,
+                       causal_conv_diagonal_least_squares_fit,
+                       causal_readout_least_squares_fit,
+                       conv_hierarchical_alignment_report,
+                       layerwise_causal_feedback_observation)
+from sdil.data import DATA_DIR, get_cifar_image_splits
+
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def provenance():
+    def run(command):
+        return subprocess.run(
+            command, cwd=ROOT, check=True, capture_output=True,
+            text=True).stdout.strip()
+    return {
+        "git_commit": run(["git", "rev-parse", "HEAD"]),
+        "git_tracked_dirty": bool(run(
+            ["git", "status", "--porcelain", "--untracked-files=no"])),
+    }
+
+
+def summarize_alignment(report):
+    values = report["teaching_negative_gradient_cosine"]
+    early = max(1, len(values) // 3)
+    ratios = report["feedback_forward_norm_ratio"]
+    cosines = report["feedback_forward_cosine"]
+    return {
+        "per_layer": values,
+        "early_third_alignment": sum(values[:early]) / early,
+        "all_layer_alignment": sum(values) / len(values),
+        "mean_feedback_forward_cosine": sum(cosines) / len(cosines),
+        "min_feedback_forward_norm_ratio": min(ratios),
+        "max_feedback_forward_norm_ratio": max(ratios),
+        "feedback_forward_cosine": cosines,
+        "feedback_forward_norm_ratio": ratios,
+    }
+
+
+def forward_state(net):
+    return [value.clone() for value in (
+        net.W + net.gamma + net.beta + net.running_mean + net.running_var
+        + net.mW + net.mgamma + net.mbeta
+        + [net.W_out, net.b_out, net.mW_out, net.mb_out])]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--data_dir", default=DATA_DIR)
+    parser.add_argument("--out", default="results/oral_a_v6_calibration/result.json")
+    args = parser.parse_args()
+    settings = {
+        "depth": 20, "width": 16, "seed": 0, "loader_seed": 0,
+        "batch_size": 128, "train_limit": 10000,
+        "val_examples": 5000, "split_seed": 2027,
+        "normalization": "batchnorm", "residual_scale": 1.0,
+        "feedback_scale": 1.0, "sigma": 0.01,
+        "perturb_seed": 5000, "events_per_stage": 20,
+        "readout_relative_ridge": 1e-6,
+        "conv_diagonal_relative_ridge": 1e-3,
+        "alignment_probe": 64, "calibration_augmentation": False,
+    }
+    torch.manual_seed(settings["seed"])
+    if str(args.device).startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA requested but unavailable")
+        torch.cuda.manual_seed_all(settings["seed"])
+        torch.cuda.reset_peak_memory_stats(torch.device(args.device))
+    train, _, _, input_shape, n_out, split = get_cifar_image_splits(
+        batch_size=settings["batch_size"], data_dir=args.data_dir,
+        device=args.device, train_limit=settings["train_limit"],
+        val_examples=settings["val_examples"], split_seed=settings["split_seed"],
+        loader_seed=settings["loader_seed"], augment_train=False)
+    if input_shape != (3, 32, 32) or n_out != 10:
+        raise AssertionError("unexpected CIFAR dimensions")
+    net = CIFARHierarchicalFAResNet(
+        depth=settings["depth"], base_width=settings["width"],
+        n_classes=10, device=args.device, seed=settings["seed"],
+        residual_scale=settings["residual_scale"],
+        normalization=settings["normalization"],
+        feedback_scale=settings["feedback_scale"])
+    audit_x = train.x[:settings["alignment_probe"]]
+    audit_y = train.y[:settings["alignment_probe"]]
+    fixed = summarize_alignment(
+        conv_hierarchical_alignment_report(net, audit_x, audit_y))
+    state_before = forward_state(net)
+    generator = torch.Generator(device=torch.device(args.device)).manual_seed(
+        settings["perturb_seed"])
+    events = 0
+
+    def collect(edge_index):
+        nonlocal events
+        observations = []
+        for event_index in range(settings["events_per_stage"]):
+            start_index = event_index * settings["batch_size"]
+            stop_index = start_index + settings["batch_size"]
+            x = train.x[start_index:stop_index]
+            y = train.y[start_index:stop_index]
+            clean = net.forward(
+                x, return_cache=True, training=False, update_stats=False)
+            signal = (torch.softmax(clean["logits"], dim=1)
+                      - F.one_hot(y, net.n_classes).to(clean["logits"].dtype))
+            observation = layerwise_causal_feedback_observation(
+                net, x, y, clean, signal, edge_index=edge_index,
+                sigma=settings["sigma"], generator=generator)
+            # These audit tensors are not inputs to either local fit.
+            observation.pop("direction")
+            observation.pop("directional")
+            observations.append(observation)
+            events += 1
+        return observations
+
+    if str(args.device).startswith("cuda"):
+        torch.cuda.synchronize(torch.device(args.device))
+    start = time.time()
+    stages = []
+    readout_observations = collect(None)
+    readout_fit = causal_readout_least_squares_fit(
+        net, readout_observations,
+        relative_ridge=settings["readout_relative_ridge"])
+    stages.append({"kind": "readout", **readout_fit})
+    print(json.dumps(stages[-1]), flush=True)
+    del readout_observations
+    for index in reversed(range(1, len(net.Q))):
+        observations = collect(index)
+        fit = causal_conv_diagonal_least_squares_fit(
+            net, observations,
+            relative_ridge=settings["conv_diagonal_relative_ridge"])
+        stages.append({"kind": "convolution", **fit})
+        print(json.dumps(stages[-1]), flush=True)
+        del observations
+    if str(args.device).startswith("cuda"):
+        torch.cuda.synchronize(torch.device(args.device))
+    wall_seconds = time.time() - start
+    state_after = forward_state(net)
+    forward_state_max_difference = max(
+        float((before - after).abs().max())
+        for before, after in zip(state_before, state_after))
+    learned = summarize_alignment(
+        conv_hierarchical_alignment_report(net, audit_x, audit_y))
+
+    batch = settings["batch_size"]
+    queries = 2 * events
+    observations_count = events * batch
+    clean_forward_examples = events * batch
+    perturbation_forward_examples = queries * batch
+    # Teaching and diagonal-correlation accounting is a conservative upper
+    # bound: every event is charged three complete feedback traversals.
+    feedback_work = 3 * events * batch * net.apical_macs_per_example
+    work = {
+        "stages": len(stages), "edge_events": events,
+        "logical_batch_loss_queries": queries,
+        "per_example_causal_observations": observations_count,
+        "per_example_cross_entropy_terms": 2 * observations_count,
+        "clean_forward_examples": clean_forward_examples,
+        "perturbation_forward_examples": perturbation_forward_examples,
+        "forward_macs": ((clean_forward_examples + perturbation_forward_examples)
+                         * net.forward_macs_per_example),
+        "feedback_fit_macs_conservative_estimate": feedback_work,
+    }
+    work["total_macs_conservative_estimate"] = (
+        work["forward_macs"] + feedback_work)
+    finite_values = [
+        fixed["early_third_alignment"], fixed["all_layer_alignment"],
+        learned["early_third_alignment"], learned["all_layer_alignment"],
+        learned["min_feedback_forward_norm_ratio"],
+        learned["max_feedback_forward_norm_ratio"],
+    ]
+    for stage in stages:
+        finite_values.extend(value for value in stage.values()
+                             if isinstance(value, float))
+    output = {
+        "schema_version": 1,
+        "protocol": "oral_a_v6_stagewise_whitened_causal_capture_v1",
+        "settings": settings, "provenance": provenance(), "split": split,
+        "architecture": {
+            "family": "CIFAR 6n+2 ResNet, option-A shortcuts",
+            "forward_parameters": net.n_forward_parameters,
+            "adaptive_feedback_parameters": net.n_fixed_feedback_parameters,
+            "forward_macs_per_example": net.forward_macs_per_example,
+            "feedback_macs_per_example": net.apical_macs_per_example,
+        },
+        "method_audit": {
+            "stage_order": ["readout"] + list(reversed(range(1, len(net.Q)))),
+            "forward_weight_reads_in_feedback_fit": 0,
+            "reverse_mode_learning_operations": 0,
+            "causal_query_normalization_state": "evaluation_running_statistics",
+            "ordinary_task_normalization_state": "not_run_forward_frozen",
+            "forward_state_max_absolute_difference": (
+                forward_state_max_difference),
+        },
+        "fixed_hfa": fixed, "learned_scib": learned,
+        "stage_fits": stages, "work": work, "wall_seconds": wall_seconds,
+        "finite": all(math.isfinite(value) for value in finite_values),
+        "test_examples_touched": 0, "validation_endpoints_observed": 0,
+        "hardware": {
+            "device": str(args.device), "torch_version": torch.__version__,
+            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
+            "cuda_device_name": (torch.cuda.get_device_name(torch.device(args.device))
+                                 if str(args.device).startswith("cuda") else None),
+            "peak_memory_allocated_bytes": (
+                torch.cuda.max_memory_allocated(torch.device(args.device))
+                if str(args.device).startswith("cuda") else None),
+        },
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
+    with open(args.out, "w") as handle:
+        json.dump(output, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+    print(json.dumps({
+        "fixed_hfa": fixed, "learned_scib": learned, "work": work,
+        "finite": output["finite"], "wall_seconds": wall_seconds,
+        "out": args.out,
+    }, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()
+