"""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 ( BatchedLocalAffineDebiaser, 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 < 0: raise ValueError("neutral cadence must be nonnegative") 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.neutral_cadence > 0 and 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) clean_square = sum(float(value.square().sum()) for value in clean) bias_square = sum(float(value.square().sum()) for value in bias) residual_square = sum( float(value.square().sum()) for value in residual_bias) clean_rms = (clean_square / total_elements) ** 0.5 bias_rms = (bias_square / total_elements) ** 0.5 residual_bias_rms = (residual_square / total_elements) ** 0.5 self.last_diagnostics = { "step": self.steps, "clean_rms": clean_rms, "bias_rms": bias_rms, "residual_bias_rms": residual_bias_rms, "bias_to_clean_rms": bias_rms / max(clean_rms, 1e-30), "residual_to_clean_rms": residual_bias_rms / max(clean_rms, 1e-30), "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) class RainLayerStateCorrector: """Correct a structured per-neuron bias before Rain's local EP update.""" MODES = RainGradientCorrector.MODES def __init__( self, *, mode: str, bias_ratio: float, predictor_rate: float = 0.2, calibration_steps: int = 1, bias_normalization: str = "clean_difference", seed: int = 1729, ) -> None: if mode not in self.MODES: raise ValueError(f"unrecognized correction mode {mode}") if bias_ratio < 0.0: raise ValueError("bias ratio must be nonnegative") if calibration_steps < 0: raise ValueError("calibration steps must be nonnegative") if bias_normalization not in {"clean_difference", "first_state"}: raise ValueError("unrecognized layer-bias normalization") self.mode = mode self.bias_ratio = bias_ratio self.predictor_rate = predictor_rate self.calibration_steps = calibration_steps self.bias_normalization = bias_normalization self.seed = seed self.steps = 0 self.last_diagnostics: dict[str, float | int] = {} self.debiaser: BatchedLocalAffineDebiaser | None = None self._metadata: list[tuple[Tensor, Tensor, Tensor]] | None = None self._noise_generators: list[torch.Generator] | None = None @torch.no_grad() def _initialize( self, first_states: list[Tensor], clean_differences: list[Tensor] ) -> None: self._metadata = [] for index, (first, clean) in enumerate(zip( first_states, clean_differences )): scale = first.square().mean(dim=0).sqrt().clamp_min(1e-6) flat_index = torch.arange( scale.numel(), dtype=first.dtype, device=first.device ).reshape(scale.shape) phase = flat_index + float(self.seed + 97 * index) offset = torch.where( torch.remainder(phase, 2.0) < 1.0, torch.full_like(scale, -0.5), torch.full_like(scale, 0.5), ) slope = 0.5 + torch.remainder(phase * 0.61803398875, 1.0) basis = torch.tanh(first / scale) source = offset.unsqueeze(0) + slope.unsqueeze(0) * basis source_rms = source.square().mean().sqrt().clamp_min(1e-30) reference_rms = ( clean.square().mean().sqrt() if self.bias_normalization == "clean_difference" else first.square().mean().sqrt() ) gain = self.bias_ratio * reference_rms / source_rms self._metadata.append( (scale, gain * offset, gain * slope) ) self.debiaser = BatchedLocalAffineDebiaser( first_states, feature_centers=[0.0] * len(first_states), feature_scales=[1.0] * len(first_states), affine=self.mode == "innovation", ) @torch.no_grad() def _measure( self, first_states: list[Tensor] ) -> tuple[list[Tensor], list[Tensor]]: bases = [] biases = [] for first, (scale, offset, slope) in zip( first_states, self._metadata ): basis = torch.tanh(first / scale) bases.append(basis) biases.append( offset.unsqueeze(0) + slope.unsqueeze(0) * basis) return bases, biases @torch.no_grad() def apply( self, layers_first: dict[str, Tensor], layers_second: dict[str, Tensor], layer_names: list[str], ) -> dict[str, Tensor]: if self.mode == "clean": return dict(layers_second) first = [layers_first[name] for name in layer_names] second = [layers_second[name] for name in layer_names] if any(value.requires_grad for value in first + second): raise ValueError("Rain layer adapter received a requires-grad tensor") clean = [after - before for before, after in zip(first, second)] if self._metadata is None: self._initialize(first, clean) bases, bias = self._measure(first) 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.steps < self.calibration_steps: 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 clean) clean_square = sum(float(value.square().sum()) for value in clean) bias_square = sum(float(value.square().sum()) for value in bias) residual_square = sum( float(value.square().sum()) for value in residual_bias) clean_rms = (clean_square / total_elements) ** 0.5 bias_rms = (bias_square / total_elements) ** 0.5 residual_rms = (residual_square / total_elements) ** 0.5 self.last_diagnostics = { "step": self.steps, "clean_state_difference_rms": clean_rms, "bias_state_difference_rms": bias_rms, "residual_state_difference_rms": residual_rms, "bias_to_clean_state_difference_rms": ( bias_rms / max(clean_rms, 1e-30)), "residual_to_clean_state_difference_rms": ( residual_rms / max(clean_rms, 1e-30)), "neutral_observations": ( 0 if self.debiaser is None else self.debiaser.neutral_observations), } used_second = dict(layers_second) for name, before, difference in zip(layer_names, first, corrected): used_second[name] = before + difference self.steps += 1 return used_second def attach_layer_to_rain_estimator( estimator, corrector: RainLayerStateCorrector ): """Patch Rain immediately before its hand-written local parameter rule.""" layer_names = [layer.name for layer in estimator._layers[1:]] @torch.no_grad() def corrected_standard_param_grads(self, layers_first, layers_second): used_second = corrector.apply( layers_first, layers_second, layer_names) 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 = used_second[layer.name] grads_second = [updater.grad() for updater in self._param_updaters] denominator = self._second_nudging - self._first_nudging return [ (second - first) / denominator for first, second in zip(grads_first, grads_second) ] estimator._standard_param_grads = MethodType( corrected_standard_param_grads, estimator) estimator.sdil_corrector = corrector return estimator