"""Convolutional local-learning primitives for CIFAR residual networks. The forward topology is the standard CIFAR ``6n+2`` basic-block family with option-A identity shortcuts. It supports canonical BatchNorm as well as a normalization-free ablation. BatchNorm's cross-example Jacobian is evaluated inside the current layer only; the synaptic update still never reads a downstream weight. Normalization-free networks use an explicit residual multiplier, which is included in every local eligibility calculation. Forward parameters are plain tensors. The local rule uses only the stored presynaptic activation, a postsynaptic ReLU gate, and a teaching vector at the same hidden population. ``torch.nn.grad.conv2d_weight`` evaluates their local correlation efficiently; it does not traverse a reverse-mode graph or access downstream weights. Autograd is confined to ``bp_step`` and diagnostic smoke tests for the exact BP comparator. """ from dataclasses import dataclass import math import torch import torch.nn.functional as F @dataclass(frozen=True) class ConvLayerSpec: """Static metadata for one locally updated convolution.""" name: str stride: int padding: int hidden_shape: tuple branch_scale: float class CIFARLocalResNet: """CIFAR ResNet with explicit local eligibilities. ``depth`` must satisfy ``depth = 6n + 2``. Hidden populations are defined after the stem ReLU, after every block's first ReLU, and after every block output ReLU. Consequently there is exactly one teaching population per convolution, including a branch-scale factor for each second convolution. """ def __init__(self, depth=20, base_width=16, n_classes=10, device="cpu", dtype=torch.float32, seed=0, weight_scale=1.0, residual_scale=None, normalization="none", bn_momentum=0.1, bn_eps=1e-5): if depth < 8 or (depth - 2) % 6: raise ValueError(f"CIFAR ResNet depth must be 6n+2 and >=8, got {depth}") if base_width <= 0: raise ValueError(f"base_width must be positive, got {base_width}") self.depth = int(depth) self.blocks_per_stage = (depth - 2) // 6 self.base_width = int(base_width) self.n_classes = int(n_classes) self.device = str(device) self.dtype = dtype if normalization not in ("none", "batchnorm"): raise ValueError(f"unknown normalization: {normalization}") self.normalization = normalization self.bn_momentum = float(bn_momentum) self.bn_eps = float(bn_eps) if not 0.0 < self.bn_momentum <= 1.0 or self.bn_eps <= 0: raise ValueError("invalid BatchNorm momentum/epsilon") self.n_blocks = 3 * self.blocks_per_stage self.residual_scale = (1.0 / math.sqrt(self.n_blocks) if residual_scale is None else float(residual_scale)) if not self.residual_scale > 0: raise ValueError("residual_scale must be positive") generator = torch.Generator(device="cpu").manual_seed(seed) self.W = [] self.layer_specs = [] self.blocks = [] self.gamma = [] self.beta = [] self.running_mean = [] self.running_var = [] def add_conv(name, in_channels, out_channels, stride, hidden_shape, branch_scale=1.0): fan_in = 9 * in_channels weight = (torch.randn( out_channels, in_channels, 3, 3, generator=generator) * (weight_scale * math.sqrt(2.0 / fan_in))) self.W.append(weight.to(device=device, dtype=dtype)) if normalization == "batchnorm": self.gamma.append(torch.ones(out_channels, device=device, dtype=dtype)) self.beta.append(torch.zeros(out_channels, device=device, dtype=dtype)) self.running_mean.append(torch.zeros( out_channels, device=device, dtype=dtype)) self.running_var.append(torch.ones( out_channels, device=device, dtype=dtype)) self.layer_specs.append(ConvLayerSpec( name=name, stride=stride, padding=1, hidden_shape=tuple(hidden_shape), branch_scale=float(branch_scale))) return len(self.W) - 1 channels = base_width spatial = 32 stem = add_conv("stem", 3, channels, 1, (channels, spatial, spatial)) if stem != 0: raise AssertionError("stem must be convolution zero") for stage, out_channels in enumerate( (base_width, 2 * base_width, 4 * base_width)): for block in range(self.blocks_per_stage): stride = 2 if stage > 0 and block == 0 else 1 if stride == 2: spatial //= 2 first = add_conv( f"stage{stage + 1}.block{block + 1}.conv1", channels, out_channels, stride, (out_channels, spatial, spatial)) second = add_conv( f"stage{stage + 1}.block{block + 1}.conv2", out_channels, out_channels, 1, (out_channels, spatial, spatial), self.residual_scale) self.blocks.append({ "first": first, "second": second, "in_channels": channels, "out_channels": out_channels, "stride": stride, }) channels = out_channels if len(self.W) != depth - 1: raise AssertionError( f"expected {depth - 1} convolutions, constructed {len(self.W)}") self.W_out = (torch.randn(n_classes, channels, generator=generator) / math.sqrt(channels)).to(device=device, dtype=dtype) self.b_out = torch.zeros(n_classes, device=device, dtype=dtype) self.mW = [torch.zeros_like(weight) for weight in self.W] self.mW_out = torch.zeros_like(self.W_out) self.mb_out = torch.zeros_like(self.b_out) self.mgamma = [torch.zeros_like(value) for value in self.gamma] self.mbeta = [torch.zeros_like(value) for value in self.beta] @property def hidden_shapes(self): return [spec.hidden_shape for spec in self.layer_specs] @property def n_hidden(self): return len(self.layer_specs) @property def n_forward_parameters(self): return (sum(weight.numel() for weight in self.W) + sum(value.numel() for value in self.gamma) + sum(value.numel() for value in self.beta) + self.W_out.numel() + self.b_out.numel()) @property def forward_macs_per_example(self): """Multiply-accumulates in convolutions plus the linear readout.""" total = 0 for weight, spec in zip(self.W, self.layer_specs): out_channels, in_channels, kh, kw = weight.shape _, height, width = spec.hidden_shape total += out_channels * height * width * in_channels * kh * kw total += self.W_out.numel() return int(total) @staticmethod def _option_a_shortcut(x, out_channels, stride): """Original CIFAR ResNet identity shortcut with striding/zero padding.""" if stride == 2: x = x[:, :, ::2, ::2] in_channels = x.shape[1] if in_channels == out_channels: return x if in_channels > out_channels: raise ValueError("option-A shortcut cannot reduce channel count") missing = out_channels - in_channels before = missing // 2 after = missing - before chunks = [] if before: chunks.append(x.new_zeros(x.shape[0], before, x.shape[2], x.shape[3])) chunks.append(x) if after: chunks.append(x.new_zeros(x.shape[0], after, x.shape[2], x.shape[3])) return torch.cat(chunks, dim=1) def _inject(self, value, perturbations, index): if perturbations is None: return value perturbation = perturbations[index] if tuple(perturbation.shape) != tuple(value.shape): raise ValueError( f"perturbation {index} shape {tuple(perturbation.shape)} " f"does not match hidden value {tuple(value.shape)}") return value + perturbation def _normalize(self, index, value, training, update_stats): if self.normalization == "none": return value, None axes = (0, 2, 3) if training: mean = value.mean(dim=axes) variance = value.var(dim=axes, unbiased=False) if update_stats: with torch.no_grad(): count = value.numel() // value.shape[1] unbiased = variance * count / max(1, count - 1) self.running_mean[index].lerp_(mean.detach(), self.bn_momentum) self.running_var[index].lerp_(unbiased.detach(), self.bn_momentum) else: mean = self.running_mean[index] variance = self.running_var[index] inverse_std = torch.rsqrt(variance + self.bn_eps) normalized = ((value - mean[None, :, None, None]) * inverse_std[None, :, None, None]) output = (self.gamma[index][None, :, None, None] * normalized + self.beta[index][None, :, None, None]) cache = { "normalized": normalized, "inverse_std": inverse_std, "training": bool(training), } return output, cache def forward(self, x, perturbations=None, return_cache=False, training=False, update_stats=False): if x.ndim != 4 or tuple(x.shape[1:]) != (3, 32, 32): raise ValueError(f"expected CIFAR NCHW input, got {tuple(x.shape)}") if perturbations is not None and len(perturbations) != self.n_hidden: raise ValueError( f"expected {self.n_hidden} perturbations, got {len(perturbations)}") hiddens = [] caches = [] pre = x u = F.conv2d(pre, self.W[0], stride=1, padding=1) normalized, norm_cache = self._normalize(0, u, training, update_stats) h_clean = F.relu(normalized) hiddens.append(h_clean) if return_cache: caches.append({"pre": pre, "gate": normalized > 0, "normalization": norm_cache}) h = self._inject(h_clean, perturbations, 0) for block in self.blocks: first = block["first"] second = block["second"] shortcut = self._option_a_shortcut( h, block["out_channels"], block["stride"]) pre_first = h u_first = F.conv2d( pre_first, self.W[first], stride=block["stride"], padding=1) normalized_first, first_norm_cache = self._normalize( first, u_first, training, update_stats) first_clean = F.relu(normalized_first) hiddens.append(first_clean) if return_cache: caches.append({"pre": pre_first, "gate": normalized_first > 0, "normalization": first_norm_cache}) first_value = self._inject(first_clean, perturbations, first) u_second = F.conv2d(first_value, self.W[second], stride=1, padding=1) normalized_second, second_norm_cache = self._normalize( second, u_second, training, update_stats) block_pre = shortcut + self.residual_scale * normalized_second block_clean = F.relu(block_pre) hiddens.append(block_clean) if return_cache: caches.append({"pre": first_value, "gate": block_pre > 0, "normalization": second_norm_cache}) h = self._inject(block_clean, perturbations, second) features = h.mean(dim=(2, 3)) logits = features @ self.W_out.t() + self.b_out result = {"logits": logits, "features": features, "hiddens": hiddens} if return_cache: if len(caches) != self.n_hidden: raise AssertionError("cache/hidden layer mismatch") result["caches"] = caches return result def logits(self, x): return self.forward(x)["logits"] def _normalization_backward(self, index, delta, cache): """Local BatchNorm Jacobian-vector product and affine directions.""" if self.normalization == "none": return delta, None, None normalized = cache["normalized"] gamma_direction = (delta * normalized).sum(dim=(0, 2, 3)) beta_direction = delta.sum(dim=(0, 2, 3)) scaled = delta * self.gamma[index][None, :, None, None] inverse_std = cache["inverse_std"][None, :, None, None] if cache["training"]: count = delta.shape[0] * delta.shape[2] * delta.shape[3] summed = scaled.sum(dim=(0, 2, 3), keepdim=True) projected = (scaled * normalized).sum( dim=(0, 2, 3), keepdim=True) input_delta = (inverse_std / count) * ( count * scaled - summed - normalized * projected) else: input_delta = inverse_std * scaled return input_delta, gamma_direction, beta_direction def local_ascent_directions(self, teaching, output_error, forward): """Return forward-parameter descent directions from local signals. ``teaching[l][i]`` represents the per-example ``-d ell_i/dh_l``. Each convolutional direction averages the exact local Jacobian-vector products using only that population's cache. The output error is the per-example ``d ell_i/dlogits`` and therefore receives an explicit minus sign. """ if len(teaching) != self.n_hidden: raise ValueError(f"expected {self.n_hidden} teaching tensors") caches = forward.get("caches") if caches is None: raise ValueError("local directions require a cached forward pass") batch = output_error.shape[0] directions = [] gamma_directions = [] beta_directions = [] with torch.no_grad(): for index, (signal, cache, spec, weight) in enumerate(zip( teaching, caches, self.layer_specs, self.W)): if tuple(signal.shape[1:]) != spec.hidden_shape: raise ValueError( f"teaching {index} has {tuple(signal.shape[1:])}, " f"expected {spec.hidden_shape}") post_norm_delta = (signal * cache["gate"].to(signal.dtype) * spec.branch_scale) delta, gamma_direction, beta_direction = self._normalization_backward( index, post_norm_delta, cache["normalization"]) direction = torch.nn.grad.conv2d_weight( cache["pre"].detach(), weight.shape, delta.detach(), stride=spec.stride, padding=spec.padding) directions.append(direction / batch) if gamma_direction is not None: gamma_directions.append(gamma_direction / batch) beta_directions.append(beta_direction / batch) output_weight = -(output_error.t() @ forward["features"].detach()) / batch output_bias = -output_error.mean(dim=0) return (directions, gamma_directions, beta_directions, output_weight, output_bias) def apply_ascent(self, directions, output_weight, output_bias, eta_hidden, eta_output=None, momentum=0.0, weight_decay=0.0, gamma_directions=None, beta_directions=None): """Apply simultaneously computed directions with optional momentum.""" if len(directions) != len(self.W): raise ValueError("one direction is required for every convolution") eta_output = eta_hidden if eta_output is None else eta_output gamma_directions = [] if gamma_directions is None else gamma_directions beta_directions = [] if beta_directions is None else beta_directions if self.normalization == "batchnorm" and not ( len(gamma_directions) == len(beta_directions) == len(self.W)): raise ValueError("BatchNorm directions must cover every convolution") with torch.no_grad(): for index, (weight, direction) in enumerate(zip(self.W, directions)): update = direction - weight_decay * weight if momentum: self.mW[index].mul_(momentum).add_(update) update = self.mW[index] weight.add_(update, alpha=eta_hidden) for index, (gamma_direction, beta_direction) in enumerate(zip( gamma_directions, beta_directions)): if momentum: self.mgamma[index].mul_(momentum).add_(gamma_direction) self.mbeta[index].mul_(momentum).add_(beta_direction) gamma_direction = self.mgamma[index] beta_direction = self.mbeta[index] self.gamma[index].add_(gamma_direction, alpha=eta_hidden) self.beta[index].add_(beta_direction, alpha=eta_hidden) out_update = output_weight - weight_decay * self.W_out if momentum: self.mW_out.mul_(momentum).add_(out_update) self.mb_out.mul_(momentum).add_(output_bias) out_update = self.mW_out output_bias = self.mb_out self.W_out.add_(out_update, alpha=eta_output) self.b_out.add_(output_bias, alpha=eta_output) def bp_step(self, x, y, eta, momentum=0.0, weight_decay=0.0): """Exact-backprop comparator on the identical forward architecture.""" parameters = self.W + self.gamma + self.beta + [self.W_out, self.b_out] for parameter in parameters: parameter.requires_grad_(True) loss = F.cross_entropy( self.forward(x, training=True, update_stats=True)["logits"], y) gradients = torch.autograd.grad(loss, parameters) n_conv = len(self.W) with torch.no_grad(): conv_directions = [-gradient for gradient in gradients[:n_conv]] if self.normalization == "batchnorm": gamma_directions = [ -gradient for gradient in gradients[n_conv:2 * n_conv]] beta_directions = [ -gradient for gradient in gradients[2 * n_conv:3 * n_conv]] else: gamma_directions = [] beta_directions = [] output_weight = -gradients[-2] output_bias = -gradients[-1] self.apply_ascent( conv_directions, output_weight, output_bias, eta, momentum=momentum, weight_decay=weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) for parameter in parameters: parameter.requires_grad_(False) return float(loss.detach()) class CIFARHierarchicalFAResNet(CIFARLocalResNet): """Residual-DAG feedback alignment with independent convolutional weights. Feedback follows the actual child edges of the forward residual graph and uses locally available ReLU/BatchNorm Jacobians, but every convolutional feedback tensor is initialized independently and never reads its forward counterpart. This is a baseline and an infrastructure step for learned hierarchical dendritic feedback, not an SDIL novelty claim. """ def __init__(self, *args, feedback_seed=None, feedback_scale=1.0, **kwargs): model_seed = kwargs.get("seed", 0) super().__init__(*args, **kwargs) if feedback_scale <= 0: raise ValueError("feedback_scale must be positive") generator = torch.Generator(device="cpu").manual_seed( model_seed + 30011 if feedback_seed is None else feedback_seed) self.Q = [] for weight in self.W: fan_in = weight.shape[1] * weight.shape[2] * weight.shape[3] value = (torch.randn(weight.shape, generator=generator) * (feedback_scale * math.sqrt(2.0 / fan_in))) self.Q.append(value.to(device=weight.device, dtype=weight.dtype)) channels = self.W_out.shape[1] self.R_out = (torch.randn( channels, self.n_classes, generator=generator) * (feedback_scale / math.sqrt(channels))).to( device=self.W_out.device, dtype=self.W_out.dtype) @property def n_fixed_feedback_parameters(self): # Q[0] maps the stem to pixels and is not used for hidden credit. return (sum(value.numel() for value in self.Q[1:]) + self.R_out.numel()) @property def apical_macs_per_example(self): conv = 0 for weight, spec in zip(self.Q[1:], self.layer_specs[1:]): out_channels, in_channels, kh, kw = weight.shape _, height, width = spec.hidden_shape conv += out_channels * height * width * in_channels * kh * kw return int(conv + self.R_out.numel()) @staticmethod def _option_a_shortcut_transpose(value, in_channels, stride, output_shape): """Adjoint of the parameter-free option-A shortcut.""" out_channels = value.shape[1] if in_channels > out_channels: raise ValueError("option-A transpose cannot recover reduced channels") missing = out_channels - in_channels before = missing // 2 selected = value[:, before:before + in_channels] if stride == 1: if tuple(selected.shape) != tuple(output_shape): raise ValueError("shortcut transpose shape mismatch") return selected result = value.new_zeros(output_shape) result[:, :, ::2, ::2] = selected return result @torch.no_grad() def hierarchical_teaching(self, output_signal, forward, return_edge_contexts=False): """Propagate teaching fields through independent feedback convolutions. When requested, ``edge_contexts[l]`` is the local child field consumed by ``Q[l]`` and ``recipients[l]`` is the hidden population receiving that transposed-convolution contribution. These values expose the sufficient statistics for causal feedback calibration without reading any forward tensor. """ caches = forward.get("caches") hiddens = forward.get("hiddens") if caches is None or hiddens is None: raise ValueError("hierarchical feedback requires cached hidden states") if len(hiddens) != self.n_hidden: raise ValueError("hierarchical hidden population mismatch") teaching = [torch.zeros_like(value) for value in hiddens] edge_contexts = [None for _ in self.Q] recipients = [None for _ in self.Q] spatial = hiddens[-1].shape[2] * hiddens[-1].shape[3] teaching[-1].copy_( (output_signal @ self.R_out.t())[:, :, None, None] / spatial) for block in reversed(self.blocks): first = block["first"] second = block["second"] parent = first - 1 second_gate = caches[second]["gate"].to(teaching[second].dtype) second_delta = teaching[second] * second_gate branch_delta, _, _ = self._normalization_backward( second, second_delta * self.residual_scale, caches[second]["normalization"]) edge_contexts[second] = branch_delta recipients[second] = first teaching[first].add_(F.conv_transpose2d( branch_delta, self.Q[second], stride=1, padding=1)) teaching[parent].add_(self._option_a_shortcut_transpose( second_delta, block["in_channels"], block["stride"], hiddens[parent].shape)) first_gate = caches[first]["gate"].to(teaching[first].dtype) first_delta = teaching[first] * first_gate first_delta, _, _ = self._normalization_backward( first, first_delta, caches[first]["normalization"]) edge_contexts[first] = first_delta recipients[first] = parent teaching[parent].add_(F.conv_transpose2d( first_delta, self.Q[first], stride=block["stride"], padding=1, output_padding=block["stride"] - 1)) if return_edge_contexts: if any(value is None for value in edge_contexts[1:]): raise AssertionError("hierarchical edge context is incomplete") return teaching, edge_contexts, recipients return teaching class CIFARKPResNet(CIFARHierarchicalFAResNet): """Modified Kolen--Pollack reciprocal-feedback baseline. Forward and reciprocal synapses receive the same two locally available activity factors and independently form equal-shaped correlation updates. Matching optimizer and decay dynamics then suppress their initial difference without reading or copying either synaptic weight. This is the inherited Akrout et al. baseline, not an SDIL contribution. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.mQ = [torch.zeros_like(value) for value in self.Q] self.mR_out = torch.zeros_like(self.R_out) @torch.no_grad() def reciprocal_feedback_directions(self, teaching, output_error, forward): """Recompute reciprocal updates from local activities only. Q is stored in forward-kernel orientation because ``conv_transpose2d`` applies its transpose during feedback. Consequently its local KP correlation has the same stored orientation as the corresponding W correlation. R is stored as ``-W_out.T`` under this runner's teaching sign convention. """ if len(teaching) != self.n_hidden: raise ValueError("KP teaching must cover every hidden population") caches = forward.get("caches") if caches is None: raise ValueError("KP feedback updates require local forward caches") batch = output_error.shape[0] directions = [None] for index in range(1, len(self.Q)): signal = teaching[index] cache = caches[index] spec = self.layer_specs[index] post_norm_delta = (signal * cache["gate"].to(signal.dtype) * spec.branch_scale) delta, _, _ = self._normalization_backward( index, post_norm_delta, cache["normalization"]) direction = torch.nn.grad.conv2d_weight( cache["pre"].detach(), self.Q[index].shape, delta.detach(), stride=spec.stride, padding=spec.padding) directions.append(direction / batch) readout_direction = ( forward["features"].detach().t() @ output_error) / batch return directions, readout_direction @torch.no_grad() def apply_reciprocal_ascent(self, directions, readout_direction, eta_hidden, eta_output=None, momentum=0.0, weight_decay=0.0): """Apply the independently formed KP feedback-synapse updates.""" if len(directions) != len(self.Q) or directions[0] is not None: raise ValueError("KP directions must cover Q[1:] only") eta_output = eta_hidden if eta_output is None else eta_output for index in range(1, len(self.Q)): update = directions[index] - weight_decay * self.Q[index] if momentum: self.mQ[index].mul_(momentum).add_(update) update = self.mQ[index] self.Q[index].add_(update, alpha=eta_hidden) update = readout_direction - weight_decay * self.R_out if momentum: self.mR_out.mul_(momentum).add_(update) update = self.mR_out self.R_out.add_(update, alpha=eta_output) class CIFARKPMixedTrafficResNet(CIFARKPResNet): """KP credit with a per-unit mixed apical compartment and neutral predictor. The reciprocal KP pathway supplies the instructional field. Fixed soma-coupled traffic is added locally at every hidden population, while a diagonal affine predictor is fitted only to instruction-off observations. Which of raw, norm-matched raw, or innovation drives plasticity is selected by the runner; every condition still computes and trains the predictor. """ def __init__(self, *args, traffic_seed=4000, **kwargs): super().__init__(*args, **kwargs) generator = torch.Generator(device="cpu").manual_seed(traffic_seed) self.B_traffic = [] self.traffic_gain = [] self.P_traffic = [] self.P_traffic_bias = [] for shape in self.hidden_shapes: coefficient = torch.exp( 0.25 * torch.randn(shape, generator=generator)) self.B_traffic.append(coefficient.to( device=self.device, dtype=self.dtype)) self.traffic_gain.append(torch.zeros( (), device=self.device, dtype=self.dtype)) self.P_traffic.append(torch.zeros( shape, device=self.device, dtype=self.dtype)) self.P_traffic_bias.append(torch.zeros( shape, device=self.device, dtype=self.dtype)) self.traffic_rule = None @property def n_predictor_parameters(self): return (sum(value.numel() for value in self.P_traffic) + sum(value.numel() for value in self.P_traffic_bias)) @property def n_fixed_traffic_coefficients(self): return sum(value.numel() for value in self.B_traffic) @property def n_apical_parameters(self): # Reciprocal Q/R parameters are logged separately as adaptive feedback. return self.n_predictor_parameters @property def mixed_units_per_example(self): return sum(math.prod(shape) for shape in self.hidden_shapes) def mixed_elementwise_ops_per_example(self, rule=None): """Transparent arithmetic count beyond convolution/linear MACs. Six operations per unit form traffic, affine prediction, raw, and innovation in the executed tensor expression. Norm matching additionally charges squares/reductions, rescaling, and scalar norm arithmetic conservatively as five per unit. """ rule = self.traffic_rule if rule is None else rule if rule not in ("raw", "matched", "innovation"): raise ValueError(f"unknown mixed-traffic rule: {rule}") multiplier = 11 if rule == "matched" else 6 return multiplier * self.mixed_units_per_example @property def predictor_elementwise_ops_per_example(self): # Conservative count for traffic/prediction, centering, reductions, # normalized moments, residual power, and affine parameter updates. return 24 * self.mixed_units_per_example @property def traffic_calibration_elementwise_ops_per_example(self): # Two RMS measurements, gain application, and achieved-ratio audit. return 8 * self.mixed_units_per_example @property def predictor_audit_elementwise_ops_per_example(self): # Traffic/prediction/residual construction and two power reductions. return 10 * self.mixed_units_per_example @property def neutral_projection_elementwise_ops_per_example(self): """Conservative arithmetic count for fast local affine projection. The count covers local means, centering, variance/covariance, coefficient formation, affine reconstruction, and subtraction. It is deliberately separate from the slow-predictor cost because the two mechanisms operate on different timescales and observations. """ return 18 * self.mixed_units_per_example @torch.no_grad() def traffic_fields(self, hiddens): if len(hiddens) != self.n_hidden: raise ValueError("traffic requires every somatic population") return [gain * coefficient * hidden for gain, coefficient, hidden in zip( self.traffic_gain, self.B_traffic, hiddens)] @torch.no_grad() def calibrate_traffic_gain(self, instruction, hiddens, target_ratio): """Fix one gain per layer from an initialization-only training prefix.""" if target_ratio <= 0: raise ValueError("traffic ratio must be positive") if not (len(instruction) == len(hiddens) == self.n_hidden): raise ValueError("traffic calibration must cover every population") instruction_rms = [] unscaled_traffic_rms = [] gains = [] realized = [] for signal, hidden, coefficient, gain in zip( instruction, hiddens, self.B_traffic, self.traffic_gain): signal_scale = signal.square().mean().sqrt() traffic_scale = (coefficient * hidden).square().mean().sqrt() if float(signal_scale) <= 0 or float(traffic_scale) <= 0: raise ValueError("traffic calibration encountered a zero RMS") value = target_ratio * signal_scale / traffic_scale gain.copy_(value) achieved = (gain * coefficient * hidden).square().mean().sqrt() instruction_rms.append(float(signal_scale)) unscaled_traffic_rms.append(float(traffic_scale)) gains.append(float(gain)) realized.append(float(achieved / signal_scale)) return { "target_ratio": float(target_ratio), "instruction_rms": instruction_rms, "unscaled_traffic_rms": unscaled_traffic_rms, "traffic_gain": gains, "realized_traffic_instruction_rms_ratio": realized, } @torch.no_grad() def neutral_residual_projection(self, hiddens, min_variance=1e-12): """Project instruction-off residuals away from current local soma. This is a fast stability controller, not a task-period predictor update. Each cell observes its soma and ordinary apical traffic with instruction absent, fits only the affine component of the *remaining* neutral residual, and returns the orthogonal remainder. No teaching signal, label, loss, forward weight, or feedback weight enters the projection. Repeating the projection on each task minibatch adapts its coefficient to the current local covariance. For diagonal affine traffic it nulls the multiplicative residual mode instead of placing that mode at a fixed negative margin whose stability depends on a changing input covariance. """ if min_variance < 0: raise ValueError("minimum projection variance must be nonnegative") traffic = self.traffic_fields(hiddens) stabilized = [] pre_power = 0.0 post_power = 0.0 traffic_power = 0.0 maximum_pre_slope = 0.0 maximum_post_slope = 0.0 maximum_positive_post_slope = 0.0 minimum_post_slope = 0.0 maximum_correction_slope = 0.0 for hidden, target, slope, bias in zip( hiddens, traffic, self.P_traffic, self.P_traffic_bias): neutral = target - (slope * hidden + bias) hidden_mean = hidden.mean(dim=0) neutral_mean = neutral.mean(dim=0) centered_h = hidden - hidden_mean centered_neutral = neutral - neutral_mean variance = centered_h.square().mean(dim=0) covariance = (centered_h * centered_neutral).mean(dim=0) active = variance > min_variance correction_slope = torch.where( active, covariance / variance.clamp_min(min_variance), torch.zeros_like(variance)) remainder = (centered_neutral - correction_slope * centered_h) # Audit the coefficient left after projection using the same local # sufficient statistics. Inactive cells are constant across the # batch and have already been centered to zero. remainder_mean = remainder.mean(dim=0) centered_remainder = remainder - remainder_mean post_covariance = (centered_h * centered_remainder).mean(dim=0) post_slope = torch.where( active, post_covariance / variance.clamp_min(min_variance), torch.zeros_like(variance)) pre_slope = torch.where( active, covariance / variance.clamp_min(min_variance), torch.zeros_like(variance)) maximum_pre_slope = max( maximum_pre_slope, float(pre_slope.abs().max())) maximum_post_slope = max( maximum_post_slope, float(post_slope.abs().max())) maximum_positive_post_slope = max( maximum_positive_post_slope, float(post_slope.max().clamp_min(0.0))) minimum_post_slope = min( minimum_post_slope, float(post_slope.min())) maximum_correction_slope = max( maximum_correction_slope, float(correction_slope.abs().max())) pre_power += float(neutral.square().sum()) post_power += float(remainder.square().sum()) traffic_power += float(target.square().sum()) stabilized.append(remainder) if traffic_power <= 0: raise ValueError("neutral projection requires nonzero traffic") return stabilized, { "pre_projection_traffic_rms_ratio": math.sqrt( pre_power / traffic_power), "post_projection_traffic_rms_ratio": math.sqrt( post_power / traffic_power), "max_absolute_pre_projection_soma_slope": maximum_pre_slope, "max_absolute_post_projection_soma_slope": maximum_post_slope, "max_positive_post_projection_soma_slope": ( maximum_positive_post_slope), "min_post_projection_soma_slope": minimum_post_slope, "max_absolute_correction_slope": maximum_correction_slope, "observations": int(hiddens[0].shape[0]), "instruction_observations": 0, } @torch.no_grad() def mixed_apical_components(self, instruction, hiddens, rule, compute_matched=False, neutral_projection=False): """Return used, raw, innovation, matched, and traffic fields. When ``neutral_projection`` is enabled, every rule pays for and reports the same instruction-off projection. Only the innovation rule applies the subtractive direction. Raw keeps the unmodified apical vector, while matched keeps its direction and borrows only the projected innovation norm. This makes raw and matched strict sham controls for the subtractive operation without changing their local information or observation budget. """ if rule not in ("raw", "matched", "innovation"): raise ValueError(f"unknown mixed-traffic rule: {rule}") if not (len(instruction) == len(hiddens) == self.n_hidden): raise ValueError("mixed apical inputs must cover every population") traffic = self.traffic_fields(hiddens) projected_neutral = None projection_report = None if neutral_projection: projected_neutral, projection_report = ( self.neutral_residual_projection(hiddens)) raw = [] innovation = [] matched = [] if (rule == "matched" or compute_matched) else None for index, (signal, hidden, ordinary, slope, bias) in enumerate(zip( instruction, hiddens, traffic, self.P_traffic, self.P_traffic_bias)): apical = signal + ordinary if projected_neutral is None: residual = apical - (slope * hidden + bias) else: residual = signal + projected_neutral[index] raw.append(apical) innovation.append(residual) if matched is not None: raw_norm = apical.flatten(1).norm(dim=1).clamp_min(1e-30) residual_norm = residual.flatten(1).norm(dim=1) scale = (residual_norm / raw_norm).reshape( residual.shape[0], *([1] * (residual.ndim - 1))) matched.append(scale * apical) choices = {"raw": raw, "matched": matched, "innovation": innovation} return { "used": choices[rule], "raw": raw, "innovation": innovation, "matched": matched, "traffic": traffic, "instruction": instruction, "neutral_projection": projection_report, } @torch.no_grad() def predictor_step(self, hiddens, eta): """Instruction-off normalized-LMS update from local soma/traffic pairs.""" if not 0.0 < eta <= 1.0: raise ValueError("predictor learning rate must lie in (0, 1]") traffic = self.traffic_fields(hiddens) squared_error = 0.0 units = 0 for hidden, target, slope, bias in zip( hiddens, traffic, self.P_traffic, self.P_traffic_bias): residual = target - slope * hidden - bias centered_h = hidden - hidden.mean(dim=0) centered_r = residual - residual.mean(dim=0) variance = centered_h.square().mean(dim=0) slope.add_((centered_r * centered_h).mean(dim=0) / (variance + 1e-12), alpha=eta) bias.add_(residual.mean(dim=0), alpha=eta) squared_error += float(residual.square().sum()) units += residual.numel() return squared_error / units @torch.no_grad() def predictor_closed_form_fit(self, hiddens, min_variance=1e-12, stability_margin=0.0): """Fit the local affine neutral relation by per-cell least squares. Each spatial cell uses only its own soma and instruction-off apical observations across the supplied batch. Cells with no observed soma variance receive a zero slope and their local target mean as bias. """ if min_variance < 0: raise ValueError("minimum predictor variance must be nonnegative") if stability_margin < 0: raise ValueError("predictor stability margin must be nonnegative") traffic = self.traffic_fields(hiddens) residual_power = 0.0 traffic_power = 0.0 maximum_residual_slope = 0.0 maximum_positive_residual_slope = 0.0 minimum_residual_slope = 0.0 maximum_applied_margin = 0.0 squared_error = 0.0 units = 0 for hidden, target, slope, bias in zip( hiddens, traffic, self.P_traffic, self.P_traffic_bias): hidden_mean = hidden.mean(dim=0) target_mean = target.mean(dim=0) centered_h = hidden - hidden_mean centered_target = target - target_mean variance = centered_h.square().mean(dim=0) covariance = (centered_h * centered_target).mean(dim=0) active = variance > min_variance fitted_slope = torch.where( active, covariance / variance.clamp_min(min_variance), torch.zeros_like(variance)) applied_margin = stability_margin * (1.0 + fitted_slope.abs()) stabilized_slope = fitted_slope + applied_margin fitted_bias = target_mean - stabilized_slope * hidden_mean slope.copy_(stabilized_slope) bias.copy_(fitted_bias) residual = target - slope * hidden - bias residual_covariance = (centered_h * ( residual - residual.mean(dim=0))).mean(dim=0) residual_slope = torch.where( active, residual_covariance / variance.clamp_min(min_variance), torch.zeros_like(variance)) maximum_residual_slope = max( maximum_residual_slope, float(residual_slope.abs().max())) maximum_positive_residual_slope = max( maximum_positive_residual_slope, float(residual_slope.max().clamp_min(0.0))) minimum_residual_slope = min( minimum_residual_slope, float(residual_slope.min())) maximum_applied_margin = max( maximum_applied_margin, float(applied_margin.max())) power = float(residual.square().sum()) residual_power += power traffic_power += float(target.square().sum()) squared_error += power units += residual.numel() if traffic_power <= 0: raise ValueError("closed-form predictor fit requires nonzero traffic") return { "mse": squared_error / units, "residual_traffic_rms_ratio": math.sqrt( residual_power / traffic_power), "max_absolute_residual_soma_slope": maximum_residual_slope, "max_positive_residual_soma_slope": ( maximum_positive_residual_slope), "min_residual_soma_slope": minimum_residual_slope, "max_applied_stability_margin": maximum_applied_margin, "stability_margin": float(stability_margin), "observations": int(hiddens[0].shape[0]), } @torch.no_grad() def predictor_traffic_residual_rms_ratio(self, hiddens): traffic = self.traffic_fields(hiddens) residual_power = 0.0 traffic_power = 0.0 for hidden, target, slope, bias in zip( hiddens, traffic, self.P_traffic, self.P_traffic_bias): residual = target - slope * hidden - bias residual_power += float(residual.square().sum()) traffic_power += float(target.square().sum()) if traffic_power <= 0: raise ValueError("predictor audit requires nonzero traffic") return math.sqrt(residual_power / traffic_power) @torch.no_grad() def hierarchical_feedback_tracking_report(net): """Cheap parameter-space tracking diagnostics; never used for learning.""" if not isinstance(net, CIFARHierarchicalFAResNet): raise TypeError("feedback tracking requires a hierarchical network") 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().clamp_min(1e-30)) for feedback, target in pairs] relative_errors = [float( (feedback - target).norm() / target.norm().clamp_min(1e-30)) for feedback, target in pairs] return { "feedback_forward_cosine": cosines, "feedback_forward_norm_ratio": norm_ratios, "feedback_forward_relative_error": relative_errors, "mean_feedback_forward_cosine": sum(cosines) / len(cosines), "mean_feedback_forward_relative_error": ( sum(relative_errors) / len(relative_errors)), } @torch.no_grad() def hierarchical_parameter_subspace_calibration( net, x, y, clean_forward, output_signal, sigma=1e-2, n_directions=1, eta=1e-3, generator=None, return_diagnostics=False): """Calibrate the residual-DAG feedback maps with two causal queries. A Rademacher tensor is drawn in every Q/R parameter space. Each tensor is applied only to its locally available child field, producing one candidate intervention at the edge's parent population. All independent candidate fields are injected in the same antithetic pair. Multiplying the scalar loss derivative back into each random tensor gives an unbiased estimate of that edge's causal target moment; subtracting the current predicted moment is the exact local squared-field delta rule. No forward weight or reverse differentiation is used. """ if not isinstance(net, CIFARHierarchicalFAResNet): raise TypeError("hierarchical calibration requires a hierarchical net") if sigma <= 0 or n_directions < 1 or eta < 0: raise ValueError("invalid hierarchical calibration hyperparameters") if generator is None: generator = torch.Generator(device=x.device).manual_seed(0) hiddens = clean_forward.get("hiddens") if hiddens is None or len(hiddens) != net.n_hidden: raise ValueError("clean forward does not match hierarchical populations") teaching, contexts, recipients = net.hierarchical_teaching( output_signal, clean_forward, return_edge_contexts=True) batch = x.shape[0] target_q = [torch.zeros_like(value) if index else None for index, value in enumerate(net.Q)] target_r = torch.zeros_like(net.R_out) diagnostic_directions = [] diagnostic_derivatives = [] for _ in range(n_directions): random_q = [None] random_q.extend([ torch.empty_like(value).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) for value in net.Q[1:] ]) random_r = torch.empty_like(net.R_out).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) interventions = [torch.zeros_like(value) for value in hiddens] spatial_out = hiddens[-1].shape[2] * hiddens[-1].shape[3] interventions[-1].add_( (output_signal @ random_r.t())[:, :, None, None] / spatial_out) for index in range(1, len(net.Q)): spec = net.layer_specs[index] recipient = recipients[index] contribution = F.conv_transpose2d( contexts[index], random_q[index], stride=spec.stride, padding=spec.padding, output_padding=spec.stride - 1) if contribution.shape != interventions[recipient].shape: raise AssertionError("hierarchical intervention shape mismatch") interventions[recipient].add_(contribution) plus = F.cross_entropy(net.forward( x, perturbations=[sigma * value for value in interventions], training=True, update_stats=False)["logits"], y) minus = F.cross_entropy(net.forward( x, perturbations=[-sigma * value for value in interventions], training=True, update_stats=False)["logits"], y) # The parameter maps are batch-shared. Scale the mean-loss derivative # into a summed-loss hidden signal, then average its moment over B and # recipient spatial sites to keep one eta meaningful across stages. directional = (plus - minus) * batch / (2.0 * sigma) target_r.add_(random_r, alpha=-float(directional) / ( batch * n_directions)) for index in range(1, len(net.Q)): recipient = recipients[index] spatial = (hiddens[recipient].shape[2] * hiddens[recipient].shape[3]) target_q[index].add_( random_q[index], alpha=-float(directional) / ( batch * spatial * n_directions)) if return_diagnostics: diagnostic_directions.append({ "hidden": interventions, "Q": random_q, "R": random_r}) diagnostic_derivatives.append({ "scaled_directional": directional, "coupling": "summed_batch_objective", }) prediction_q = [None] for index in range(1, len(net.Q)): recipient = recipients[index] spec = net.layer_specs[index] spatial = (hiddens[recipient].shape[2] * hiddens[recipient].shape[3]) prediction_q.append(torch.nn.grad.conv2d_weight( teaching[recipient], net.Q[index].shape, contexts[index], stride=spec.stride, padding=spec.padding) / (batch * spatial)) prediction_r = (teaching[-1].mean(dim=(2, 3)).t() @ output_signal) / batch errors_q = [None] errors_q.extend([ target_q[index] - prediction_q[index] for index in range(1, len(net.Q)) ]) error_r = target_r - prediction_r for index in range(1, len(net.Q)): net.Q[index].add_(errors_q[index], alpha=eta) net.R_out.add_(error_r, alpha=eta) target_power = float(target_r.square().sum()) prediction_power = float(prediction_r.square().sum()) error_power = float(error_r.square().sum()) dot = float((target_r * prediction_r).sum()) parameters = target_r.numel() for index in range(1, len(net.Q)): target_power += float(target_q[index].square().sum()) prediction_power += float(prediction_q[index].square().sum()) error_power += float(errors_q[index].square().sum()) dot += float((target_q[index] * prediction_q[index]).sum()) parameters += target_q[index].numel() denominator = math.sqrt(target_power * prediction_power) calibration = { "calibration_mse": error_power / parameters, "target_power": target_power / parameters, "prediction_target_cosine": dot / denominator if denominator else 0.0, "parameter_update_rms": math.sqrt(error_power / parameters), } if return_diagnostics: return calibration, { "directions": diagnostic_directions, "directional_derivatives": diagnostic_derivatives, "teaching": teaching, "contexts": contexts, "recipients": recipients, "target_Q": target_q, "prediction_Q": prediction_q, "target_R": target_r, "prediction_R": prediction_r, } return calibration @torch.no_grad() def hierarchical_mirror_observations(net, batch_size=1, noise_std=1.0, generator=None): """Generate local bias-blocked probe/child-response pairs. This is the observation phase of a normalized weight-mirror baseline. A parent population emits independent Gaussian noise and the ordinary forward synapses generate the child's preactivation response. The return value contains activities only; the subsequent feedback update has no access to forward parameters. """ if not isinstance(net, CIFARHierarchicalFAResNet): raise TypeError("mirror observations require a hierarchical net") if batch_size < 1 or noise_std <= 0: raise ValueError("invalid mirror observation hyperparameters") if generator is None: generator = torch.Generator(device=net.W_out.device).manual_seed(0) conv = [None] for index in range(1, len(net.W)): weight = net.W[index] spec = net.layer_specs[index] _, in_channels, _, _ = weight.shape _, out_height, out_width = spec.hidden_shape input_shape = ( batch_size, in_channels, out_height * spec.stride, out_width * spec.stride) probe = torch.randn( input_shape, generator=generator, device=weight.device, dtype=weight.dtype).mul_(noise_std) response = F.conv2d( probe, weight, stride=spec.stride, padding=spec.padding) conv.append((probe, response)) # The dense readout has no spatial sample multiplicity. Use a local probe # population large enough to estimate its channel covariance accurately; # its MAC cost is recorded explicitly by the runner. readout_batch = max( batch_size, 16 * net.W_out.shape[1], 16 * net.W_out.shape[0]) readout_probe = torch.randn( readout_batch, net.W_out.shape[1], generator=generator, device=net.W_out.device, dtype=net.W_out.dtype).mul_(noise_std) readout_response = readout_probe @ net.W_out.t() return { "conv": conv, "readout": (readout_probe, readout_response), "noise_variance": noise_std ** 2, "conv_batch_size": batch_size, "readout_batch_size": readout_batch, } @torch.no_grad() def normalized_response_mirror_update(net, observations, eta=0.1): """Update Q/R from local probe/response pairs without reading W. For white parent noise ``z`` and a bias-blocked child response ``u=Wz``, the normalized local correlation is an unbiased estimator of ``W``. An exponential delta rule makes Q track that estimate. This is a stabilized weight-estimation baseline inherited from weight-mirror work, not an SDIL contribution. """ if not isinstance(net, CIFARHierarchicalFAResNet): raise TypeError("mirror update requires a hierarchical net") if not 0.0 < eta <= 1.0: raise ValueError("mirror eta must be in (0, 1]") variance = float(observations.get("noise_variance", 0.0)) conv = observations.get("conv") if variance <= 0 or conv is None or len(conv) != len(net.Q): raise ValueError("invalid mirror observations") estimates = [None] update_power = 0.0 estimate_power = 0.0 parameters = 0 for index in range(1, len(net.Q)): probe, response = conv[index] spec = net.layer_specs[index] if response.shape[1] != net.Q[index].shape[0]: raise ValueError("mirror child response channel mismatch") correlation = torch.nn.grad.conv2d_weight( probe, net.Q[index].shape, response, stride=spec.stride, padding=spec.padding) counts = torch.nn.grad.conv2d_weight( torch.ones_like(probe), net.Q[index].shape, torch.ones_like(response), stride=spec.stride, padding=spec.padding) estimate = correlation / (variance * counts.clamp_min(1.0)) update = estimate - net.Q[index] net.Q[index].add_(update, alpha=eta) estimates.append(estimate) update_power += float(update.square().sum()) estimate_power += float(estimate.square().sum()) parameters += estimate.numel() readout_probe, readout_response = observations["readout"] estimate_readout = -(readout_probe.t() @ readout_response) / ( variance * readout_probe.shape[0]) if estimate_readout.shape != net.R_out.shape: raise ValueError("mirror readout response shape mismatch") update_readout = estimate_readout - net.R_out net.R_out.add_(update_readout, alpha=eta) update_power += float(update_readout.square().sum()) estimate_power += float(estimate_readout.square().sum()) parameters += estimate_readout.numel() return { "mirror_update_rms": math.sqrt(update_power / parameters), "mirror_estimate_rms": math.sqrt(estimate_power / parameters), "conv_batch_size": int(observations["conv_batch_size"]), "readout_batch_size": int(observations["readout_batch_size"]), }, {"Q": estimates, "R": estimate_readout} @torch.no_grad() def normalized_residual_mirror_update(net, observations, eta=0.1): """Local normalized LMS on child-response prediction residuals. Unlike estimate-then-average mirroring, the stochastic update vanishes for every probe when Q exactly matches W. It therefore removes the stationary estimator noise that compounds through a deep feedback chain. The update still consumes only fixed probe/response observations and Q/R. """ if not isinstance(net, CIFARHierarchicalFAResNet): raise TypeError("residual mirror update requires a hierarchical net") if not 0.0 < eta <= 1.0: raise ValueError("residual mirror eta must be in (0, 1]") variance = float(observations.get("noise_variance", 0.0)) conv = observations.get("conv") if variance <= 0 or conv is None or len(conv) != len(net.Q): raise ValueError("invalid residual mirror observations") update_power = 0.0 residual_power = 0.0 response_power = 0.0 parameters = 0 response_units = 0 for index in range(1, len(net.Q)): probe, response = conv[index] spec = net.layer_specs[index] prediction = F.conv2d( probe, net.Q[index], stride=spec.stride, padding=spec.padding) residual = response - prediction correlation = torch.nn.grad.conv2d_weight( probe, net.Q[index].shape, residual, stride=spec.stride, padding=spec.padding) counts = torch.nn.grad.conv2d_weight( torch.ones_like(probe), net.Q[index].shape, torch.ones_like(response), stride=spec.stride, padding=spec.padding) update = correlation / (variance * counts.clamp_min(1.0)) net.Q[index].add_(update, alpha=eta) update_power += float(update.square().sum()) residual_power += float(residual.square().sum()) response_power += float(response.square().sum()) parameters += update.numel() response_units += response.numel() readout_probe, readout_response = observations["readout"] readout_prediction = -(readout_probe @ net.R_out) readout_residual = readout_response - readout_prediction readout_update = -(readout_probe.t() @ readout_residual) / ( variance * readout_probe.shape[0]) net.R_out.add_(readout_update, alpha=eta) update_power += float(readout_update.square().sum()) residual_power += float(readout_residual.square().sum()) response_power += float(readout_response.square().sum()) parameters += readout_update.numel() response_units += readout_response.numel() return { "mirror_update_rms": math.sqrt(update_power / parameters), "mirror_response_residual_rms": math.sqrt( residual_power / response_units), "mirror_response_residual_fraction": math.sqrt( residual_power / max(response_power, 1e-300)), "conv_batch_size": int(observations["conv_batch_size"]), "readout_batch_size": int(observations["readout_batch_size"]), } @torch.no_grad() def normalized_response_mirror_step( net, batch_size=1, noise_std=1.0, eta=0.1, generator=None): observations = hierarchical_mirror_observations( net, batch_size=batch_size, noise_std=noise_std, generator=generator) metrics, _ = normalized_response_mirror_update( net, observations, eta=eta) return metrics @torch.no_grad() def normalized_residual_mirror_step( net, batch_size=1, noise_std=1.0, eta=0.1, generator=None): observations = hierarchical_mirror_observations( net, batch_size=batch_size, noise_std=noise_std, generator=generator) return normalized_residual_mirror_update(net, observations, eta=eta) def conv_hierarchical_step(net, x, y, config): """One hierarchical-FA update using no reverse-mode graph or weight transport.""" config.validate() with torch.no_grad(): forward = net.forward( x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) teaching = net.hierarchical_teaching(output_error, forward) total_units = sum(value.numel() for value in teaching) teaching_rms = math.sqrt( sum(float(value.square().sum()) for value in teaching) / total_units) (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( teaching, output_error, forward) net.apply_ascent( directions, output_weight, output_bias, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) return { "loss": float(loss), "did_perturb": False, "calibration": None, "predictor_mse": None, "teaching_rms": teaching_rms, "raw_apical_rms": teaching_rms, "innovation_rms": teaching_rms, } def conv_kolen_pollack_step(net, x, y, config): """One modified-KP update with independently computed reciprocal changes.""" if not isinstance(net, CIFARKPResNet): raise TypeError("KP step requires CIFARKPResNet") config.validate() with torch.no_grad(): forward = net.forward( x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) teaching = net.hierarchical_teaching(output_error, forward) total_units = sum(value.numel() for value in teaching) teaching_rms = math.sqrt( sum(float(value.square().sum()) for value in teaching) / total_units) (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)) # Both parameter sets consume their own local correlation calculation. # Apply only after every direction has been formed, preserving a # simultaneous update with no weight or weight-change read across paths. net.apply_reciprocal_ascent( reciprocal_directions, reciprocal_readout, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay) net.apply_ascent( directions, output_weight, output_bias, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) return { "loss": float(loss), "did_perturb": False, "calibration": None, "predictor_mse": None, "teaching_rms": teaching_rms, "raw_apical_rms": teaching_rms, "innovation_rms": teaching_rms, } def conv_kp_mixed_traffic_step(net, x, y, config, step, rule, predictor_every, neutral_projection=False): """One reciprocal-KP update using a selected mixed-apical signal.""" if not isinstance(net, CIFARKPMixedTrafficResNet): raise TypeError("mixed-traffic KP step requires CIFARKPMixedTrafficResNet") if predictor_every < 0: raise ValueError("predictor cadence must be nonnegative") config.validate() with torch.no_grad(): forward = net.forward( x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) instruction = net.hierarchical_teaching(output_error, forward) components = net.mixed_apical_components( instruction, forward["hiddens"], rule, neutral_projection=neutral_projection) used = components["used"] total_units = sum(value.numel() for value in used) def rms(values): return math.sqrt( sum(float(value.square().sum()) for value in values) / total_units) (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( used, output_error, forward) reciprocal_directions, reciprocal_readout = ( net.reciprocal_feedback_directions(used, output_error, forward)) # Form both local correlations before either parameter path changes. net.apply_reciprocal_ascent( reciprocal_directions, reciprocal_readout, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay) net.apply_ascent( directions, output_weight, output_bias, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) did_predictor_update = ( predictor_every > 0 and step % predictor_every == 0) predictor_mse = None if did_predictor_update: predictor_mse = net.predictor_step( forward["hiddens"], config.eta_P) return { "loss": float(loss), "did_perturb": False, "calibration": None, "did_predictor_update": did_predictor_update, "predictor_mse": predictor_mse, "teaching_rms": rms(used), "instruction_rms": rms(components["instruction"]), "raw_apical_rms": rms(components["raw"]), "innovation_rms": rms(components["innovation"]), "traffic_rms": rms(components["traffic"]), "neutral_projection": components["neutral_projection"], } def conv_learned_hierarchical_step(net, x, y, config, step, generator=None): """One task update with optional causal calibration of hierarchical Q/R.""" config.validate() with torch.no_grad(): forward = net.forward( x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) teaching = net.hierarchical_teaching(output_error, forward) total_units = sum(value.numel() for value in teaching) teaching_rms = math.sqrt( sum(float(value.square().sum()) for value in teaching) / total_units) did_perturb = config.learn_A and step % config.pert_every == 0 calibration = None if did_perturb: calibration = hierarchical_parameter_subspace_calibration( net, x, y, forward, output_error, sigma=config.pert_sigma, n_directions=config.pert_directions, eta=config.eta_A, generator=generator) (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( teaching, output_error, forward) net.apply_ascent( directions, output_weight, output_bias, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) return { "loss": float(loss), "did_perturb": did_perturb, "calibration": calibration, "predictor_mse": None, "teaching_rms": teaching_rms, "raw_apical_rms": teaching_rms, "innovation_rms": teaching_rms, } def conv_hierarchical_alignment_report(net, x, y): """Audit hierarchical teaching against exact hidden gradients.""" 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) gradients = torch.autograd.grad( F.cross_entropy(forward["logits"], y), forward["hiddens"]) batch = x.shape[0] negative_gradients = [-batch * value.detach() for value in gradients] with torch.no_grad(): output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, net.n_classes).to(forward["logits"].dtype)) teaching = net.hierarchical_teaching(output_error, forward) values = [float(F.cosine_similarity( left.flatten(1), right.flatten(1), dim=1).mean()) for left, right in zip(teaching, negative_gradients)] for parameter in parameters: parameter.requires_grad_(False) report = { "normalization_state": "training_batch_stats_without_running_update", "teaching_negative_gradient_cosine": values, "raw_negative_gradient_cosine": values, "innovation_negative_gradient_cosine": values, } report.update(hierarchical_feedback_tracking_report(net)) return report def conv_kp_mixed_traffic_alignment_report( net, x, y, rule, neutral_projection=False): """Same-state audit of instruction/raw/innovation/matched directions.""" if not isinstance(net, CIFARKPMixedTrafficResNet): raise TypeError("mixed-traffic audit requires CIFARKPMixedTrafficResNet") 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) gradients = torch.autograd.grad( F.cross_entropy(forward["logits"], y), forward["hiddens"]) batch = x.shape[0] negative_gradients = [-batch * value.detach() for value in gradients] with torch.no_grad(): output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, net.n_classes).to(forward["logits"].dtype)) instruction = net.hierarchical_teaching(output_error, forward) components = net.mixed_apical_components( instruction, forward["hiddens"], rule, compute_matched=True, neutral_projection=neutral_projection) def align(values): return [float(F.cosine_similarity( left.flatten(1), right.flatten(1), dim=1).mean()) for left, right in zip(values, negative_gradients)] 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) target_norm = innovation_flat.norm(dim=1) norm_errors.append(float(( (matched_flat.norm(dim=1) - target_norm).abs() / target_norm.clamp_min(1e-30)).max())) raw_match_cosine = F.cosine_similarity( raw_flat, matched_flat, dim=1) direction_errors.append(float((raw_match_cosine - 1.0).abs().max())) instruction_power = sum(float(value.square().sum()) for value in components["instruction"]) traffic_power = sum(float(value.square().sum()) for value in components["traffic"]) report = { "normalization_state": "training_batch_stats_without_running_update", "teaching_negative_gradient_cosine": align(components["used"]), "used_negative_gradient_cosine": align(components["used"]), "instruction_negative_gradient_cosine": align( components["instruction"]), "raw_negative_gradient_cosine": align(components["raw"]), "innovation_negative_gradient_cosine": align( components["innovation"]), "matched_negative_gradient_cosine": align(components["matched"]), "max_norm_match_relative_error": max(norm_errors), "max_norm_match_direction_error": max(direction_errors), "traffic_instruction_rms_ratio": math.sqrt( traffic_power / instruction_power), "predictor_traffic_residual_rms_ratio": ( net.predictor_traffic_residual_rms_ratio(forward["hiddens"])), "neutral_projection": components["neutral_projection"], } report.update(hierarchical_feedback_tracking_report(net)) for parameter in parameters: parameter.requires_grad_(False) return report class CIFARSDILResNet(CIFARLocalResNet): """CIFAR local ResNet with per-unit apical vectorizers and predictors. ``spatial_template`` gives every feature unit a class-error vectorizer. ``channel_gated`` instead shares class coefficients across position and obtains spatially heterogeneous credit through a local ``tanh(h)`` gate. The latter respects convolutional translation sharing and sharply reduces feedback parameters. The predictor remains Harnett-faithful and diagonal: each unit fits its own affine soma--apical relation. """ def __init__(self, *args, a_scale=1.0, apical_seed=None, vectorizer_mode="spatial_template", **kwargs): model_seed = kwargs.get("seed", 0) super().__init__(*args, **kwargs) generator = torch.Generator(device="cpu").manual_seed( model_seed + 10007 if apical_seed is None else apical_seed) nuisance_generator = torch.Generator(device="cpu").manual_seed( model_seed + 20011 if apical_seed is None else apical_seed + 1) if vectorizer_mode not in ("spatial_template", "channel_gated"): raise ValueError(f"unknown convolutional vectorizer: {vectorizer_mode}") self.vectorizer_mode = vectorizer_mode self.A = [] self.A_gate = [] self.P = [] self.P_bias = [] self.Bnuis = [] for channels, height, width in self.hidden_shapes: units = channels * height * width # A global-average readout makes early per-unit gradients shrink # approximately as 1/(H*W). This scale keeps fixed-DFA controls # finite while learned A remains free to change its gain. std = a_scale / (height * width * math.sqrt(self.n_classes)) vectorizer_units = units if vectorizer_mode == "spatial_template" else channels self.A.append((torch.randn( vectorizer_units, self.n_classes, generator=generator) * std).to(device=self.device, dtype=self.dtype)) if vectorizer_mode == "channel_gated": self.A_gate.append(torch.zeros( channels, self.n_classes, device=self.device, dtype=self.dtype)) shape = (channels, height, width) self.P.append(torch.zeros(shape, device=self.device, dtype=self.dtype)) self.P_bias.append(torch.zeros(shape, device=self.device, dtype=self.dtype)) self.Bnuis.append(torch.exp( 0.25 * torch.randn(shape, generator=nuisance_generator) ).to(device=self.device, dtype=self.dtype)) @property def n_vectorizer_parameters(self): return (sum(value.numel() for value in self.A) + sum(value.numel() for value in self.A_gate)) @property def n_predictor_parameters(self): return (sum(value.numel() for value in self.P) + sum(value.numel() for value in self.P_bias)) @property def n_apical_parameters(self): return self.n_vectorizer_parameters + self.n_predictor_parameters @property def n_fixed_traffic_coefficients(self): return sum(value.numel() for value in self.Bnuis) @property def apical_macs_per_example(self): """MACs for projecting one class-error vector to all hidden units.""" if self.vectorizer_mode == "spatial_template": return sum(value.numel() for value in self.A) projection = sum(value.numel() for value in self.A + self.A_gate) gating = sum(math.prod(shape) for shape in self.hidden_shapes) return projection + gating def instruction(self, index, output_signal, hidden): shape = self.hidden_shapes[index] if self.vectorizer_mode == "spatial_template": return (output_signal @ self.A[index].t()).reshape( output_signal.shape[0], *shape) base = (output_signal @ self.A[index].t())[:, :, None, None] gate = (output_signal @ self.A_gate[index].t())[:, :, None, None] return base + torch.tanh(hidden) * gate def apical_components(self, output_signal, hiddens, nuisance_scale=0.0, use_residual=True): """Return teaching, raw apical, and innovation at every population.""" if len(hiddens) != self.n_hidden: raise ValueError("one somatic state is required per apical population") teaching = [] raw_apical = [] innovations = [] for index, hidden in enumerate(hiddens): instruction = self.instruction(index, output_signal, hidden) traffic = nuisance_scale * self.Bnuis[index] * hidden raw = instruction + traffic baseline = self.P[index] * hidden + self.P_bias[index] innovation = raw - baseline teaching.append(innovation if use_residual else raw) raw_apical.append(raw) innovations.append(innovation) return teaching, raw_apical, innovations @torch.no_grad() def predictor_step(self, hiddens, eta, nuisance_scale): """Neutral-period normalized LMS fit to soma-predictable traffic.""" squared_error = 0.0 units = 0 for index, hidden in enumerate(hiddens): target = nuisance_scale * self.Bnuis[index] * hidden residual = target - self.P[index] * hidden - self.P_bias[index] centered_h = hidden - hidden.mean(dim=0) centered_r = residual - residual.mean(dim=0) variance = centered_h.square().mean(dim=0) self.P[index].add_( (centered_r * centered_h).mean(dim=0) / (variance + 1e-6), alpha=eta) self.P_bias[index].add_(residual.mean(dim=0), alpha=eta) squared_error += float(residual.square().sum()) units += residual.numel() return squared_error / units @torch.no_grad() def calibrate_apical(self, output_signal, hiddens, predicted_teaching, targets, eta): """Local delta rule fitting innovation to causal perturbation targets.""" if not (len(hiddens) == len(predicted_teaching) == len(targets) == self.n_hidden): raise ValueError("calibration lists must cover every hidden population") batch = output_signal.shape[0] before_error = 0.0 target_power = 0.0 dot = 0.0 prediction_power = 0.0 for index, (hidden, prediction, target) in enumerate(zip( hiddens, predicted_teaching, targets)): error = target - prediction flat_error = error.flatten(1) if self.vectorizer_mode == "spatial_template": self.A[index].add_( flat_error.t() @ output_signal / batch, alpha=eta) else: spatial_error = error.mean(dim=(2, 3)) gated_error = (error * torch.tanh(hidden)).mean(dim=(2, 3)) self.A[index].add_( spatial_error.t() @ output_signal / batch, alpha=eta) self.A_gate[index].add_( gated_error.t() @ output_signal / batch, alpha=eta) before_error += float(error.square().sum()) target_power += float(target.square().sum()) prediction_power += float(prediction.square().sum()) dot += float((target * prediction).sum()) denominator = math.sqrt(target_power * prediction_power) return { "calibration_mse": before_error / sum( target.numel() for target in targets), "target_power": target_power / sum(target.numel() for target in targets), "prediction_target_cosine": dot / denominator if denominator else 0.0, } @torch.no_grad() def simultaneous_conv_node_perturbation(net, x, y, clean_forward, sigma=1e-2, n_directions=1, generator=None, return_diagnostics=False): """Forward-only antithetic targets for all convolutional populations. Independent Rademacher interventions are injected into every hidden map in the same plus/minus evaluations. Cross-layer interference is zero mean and is handled by the variance theorem in ``THEORY.md``. The antithetic trials are evaluated as separate B-sized batches: concatenating them would couple their BatchNorm statistics and change the intervention being estimated. """ if sigma <= 0: raise ValueError("perturbation sigma must be positive") if n_directions < 1: raise ValueError("n_directions must be positive") if len(clean_forward["hiddens"]) != net.n_hidden: raise ValueError("clean forward does not match network hidden populations") if generator is None: generator = torch.Generator(device=x.device).manual_seed(0) targets = [torch.zeros_like(hidden) for hidden in clean_forward["hiddens"]] diagnostic_directions = [] diagnostic_derivatives = [] for _ in range(n_directions): directions = [] for hidden in clean_forward["hiddens"]: direction = torch.empty_like(hidden).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) directions.append(direction) # Build one signed intervention at a time and retain only its scalar # losses. Keeping both complete hidden dictionaries would needlessly # double peak memory at ResNet-56. plus = F.cross_entropy(net.forward( x, perturbations=[sigma * direction for direction in directions], training=True, update_stats=False)["logits"], y, reduction="none") minus = F.cross_entropy(net.forward( x, perturbations=[-sigma * direction for direction in directions], training=True, update_stats=False)["logits"], y, reduction="none") if net.normalization == "batchnorm": # BN couples examples. The per-example loss difference is not a # valid node-perturbation target because ell_i also responds to # xi_j for j != i. The scalar batch objective is valid; multiplying # its derivative by B recovers the derivative of the summed loss, # matching the per-example signal convention of the local update. batch_directional = (plus.mean() - minus.mean()) / (2.0 * sigma) directional = batch_directional.mul(x.shape[0]).expand(x.shape[0]) else: batch_directional = None directional = (plus - minus) / (2.0 * sigma) for index, direction in enumerate(directions): expand = directional.reshape( directional.shape[0], *([1] * (direction.ndim - 1))) targets[index].add_(-expand * direction / n_directions) if return_diagnostics: diagnostic_directions.append(directions) diagnostic_derivatives.append({ "scaled_directional": directional, "batch_mean_directional": batch_directional, "coupling": ("batch_objective" if batch_directional is not None else "per_example_objective"), }) if return_diagnostics: return targets, { "directions": diagnostic_directions, "directional_derivatives": diagnostic_derivatives, } return targets @torch.no_grad() def channel_subspace_apical_calibration( net, x, y, clean_forward, output_signal, sigma=1e-2, n_directions=1, eta=1e-3, generator=None, return_diagnostics=False): """Calibrate channel-gated feedback in its representable causal subspace. The legacy estimator perturbs every spatial unit independently, estimates a full hidden target, and only then averages that target into the shared channel coefficients. With K=1, most of its variance lies outside the vectorizer's representable subspace. Here each intervention is instead ``(z_base + tanh(h) z_gate) / sqrt(2)`` with one Rademacher coefficient per example and channel. If ``D`` is the antithetic loss derivative and ``S`` is the number of spatial sites, ``-sqrt(2) D z/S`` is an unbiased estimate of the corresponding negative-gradient moment. Cross-example and cross-layer terms remain zero mean. Subtracting the predicted moments gives exactly the expected local delta-rule update that full unit targets would produce, without first estimating directions outside the feedback model's representable subspace. This is still forward-only causal calibration: it consumes the same two scalar loss queries per direction, never differentiates through the network, and updates only the local A/A_gate tensors. """ if getattr(net, "vectorizer_mode", None) != "channel_gated": raise ValueError("channel-subspace calibration requires channel_gated A") if sigma <= 0 or n_directions < 1 or eta < 0: raise ValueError("invalid channel-subspace calibration hyperparameters") if len(clean_forward["hiddens"]) != net.n_hidden: raise ValueError("clean forward does not match network hidden populations") if generator is None: generator = torch.Generator(device=x.device).manual_seed(0) batch = x.shape[0] target_base = [torch.zeros( batch, hidden.shape[1], device=hidden.device, dtype=hidden.dtype) for hidden in clean_forward["hiddens"]] target_gate = [torch.zeros_like(value) for value in target_base] diagnostic_directions = [] diagnostic_derivatives = [] inverse_sqrt_two = 1.0 / math.sqrt(2.0) for _ in range(n_directions): base_random = [] gate_random = [] directions = [] for hidden in clean_forward["hiddens"]: shape = (batch, hidden.shape[1]) base = torch.empty( shape, device=hidden.device, 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) direction = (base[:, :, None, None] + torch.tanh(hidden) * gate[:, :, None, None]) direction.mul_(inverse_sqrt_two) base_random.append(base) gate_random.append(gate) directions.append(direction) plus = F.cross_entropy(net.forward( x, perturbations=[sigma * value for value in directions], training=True, update_stats=False)["logits"], y, reduction="none") minus = F.cross_entropy(net.forward( x, perturbations=[-sigma * value for value in directions], training=True, update_stats=False)["logits"], y, reduction="none") if net.normalization == "batchnorm": batch_directional = (plus.mean() - minus.mean()) / (2.0 * sigma) directional = batch_directional.mul(batch).expand(batch) else: batch_directional = None directional = (plus - minus) / (2.0 * sigma) for index, (hidden, base, gate) in enumerate(zip( clean_forward["hiddens"], base_random, gate_random)): spatial = hidden.shape[2] * hidden.shape[3] scale = -math.sqrt(2.0) / (spatial * n_directions) target_base[index].add_(directional[:, None] * base, alpha=scale) target_gate[index].add_(directional[:, None] * gate, alpha=scale) if return_diagnostics: diagnostic_directions.append({ "hidden": directions, "base": base_random, "gate": gate_random}) diagnostic_derivatives.append({ "scaled_directional": directional, "batch_mean_directional": batch_directional, "coupling": ("batch_objective" if batch_directional is not None else "per_example_objective"), }) before_error = 0.0 target_power = 0.0 prediction_power = 0.0 dot = 0.0 update_power = 0.0 coefficients = 0 for index, (base_target, gate_target) in enumerate(zip( target_base, target_gate)): base_coefficient = output_signal @ net.A[index].t() gate_coefficient = output_signal @ net.A_gate[index].t() gate = torch.tanh(clean_forward["hiddens"][index]) gate_mean = gate.mean(dim=(2, 3)) gate_second_moment = gate.square().mean(dim=(2, 3)) # For prediction b + tanh(h) g, these are its inner products with # the two representable basis fields. The errors are therefore the # exact stochastic gradients of full-field squared prediction error. base_prediction = base_coefficient + gate_mean * gate_coefficient gate_prediction = (gate_mean * base_coefficient + gate_second_moment * gate_coefficient) base_error = base_target - base_prediction gate_error = gate_target - gate_prediction base_update = base_error.t() @ output_signal / batch gate_update = gate_error.t() @ output_signal / batch net.A[index].add_(base_update, alpha=eta) net.A_gate[index].add_(gate_update, alpha=eta) for prediction, target, error in ( (base_prediction, base_target, base_error), (gate_prediction, gate_target, gate_error)): before_error += float(error.square().sum()) target_power += float(target.square().sum()) prediction_power += float(prediction.square().sum()) dot += float((prediction * target).sum()) coefficients += target.numel() update_power += float(base_update.square().sum() + gate_update.square().sum()) denominator = math.sqrt(target_power * prediction_power) calibration = { "calibration_mse": before_error / coefficients, "target_power": target_power / coefficients, "prediction_target_cosine": dot / denominator if denominator else 0.0, "parameter_update_rms": math.sqrt( update_power / max(1, net.n_vectorizer_parameters)), } if return_diagnostics: return calibration, { "directions": diagnostic_directions, "directional_derivatives": diagnostic_derivatives, "target_base": target_base, "target_gate": target_gate, } return calibration @torch.no_grad() def vectorizer_subspace_apical_calibration( net, x, y, clean_forward, output_signal, sigma=1e-2, n_directions=1, eta=1e-3, generator=None, return_diagnostics=False): """Estimate the causal A/G delta rule directly in parameter space. Channel-subspace calibration first estimates one causal coefficient target per example/channel and then regresses those targets on ``output_signal``. This estimator instead draws Rademacher matrices with the exact shapes of A and A_gate. Their induced hidden intervention already contains the output context, so the antithetic scalar directly estimates the matrix moments ``mean(q c^T)`` and ``mean(q tanh(h) c^T)`` required by the local vectorizer delta rule. It uses the same two loss queries per direction and removes variance in coefficient directions that the shared A/G maps cannot represent. """ if getattr(net, "vectorizer_mode", None) != "channel_gated": raise ValueError("vectorizer-subspace calibration requires channel_gated A") if sigma <= 0 or n_directions < 1 or eta < 0: raise ValueError("invalid vectorizer-subspace calibration hyperparameters") if len(clean_forward["hiddens"]) != net.n_hidden: raise ValueError("clean forward does not match network hidden populations") if generator is None: generator = torch.Generator(device=x.device).manual_seed(0) batch = x.shape[0] target_base = [torch.zeros_like(value) for value in net.A] target_gate = [torch.zeros_like(value) for value in net.A_gate] diagnostic_directions = [] diagnostic_derivatives = [] inverse_sqrt_two = 1.0 / math.sqrt(2.0) for _ in range(n_directions): base_random = [] gate_random = [] directions = [] for hidden, base_map, gate_map in zip( clean_forward["hiddens"], net.A, net.A_gate): base = torch.empty_like(base_map).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) gate = torch.empty_like(gate_map).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) base_field = output_signal @ base.t() gate_field = output_signal @ gate.t() direction = (base_field[:, :, None, None] + torch.tanh(hidden) * gate_field[:, :, None, None]) direction.mul_(inverse_sqrt_two) base_random.append(base) gate_random.append(gate) directions.append(direction) plus = F.cross_entropy(net.forward( x, perturbations=[sigma * value for value in directions], training=True, update_stats=False)["logits"], y) minus = F.cross_entropy(net.forward( x, perturbations=[-sigma * value for value in directions], training=True, update_stats=False)["logits"], y) # A/G are shared across the minibatch, so their sufficient statistic is # the derivative of the summed loss even when examples are uncoupled. directional = (plus - minus) * batch / (2.0 * sigma) for index, (hidden, base, gate) in enumerate(zip( clean_forward["hiddens"], base_random, gate_random)): spatial = hidden.shape[2] * hidden.shape[3] scale = -math.sqrt(2.0) / ( batch * spatial * n_directions) target_base[index].add_(directional * base, alpha=scale) target_gate[index].add_(directional * gate, alpha=scale) if return_diagnostics: diagnostic_directions.append({ "hidden": directions, "base": base_random, "gate": gate_random}) diagnostic_derivatives.append({ "scaled_directional": directional, "coupling": "summed_batch_objective", }) before_error = 0.0 target_power = 0.0 prediction_power = 0.0 dot = 0.0 update_power = 0.0 coefficients = 0 for index, (base_target, gate_target) in enumerate(zip( target_base, target_gate)): base_coefficient = output_signal @ net.A[index].t() gate_coefficient = output_signal @ net.A_gate[index].t() gate = torch.tanh(clean_forward["hiddens"][index]) gate_mean = gate.mean(dim=(2, 3)) gate_second_moment = gate.square().mean(dim=(2, 3)) base_moment = base_coefficient + gate_mean * gate_coefficient gate_moment = (gate_mean * base_coefficient + gate_second_moment * gate_coefficient) base_prediction = base_moment.t() @ output_signal / batch gate_prediction = gate_moment.t() @ output_signal / batch base_error = base_target - base_prediction gate_error = gate_target - gate_prediction net.A[index].add_(base_error, alpha=eta) net.A_gate[index].add_(gate_error, alpha=eta) for prediction, target, error in ( (base_prediction, base_target, base_error), (gate_prediction, gate_target, gate_error)): before_error += float(error.square().sum()) target_power += float(target.square().sum()) prediction_power += float(prediction.square().sum()) dot += float((prediction * target).sum()) update_power += float(error.square().sum()) coefficients += target.numel() denominator = math.sqrt(target_power * prediction_power) calibration = { "calibration_mse": before_error / coefficients, "target_power": target_power / coefficients, "prediction_target_cosine": dot / denominator if denominator else 0.0, "parameter_update_rms": math.sqrt(update_power / coefficients), } if return_diagnostics: return calibration, { "directions": diagnostic_directions, "directional_derivatives": diagnostic_derivatives, "target_base": target_base, "target_gate": target_gate, } return calibration @dataclass class ConvSDILConfig: eta: float = 0.01 eta_output: float = None eta_A: float = 0.01 eta_P: float = 0.01 momentum: float = 0.9 weight_decay: float = 5e-4 learn_A: bool = True learn_P: bool = False use_residual: bool = True nuisance_scale: float = 0.0 pert_sigma: float = 1e-2 pert_every: int = 4 pert_directions: int = 1 apical_calibration_mode: str = "unit_targets" direct_node_perturbation: bool = False def validate(self): if self.eta <= 0 or (self.eta_output is not None and self.eta_output <= 0): raise ValueError("forward learning rates must be positive") if self.eta_A < 0 or self.eta_P < 0: raise ValueError("apical learning rates must be nonnegative") if self.pert_every < 1 or self.pert_directions < 1: raise ValueError("perturbation cadence/directions must be positive") if self.apical_calibration_mode not in ( "unit_targets", "channel_subspace", "vectorizer_subspace", "hierarchical_parameter_subspace"): raise ValueError("unknown apical calibration mode") if self.direct_node_perturbation and self.pert_every != 1: raise ValueError("direct node perturbation requires a target every step") if (self.direct_node_perturbation and self.apical_calibration_mode != "unit_targets"): raise ValueError("direct node perturbation requires unit targets") def conv_local_step(net, x, y, config, step, generator=None): """One DFA/learned-feedback/direct-NP minibatch update without autograd.""" config.validate() with torch.no_grad(): forward = net.forward( x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) teaching, raw, innovations = net.apical_components( output_error, forward["hiddens"], config.nuisance_scale, config.use_residual) total_units = sum(value.numel() for value in teaching) teaching_rms = math.sqrt( sum(float(value.square().sum()) for value in teaching) / total_units) raw_apical_rms = math.sqrt( sum(float(value.square().sum()) for value in raw) / total_units) innovation_rms = math.sqrt( sum(float(value.square().sum()) for value in innovations) / total_units) del raw, innovations did_perturb = ((config.learn_A or config.direct_node_perturbation) and step % config.pert_every == 0) targets = None calibration = None if did_perturb: if config.apical_calibration_mode == "channel_subspace": calibration = channel_subspace_apical_calibration( net, x, y, forward, output_error, sigma=config.pert_sigma, n_directions=config.pert_directions, eta=config.eta_A, generator=generator) elif config.apical_calibration_mode == "vectorizer_subspace": calibration = vectorizer_subspace_apical_calibration( net, x, y, forward, output_error, sigma=config.pert_sigma, n_directions=config.pert_directions, eta=config.eta_A, generator=generator) else: targets = simultaneous_conv_node_perturbation( net, x, y, forward, sigma=config.pert_sigma, n_directions=config.pert_directions, generator=generator) weight_teaching = targets if config.direct_node_perturbation else teaching if weight_teaching is None: raise RuntimeError("direct perturbation target is unavailable") (directions, gamma_directions, beta_directions, output_weight, output_bias) = net.local_ascent_directions( weight_teaching, output_error, forward) net.apply_ascent( directions, output_weight, output_bias, eta_hidden=config.eta, eta_output=config.eta_output, momentum=config.momentum, weight_decay=config.weight_decay, gamma_directions=gamma_directions, beta_directions=beta_directions) if (did_perturb and config.learn_A and config.apical_calibration_mode == "unit_targets"): calibration = net.calibrate_apical( output_error, forward["hiddens"], teaching, targets, config.eta_A) predictor_mse = None if config.learn_P: predictor_mse = net.predictor_step( forward["hiddens"], config.eta_P, config.nuisance_scale) return { "loss": float(loss), "did_perturb": did_perturb, "calibration": calibration, "predictor_mse": predictor_mse, "teaching_rms": teaching_rms, "raw_apical_rms": raw_apical_rms, "innovation_rms": innovation_rms, } @torch.no_grad() def conv_apical_calibration_step(net, x, y, config, generator=None): """Fit A from one causal intervention event while forward weights stay fixed.""" config.validate() if not config.learn_A: raise ValueError("apical-only calibration requires learn_A=True") forward = net.forward( x, return_cache=False, training=True, update_stats=False) logits = forward["logits"] output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) teaching, raw, innovations = net.apical_components( output_error, forward["hiddens"], config.nuisance_scale, config.use_residual) del raw, innovations if config.apical_calibration_mode == "channel_subspace": calibration = channel_subspace_apical_calibration( net, x, y, forward, output_error, sigma=config.pert_sigma, n_directions=config.pert_directions, eta=config.eta_A, generator=generator) elif config.apical_calibration_mode == "vectorizer_subspace": calibration = vectorizer_subspace_apical_calibration( net, x, y, forward, output_error, sigma=config.pert_sigma, n_directions=config.pert_directions, eta=config.eta_A, generator=generator) else: targets = simultaneous_conv_node_perturbation( net, x, y, forward, sigma=config.pert_sigma, n_directions=config.pert_directions, generator=generator) calibration = net.calibrate_apical( output_error, forward["hiddens"], teaching, targets, config.eta_A) return float(F.cross_entropy(logits, y)), calibration def conv_alignment_report(net, x, y, config): """Measure apical alignment to exact hidden gradients; never used to learn.""" 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, training=True, update_stats=False) gradients = torch.autograd.grad( F.cross_entropy(forward["logits"], y), forward["hiddens"]) batch = x.shape[0] negative_gradients = [-batch * gradient.detach() for gradient in gradients] with torch.no_grad(): output_error = (torch.softmax(forward["logits"], dim=1) - F.one_hot(y, net.n_classes).to(forward["logits"].dtype)) teaching, raw, innovations = net.apical_components( output_error, [value.detach() for value in forward["hiddens"]], config.nuisance_scale, config.use_residual) def cosine(left, right): left = left.flatten(1) right = right.flatten(1) return float(F.cosine_similarity(left, right, dim=1).mean()) report = { "normalization_state": "training_batch_stats_without_running_update", "teaching_negative_gradient_cosine": [ cosine(left, right) for left, right in zip(teaching, negative_gradients)], "raw_negative_gradient_cosine": [ cosine(left, right) for left, right in zip(raw, negative_gradients)], "innovation_negative_gradient_cosine": [ cosine(left, right) for left, right in zip(innovations, negative_gradients)], } for parameter in parameters: parameter.requires_grad_(False) return report @torch.no_grad() def evaluate_conv(net, loader): correct = 0 total = 0 total_loss = 0.0 for x, y in loader: logits = net.logits(x) total_loss += F.cross_entropy(logits, y, reduction="sum").item() correct += (logits.argmax(dim=1) == y).sum().item() total += y.numel() return correct / total, total_loss / total