experiment: implement causally whitened feedback fits

1 files changed, 97 insertions, 0 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index 7f4d144..6b97b3c 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -11,6 +11,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
                        CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet,
                        ConvSDILConfig,
+                       causal_conv_diagonal_least_squares_fit,
+                       causal_readout_least_squares_fit,
                        channel_subspace_apical_calibration,
                        conv_hierarchical_step, conv_kolen_pollack_step,
                        conv_kp_mixed_traffic_step,
@@ -841,6 +843,100 @@ def layerwise_causal_bootstrap_checks():
     }
 
 
+def stagewise_whitened_causal_checks():
+    """Audit closed/diagonal local fits and their forward independence."""
+    torch.manual_seed(125)
+    left = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=126, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    right = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=126, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    features = torch.randn(128, 10, dtype=torch.float64)
+    known_readout = torch.randn_like(left.R_out)
+    targets = features @ known_readout.t()
+    readout_observations = [{
+        "kind": "readout", "features": features, "target": targets}]
+    for weight in right.W:
+        weight.add_(torch.randn_like(weight))
+    right.W_out.add_(torch.randn_like(right.W_out))
+    causal_readout_least_squares_fit(left, readout_observations)
+    causal_readout_least_squares_fit(right, readout_observations)
+    readout_locality = float((left.R_out - right.R_out).abs().max())
+    assert readout_locality < 1e-12
+
+    edge = len(left.Q) - 1
+    left.Q[edge].zero_()
+    right.Q[edge].zero_()
+    spec = left.layer_specs[edge]
+    known = torch.randn_like(left.Q[edge]) * 0.2
+    observations = []
+    for _ in range(8):
+        context = torch.randn(
+            16, known.shape[0], spec.hidden_shape[1], spec.hidden_shape[2],
+            dtype=torch.float64)
+        target = F.conv_transpose2d(
+            context, known, stride=spec.stride, padding=spec.padding,
+            output_padding=spec.stride - 1)
+        observations.append({
+            "kind": "convolution", "edge_index": edge,
+            "target": target, "prediction": torch.zeros_like(target),
+            "context": context, "stride": spec.stride,
+            "padding": spec.padding,
+        })
+    before_q = [value.clone() for value in left.Q]
+    causal_conv_diagonal_least_squares_fit(left, observations)
+    causal_conv_diagonal_least_squares_fit(right, observations)
+    conv_locality = max(float((a - b).abs().max()) for a, b in zip(
+        left.Q, right.Q))
+    assert conv_locality < 1e-12
+    changed = [not torch.equal(a, b) for a, b in zip(before_q, left.Q)]
+    assert changed == [index == edge for index in range(len(left.Q))]
+    heldout = torch.randn(
+        32, known.shape[0], spec.hidden_shape[1], spec.hidden_shape[2],
+        dtype=torch.float64)
+    heldout_target = F.conv_transpose2d(
+        heldout, known, stride=spec.stride, padding=spec.padding,
+        output_padding=spec.stride - 1)
+    heldout_before = float(heldout_target.square().mean())
+    heldout_after = float((heldout_target - F.conv_transpose2d(
+        heldout, left.Q[edge], stride=spec.stride, padding=spec.padding,
+        output_padding=spec.stride - 1)).square().mean())
+    heldout_ratio = heldout_after / heldout_before
+    assert heldout_ratio < 0.20
+
+    stage = CIFARHierarchicalFAResNet(
+        depth=8, base_width=2, seed=127, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0)
+    x = torch.randn(2, 3, 32, 32, dtype=torch.float64)
+    y = torch.tensor([2, 7])
+    generator = torch.Generator().manual_seed(128)
+
+    def observe(selected):
+        clean = stage.forward(
+            x, return_cache=True, training=False, update_stats=False)
+        signal = (torch.softmax(clean["logits"], dim=1)
+                  - F.one_hot(y, 10).to(torch.float64))
+        return layerwise_causal_feedback_observation(
+            stage, x, y, clean, signal, edge_index=selected,
+            sigma=1e-3, generator=generator)
+
+    causal_readout_least_squares_fit(stage, [observe(None)])
+    stage_metrics = []
+    for index in reversed(range(1, len(stage.Q))):
+        stage_metrics.append(causal_conv_diagonal_least_squares_fit(
+            stage, [observe(index)]))
+    assert all(torch.isfinite(value).all()
+               for value in stage.Q[1:] + [stage.R_out])
+    return {
+        "causal_whitened_readout_forward_independence_error": readout_locality,
+        "causal_whitened_conv_forward_independence_error": conv_locality,
+        "causal_whitened_heldout_mse_ratio": heldout_ratio,
+        "causal_whitened_tiny_max_update_rms": max(
+            value["parameter_update_rms"] for value in stage_metrics),
+    }
+
+
 def normalized_response_mirror_checks():
     """Audit local response estimation and absence of W access in the update."""
     net = CIFARHierarchicalFAResNet(
@@ -1391,6 +1487,7 @@ def main():
     report.update(hierarchical_feedback_checks())
     report.update(hierarchical_parameter_calibration_checks())
     report.update(layerwise_causal_bootstrap_checks())
+    report.update(stagewise_whitened_causal_checks())
     report.update(normalized_response_mirror_checks())
     report.update(kolen_pollack_checks())
     report.update(kp_mixed_traffic_checks())