#!/usr/bin/env python3 """Prove the convolutional local eligibility matches exact BP when instructed.""" import math import os 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 (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet, CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, channel_subspace_apical_calibration, conv_hierarchical_step, conv_kolen_pollack_step, conv_kp_mixed_traffic_step, conv_local_step, hierarchical_mirror_observations, hierarchical_parameter_subspace_calibration, normalized_residual_mirror_update, normalized_response_mirror_update, simultaneous_conv_node_perturbation, vectorizer_subspace_apical_calibration) def architecture_checks(): expected = { 8: (7, 74810), 20: (19, 268346), 32: (31, 461882), 56: (55, 848954), } for depth, (hidden, parameters) in expected.items(): net = CIFARLocalResNet(depth=depth) assert net.n_hidden == hidden assert net.n_forward_parameters == parameters assert len(net.blocks) == 3 * ((depth - 2) // 6) assert len(net.W) + 1 == depth batchnorm_parameters = {8: 75290, 20: 269722, 32: 464154, 56: 853018} for depth, parameters in batchnorm_parameters.items(): net = CIFARLocalResNet( depth=depth, normalization="batchnorm", residual_scale=1.0) assert net.n_forward_parameters == parameters for depth in (7, 9, 21): try: CIFARLocalResNet(depth=depth) except ValueError: pass else: raise AssertionError(f"invalid depth {depth} was accepted") x = torch.arange(2 * 3 * 8 * 8, dtype=torch.float32).reshape(2, 3, 8, 8) shortcut = CIFARLocalResNet._option_a_shortcut(x, 6, 2) assert tuple(shortcut.shape) == (2, 6, 4, 4) assert torch.count_nonzero(shortcut[:, 0]) == 0 assert torch.count_nonzero(shortcut[:, -2:]) == 0 assert torch.equal(shortcut[:, 1:4], x[:, :, ::2, ::2]) def exact_local_gradient_check(): torch.manual_seed(123) batch = 3 x = torch.randn(batch, 3, 32, 32) y = torch.tensor([0, 3, 8]) local = CIFARLocalResNet(depth=8, base_width=4, seed=19) bp = CIFARLocalResNet(depth=8, base_width=4, seed=19) parameters = local.W + [local.W_out, local.b_out] for parameter in parameters: parameter.requires_grad_(True) forward = local.forward(x, return_cache=True) for hidden in forward["hiddens"]: hidden.retain_grad() loss = F.cross_entropy(forward["logits"], y) loss.backward() # .backward() differentiated a batch-mean loss. Multiplying hidden grads # by B recovers the per-example convention consumed by the local rule. teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, local.n_classes)) (directions, gamma_directions, beta_directions, out_direction, bias_direction) = local.local_ascent_directions( teaching, output_error, forward) assert gamma_directions == beta_directions == [] relative_errors = [] for direction, parameter in zip(directions, local.W): absolute = (direction + parameter.grad).abs().max() scale = parameter.grad.abs().max().clamp_min(1e-12) relative_errors.append(float(absolute / scale)) output_abs = float((out_direction + local.W_out.grad).abs().max()) bias_abs = float((bias_direction + local.b_out.grad).abs().max()) assert max(relative_errors) < 3e-5 assert output_abs < 2e-6 and bias_abs < 2e-6 for parameter in parameters: parameter.requires_grad_(False) eta = 0.017 local.apply_ascent(directions, out_direction, bias_direction, eta) bp_loss = bp.bp_step(x, y, eta) assert abs(float(loss.detach()) - bp_loss) < 1e-7 parameter_differences = [ float((left - right).abs().max()) for left, right in zip( local.W + [local.W_out, local.b_out], bp.W + [bp.W_out, bp.b_out])] assert max(parameter_differences) < 2e-7 return { "max_relative_local_gradient_error": max(relative_errors), "output_absolute_error": output_abs, "post_update_parameter_max_error": max(parameter_differences), } def exact_batchnorm_local_gradient_check(): torch.manual_seed(31) batch = 4 x = torch.randn(batch, 3, 32, 32) y = torch.tensor([0, 1, 2, 3]) common = dict( depth=8, base_width=2, seed=29, normalization="batchnorm", residual_scale=1.0) local = CIFARLocalResNet(**common) bp = CIFARLocalResNet(**common) parameters = local.W + local.gamma + local.beta + [local.W_out, local.b_out] for parameter in parameters: parameter.requires_grad_(True) forward = local.forward( x, return_cache=True, training=True, update_stats=True) for hidden in forward["hiddens"]: hidden.retain_grad() loss = F.cross_entropy(forward["logits"], y) loss.backward() teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, 10)) (directions, gamma_directions, beta_directions, out_direction, bias_direction) = local.local_ascent_directions( teaching, output_error, forward) groups = ( (directions, local.W), (gamma_directions, local.gamma), (beta_directions, local.beta), ) relative_errors = [] for direction_group, parameter_group in groups: for direction, parameter in zip(direction_group, parameter_group): absolute = (direction + parameter.grad).abs().max() relative_errors.append(float( absolute / parameter.grad.abs().max().clamp_min(1e-12))) assert max(relative_errors) < 3e-5 for parameter in parameters: parameter.requires_grad_(False) eta = 0.013 local.apply_ascent( directions, out_direction, bias_direction, eta, gamma_directions=gamma_directions, beta_directions=beta_directions) bp.bp_step(x, y, eta) parameter_differences = [ float((left - right).abs().max()) for left, right in zip( local.W + local.gamma + local.beta + [local.W_out, local.b_out], bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])] running_differences = [ float((left - right).abs().max()) for left, right in zip( local.running_mean + local.running_var, bp.running_mean + bp.running_var)] assert max(parameter_differences) < 2e-7 assert max(running_differences) == 0.0 running_before = [value.clone() for value in local.running_mean + local.running_var] clean = local.forward(x, training=True, update_stats=False) simultaneous_conv_node_perturbation( local, x, y, clean, sigma=1e-3, n_directions=1, generator=torch.Generator(device="cpu").manual_seed(9)) assert all(torch.equal(before, after) for before, after in zip( running_before, local.running_mean + local.running_var)) assert torch.equal(local.logits(x), local.logits(x)) return { "batchnorm_max_relative_local_gradient_error": max(relative_errors), "batchnorm_post_update_parameter_max_error": max(parameter_differences), } def perturbation_checks(): net = CIFARLocalResNet(depth=8, base_width=4, seed=3) x = torch.randn(2, 3, 32, 32) clean = net.forward(x) perturbations = [torch.zeros_like(hidden) for hidden in clean["hiddens"]] perturbed = net.forward(x, perturbations=perturbations) assert torch.equal(clean["logits"], perturbed["logits"]) perturbations[0] = torch.ones_like(perturbations[0]) * 0.01 changed = net.forward(x, perturbations=perturbations) assert not torch.equal(clean["logits"], changed["logits"]) try: net.forward(x, perturbations=perturbations[:-1]) except ValueError: pass else: raise AssertionError("short perturbation list was accepted") def perturbation_estimator_check(): """Antithetic finite differences equal the simultaneous hidden JVP.""" torch.manual_seed(3) batch = 2 net = CIFARSDILResNet( depth=8, base_width=2, seed=4, dtype=torch.float64) x = torch.randn(batch, 3, 32, 32, dtype=torch.float64) y = torch.tensor([2, 7]) parameters = net.W + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) clean = net.forward(x, return_cache=True) for hidden in clean["hiddens"]: hidden.retain_grad() F.cross_entropy(clean["logits"], y).backward() generator = torch.Generator(device="cpu").manual_seed(99) targets, diagnostics = simultaneous_conv_node_perturbation( net, x, y, clean, sigma=1e-6, n_directions=1, generator=generator, return_diagnostics=True) directions = diagnostics["directions"][0] derivative = diagnostics["directional_derivatives"][0] assert derivative["coupling"] == "per_example_objective" finite_difference = derivative["scaled_directional"] exact = sum( (batch * hidden.grad * direction).flatten(1).sum(dim=1) for hidden, direction in zip(clean["hiddens"], directions)) relative = (finite_difference - exact).abs() / exact.abs().clamp_min(1e-12) assert float(relative.max()) < 2e-6 for target, direction in zip(targets, directions): expected = -finite_difference[:, None, None, None] * direction assert torch.equal(target, expected) for parameter in parameters: parameter.requires_grad_(False) no_norm_relative = float(relative.max()) batch = 3 batchnorm = CIFARSDILResNet( depth=8, base_width=2, seed=14, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0) xb = torch.randn(batch, 3, 32, 32, dtype=torch.float64) yb = torch.tensor([1, 4, 9]) parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta + [batchnorm.W_out, batchnorm.b_out]) for parameter in parameters: parameter.requires_grad_(True) clean = batchnorm.forward(xb, training=True, update_stats=False) for hidden in clean["hiddens"]: hidden.retain_grad() F.cross_entropy(clean["logits"], yb).backward() targets, diagnostics = simultaneous_conv_node_perturbation( batchnorm, xb, yb, clean, sigma=1e-6, n_directions=1, generator=torch.Generator(device="cpu").manual_seed(101), return_diagnostics=True) directions = diagnostics["directions"][0] derivative = diagnostics["directional_derivatives"][0] assert derivative["coupling"] == "batch_objective" scaled = derivative["scaled_directional"] exact_sum_directional = sum( float((batch * hidden.grad * direction).sum()) for hidden, direction in zip(clean["hiddens"], directions)) batchnorm_relative = abs(float(scaled[0]) - exact_sum_directional) / max( abs(exact_sum_directional), 1e-12) assert batchnorm_relative < 2e-6 assert torch.equal(scaled, scaled[:1].expand_as(scaled)) for target, direction in zip(targets, directions): expected = -scaled[:, None, None, None] * direction assert torch.equal(target, expected) for parameter in parameters: parameter.requires_grad_(False) return { "perturbation_jvp_max_relative_error": no_norm_relative, "batchnorm_batch_objective_jvp_relative_error": batchnorm_relative, } def channel_subspace_estimator_check(): """The structured estimator targets representable base/gate moments.""" torch.manual_seed(71) batch = 3 hiddens = [ torch.randn(batch, 2, 4, 4, dtype=torch.float64), torch.randn(batch, 3, 2, 2, dtype=torch.float64), ] negative_gradients = [torch.randn_like(value) for value in hiddens] estimated_base = [torch.zeros( batch, value.shape[1], dtype=value.dtype) for value in hiddens] estimated_gate = [torch.zeros_like(value) for value in estimated_base] generator = torch.Generator(device="cpu").manual_seed(211) directions = 4096 inverse_sqrt_two = 1.0 / (2.0 ** 0.5) for _ in range(directions): hidden_directions = [] base_random = [] gate_random = [] for hidden in hiddens: shape = (batch, hidden.shape[1]) base = torch.empty(shape, dtype=hidden.dtype).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) gate = torch.empty_like(base).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) hidden_directions.append(( base[:, :, None, None] + torch.tanh(hidden) * gate[:, :, None, None]) * inverse_sqrt_two) base_random.append(base) gate_random.append(gate) # The exact loss derivative uses g=-negative_gradient and remains # per-example without BatchNorm coupling. directional = -sum( (gradient * direction).flatten(1).sum(dim=1) for gradient, direction in zip( negative_gradients, hidden_directions)) for index, (hidden, base, gate) in enumerate(zip( hiddens, base_random, gate_random)): spatial = hidden.shape[2] * hidden.shape[3] scale = -(2.0 ** 0.5) / (spatial * directions) estimated_base[index].add_(directional[:, None] * base, alpha=scale) estimated_gate[index].add_(directional[:, None] * gate, alpha=scale) exact_base = [value.mean(dim=(2, 3)) for value in negative_gradients] exact_gate = [(value * torch.tanh(hidden)).mean(dim=(2, 3)) for value, hidden in zip(negative_gradients, hiddens)] estimated = torch.cat([ value.flatten() for pair in zip(estimated_base, estimated_gate) for value in pair]) exact = torch.cat([ value.flatten() for pair in zip(exact_base, exact_gate) for value in pair]) cosine = float(F.cosine_similarity(estimated, exact, dim=0)) norm_ratio = float(estimated.norm() / exact.norm()) assert cosine > 0.985 assert 0.90 < norm_ratio < 1.10 # The executable antithetic implementation must match the same structured # directional derivative, not an autograd surrogate. # Use a fixed nondegenerate point. Width-one, zero-bias ReLU networks can # contain structurally exact-zero preactivations, where central differences # and PyTorch's chosen subgradient need not agree even as sigma -> 0. torch.manual_seed(3) net = CIFARSDILResNet( depth=8, base_width=2, seed=72, dtype=torch.float64, vectorizer_mode="channel_gated") x = torch.randn(2, 3, 32, 32, dtype=torch.float64) y = torch.tensor([2, 8]) parameters = net.W + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) clean = net.forward(x, return_cache=True) for hidden in clean["hiddens"]: hidden.retain_grad() loss = F.cross_entropy(clean["logits"], y) loss.backward() output_signal = (torch.softmax(clean["logits"].detach(), dim=1) - F.one_hot(y, 10).to(torch.float64)) _, diagnostics = channel_subspace_apical_calibration( net, x, y, clean, output_signal, sigma=1e-6, n_directions=1, eta=0.0, generator=torch.Generator(device="cpu").manual_seed(307), return_diagnostics=True) hidden_direction = diagnostics["directions"][0]["hidden"] finite_difference = diagnostics["directional_derivatives"][0][ "scaled_directional"] exact_directional = sum( (x.shape[0] * hidden.grad * direction).flatten(1).sum(dim=1) for hidden, direction in zip(clean["hiddens"], hidden_direction)) relative = ((finite_difference - exact_directional).abs() / exact_directional.abs().clamp_min(1e-12)) assert float(relative.max()) < 2e-6 for parameter in parameters: parameter.requires_grad_(False) # The local A/G update must equal the full-field delta rule after replacing # only its two target moments with the structured causal estimates. eta = 0.0023 before_a = [value.clone() for value in net.A] before_g = [value.clone() for value in net.A_gate] expected_a = [] expected_g = [] for index, hidden in enumerate(clean["hiddens"]): base = output_signal @ before_a[index].t() gate_coefficient = output_signal @ before_g[index].t() gate = torch.tanh(hidden.detach()) mean = gate.mean(dim=(2, 3)) second = gate.square().mean(dim=(2, 3)) base_error = diagnostics["target_base"][index] - ( base + mean * gate_coefficient) gate_error = diagnostics["target_gate"][index] - ( mean * base + second * gate_coefficient) expected_a.append(before_a[index] + eta * ( base_error.t() @ output_signal / x.shape[0])) expected_g.append(before_g[index] + eta * ( gate_error.t() @ output_signal / x.shape[0])) channel_subspace_apical_calibration( net, x, y, clean, output_signal, sigma=1e-6, n_directions=1, eta=eta, generator=torch.Generator(device="cpu").manual_seed(307)) update_error = max(float((actual - expected).abs().max()) for actual, expected in zip( net.A + net.A_gate, expected_a + expected_g)) assert update_error < 1e-14 return { "channel_subspace_moment_cosine": cosine, "channel_subspace_moment_norm_ratio": norm_ratio, "channel_subspace_jvp_relative_error": float(relative.max()), "channel_subspace_delta_rule_absolute_error": update_error, } def vectorizer_subspace_estimator_check(): """Direct A/G perturbations are unbiased and lower variance at batch 128.""" torch.manual_seed(81) batch = 128 output_dim = 5 hiddens = [ torch.randn(batch, 2, 4, 4, dtype=torch.float64), torch.randn(batch, 3, 2, 2, dtype=torch.float64), ] negative_gradients = [torch.randn_like(value) for value in hiddens] output_signal = torch.randn(batch, output_dim, dtype=torch.float64) exact_pairs = [] for hidden, target in zip(hiddens, negative_gradients): exact_pairs.extend([ target.mean(dim=(2, 3)).t() @ output_signal / batch, (target * torch.tanh(hidden)).mean(dim=(2, 3)).t() @ output_signal / batch, ]) exact = torch.cat([value.flatten() for value in exact_pairs]) coefficient_sum = torch.zeros_like(exact) vectorizer_sum = torch.zeros_like(exact) coefficient_mse = 0.0 vectorizer_mse = 0.0 directions = 2048 generator = torch.Generator(device="cpu").manual_seed(401) inverse_sqrt_two = 1.0 / (2.0 ** 0.5) for _ in range(directions): random_coefficients = [] hidden_directions = [] for hidden in hiddens: shape = (batch, hidden.shape[1]) base = torch.empty(shape, dtype=hidden.dtype).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) gate = torch.empty_like(base).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) random_coefficients.append((base, gate)) hidden_directions.append(( base[:, :, None, None] + torch.tanh(hidden) * gate[:, :, None, None]) * inverse_sqrt_two) directional = -sum((target * direction).sum() for target, direction in zip( negative_gradients, hidden_directions)) coefficient_values = [] for hidden, (base, gate) in zip(hiddens, random_coefficients): spatial = hidden.shape[2] * hidden.shape[3] base_target = -(2.0 ** 0.5) * directional * base / spatial gate_target = -(2.0 ** 0.5) * directional * gate / spatial coefficient_values.extend([ base_target.t() @ output_signal / batch, gate_target.t() @ output_signal / batch, ]) coefficient_sample = torch.cat( [value.flatten() for value in coefficient_values]) random_matrices = [] hidden_directions = [] for hidden in hiddens: shape = (hidden.shape[1], output_dim) base = torch.empty(shape, dtype=hidden.dtype).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) gate = torch.empty_like(base).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) base_field = output_signal @ base.t() gate_field = output_signal @ gate.t() random_matrices.append((base, gate)) hidden_directions.append(( base_field[:, :, None, None] + torch.tanh(hidden) * gate_field[:, :, None, None]) * inverse_sqrt_two) directional = -sum((target * direction).sum() for target, direction in zip( negative_gradients, hidden_directions)) vectorizer_values = [] for hidden, (base, gate) in zip(hiddens, random_matrices): spatial = hidden.shape[2] * hidden.shape[3] scale = -(2.0 ** 0.5) * directional / (batch * spatial) vectorizer_values.extend([scale * base, scale * gate]) vectorizer_sample = torch.cat( [value.flatten() for value in vectorizer_values]) coefficient_sum.add_(coefficient_sample) vectorizer_sum.add_(vectorizer_sample) coefficient_mse += float((coefficient_sample - exact).square().mean()) vectorizer_mse += float((vectorizer_sample - exact).square().mean()) coefficient_mean = coefficient_sum / directions vectorizer_mean = vectorizer_sum / directions vectorizer_cosine = float(F.cosine_similarity( vectorizer_mean, exact, dim=0)) vectorizer_norm_ratio = float(vectorizer_mean.norm() / exact.norm()) variance_ratio = vectorizer_mse / coefficient_mse assert vectorizer_cosine > 0.95 assert 0.90 < vectorizer_norm_ratio < 1.10 assert variance_ratio < 0.25 # Match the executable forward-only derivative and its exact A/G update. torch.manual_seed(3) net = CIFARSDILResNet( depth=8, base_width=2, seed=82, dtype=torch.float64, vectorizer_mode="channel_gated") x = torch.randn(2, 3, 32, 32, dtype=torch.float64) y = torch.tensor([1, 6]) parameters = net.W + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) clean = net.forward(x, return_cache=True) for hidden in clean["hiddens"]: hidden.retain_grad() F.cross_entropy(clean["logits"], y).backward() output_error = (torch.softmax(clean["logits"].detach(), dim=1) - F.one_hot(y, 10).to(torch.float64)) _, diagnostics = vectorizer_subspace_apical_calibration( net, x, y, clean, output_error, sigma=1e-6, n_directions=1, eta=0.0, generator=torch.Generator(device="cpu").manual_seed(503), return_diagnostics=True) finite_difference = diagnostics["directional_derivatives"][0][ "scaled_directional"] exact_directional = sum( (x.shape[0] * hidden.grad * direction).sum() for hidden, direction in zip( clean["hiddens"], diagnostics["directions"][0]["hidden"])) jvp_relative = float((finite_difference - exact_directional).abs() / exact_directional.abs().clamp_min(1e-12)) assert jvp_relative < 2e-6 for parameter in parameters: parameter.requires_grad_(False) eta = 0.0017 before_a = [value.clone() for value in net.A] before_g = [value.clone() for value in net.A_gate] expected_a = [] expected_g = [] for index, hidden in enumerate(clean["hiddens"]): base = output_error @ before_a[index].t() gate_coefficient = output_error @ before_g[index].t() gate = torch.tanh(hidden.detach()) mean = gate.mean(dim=(2, 3)) second = gate.square().mean(dim=(2, 3)) base_prediction = (base + mean * gate_coefficient).t() @ output_error / 2 gate_prediction = ( mean * base + second * gate_coefficient).t() @ output_error / 2 expected_a.append(before_a[index] + eta * ( diagnostics["target_base"][index] - base_prediction)) expected_g.append(before_g[index] + eta * ( diagnostics["target_gate"][index] - gate_prediction)) vectorizer_subspace_apical_calibration( net, x, y, clean, output_error, sigma=1e-6, n_directions=1, eta=eta, generator=torch.Generator(device="cpu").manual_seed(503)) update_error = max(float((actual - expected).abs().max()) for actual, expected in zip( net.A + net.A_gate, expected_a + expected_g)) assert update_error < 1e-14 batchnorm = CIFARSDILResNet( depth=8, base_width=2, seed=83, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0, vectorizer_mode="channel_gated") xb = torch.randn(3, 3, 32, 32, dtype=torch.float64) yb = torch.tensor([0, 4, 9]) parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta + [batchnorm.W_out, batchnorm.b_out]) for parameter in parameters: parameter.requires_grad_(True) clean_b = batchnorm.forward(xb, training=True, update_stats=False) for hidden in clean_b["hiddens"]: hidden.retain_grad() F.cross_entropy(clean_b["logits"], yb).backward() output_b = (torch.softmax(clean_b["logits"].detach(), dim=1) - F.one_hot(yb, 10).to(torch.float64)) _, diagnostics_b = vectorizer_subspace_apical_calibration( batchnorm, xb, yb, clean_b, output_b, sigma=1e-6, n_directions=1, eta=0.0, generator=torch.Generator(device="cpu").manual_seed(509), return_diagnostics=True) finite_b = diagnostics_b["directional_derivatives"][0][ "scaled_directional"] exact_b = sum( (xb.shape[0] * hidden.grad * direction).sum() for hidden, direction in zip( clean_b["hiddens"], diagnostics_b["directions"][0]["hidden"])) batchnorm_relative = float( (finite_b - exact_b).abs() / exact_b.abs().clamp_min(1e-12)) assert batchnorm_relative < 2e-6 for parameter in parameters: parameter.requires_grad_(False) return { "vectorizer_subspace_mean_cosine": vectorizer_cosine, "vectorizer_subspace_mean_norm_ratio": vectorizer_norm_ratio, "vectorizer_vs_coefficient_mse_ratio": variance_ratio, "vectorizer_subspace_jvp_relative_error": jvp_relative, "vectorizer_subspace_batchnorm_jvp_relative_error": batchnorm_relative, "vectorizer_subspace_delta_rule_absolute_error": update_error, } def hierarchical_feedback_checks(): """The residual feedback graph becomes exact only under an audit copy.""" torch.manual_seed(91) exact = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=92, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0) exact.Q = [value.clone() for value in exact.W] exact.R_out.copy_(-exact.W_out.t()) x = torch.randn(3, 3, 32, 32, dtype=torch.float64) y = torch.tensor([1, 5, 8]) parameters = (exact.W + exact.gamma + exact.beta + [exact.W_out, exact.b_out]) for parameter in parameters: parameter.requires_grad_(True) forward = exact.forward(x, return_cache=True, training=True, update_stats=False) gradients = torch.autograd.grad( F.cross_entropy(forward["logits"], y), forward["hiddens"]) output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, 10).to(torch.float64)) teaching = exact.hierarchical_teaching(output_error, forward) relative = [float((signal + x.shape[0] * gradient).abs().max() / gradient.abs().max().clamp_min(1e-30) / x.shape[0]) for signal, gradient in zip(teaching, gradients)] assert max(relative) < 2e-12 for parameter in parameters: parameter.requires_grad_(False) # With independent feedback the forward model is bitwise unchanged, while # changing only the feedback seed changes Q/R. This is the actual baseline. left = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=93, feedback_seed=1001) right = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=93, feedback_seed=1002) assert all(torch.equal(a, b) for a, b in zip( left.W + [left.W_out], right.W + [right.W_out])) assert any(not torch.equal(a, b) for a, b in zip( left.Q + [left.R_out], right.Q + [right.R_out])) # The local update must match exact BP when feedback is explicitly copied # in this audit-only comparator. Actual HFA never performs this copy. local = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=94, normalization="batchnorm", residual_scale=1.0) bp = CIFARLocalResNet( depth=8, base_width=2, seed=94, normalization="batchnorm", residual_scale=1.0) local.Q = [value.clone() for value in local.W] local.R_out.copy_(-local.W_out.t()) xf = torch.randn(4, 3, 32, 32) yf = torch.tensor([0, 2, 5, 9]) eta = 0.013 conv_hierarchical_step( local, xf, yf, ConvSDILConfig( eta=eta, eta_output=eta, eta_A=0.0, momentum=0.0, weight_decay=0.0, learn_A=False)) bp.bp_step(xf, yf, eta, momentum=0.0, weight_decay=0.0) parameter_error = max(float((a - b).abs().max()) for a, b in zip( local.W + local.gamma + local.beta + [local.W_out, local.b_out], bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])) running_error = max(float((a - b).abs().max()) for a, b in zip( local.running_mean + local.running_var, bp.running_mean + bp.running_var)) assert parameter_error < 2e-7 assert running_error == 0.0 return { "hierarchical_symmetric_hidden_relative_error": max(relative), "hierarchical_symmetric_update_absolute_error": parameter_error, "hierarchical_feedback_to_forward_mac_ratio": ( local.apical_macs_per_example / local.forward_macs_per_example), } def hierarchical_parameter_calibration_checks(): """Audit the causal JVP and the exact local Q/R delta-rule moments.""" torch.manual_seed(109) net = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=110, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0) x = torch.randn(3, 3, 32, 32, dtype=torch.float64) y = torch.tensor([0, 4, 7]) parameters = net.W + net.gamma + net.beta + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) forward = net.forward( x, return_cache=True, training=True, update_stats=False) loss = F.cross_entropy(forward["logits"], y) hidden_gradients = torch.autograd.grad(loss, forward["hiddens"]) output_signal = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, 10).to(torch.float64)) _, diagnostic = hierarchical_parameter_subspace_calibration( net, x, y, forward, output_signal, sigma=1e-5, n_directions=1, eta=0.0, generator=torch.Generator().manual_seed(111), return_diagnostics=True) directions = diagnostic["directions"][0]["hidden"] exact_directional = x.shape[0] * sum( (gradient * direction).sum() for gradient, direction in zip(hidden_gradients, directions)) estimated_directional = diagnostic[ "directional_derivatives"][0]["scaled_directional"] jvp_relative = float((estimated_directional - exact_directional).abs() / exact_directional.abs().clamp_min(1e-30)) assert jvp_relative < 3e-7 for parameter in parameters: parameter.requires_grad_(False) # Under an audit-only symmetric copy, the hierarchical field is the exact # negative gradient. Consequently every local Q/R predicted moment equals # its exact causal regression target, including option-A shortcut terms. exact = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=112, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0) exact.Q = [value.clone() for value in exact.W] exact.R_out.copy_(-exact.W_out.t()) for parameter in exact.W + exact.gamma + exact.beta + [ exact.W_out, exact.b_out]: parameter.requires_grad_(True) clean = exact.forward( x, return_cache=True, training=True, update_stats=False) gradients = torch.autograd.grad( F.cross_entropy(clean["logits"], y), clean["hiddens"]) negative = [-x.shape[0] * value.detach() for value in gradients] signal = (torch.softmax(clean["logits"].detach(), dim=1) - F.one_hot(y, 10).to(torch.float64)) teaching, contexts, recipients = exact.hierarchical_teaching( signal, clean, return_edge_contexts=True) numerator = 0.0 denominator = 0.0 for index in range(1, len(exact.Q)): recipient = recipients[index] spec = exact.layer_specs[index] spatial = (negative[recipient].shape[2] * negative[recipient].shape[3]) target = torch.nn.grad.conv2d_weight( negative[recipient], exact.Q[index].shape, contexts[index], stride=spec.stride, padding=spec.padding) / (x.shape[0] * spatial) prediction = torch.nn.grad.conv2d_weight( teaching[recipient], exact.Q[index].shape, contexts[index], stride=spec.stride, padding=spec.padding) / (x.shape[0] * spatial) numerator += float((target - prediction).square().sum()) denominator += float(target.square().sum()) target_r = negative[-1].mean(dim=(2, 3)).t() @ signal / x.shape[0] prediction_r = teaching[-1].mean(dim=(2, 3)).t() @ signal / x.shape[0] numerator += float((target_r - prediction_r).square().sum()) denominator += float(target_r.square().sum()) delta_rule_relative = math.sqrt(numerator / max(denominator, 1e-300)) assert delta_rule_relative < 2e-12 return { "hierarchical_parameter_subspace_jvp_relative_error": jvp_relative, "hierarchical_parameter_delta_rule_relative_error": delta_rule_relative, } def normalized_response_mirror_checks(): """Audit local response estimation and absence of W access in the update.""" net = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=121, dtype=torch.float64, normalization="batchnorm") observations = hierarchical_mirror_observations( net, batch_size=16, noise_std=1.0, generator=torch.Generator().manual_seed(122)) metrics, _ = normalized_response_mirror_update( net, observations, eta=1.0) pairs = list(zip(net.Q[1:], net.W[1:])) + [ (net.R_out, -net.W_out.t())] cosines = [float(F.cosine_similarity( feedback.flatten(), target.flatten(), dim=0)) for feedback, target in pairs] norm_ratios = [float(feedback.norm() / target.norm()) for feedback, target in pairs] assert sum(cosines) / len(cosines) > 0.985 assert min(cosines) > 0.95 assert min(norm_ratios) > 0.90 and max(norm_ratios) < 1.10 # The update consumes observations only. Changing every forward parameter # after those observations were generated must not change the Q/R update. left = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=123, dtype=torch.float64) right = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=123, dtype=torch.float64) shared_observations = hierarchical_mirror_observations( left, batch_size=2, generator=torch.Generator().manual_seed(124)) for value in right.W + [right.W_out]: value.normal_(generator=torch.Generator().manual_seed(value.numel())) normalized_response_mirror_update(left, shared_observations, eta=0.2) normalized_response_mirror_update(right, shared_observations, eta=0.2) independence_error = max(float((a - b).abs().max()) for a, b in zip( left.Q[1:] + [left.R_out], right.Q[1:] + [right.R_out])) assert independence_error == 0.0 # Residual-response LMS has a per-observation exact fixed point: its update # is zero, not merely zero in expectation, when Q/R match the forward maps. fixed = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=125, dtype=torch.float64) fixed.Q = [value.clone() for value in fixed.W] fixed.R_out.copy_(-fixed.W_out.t()) fixed_observations = hierarchical_mirror_observations( fixed, batch_size=2, generator=torch.Generator().manual_seed(126)) fixed_metrics = normalized_residual_mirror_update( fixed, fixed_observations, eta=1.0) assert fixed_metrics["mirror_update_rms"] < 1e-14 assert fixed_metrics["mirror_response_residual_fraction"] < 1e-14 residual_left = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=127, dtype=torch.float64) residual_right = CIFARHierarchicalFAResNet( depth=8, base_width=2, seed=127, dtype=torch.float64) residual_observations = hierarchical_mirror_observations( residual_left, batch_size=2, generator=torch.Generator().manual_seed(128)) for value in residual_right.W + [residual_right.W_out]: value.normal_(generator=torch.Generator().manual_seed(value.numel() + 1)) normalized_residual_mirror_update( residual_left, residual_observations, eta=0.2) normalized_residual_mirror_update( residual_right, residual_observations, eta=0.2) residual_independence_error = max(float((a - b).abs().max()) for a, b in zip( residual_left.Q[1:] + [residual_left.R_out], residual_right.Q[1:] + [residual_right.R_out])) assert residual_independence_error == 0.0 return { "mirror_estimate_mean_forward_cosine": sum(cosines) / len(cosines), "mirror_estimate_min_forward_cosine": min(cosines), "mirror_estimate_min_norm_ratio": min(norm_ratios), "mirror_estimate_max_norm_ratio": max(norm_ratios), "mirror_update_forward_parameter_independence_error": independence_error, "mirror_update_rms": metrics["mirror_update_rms"], "residual_mirror_exact_fixed_point_update_rms": fixed_metrics[ "mirror_update_rms"], "residual_mirror_exact_fixed_point_fraction": fixed_metrics[ "mirror_response_residual_fraction"], "residual_mirror_forward_parameter_independence_error": ( residual_independence_error), } def kolen_pollack_checks(): """KP's reciprocal correlations are local and preserve exact symmetry.""" torch.manual_seed(127) common = dict( depth=8, base_width=2, seed=53, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0) net = CIFARKPResNet(**common) for index in range(1, len(net.Q)): net.Q[index].copy_(net.W[index]) net.R_out.copy_(-net.W_out.t()) x = torch.randn(3, 3, 32, 32, dtype=torch.float64) y = torch.tensor([1, 4, 7]) forward = net.forward( x, return_cache=True, training=True, update_stats=False) output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, 10).to(torch.float64)) teaching = net.hierarchical_teaching(output_error, forward) (forward_directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( teaching, output_error, forward) reciprocal_directions, reciprocal_readout = ( net.reciprocal_feedback_directions( teaching, output_error, forward)) direction_error = max([ float((left - right).abs().max()) for left, right in zip(forward_directions[1:], reciprocal_directions[1:]) ] + [float((reciprocal_readout + output_weight.t()).abs().max())]) assert direction_error < 1e-14 # Once local activities have been observed, neither forward nor feedback # parameter values may alter the independently formed reciprocal update. before = [value.clone() for value in reciprocal_directions[1:]] before_readout = reciprocal_readout.clone() for value in net.W + net.Q: value.add_(torch.randn_like(value)) net.W_out.add_(torch.randn_like(net.W_out)) net.R_out.add_(torch.randn_like(net.R_out)) independent_directions, independent_readout = ( net.reciprocal_feedback_directions( teaching, output_error, forward)) independence_error = max([ float((left - right).abs().max()) for left, right in zip(before, independent_directions[1:]) ] + [float((before_readout - independent_readout).abs().max())]) assert independence_error == 0.0 # A fresh symmetric state must remain symmetric under two momentum steps. net = CIFARKPResNet(**common) for index in range(1, len(net.Q)): net.Q[index].copy_(net.W[index]) net.R_out.copy_(-net.W_out.t()) for _ in range(2): forward = net.forward( x, return_cache=True, training=True, update_stats=False) output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, 10).to(torch.float64)) teaching = net.hierarchical_teaching(output_error, forward) (forward_directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( teaching, output_error, forward) reciprocal_directions, reciprocal_readout = ( net.reciprocal_feedback_directions( teaching, output_error, forward)) net.apply_reciprocal_ascent( reciprocal_directions, reciprocal_readout, eta_hidden=0.013, eta_output=0.017, momentum=0.9, weight_decay=1e-4) net.apply_ascent( forward_directions, output_weight, output_bias, eta_hidden=0.013, eta_output=0.017, momentum=0.9, weight_decay=1e-4, gamma_directions=gamma_directions, beta_directions=beta_directions) symmetry_error = max([ float((net.Q[index] - net.W[index]).abs().max()) for index in range(1, len(net.Q)) ] + [float((net.R_out + net.W_out.t()).abs().max())]) assert symmetry_error < 1e-14 # Exercise the public training step and ensure it remains graph-free. result = conv_kolen_pollack_step( net, x, y, ConvSDILConfig( eta=1e-3, eta_output=1e-3, momentum=0.9, weight_decay=1e-4, learn_A=False, learn_P=False)) assert math.isfinite(result["loss"]) assert all(not value.requires_grad for value in net.W + net.Q + [net.W_out, net.R_out, net.b_out]) return { "kp_local_direction_absolute_error": direction_error, "kp_forward_parameter_independence_error": independence_error, "kp_symmetric_update_absolute_error": symmetry_error, } def kp_mixed_traffic_checks(): """Mixed traffic isolates subtraction from norm and preserves KP locality.""" torch.manual_seed(211) common = dict( depth=8, base_width=2, seed=67, dtype=torch.float64, normalization="batchnorm", residual_scale=1.0, traffic_seed=4000) net = CIFARKPMixedTrafficResNet(**common) x = torch.randn(5, 3, 32, 32, dtype=torch.float64) y = torch.tensor([0, 2, 4, 6, 8]) forward = net.forward( x, return_cache=True, training=True, update_stats=False) output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, 10).to(torch.float64)) instruction = net.hierarchical_teaching(output_error, forward) # Zero traffic and zero predictor collapse all three rules to clean KP. zero_errors = [] for rule in ("raw", "matched", "innovation"): components = net.mixed_apical_components( instruction, forward["hiddens"], rule) zero_errors.extend(float((left - right).abs().max()) for left, right in zip(components["used"], instruction)) assert max(zero_errors) == 0.0 calibration = net.calibrate_traffic_gain( instruction, forward["hiddens"], target_ratio=4.0) ratio_error = max(abs(value - 4.0) for value in calibration["realized_traffic_instruction_rms_ratio"]) assert ratio_error < 1e-12 # An exact per-unit predictor removes all predictable traffic. for slope, gain, coefficient in zip( net.P_traffic, net.traffic_gain, net.B_traffic): slope.copy_(gain * coefficient) exact = net.mixed_apical_components( instruction, forward["hiddens"], "innovation") exact_predictor_error = max(float((left - right).abs().max()) for left, right in zip( exact["innovation"], instruction)) assert exact_predictor_error < 1e-14 for slope, bias in zip(net.P_traffic, net.P_traffic_bias): slope.zero_() bias.zero_() closed_form = net.predictor_closed_form_fit(forward["hiddens"]) fitted = net.mixed_apical_components( instruction, forward["hiddens"], "innovation") closed_form_error = max(float((left - right).abs().max()) for left, right in zip( fitted["innovation"], instruction)) assert closed_form_error < 1e-14 assert closed_form["residual_traffic_rms_ratio"] < 1e-14 assert closed_form["max_absolute_residual_soma_slope"] < 1e-14 for slope, bias in zip(net.P_traffic, net.P_traffic_bias): slope.zero_() bias.zero_() stable_fit = net.predictor_closed_form_fit( forward["hiddens"], stability_margin=1e-3) assert stable_fit["max_positive_residual_soma_slope"] < 1e-14 assert stable_fit["min_residual_soma_slope"] < -9e-4 assert stable_fit["max_applied_stability_margin"] >= 1e-3 # A deliberately inaccurate slow predictor leaves an affine neutral mode. # The fast controller must remove that mode using only paired neutral # soma/traffic observations, without changing the predictor parameters or # reading the task instruction during its coefficient fit. frozen_before_projection = [value.clone() for value in net.P_traffic + net.P_traffic_bias] projected = net.mixed_apical_components( instruction, forward["hiddens"], "innovation", neutral_projection=True) projection = projected["neutral_projection"] projected_instruction_error = max(float((left - right).abs().max()) for left, right in zip( projected["innovation"], instruction)) assert projected_instruction_error < 1e-14 assert projection["post_projection_traffic_rms_ratio"] < 1e-14 assert projection["max_absolute_post_projection_soma_slope"] < 1e-14 assert projection["instruction_observations"] == 0 assert all(torch.equal(before, after) for before, after in zip( frozen_before_projection, net.P_traffic + net.P_traffic_bias)) # Raw and matched controls pay for the identical neutral projection but do # not apply its subtractive direction. Raw must remain exactly the mixed # apical vector. Matched may borrow only projected innovation's norm. projected_raw = net.mixed_apical_components( instruction, forward["hiddens"], "raw", neutral_projection=True) projected_matched = net.mixed_apical_components( instruction, forward["hiddens"], "matched", neutral_projection=True) sham_raw_error = max(float((left - right).abs().max()) for left, right in zip( projected_raw["used"], projected_raw["raw"])) sham_projection_report_error = max( abs(projected_raw["neutral_projection"][key] - projected_matched["neutral_projection"][key]) for key in ( "pre_projection_traffic_rms_ratio", "post_projection_traffic_rms_ratio", "max_absolute_pre_projection_soma_slope", "max_absolute_post_projection_soma_slope", "max_positive_post_projection_soma_slope", "min_post_projection_soma_slope", "max_absolute_correction_slope", )) sham_matched_norm_errors = [] sham_matched_direction_errors = [] for raw, innovation, matched in zip( projected_matched["raw"], projected_matched["innovation"], projected_matched["matched"]): raw_flat = raw.flatten(1) innovation_flat = innovation.flatten(1) matched_flat = matched.flatten(1) sham_matched_norm_errors.append(float(( (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs() / innovation_flat.norm(dim=1).clamp_min(1e-30)).max())) sham_matched_direction_errors.append(float((F.cosine_similarity( raw_flat, matched_flat, dim=1) - 1.0).abs().max())) assert sham_raw_error == 0.0 assert sham_projection_report_error == 0.0 assert max(sham_matched_norm_errors) < 1e-12 assert max(sham_matched_direction_errors) < 1e-12 assert projected_raw["neutral_projection"]["instruction_observations"] == 0 assert all(torch.equal(before, after) for before, after in zip( frozen_before_projection, net.P_traffic + net.P_traffic_bias)) for slope, bias in zip(net.P_traffic, net.P_traffic_bias): slope.zero_() bias.zero_() components = net.mixed_apical_components( instruction, forward["hiddens"], "matched") norm_errors = [] direction_errors = [] for raw, innovation, matched in zip( components["raw"], components["innovation"], components["matched"]): raw_flat = raw.flatten(1) innovation_flat = innovation.flatten(1) matched_flat = matched.flatten(1) norm_errors.append(float(( (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs() / innovation_flat.norm(dim=1).clamp_min(1e-30)).max())) direction_errors.append(float((F.cosine_similarity( raw_flat, matched_flat, dim=1) - 1.0).abs().max())) assert max(norm_errors) < 1e-12 assert max(direction_errors) < 1e-12 # All used signals retain equal independently recomputed local KP products. correlation_errors = [] for rule in ("raw", "matched", "innovation"): used = net.mixed_apical_components( instruction, forward["hiddens"], rule)["used"] forward_directions, _, _, output_weight, _ = ( net.local_ascent_directions(used, output_error, forward)) reciprocal, reciprocal_readout = net.reciprocal_feedback_directions( used, output_error, forward) correlation_errors.extend(float((left - right).abs().max()) for left, right in zip( forward_directions[1:], reciprocal[1:])) correlation_errors.append(float( (reciprocal_readout + output_weight.t()).abs().max())) assert max(correlation_errors) < 1e-14 # Predictor plasticity consumes only the supplied soma/traffic pair: once # those are fixed, changing every forward/feedback weight has no effect. left = CIFARKPMixedTrafficResNet(**common) right = CIFARKPMixedTrafficResNet(**common) for left_gain, right_gain, source in zip( left.traffic_gain, right.traffic_gain, net.traffic_gain): left_gain.copy_(source) right_gain.copy_(source) fixed_hiddens = [value.detach().clone() for value in forward["hiddens"]] for value in right.W + right.Q + [right.W_out, right.R_out]: value.add_(torch.randn_like(value)) left.predictor_step(fixed_hiddens, eta=0.1) right.predictor_step(fixed_hiddens, eta=0.1) predictor_independence_error = max(float((a - b).abs().max()) for a, b in zip( left.P_traffic + left.P_traffic_bias, right.P_traffic + right.P_traffic_bias)) assert predictor_independence_error == 0.0 closed_left = CIFARKPMixedTrafficResNet(**common) closed_right = CIFARKPMixedTrafficResNet(**common) for left_gain, right_gain, source in zip( closed_left.traffic_gain, closed_right.traffic_gain, net.traffic_gain): left_gain.copy_(source) right_gain.copy_(source) for value in (closed_right.W + closed_right.Q + [closed_right.W_out, closed_right.R_out]): value.add_(torch.randn_like(value)) closed_left.predictor_closed_form_fit(fixed_hiddens) closed_right.predictor_closed_form_fit(fixed_hiddens) closed_form_independence_error = max(float((a - b).abs().max()) for a, b in zip( closed_left.P_traffic + closed_left.P_traffic_bias, closed_right.P_traffic + closed_right.P_traffic_bias)) assert closed_form_independence_error == 0.0 net.traffic_rule = "innovation" result = conv_kp_mixed_traffic_step( net, x, y, ConvSDILConfig( eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0, weight_decay=0.0, learn_A=False, learn_P=True), step=0, rule="innovation", predictor_every=16) assert math.isfinite(result["loss"]) and result["did_predictor_update"] frozen_predictor = [value.clone() for value in net.P_traffic + net.P_traffic_bias] frozen_result = conv_kp_mixed_traffic_step( net, x, y, ConvSDILConfig( eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0, weight_decay=0.0, learn_A=False, learn_P=True), step=1, rule="innovation", predictor_every=0) assert math.isfinite(frozen_result["loss"]) assert not frozen_result["did_predictor_update"] assert all(torch.equal(before, after) for before, after in zip( frozen_predictor, net.P_traffic + net.P_traffic_bias)) assert all(not value.requires_grad for value in net.W + net.Q + net.P_traffic + net.P_traffic_bias + [net.W_out, net.R_out, net.b_out]) projected_result = conv_kp_mixed_traffic_step( net, x, y, ConvSDILConfig( eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0, weight_decay=0.0, learn_A=False, learn_P=True), step=2, rule="innovation", predictor_every=0, neutral_projection=True) assert math.isfinite(projected_result["loss"]) assert projected_result["neutral_projection"] is not None assert net.mixed_elementwise_ops_per_example("matched") > ( net.mixed_elementwise_ops_per_example("raw")) return { "kp_traffic_zero_limit_error": max(zero_errors), "kp_traffic_ratio_error": ratio_error, "kp_traffic_exact_predictor_error": exact_predictor_error, "kp_traffic_closed_form_predictor_error": closed_form_error, "kp_traffic_closed_form_residual_ratio": closed_form[ "residual_traffic_rms_ratio"], "kp_traffic_closed_form_residual_slope": closed_form[ "max_absolute_residual_soma_slope"], "kp_traffic_projected_instruction_error": projected_instruction_error, "kp_traffic_projected_residual_ratio": projection[ "post_projection_traffic_rms_ratio"], "kp_traffic_projected_residual_slope": projection[ "max_absolute_post_projection_soma_slope"], "kp_traffic_sham_raw_error": sham_raw_error, "kp_traffic_sham_projection_report_error": ( sham_projection_report_error), "kp_traffic_sham_matched_norm_error": max( sham_matched_norm_errors), "kp_traffic_sham_matched_direction_error": max( sham_matched_direction_errors), "kp_traffic_matched_norm_error": max(norm_errors), "kp_traffic_matched_direction_error": max(direction_errors), "kp_traffic_reciprocal_correlation_error": max(correlation_errors), "kp_traffic_predictor_parameter_independence_error": ( predictor_independence_error), "kp_traffic_closed_form_parameter_independence_error": ( closed_form_independence_error), } def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) x = torch.randn(8, 3, 32, 32) y = torch.arange(8) % 10 clean = net.forward(x, return_cache=True) output_signal = (torch.softmax(clean["logits"], dim=1) - F.one_hot(y, 10)) prediction, _, _ = net.apical_components( output_signal, clean["hiddens"], use_residual=True) targets = [torch.randn_like(value) * 0.01 for value in prediction] before = sum(float((target - value).square().sum()) for target, value in zip(targets, prediction)) net.calibrate_apical( output_signal, clean["hiddens"], prediction, targets, eta=0.1) after_prediction, _, _ = net.apical_components( output_signal, clean["hiddens"], use_residual=True) after = sum(float((target - value).square().sum()) for target, value in zip(targets, after_prediction)) assert after < before gated = CIFARSDILResNet( depth=8, base_width=2, seed=6, vectorizer_mode="channel_gated") gated_clean = gated.forward(x) gated_signal = (torch.softmax(gated_clean["logits"], dim=1) - F.one_hot(y, 10)) gated_prediction, _, _ = gated.apical_components( gated_signal, gated_clean["hiddens"], use_residual=True) gated_targets = [torch.randn_like(value) * 0.01 for value in gated_prediction] gated_before = sum(float((target - value).square().sum()) for target, value in zip(gated_targets, gated_prediction)) gated.calibrate_apical( gated_signal, gated_clean["hiddens"], gated_prediction, gated_targets, eta=0.1) gated_after_prediction, _, _ = gated.apical_components( gated_signal, gated_clean["hiddens"], use_residual=True) gated_after = sum(float((target - value).square().sum()) for target, value in zip(gated_targets, gated_after_prediction)) assert gated_after < gated_before shifted_hidden = [torch.roll(value, shifts=(3, -2), dims=(2, 3)) for value in gated_clean["hiddens"]] shifted_instruction, _, _ = gated.apical_components( gated_signal, shifted_hidden, use_residual=True) original_instruction, _, _ = gated.apical_components( gated_signal, gated_clean["hiddens"], use_residual=True) assert all(torch.allclose( shifted, torch.roll(original, shifts=(3, -2), dims=(2, 3))) for shifted, original in zip(shifted_instruction, original_instruction)) spatial_56 = CIFARSDILResNet(depth=56, vectorizer_mode="spatial_template") gated_56 = CIFARSDILResNet(depth=56, vectorizer_mode="channel_gated") assert spatial_56.n_vectorizer_parameters == 5_324_800 assert gated_56.n_vectorizer_parameters == 40_640 predictor_net = CIFARSDILResNet(depth=8, base_width=2, seed=8) hiddens = [torch.randn(64, *shape) for shape in predictor_net.hidden_shapes] initial = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5) final = initial for _ in range(60): final = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5) assert final < initial * 1e-3 weights_before = [weight.clone() for weight in net.W] result = conv_local_step( net, x[:2], y[:2], ConvSDILConfig( eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0, pert_every=1), step=0, generator=torch.Generator(device="cpu").manual_seed(7)) assert result["did_perturb"] and result["calibration"] is not None assert torch.isfinite(torch.tensor(list( value for key, value in result.items() if isinstance(value, float) and key != "predictor_mse"))).all() assert any(not torch.equal(before_weight, after_weight) for before_weight, after_weight in zip(weights_before, net.W)) assert all(not parameter.requires_grad for parameter in net.W + [net.W_out, net.b_out]) gated_weights_before = [value.clone() for value in gated.A + gated.A_gate] gated_result = conv_local_step( gated, x[:2], y[:2], ConvSDILConfig( eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0, pert_every=1, apical_calibration_mode="channel_subspace"), step=0, generator=torch.Generator(device="cpu").manual_seed(17)) assert gated_result["did_perturb"] assert all(torch.isfinite(torch.tensor(value)) for value in gated_result["calibration"].values()) assert any(not torch.equal(before, after) for before, after in zip( gated_weights_before, gated.A + gated.A_gate)) return {"apical_mse_ratio": after / before, "gated_apical_mse_ratio": gated_after / gated_before, "gated_vectorizer_parameter_reduction": ( spatial_56.n_vectorizer_parameters / gated_56.n_vectorizer_parameters), "predictor_mse_ratio": final / initial} def main(): architecture_checks() perturbation_checks() report = exact_local_gradient_check() report.update(exact_batchnorm_local_gradient_check()) report.update(perturbation_estimator_check()) report.update(channel_subspace_estimator_check()) report.update(vectorizer_subspace_estimator_check()) report.update(hierarchical_feedback_checks()) report.update(hierarchical_parameter_calibration_checks()) report.update(normalized_response_mirror_checks()) report.update(kolen_pollack_checks()) report.update(kp_mixed_traffic_checks()) report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") if __name__ == "__main__": main()