#!/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()