experiment: implement causally whitened feedback fits

2 files changed, 186 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())
diff --git a/sdil/conv.py b/sdil/conv.py
index 731cbab..404bca6 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -1283,6 +1283,95 @@ def layerwise_causal_feedback_update(net, observation, eta=0.1, eps=1e-12):
 
 
 @torch.no_grad()
+def causal_readout_least_squares_fit(
+        net, observations, relative_ridge=1e-6):
+    """Fit readout feedback from stored local causal observations."""
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("causal readout fit requires a hierarchical net")
+    if not observations or relative_ridge <= 0:
+        raise ValueError("invalid causal readout fit inputs")
+    if any(value.get("kind") != "readout" for value in observations):
+        raise ValueError("readout fit received a convolutional observation")
+    features = torch.cat([value["features"] for value in observations], dim=0)
+    targets = torch.cat([value["target"] for value in observations], dim=0)
+    gram = features.t() @ features
+    ridge = relative_ridge * float(torch.trace(gram)) / gram.shape[0]
+    ridge = max(ridge, torch.finfo(gram.dtype).tiny)
+    regularized = gram + ridge * torch.eye(
+        gram.shape[0], device=gram.device, dtype=gram.dtype)
+    solution = torch.linalg.solve(regularized, features.t() @ targets)
+    fitted = features @ solution
+    before = features @ net.R_out.t()
+    net.R_out.copy_(solution.t())
+    target_power = float(targets.square().sum())
+    fitted_power = float(fitted.square().sum())
+    denominator = math.sqrt(target_power * fitted_power)
+    return {
+        "examples": int(features.shape[0]),
+        "relative_ridge": float(relative_ridge),
+        "absolute_ridge": ridge,
+        "before_mse": float((targets - before).square().mean()),
+        "after_mse": float((targets - fitted).square().mean()),
+        "fit_target_cosine": (
+            float((targets * fitted).sum()) / denominator
+            if denominator else 0.0),
+        "parameter_rms": math.sqrt(float(net.R_out.square().mean())),
+    }
+
+
+@torch.no_grad()
+def causal_conv_diagonal_least_squares_fit(
+        net, observations, relative_ridge=1e-3):
+    """Apply one diagonal-whitened local fit to a selected feedback edge."""
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("causal convolutional fit requires a hierarchical net")
+    if not observations or relative_ridge <= 0:
+        raise ValueError("invalid causal convolutional fit inputs")
+    if any(value.get("kind") != "convolution" for value in observations):
+        raise ValueError("convolutional fit received a readout observation")
+    indices = {int(value["edge_index"]) for value in observations}
+    if len(indices) != 1:
+        raise ValueError("one causal fit must contain exactly one feedback edge")
+    index = indices.pop()
+    numerator = torch.zeros_like(net.Q[index])
+    diagonal = torch.zeros_like(net.Q[index])
+    before_error = 0.0
+    target_power = 0.0
+    units = 0
+    for value in observations:
+        error = value["target"] - value["prediction"]
+        context = value["context"]
+        stride = int(value["stride"])
+        padding = int(value["padding"])
+        numerator.add_(torch.nn.grad.conv2d_weight(
+            error, net.Q[index].shape, context,
+            stride=stride, padding=padding))
+        diagonal.add_(torch.nn.grad.conv2d_weight(
+            torch.ones_like(error), net.Q[index].shape, context.square(),
+            stride=stride, padding=padding))
+        before_error += float(error.square().sum())
+        target_power += float(value["target"].square().sum())
+        units += error.numel()
+    ridge = relative_ridge * float(diagonal.mean())
+    ridge = max(ridge, torch.finfo(diagonal.dtype).tiny)
+    update = numerator / (diagonal + ridge)
+    net.Q[index].add_(update)
+    return {
+        "edge": index,
+        "observations": len(observations),
+        "examples": sum(value["target"].shape[0] for value in observations),
+        "relative_ridge": float(relative_ridge),
+        "absolute_ridge": ridge,
+        "before_field_mse": before_error / units,
+        "target_rms": math.sqrt(target_power / units),
+        "parameter_update_rms": math.sqrt(float(update.square().mean())),
+        "parameter_rms": math.sqrt(float(net.Q[index].square().mean())),
+        "minimum_diagonal": float(diagonal.min()),
+        "mean_diagonal": float(diagonal.mean()),
+    }
+
+
+@torch.no_grad()
 def layerwise_causal_bootstrap_sweep(
         net, x, y, sigma=1e-2, eta=0.1, generator=None):
     """Calibrate readout then convolutional feedback in reverse DAG order."""