"""No-autograd structured-bias adapter for the Rain EP implementation. The adapter monkey-patches only the estimator's two-state parameter-gradient measurement. Rain's interaction objects already compute dense, convolutional and bias energy derivatives by explicit local tensor operations, so no reverse mode is introduced here. """ from __future__ import annotations from types import MethodType from typing import Iterable import torch from sdil.two_state_debias import LocalAffineDebiaser Tensor = torch.Tensor class LocalStructuredBias: """Fixed per-element bias affine in a local first-state measurement.""" def __init__(self, ratio: float, seed: int = 1729) -> None: if ratio < 0.0: raise ValueError("bias ratio must be nonnegative") self.ratio = ratio self.seed = seed self._metadata: list[tuple[Tensor, Tensor, Tensor]] | None = None @torch.no_grad() def _initialize(self, local_states: list[Tensor]) -> None: self._metadata = [] for tensor_index, state in enumerate(local_states): scale = state.square().mean().sqrt().clamp_min(1e-6) flat_index = torch.arange( state.numel(), dtype=state.dtype, device=state.device ).reshape(state.shape) phase = flat_index + float(self.seed + 97 * tensor_index) offset = torch.where( torch.remainder(phase, 2.0) < 1.0, torch.full_like(state, -0.5), torch.full_like(state, 0.5), ) slope = 0.5 + torch.remainder(phase * 0.61803398875, 1.0) amplitude = self.ratio * scale self._metadata.append((scale, offset, amplitude * slope)) @torch.no_grad() def measure(self, local_states: Iterable[Tensor]) -> tuple[list[Tensor], list[Tensor]]: local_states = list(local_states) if any(state.requires_grad for state in local_states): raise ValueError("local bias state must be detached") if self._metadata is None: self._initialize(local_states) if len(local_states) != len(self._metadata): raise ValueError("parameter collection changed after bias initialization") bases = [] biases = [] for state, (scale, offset, scaled_slope) in zip( local_states, self._metadata ): basis = torch.tanh(state / scale) amplitude = self.ratio * scale bias = amplitude * offset + scaled_slope * basis bases.append(basis) biases.append(bias) return bases, biases class RainGradientCorrector: """Apply clean/raw/constant/SDIL/oracle/noise measurement policies.""" MODES = { "clean", "raw", "constant", "innovation", "oracle", "same_rms_noise" } def __init__( self, *, mode: str, bias_ratio: float, predictor_rate: float = 0.05, neutral_cadence: int = 1, seed: int = 1729, ) -> None: if mode not in self.MODES: raise ValueError(f"unrecognized correction mode {mode}") if neutral_cadence < 1: raise ValueError("neutral cadence must be positive") self.mode = mode self.predictor_rate = predictor_rate self.neutral_cadence = neutral_cadence self.bias = LocalStructuredBias(bias_ratio, seed) self.debiaser: LocalAffineDebiaser | None = None self.steps = 0 self.last_diagnostics: dict[str, float | int] = {} self._noise_generators: list[torch.Generator] | None = None self.seed = seed @torch.no_grad() def _initialize_debiaser(self, templates: list[Tensor]) -> None: self.debiaser = LocalAffineDebiaser( templates, feature_centers=[0.0] * len(templates), feature_scales=[1.0] * len(templates), affine=self.mode == "innovation", ) @torch.no_grad() def observe_neutral(self, local_states: Iterable[Tensor]) -> None: """Fit the local bias field from one instruction-off observation.""" if self.mode not in {"constant", "innovation"}: raise ValueError( "neutral predictor observations require constant or innovation mode") local_states = list(local_states) if any(value.requires_grad for value in local_states): raise ValueError("Rain adapter received a requires-grad tensor") bases, bias = self.bias.measure(local_states) if self.debiaser is None: self._initialize_debiaser(local_states) self.debiaser.update_neutral(bases, bias, self.predictor_rate) @torch.no_grad() def apply(self, clean: Iterable[Tensor], local_states: Iterable[Tensor]) -> list[Tensor]: clean = list(clean) local_states = list(local_states) if any(value.requires_grad for value in clean + local_states): raise ValueError("Rain adapter received a requires-grad tensor") if self.mode == "clean": return [value.clone() for value in clean] bases, bias = self.bias.measure(local_states) measured = [value + corruption for value, corruption in zip(clean, bias)] if self.mode == "raw": corrected = measured elif self.mode == "oracle": corrected = [value.clone() for value in clean] elif self.mode == "same_rms_noise": if self._noise_generators is None: self._noise_generators = [] for index, template in enumerate(clean): generator = torch.Generator(device=template.device) generator.manual_seed(self.seed + 1009 * index) self._noise_generators.append(generator) corrected = [] for value, corruption, generator in zip( clean, bias, self._noise_generators ): noise = torch.randn( value.shape, dtype=value.dtype, device=value.device, generator=generator) noise.mul_(corruption.square().mean().sqrt()) corrected.append(value + noise) else: if self.debiaser is None: self._initialize_debiaser(clean) if self.steps % self.neutral_cadence == 0: self.debiaser.update_neutral( bases, bias, self.predictor_rate) corrected = self.debiaser.residual(bases, measured) residual_bias = [ value - target for value, target in zip(corrected, clean) ] total_elements = sum(value.numel() for value in bias) bias_square = sum(float(value.square().sum()) for value in bias) residual_square = sum( float(value.square().sum()) for value in residual_bias) self.last_diagnostics = { "step": self.steps, "bias_rms": (bias_square / total_elements) ** 0.5, "residual_bias_rms": (residual_square / total_elements) ** 0.5, "neutral_observations": ( 0 if self.debiaser is None else self.debiaser.neutral_observations ), } self.steps += 1 return corrected def attach_to_rain_estimator(estimator, corrector: RainGradientCorrector): """Replace Rain's standard two-state measurement with a corrected one.""" @torch.no_grad() def corrected_standard_param_grads(self, layers_first, layers_second): for layer in self._layers: layer.state = layers_first[layer.name] grads_first = [updater.grad() for updater in self._param_updaters] for layer in self._layers: layer.state = layers_second[layer.name] grads_second = [updater.grad() for updater in self._param_updaters] denominator = self._second_nudging - self._first_nudging clean = [ (second - first) / denominator for first, second in zip(grads_first, grads_second) ] return corrector.apply(clean, grads_first) estimator._standard_param_grads = MethodType( corrected_standard_param_grads, estimator) estimator.sdil_corrector = corrector return estimator @torch.no_grad() def observe_rain_neutral(estimator, corrector: RainGradientCorrector) -> None: """Expose one label-free Rain equilibrium to the local predictor.""" local_states = [updater.grad() for updater in estimator._param_updaters] corrector.observe_neutral(local_states)