experiment: implement causally whitened feedback fits

1 files changed, 89 insertions, 0 deletions

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."""