path: root/sdil/rain_ep_adapter.py

"""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