algorithm: estimate causal credit in vectorizer space

1 files changed, 134 insertions, 1 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 7935027..384584e 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -776,6 +776,128 @@ def channel_subspace_apical_calibration(
     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
@@ -802,7 +924,7 @@ class ConvSDILConfig:
         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"):
+                "unit_targets", "channel_subspace", "vectorizer_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")
@@ -843,6 +965,12 @@ def conv_local_step(net, x, y, config, step, generator=None):
                     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,
@@ -898,6 +1026,11 @@ def conv_apical_calibration_step(net, x, y, config, generator=None):
             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,