baseline: add normalized local response mirroring

1 files changed, 120 insertions, 0 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index ba4e980..cc373e8 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -661,6 +661,126 @@ def hierarchical_parameter_subspace_calibration(
     return calibration
 
 
+@torch.no_grad()
+def hierarchical_mirror_observations(net, batch_size=1, noise_std=1.0,
+                                     generator=None):
+    """Generate local bias-blocked probe/child-response pairs.
+
+    This is the observation phase of a normalized weight-mirror baseline.  A
+    parent population emits independent Gaussian noise and the ordinary
+    forward synapses generate the child's preactivation response.  The return
+    value contains activities only; the subsequent feedback update has no
+    access to forward parameters.
+    """
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("mirror observations require a hierarchical net")
+    if batch_size < 1 or noise_std <= 0:
+        raise ValueError("invalid mirror observation hyperparameters")
+    if generator is None:
+        generator = torch.Generator(device=net.W_out.device).manual_seed(0)
+    conv = [None]
+    for index in range(1, len(net.W)):
+        weight = net.W[index]
+        spec = net.layer_specs[index]
+        _, in_channels, _, _ = weight.shape
+        _, out_height, out_width = spec.hidden_shape
+        input_shape = (
+            batch_size, in_channels,
+            out_height * spec.stride, out_width * spec.stride)
+        probe = torch.randn(
+            input_shape, generator=generator, device=weight.device,
+            dtype=weight.dtype).mul_(noise_std)
+        response = F.conv2d(
+            probe, weight, stride=spec.stride, padding=spec.padding)
+        conv.append((probe, response))
+    # The dense readout has no spatial sample multiplicity.  Use a local probe
+    # population large enough to estimate its channel covariance accurately;
+    # its MAC cost is recorded explicitly by the runner.
+    readout_batch = max(
+        batch_size, 16 * net.W_out.shape[1], 16 * net.W_out.shape[0])
+    readout_probe = torch.randn(
+        readout_batch, net.W_out.shape[1], generator=generator,
+        device=net.W_out.device, dtype=net.W_out.dtype).mul_(noise_std)
+    readout_response = readout_probe @ net.W_out.t()
+    return {
+        "conv": conv,
+        "readout": (readout_probe, readout_response),
+        "noise_variance": noise_std ** 2,
+        "conv_batch_size": batch_size,
+        "readout_batch_size": readout_batch,
+    }
+
+
+@torch.no_grad()
+def normalized_response_mirror_update(net, observations, eta=0.1):
+    """Update Q/R from local probe/response pairs without reading W.
+
+    For white parent noise ``z`` and a bias-blocked child response ``u=Wz``,
+    the normalized local correlation is an unbiased estimator of ``W``.  An
+    exponential delta rule makes Q track that estimate.  This is a stabilized
+    weight-estimation baseline inherited from weight-mirror work, not an SDIL
+    contribution.
+    """
+    if not isinstance(net, CIFARHierarchicalFAResNet):
+        raise TypeError("mirror update requires a hierarchical net")
+    if not 0.0 < eta <= 1.0:
+        raise ValueError("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 mirror observations")
+    estimates = [None]
+    update_power = 0.0
+    estimate_power = 0.0
+    parameters = 0
+    for index in range(1, len(net.Q)):
+        probe, response = conv[index]
+        spec = net.layer_specs[index]
+        if response.shape[1] != net.Q[index].shape[0]:
+            raise ValueError("mirror child response channel mismatch")
+        correlation = torch.nn.grad.conv2d_weight(
+            probe, net.Q[index].shape, response,
+            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)
+        estimate = correlation / (variance * counts.clamp_min(1.0))
+        update = estimate - net.Q[index]
+        net.Q[index].add_(update, alpha=eta)
+        estimates.append(estimate)
+        update_power += float(update.square().sum())
+        estimate_power += float(estimate.square().sum())
+        parameters += estimate.numel()
+    readout_probe, readout_response = observations["readout"]
+    estimate_readout = -(readout_probe.t() @ readout_response) / (
+        variance * readout_probe.shape[0])
+    if estimate_readout.shape != net.R_out.shape:
+        raise ValueError("mirror readout response shape mismatch")
+    update_readout = estimate_readout - net.R_out
+    net.R_out.add_(update_readout, alpha=eta)
+    update_power += float(update_readout.square().sum())
+    estimate_power += float(estimate_readout.square().sum())
+    parameters += estimate_readout.numel()
+    return {
+        "mirror_update_rms": math.sqrt(update_power / parameters),
+        "mirror_estimate_rms": math.sqrt(estimate_power / parameters),
+        "conv_batch_size": int(observations["conv_batch_size"]),
+        "readout_batch_size": int(observations["readout_batch_size"]),
+    }, {"Q": estimates, "R": estimate_readout}
+
+
+@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(
+        net, batch_size=batch_size, noise_std=noise_std,
+        generator=generator)
+    metrics, _ = normalized_response_mirror_update(
+        net, observations, eta=eta)
+    return metrics
+
+
 def conv_hierarchical_step(net, x, y, config):
     """One hierarchical-FA update using no reverse-mode graph or weight transport."""
     config.validate()