experiment: add dynamic neutral residual projection

1 files changed, 122 insertions, 6 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index d9364a2..61fee3a 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -676,6 +676,17 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         # 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:
@@ -716,22 +727,124 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         }
 
     @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):
+                                compute_matched=False,
+                                neutral_projection=False):
         """Return used, raw, innovation, matched, and traffic fields."""
         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:
+            if rule != "innovation":
+                raise ValueError(
+                    "neutral projection is defined only for innovation")
+            projected_neutral, projection_report = (
+                self.neutral_residual_projection(hiddens))
         raw = []
         innovation = []
         matched = [] if (rule == "matched" or compute_matched) else None
-        for signal, hidden, ordinary, slope, bias in zip(
+        for index, (signal, hidden, ordinary, slope, bias) in enumerate(zip(
                 instruction, hiddens, traffic,
-                self.P_traffic, self.P_traffic_bias):
+                self.P_traffic, self.P_traffic_bias)):
             apical = signal + ordinary
-            residual = apical - (slope * hidden + bias)
+            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:
@@ -745,6 +858,7 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
             "used": choices[rule], "raw": raw, "innovation": innovation,
             "matched": matched, "traffic": traffic,
             "instruction": instruction,
+            "neutral_projection": projection_report,
         }
 
     @torch.no_grad()
@@ -1283,7 +1397,7 @@ def conv_kolen_pollack_step(net, x, y, config):
 
 
 def conv_kp_mixed_traffic_step(net, x, y, config, step, rule,
-                               predictor_every):
+                               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")
@@ -1299,7 +1413,8 @@ def conv_kp_mixed_traffic_step(net, x, y, config, step, rule,
                         - 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)
+            instruction, forward["hiddens"], rule,
+            neutral_projection=neutral_projection)
         used = components["used"]
         total_units = sum(value.numel() for value in used)
 
@@ -1338,6 +1453,7 @@ def conv_kp_mixed_traffic_step(net, x, y, config, step, rule,
             "raw_apical_rms": rms(components["raw"]),
             "innovation_rms": rms(components["innovation"]),
             "traffic_rms": rms(components["traffic"]),
+            "neutral_projection": components["neutral_projection"],
         }