experiment: add dynamic neutral residual projection

3 files changed, 163 insertions, 8 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index dbae759..cdf1ad0 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -1003,6 +1003,26 @@ def kp_mixed_traffic_checks():
     assert stable_fit["min_residual_soma_slope"] < -9e-4
     assert stable_fit["max_applied_stability_margin"] >= 1e-3
 
+    # A deliberately inaccurate slow predictor leaves an affine neutral mode.
+    # The fast controller must remove that mode using only paired neutral
+    # soma/traffic observations, without changing the predictor parameters or
+    # reading the task instruction during its coefficient fit.
+    frozen_before_projection = [value.clone() for value in
+                                net.P_traffic + net.P_traffic_bias]
+    projected = net.mixed_apical_components(
+        instruction, forward["hiddens"], "innovation",
+        neutral_projection=True)
+    projection = projected["neutral_projection"]
+    projected_instruction_error = max(float((left - right).abs().max())
+                                      for left, right in zip(
+        projected["innovation"], instruction))
+    assert projected_instruction_error < 1e-14
+    assert projection["post_projection_traffic_rms_ratio"] < 1e-14
+    assert projection["max_absolute_post_projection_soma_slope"] < 1e-14
+    assert projection["instruction_observations"] == 0
+    assert all(torch.equal(before, after) for before, after in zip(
+        frozen_before_projection, net.P_traffic + net.P_traffic_bias))
+
     for slope, bias in zip(net.P_traffic, net.P_traffic_bias):
         slope.zero_()
         bias.zero_()
@@ -1098,6 +1118,14 @@ def kp_mixed_traffic_checks():
     assert all(not value.requires_grad for value in
                net.W + net.Q + net.P_traffic + net.P_traffic_bias
                + [net.W_out, net.R_out, net.b_out])
+    projected_result = conv_kp_mixed_traffic_step(
+        net, x, y, ConvSDILConfig(
+            eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0,
+            weight_decay=0.0, learn_A=False, learn_P=True),
+        step=2, rule="innovation", predictor_every=0,
+        neutral_projection=True)
+    assert math.isfinite(projected_result["loss"])
+    assert projected_result["neutral_projection"] is not None
     assert net.mixed_elementwise_ops_per_example("matched") > (
         net.mixed_elementwise_ops_per_example("raw"))
     return {
@@ -1109,6 +1137,11 @@ def kp_mixed_traffic_checks():
             "residual_traffic_rms_ratio"],
         "kp_traffic_closed_form_residual_slope": closed_form[
             "max_absolute_residual_soma_slope"],
+        "kp_traffic_projected_instruction_error": projected_instruction_error,
+        "kp_traffic_projected_residual_ratio": projection[
+            "post_projection_traffic_rms_ratio"],
+        "kp_traffic_projected_residual_slope": projection[
+            "max_absolute_post_projection_soma_slope"],
         "kp_traffic_matched_norm_error": max(norm_errors),
         "kp_traffic_matched_direction_error": max(direction_errors),
         "kp_traffic_reciprocal_correlation_error": max(correlation_errors),
diff --git a/experiments/diagnose_kp_traffic_nonfinite.py b/experiments/diagnose_kp_traffic_nonfinite.py
index 7b7389f..d438e01 100644
--- a/experiments/diagnose_kp_traffic_nonfinite.py
+++ b/experiments/diagnose_kp_traffic_nonfinite.py
@@ -91,6 +91,7 @@ def main():
                         default="nlms20")
     parser.add_argument("--predictor_every", type=int, default=16)
     parser.add_argument("--stability_margin", type=float, default=0.0)
+    parser.add_argument("--neutral_projection", action="store_true")
     parser.add_argument("--device", default="cuda")
     parser.add_argument("--max_steps", type=int, default=352)
     parser.add_argument("--out", required=True)
@@ -176,7 +177,8 @@ def main():
             break
         result = conv_kp_mixed_traffic_step(
             net, x, y, config, step, args.rule,
-            predictor_every=args.predictor_every)
+            predictor_every=args.predictor_every,
+            neutral_projection=args.neutral_projection)
         state = network_state(net)
         bad = nonfinite_groups(state)
         for name, indices in bad.items():
@@ -210,6 +212,7 @@ def main():
             "traffic_rms": result["traffic_rms"],
             "predictor_updated": result["did_predictor_update"],
             "predictor_mse": result["predictor_mse"],
+            "neutral_projection": result["neutral_projection"],
             "parameter_state": state,
         })
         if active_bad or active_metric_names:
@@ -226,13 +229,16 @@ def main():
         if row["predictor_mse"] is not None:
             numeric.append(float(row["predictor_mse"]))
     output = {
-        "protocol": "kp_mixed_traffic_nonfinite_diagnosis_v1",
+        "protocol": ("kp_dynamic_neutral_projection_diagnosis_v1"
+                     if args.neutral_projection
+                     else "kp_mixed_traffic_nonfinite_diagnosis_v1"),
         "scope": "training_only_no_validation_or_test_evaluation",
         "provenance": provenance(),
         "rule": args.rule,
         "predictor_mode": args.predictor_mode,
         "predictor_every": args.predictor_every,
         "stability_margin": args.stability_margin,
+        "neutral_projection": args.neutral_projection,
         "max_steps": args.max_steps,
         "split": split,
         "traffic_calibration": calibration,
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"],
         }