#!/usr/bin/env python3 """Post-failure oracle audit for spatially hierarchical feedback. The audit uses exact gradients only as held-out regression targets. It asks whether the first ResNet stage can be represented by direct output feedback, ordinary downstream activation context, or a learned local map from the true error fields at the actual child nodes in the residual DAG. No parameter of the network or learning algorithm is updated. """ import argparse import json import os import subprocess import sys 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 CIFARSDILResNet from sdil.data import DATA_DIR, get_cifar_image_splits ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) EARLY_CHILDREN = ((1, 2), (2,), (3, 4), (4,), (5, 6), (6,)) 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 exact_batch(net, x, y): forward = net.forward(x, training=True, update_stats=False) gradients = torch.autograd.grad( F.cross_entropy(forward["logits"], y), forward["hiddens"]) targets = [-x.shape[0] * value.detach() for value in gradients] output_signal = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, 10).to(forward["logits"].dtype)) return ([value.detach() for value in forward["hiddens"]], targets, output_signal.detach(), forward["features"].detach()) def cosine(prediction, target): return float(F.cosine_similarity( prediction.flatten(1), target.flatten(1), dim=1).mean()) def energy_ratio(prediction, target): return float(prediction.square().sum() / target.square().sum().clamp_min(1e-30)) def solve_ridge(features, target, relative_ridge=1e-6): gram = features.t() @ features ridge = relative_ridge * gram.diag().mean().clamp_min(1e-30) return torch.linalg.solve( gram + ridge * torch.eye( gram.shape[0], dtype=gram.dtype, device=gram.device), features.t() @ target) def direct_context_oracle(output_train, output_eval, hidden_train, hidden_eval, target_train): """Best held-out A c + tanh(h) G c map at fixed initialization.""" batch, channels, height, width = target_train.shape spatial = height * width context = output_train.shape[1] train_fields = (torch.ones_like(hidden_train), torch.tanh(hidden_train)) eval_fields = (torch.ones_like(hidden_eval), torch.tanh(hidden_eval)) gram = target_train.new_zeros(channels, 2 * context, 2 * context) rhs = target_train.new_zeros(channels, 2 * context) for left in range(2): projected = (train_fields[left] * target_train).flatten(2).sum(2) rhs[:, left * context:(left + 1) * context] = torch.einsum( "bd,bc->cd", output_train, projected) for right in range(2): weights = (train_fields[left] * train_fields[right]).flatten( 2).sum(2) gram[:, left * context:(left + 1) * context, right * context:(right + 1) * context] = torch.einsum( "bd,be,bc->cde", output_train, output_train, weights) coefficients = (torch.linalg.pinv(gram) @ rhs[:, :, None]).squeeze(-1) coefficients = coefficients.reshape(channels, 2, context) prediction = target_train.new_zeros( output_eval.shape[0], channels, height, width) for field in range(2): prediction.add_( (output_eval @ coefficients[:, field].t())[:, :, None, None] * eval_fields[field]) return prediction def activation_context_oracle(output_train, output_eval, hidden_train, hidden_eval, child_train, child_eval, target_train): """Add one spatial basis from ordinary downstream activation context.""" context_train = torch.tanh(child_train).mean(dim=1, keepdim=True).expand_as( hidden_train) context_eval = torch.tanh(child_eval).mean(dim=1, keepdim=True).expand_as( hidden_eval) fields_train = (torch.ones_like(hidden_train), torch.tanh(hidden_train), context_train) fields_eval = (torch.ones_like(hidden_eval), torch.tanh(hidden_eval), context_eval) batch, channels, height, width = target_train.shape context = output_train.shape[1] n_fields = len(fields_train) gram = target_train.new_zeros( channels, n_fields * context, n_fields * context) rhs = target_train.new_zeros(channels, n_fields * context) for left in range(n_fields): projected = (fields_train[left] * target_train).flatten(2).sum(2) rhs[:, left * context:(left + 1) * context] = torch.einsum( "bd,bc->cd", output_train, projected) for right in range(n_fields): weights = (fields_train[left] * fields_train[right]).flatten( 2).sum(2) gram[:, left * context:(left + 1) * context, right * context:(right + 1) * context] = torch.einsum( "bd,be,bc->cde", output_train, output_train, weights) coefficients = (torch.linalg.pinv(gram) @ rhs[:, :, None]).squeeze(-1) coefficients = coefficients.reshape(channels, n_fields, context) prediction = target_train.new_zeros( output_eval.shape[0], channels, height, width) for field in range(n_fields): prediction.add_( (output_eval @ coefficients[:, field].t())[:, :, None, None] * fields_eval[field]) return prediction def hierarchical_features(hiddens, targets, child_indices, kernel, gated): features = [] for child in child_indices: signal = targets[child] if gated: signal = signal * (hiddens[child] > 0).to(signal.dtype) if kernel == 1: features.append(signal.permute(0, 2, 3, 1).reshape( -1, signal.shape[1])) else: features.append(F.unfold( signal, kernel_size=3, padding=1).transpose(1, 2).reshape( -1, signal.shape[1] * 9)) return torch.cat(features, dim=1).to(torch.float64) def flat_target(target): return target.permute(0, 2, 3, 1).reshape( -1, target.shape[1]).to(torch.float64) def summarize(values): return {"per_layer": values, "early_third_mean": sum(values) / len(values)} def main(): parser = argparse.ArgumentParser() parser.add_argument("--device", default="cpu") parser.add_argument("--data_dir", default=DATA_DIR) parser.add_argument("--fit_examples", type=int, default=128) parser.add_argument("--eval_examples", type=int, default=128) parser.add_argument( "--out", default="results/oral_a_hierarchical_oracle.json") args = parser.parse_args() if (args.fit_examples, args.eval_examples) != (128, 128): raise ValueError("the audited hierarchical oracle is fixed to 128/128") torch.manual_seed(0) train, _, _, _, _, split = get_cifar_image_splits( batch_size=256, data_dir=args.data_dir, device=args.device, train_limit=10000, val_examples=5000, split_seed=2027, loader_seed=0, augment_train=False) net = CIFARSDILResNet( depth=20, base_width=16, n_classes=10, device=args.device, seed=0, residual_scale=1.0, normalization="batchnorm", vectorizer_mode="channel_gated", a_scale=0.0) parameters = net.W + net.gamma + net.beta + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) fit_end = args.fit_examples eval_end = fit_end + args.eval_examples h_train, q_train, c_train, _ = exact_batch( net, train.x[:fit_end], train.y[:fit_end]) h_eval, q_eval, c_eval, _ = exact_batch( net, train.x[fit_end:eval_end], train.y[fit_end:eval_end]) for parameter in parameters: parameter.requires_grad_(False) h_train = [value.to(torch.float64) for value in h_train] h_eval = [value.to(torch.float64) for value in h_eval] q_train = [value.to(torch.float64) for value in q_train] q_eval = [value.to(torch.float64) for value in q_eval] c_train = c_train.to(torch.float64) c_eval = c_eval.to(torch.float64) reports = { "direct_output_channel_gated": [], "downstream_activation_context": [], "hierarchical_1x1": [], "hierarchical_1x1_gated": [], "hierarchical_3x3": [], "hierarchical_3x3_gated": [], } energies = {key: [] for key in reports} for layer, children in enumerate(EARLY_CHILDREN): prediction = direct_context_oracle( c_train, c_eval, h_train[layer], h_eval[layer], q_train[layer]) reports["direct_output_channel_gated"].append( cosine(prediction, q_eval[layer])) energies["direct_output_channel_gated"].append( energy_ratio(prediction, q_eval[layer])) prediction = activation_context_oracle( c_train, c_eval, h_train[layer], h_eval[layer], h_train[children[0]], h_eval[children[0]], q_train[layer]) reports["downstream_activation_context"].append( cosine(prediction, q_eval[layer])) energies["downstream_activation_context"].append( energy_ratio(prediction, q_eval[layer])) for kernel, gated, key in ( (1, False, "hierarchical_1x1"), (1, True, "hierarchical_1x1_gated"), (3, False, "hierarchical_3x3"), (3, True, "hierarchical_3x3_gated")): features_train = hierarchical_features( h_train, q_train, children, kernel, gated) features_eval = hierarchical_features( h_eval, q_eval, children, kernel, gated) weights = solve_ridge( features_train, flat_target(q_train[layer])) flat_prediction = features_eval @ weights prediction = flat_prediction.reshape( args.eval_examples, 32, 32, q_eval[layer].shape[1]).permute( 0, 3, 1, 2) reports[key].append(cosine(prediction, q_eval[layer])) energies[key].append(energy_ratio(prediction, q_eval[layer])) output = { "protocol": "oral_a_post_failure_hierarchical_oracle_v1", "provenance": provenance(), "split": split, "network": { "depth": 20, "width": 16, "seed": 0, "normalization": "batchnorm", "residual_scale": 1.0, "forward_state": "deterministic_initialization_no_updates", }, "probe": { "fit_examples": args.fit_examples, "evaluation_examples": args.eval_examples, "source": "unaugmented first training-prefix examples", "separate_batchnorm_graphs": True, "test_examples_touched": 0, "layers": list(range(6)), "residual_dag_children": [list(value) for value in EARLY_CHILDREN], }, "cosine": {key: summarize(value) for key, value in reports.items()}, "prediction_energy_over_target_energy": { key: summarize(value) for key, value in energies.items()}, "fit": { "hierarchical_relative_ridge": 1e-6, "ridge_role": "fixed numerical conditioning, not endpoint selection", }, "interpretation": { "status": "post_failure_oracle_not_a_trainable_method_or_gate", "direct_output_channel_gated": "current output-error-only family", "downstream_activation_context": ( "current family plus ordinary next-node activation map"), "hierarchical": ( "local linear feedback from exact child error fields following " "the true residual DAG; gated variants multiply by child ReLU state"), }, } 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({"cosine": output["cosine"]}, indent=2)) if __name__ == "__main__": main()