#!/usr/bin/env python3 """Deterministic equation audits for matched ResNet crossover adapters.""" import json import os import sys import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import CIFARLocalResNet from sdil.conv_crossover import CIFARDualPropResNet, CIFARPEPITAResNet def relative_error(actual, expected): numerator = torch.linalg.vector_norm(actual - expected) denominator = torch.linalg.vector_norm(expected).clamp_min(1e-30) return float(numerator / denominator) def audit_dualprop(): common = dict( depth=8, base_width=2, seed=71, normalization="batchnorm", residual_scale=1.0, dtype=torch.float64) reference = CIFARLocalResNet(**common) net = CIFARDualPropResNet( **common, alpha=0.0, dp_beta=0.1, inference_passes=2) generator = torch.Generator().manual_seed(72) image = torch.randn(3, 3, 32, 32, generator=generator, dtype=torch.float64) labels = torch.tensor([0, 3, 7]) one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64) with torch.no_grad(): reference_output = reference.forward( image, training=True, update_stats=False)["logits"] net_output = net.forward( image, training=True, update_stats=False)["logits"] clean = net.forward( image, return_cache=True, training=True, update_stats=False) plus, minus = net.infer_dual_states(image, one_hot, clean) forward_error = float(torch.max(torch.abs( reference_output - net_output))) parameters = ( net.W + net.gamma + net.beta + [net.W_out, net.b_out]) for parameter in parameters: parameter.requires_grad_(True) alpha = net.dp_alpha beta = net.dp_beta states = [ (alpha * positive + (1.0 - alpha) * negative).detach() for positive, negative in zip(plus[:-1], minus[:-1])] deltas = [ ((positive - negative) / beta).detach() for positive, negative in zip(plus, minus)] objective = image.new_zeros(()) for index in range(net.n_hidden): prediction, _ = net._node_prediction(index, states, image) objective -= torch.sum(deltas[index] * prediction) / image.shape[0] features = states[-1].mean(dim=(2, 3)) output_prediction = features @ net.W_out.t() + net.b_out objective -= torch.sum(deltas[-1] * output_prediction) / image.shape[0] gradients = torch.autograd.grad(objective, parameters) (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.dualprop_ascent_directions( image, plus, minus) actual = ( directions + gamma_directions + beta_directions + [output_weight, output_bias]) errors = [ relative_error(direction, -gradient) for direction, gradient in zip(actual, gradients)] for parameter in parameters: parameter.requires_grad_(False) if forward_error >= 1e-12 or max(errors) >= 2e-12: raise AssertionError({ "forward_error": forward_error, "direction_errors": errors, }) return { "matched_forward_max_absolute_error": forward_error, "contrastive_direction_max_relative_error": max(errors), "num_audited_parameter_tensors": len(errors), "uses_symmetric_forward_edge_transposes": True, "uses_reverse_task_loss_graph": False, } def audit_pepita(): common = dict( depth=8, base_width=2, seed=81, normalization="batchnorm", residual_scale=1.0, dtype=torch.float64) reference = CIFARLocalResNet(**common) net = CIFARPEPITAResNet( **common, projection_scale=0.05, projection_seed=82) generator = torch.Generator().manual_seed(83) image = torch.randn(3, 3, 32, 32, generator=generator, dtype=torch.float64) labels = torch.tensor([1, 4, 8]) one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64) with torch.no_grad(): reference_output = reference.forward( image, training=True, update_stats=False)["logits"] clean = net.forward( image, return_cache=True, training=True, update_stats=False) clean_error = torch.softmax(clean["logits"], dim=1) - one_hot input_error = torch.einsum( "bc,cijk->bijk", clean_error, net.input_feedback) modulated = net.forward( image + input_error, return_cache=True, training=True, update_stats=False) modulated_error = ( torch.softmax(modulated["logits"], dim=1) - one_hot) forward_error = float(torch.max(torch.abs( reference_output - clean["logits"]))) parameters = ( net.W + net.gamma + net.beta + [net.W_out, net.b_out]) for parameter in parameters: parameter.requires_grad_(True) objective = image.new_zeros(()) batch = image.shape[0] for index, (clean_hidden, modulated_hidden, cache, spec) in enumerate(zip( clean["hiddens"], modulated["hiddens"], modulated["caches"], net.layer_specs)): field = (clean_hidden - modulated_hidden).detach() convolution = torch.nn.functional.conv2d( cache["pre"].detach(), net.W[index], stride=spec.stride, padding=spec.padding) normalized, _ = net._normalize( index, convolution, training=True, update_stats=False) prediction = spec.branch_scale * normalized spatial = prediction.shape[2] * prediction.shape[3] objective += torch.sum(field * prediction) / (batch * spatial) output_prediction = ( modulated["features"].detach() @ net.W_out.t() + net.b_out) objective += torch.sum( modulated_error.detach() * output_prediction) / batch gradients = torch.autograd.grad(objective, parameters) (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.pepita_ascent_directions( clean, modulated, modulated_error) actual = ( directions + gamma_directions + beta_directions + [output_weight, output_bias]) errors = [ relative_error(direction, -gradient) for direction, gradient in zip(actual, gradients)] for parameter in parameters: parameter.requires_grad_(False) projection_limit = (6.0 / (3 * 32 * 32)) ** 0.5 * 0.05 observed_limit = float(torch.max(torch.abs(net.input_feedback))) if (forward_error >= 1e-12 or max(errors) >= 2e-12 or observed_limit > projection_limit): raise AssertionError({ "forward_error": forward_error, "direction_errors": errors, "projection_limit": projection_limit, "observed_limit": observed_limit, }) return { "matched_forward_max_absolute_error": forward_error, "local_equation_max_relative_error": max(errors), "input_projection_shape": list(net.input_feedback.shape), "input_projection_limit": projection_limit, "observed_input_projection_max": observed_limit, "uses_reverse_task_loss_graph": False, } def main(): print(json.dumps({ "dualprop": audit_dualprop(), "pepita": audit_pepita(), }, indent=2, sort_keys=True)) if __name__ == "__main__": main()