oral-a: add audited convolutional runner

1 files changed, 76 insertions, 0 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 7f805fa..996817b 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -128,6 +128,17 @@ class CIFARLocalResNet:
         return (sum(weight.numel() for weight in self.W)
                 + 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."""
@@ -327,6 +338,15 @@ class CIFARSDILResNet(CIFARLocalResNet):
                 + sum(value.numel() for value in self.P)
                 + sum(value.numel() for value in self.P_bias))
 
+    @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."""
+        return sum(value.numel() for value in self.A)
+
     def instruction(self, index, output_signal):
         shape = self.hidden_shapes[index]
         return (output_signal @ self.A[index].t()).reshape(
@@ -529,6 +549,62 @@ def conv_local_step(net, x, y, config, step, generator=None):
 
 
 @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)
+    logits = forward["logits"]
+    output_error = (torch.softmax(logits, dim=1)
+                    - F.one_hot(y, net.n_classes).to(logits.dtype))
+    teaching, _, _ = net.apical_components(
+        output_error, forward["hiddens"], config.nuisance_scale,
+        config.use_residual)
+    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, 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.W_out, net.b_out]
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    forward = net.forward(x)
+    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 = {
+            "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