From 1e4fbaf8e773509798387d19b806f6daf38f8b5c Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 13:47:35 -0500 Subject: baseline: add residual response mirroring --- sdil/conv.py | 73 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) (limited to 'sdil') diff --git a/sdil/conv.py b/sdil/conv.py index cc373e8..b3d9235 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -770,6 +770,70 @@ def normalized_response_mirror_update(net, observations, eta=0.1): }, {"Q": estimates, "R": estimate_readout} +@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): @@ -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() -- cgit v1.2.3