algorithm: calibrate feedback in representable subspace

1 files changed, 167 insertions, 10 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 279acac..7935027 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -640,6 +640,142 @@ def simultaneous_conv_node_perturbation(net, x, y, clean_forward, sigma=1e-2,
     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
+
+
 @dataclass
 class ConvSDILConfig:
     eta: float = 0.01
@@ -655,6 +791,7 @@ class ConvSDILConfig:
     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):
@@ -664,8 +801,14 @@ class ConvSDILConfig:
             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"):
+            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):
@@ -692,10 +835,18 @@ def conv_local_step(net, x, y, config, step, generator=None):
         did_perturb = ((config.learn_A or config.direct_node_perturbation)
                        and step % config.pert_every == 0)
         targets = None
+        calibration = None
         if did_perturb:
-            targets = simultaneous_conv_node_perturbation(
-                net, x, y, forward, sigma=config.pert_sigma,
-                n_directions=config.pert_directions, generator=generator)
+            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)
+            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")
@@ -708,8 +859,8 @@ def conv_local_step(net, x, y, config, step, generator=None):
             momentum=config.momentum, weight_decay=config.weight_decay,
             gamma_directions=gamma_directions,
             beta_directions=beta_directions)
-        calibration = None
-        if did_perturb and config.learn_A:
+        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
@@ -742,11 +893,17 @@ def conv_apical_calibration_step(net, x, y, config, generator=None):
         output_error, forward["hiddens"], config.nuisance_scale,
         config.use_residual)
     del raw, innovations
-    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)
+    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)
+    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