baseline: add residual response mirroring

1 files changed, 73 insertions, 0 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index cc373e8..b3d9235 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -771,6 +771,70 @@ def normalized_response_mirror_update(net, observations, eta=0.1):
 
 
 @torch.no_grad()
+def normalized_residual_mirror_update(net, observations, eta=0.1):
+    """Local normalized LMS on child-response prediction residuals.
+
+    Unlike estimate-then-average mirroring, the stochastic update vanishes for
+    every probe when Q exactly matches W.  It therefore removes the stationary
+    estimator noise that compounds through a deep feedback chain.  The update
+    still consumes only fixed probe/response observations and Q/R.
+    """
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("residual mirror update requires a hierarchical net")
+    if not 0.0 < eta <= 1.0:
+        raise ValueError("residual mirror eta must be in (0, 1]")
+    variance = float(observations.get("noise_variance", 0.0))
+    conv = observations.get("conv")
+    if variance <= 0 or conv is None or len(conv) != len(net.Q):
+        raise ValueError("invalid residual mirror observations")
+    update_power = 0.0
+    residual_power = 0.0
+    response_power = 0.0
+    parameters = 0
+    response_units = 0
+    for index in range(1, len(net.Q)):
+        probe, response = conv[index]
+        spec = net.layer_specs[index]
+        prediction = F.conv2d(
+            probe, net.Q[index], stride=spec.stride, padding=spec.padding)
+        residual = response - prediction
+        correlation = torch.nn.grad.conv2d_weight(
+            probe, net.Q[index].shape, residual,
+            stride=spec.stride, padding=spec.padding)
+        counts = torch.nn.grad.conv2d_weight(
+            torch.ones_like(probe), net.Q[index].shape,
+            torch.ones_like(response), stride=spec.stride,
+            padding=spec.padding)
+        update = correlation / (variance * counts.clamp_min(1.0))
+        net.Q[index].add_(update, alpha=eta)
+        update_power += float(update.square().sum())
+        residual_power += float(residual.square().sum())
+        response_power += float(response.square().sum())
+        parameters += update.numel()
+        response_units += response.numel()
+    readout_probe, readout_response = observations["readout"]
+    readout_prediction = -(readout_probe @ net.R_out)
+    readout_residual = readout_response - readout_prediction
+    readout_update = -(readout_probe.t() @ readout_residual) / (
+        variance * readout_probe.shape[0])
+    net.R_out.add_(readout_update, alpha=eta)
+    update_power += float(readout_update.square().sum())
+    residual_power += float(readout_residual.square().sum())
+    response_power += float(readout_response.square().sum())
+    parameters += readout_update.numel()
+    response_units += readout_response.numel()
+    return {
+        "mirror_update_rms": math.sqrt(update_power / parameters),
+        "mirror_response_residual_rms": math.sqrt(
+            residual_power / response_units),
+        "mirror_response_residual_fraction": math.sqrt(
+            residual_power / max(response_power, 1e-300)),
+        "conv_batch_size": int(observations["conv_batch_size"]),
+        "readout_batch_size": int(observations["readout_batch_size"]),
+    }
+
+
+@torch.no_grad()
 def normalized_response_mirror_step(
         net, batch_size=1, noise_std=1.0, eta=0.1, generator=None):
     observations = hierarchical_mirror_observations(
@@ -781,6 +845,15 @@ def normalized_response_mirror_step(
     return metrics
 
 
+@torch.no_grad()
+def normalized_residual_mirror_step(
+        net, batch_size=1, noise_std=1.0, eta=0.1, generator=None):
+    observations = hierarchical_mirror_observations(
+        net, batch_size=batch_size, noise_std=noise_std,
+        generator=generator)
+    return normalized_residual_mirror_update(net, observations, eta=eta)
+
+
 def conv_hierarchical_step(net, x, y, config):
     """One hierarchical-FA update using no reverse-mode graph or weight transport."""
     config.validate()