experiment: freeze KP teaching-signal ablation

1 files changed, 52 insertions, 0 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index cdf1ad0..1b18048 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -1023,6 +1023,51 @@ def kp_mixed_traffic_checks():
     assert all(torch.equal(before, after) for before, after in zip(
         frozen_before_projection, net.P_traffic + net.P_traffic_bias))
 
+    # Raw and matched controls pay for the identical neutral projection but do
+    # not apply its subtractive direction. Raw must remain exactly the mixed
+    # apical vector. Matched may borrow only projected innovation's norm.
+    projected_raw = net.mixed_apical_components(
+        instruction, forward["hiddens"], "raw",
+        neutral_projection=True)
+    projected_matched = net.mixed_apical_components(
+        instruction, forward["hiddens"], "matched",
+        neutral_projection=True)
+    sham_raw_error = max(float((left - right).abs().max())
+                         for left, right in zip(
+        projected_raw["used"], projected_raw["raw"]))
+    sham_projection_report_error = max(
+        abs(projected_raw["neutral_projection"][key]
+            - projected_matched["neutral_projection"][key])
+        for key in (
+            "pre_projection_traffic_rms_ratio",
+            "post_projection_traffic_rms_ratio",
+            "max_absolute_pre_projection_soma_slope",
+            "max_absolute_post_projection_soma_slope",
+            "max_positive_post_projection_soma_slope",
+            "min_post_projection_soma_slope",
+            "max_absolute_correction_slope",
+        ))
+    sham_matched_norm_errors = []
+    sham_matched_direction_errors = []
+    for raw, innovation, matched in zip(
+            projected_matched["raw"], projected_matched["innovation"],
+            projected_matched["matched"]):
+        raw_flat = raw.flatten(1)
+        innovation_flat = innovation.flatten(1)
+        matched_flat = matched.flatten(1)
+        sham_matched_norm_errors.append(float((
+            (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs()
+            / innovation_flat.norm(dim=1).clamp_min(1e-30)).max()))
+        sham_matched_direction_errors.append(float((F.cosine_similarity(
+            raw_flat, matched_flat, dim=1) - 1.0).abs().max()))
+    assert sham_raw_error == 0.0
+    assert sham_projection_report_error == 0.0
+    assert max(sham_matched_norm_errors) < 1e-12
+    assert max(sham_matched_direction_errors) < 1e-12
+    assert projected_raw["neutral_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_()
@@ -1142,6 +1187,13 @@ def kp_mixed_traffic_checks():
             "post_projection_traffic_rms_ratio"],
         "kp_traffic_projected_residual_slope": projection[
             "max_absolute_post_projection_soma_slope"],
+        "kp_traffic_sham_raw_error": sham_raw_error,
+        "kp_traffic_sham_projection_report_error": (
+            sham_projection_report_error),
+        "kp_traffic_sham_matched_norm_error": max(
+            sham_matched_norm_errors),
+        "kp_traffic_sham_matched_direction_error": max(
+            sham_matched_direction_errors),
         "kp_traffic_matched_norm_error": max(norm_errors),
         "kp_traffic_matched_direction_error": max(direction_errors),
         "kp_traffic_reciprocal_correlation_error": max(correlation_errors),