feat: add local online least-squares residual predictor

4 files changed, 185 insertions, 7 deletions

diff --git a/experiments/rain_ep_bias_train.py b/experiments/rain_ep_bias_train.py
index d6f1f85..70ff441 100644
--- a/experiments/rain_ep_bias_train.py
+++ b/experiments/rain_ep_bias_train.py
@@ -61,6 +61,8 @@ def parse_args() -> argparse.Namespace:
         help="released physical state-dependence report; omitted means constant B")
     parser.add_argument("--predictor-rate", type=float, default=0.1)
     parser.add_argument(
+        "--predictor-kind", choices=("nlms", "ols"), default="nlms")
+    parser.add_argument(
         "--neutral-cadence", type=int, default=1,
         help="training steps per neutral update; zero freezes after calibration")
     parser.add_argument("--calibration-batches", type=int, default=0)
@@ -247,6 +249,7 @@ def main() -> None:
             neutral_cadence=args.neutral_cadence,
             drift_ratio=args.dillavou_drift_ratio,
             empirical_profile=empirical_profile,
+            predictor_kind=args.predictor_kind,
             seed=args.seed + 1729,
         )
         attach_dillavou_to_rain_estimator(estimator, corrector)
@@ -370,6 +373,7 @@ def main() -> None:
                 or corrector.empirical_profile is None
                 else corrector.empirical_profile.as_dict()),
             "predictor_rate": args.predictor_rate,
+            "predictor_kind": args.predictor_kind,
             "neutral_cadence": args.neutral_cadence,
             "layer_calibration_steps": args.layer_calibration_steps,
             "layer_bias_normalization": args.layer_bias_normalization,
diff --git a/experiments/rain_ep_dillavou_smoke.py b/experiments/rain_ep_dillavou_smoke.py
index 39bf5fd..3d35299 100644
--- a/experiments/rain_ep_dillavou_smoke.py
+++ b/experiments/rain_ep_dillavou_smoke.py
@@ -148,6 +148,30 @@ def main() -> None:
     assert held_innovation_error < 0.05 * held_constant_error, (
         held_innovation_error, held_constant_error)
 
+    # Two distinct local neutral states identify an exactly affine field when
+    # the online sufficient-statistics predictor is selected.
+    profile_ols = DillavouUpdateCorrector(
+        mode="innovation", bias_ratio=0.2, predictor_rate=1.0,
+        calibration_steps=2, neutral_cadence=0,
+        empirical_profile=profile, predictor_kind="ols", seed=67)
+    for displacement in (-0.25, 0.25):
+        state = [
+            value + displacement * scale
+            for value, scale in zip(parameters_a, parameter_scale)
+        ]
+        profile_ols.apply(clean_a, state)
+    ols_state = [
+        value + 0.7 * scale
+        for value, scale in zip(parameters_a, parameter_scale)
+    ]
+    held_ols = profile_ols.apply(clean_b, ols_state)
+    held_ols_error = sum(
+        float((actual - target).square().sum())
+        for actual, target in zip(held_ols, clean_b)
+    )
+    assert held_ols_error < 1e-10, held_ols_error
+    assert profile_ols.debiaser.neutral_observations == 2
+
     # Integration check: the corruption is attached after Rain's hand-written
     # local EP estimator and introduces no autograd graph.
     energy, network, cost, augmented, minimizer, estimator = build_estimator(
@@ -183,6 +207,7 @@ def main() -> None:
             profile.normalized_state_variations),
         "released_profile_heldout_mse_ratio_affine_over_constant": (
             held_innovation_error / held_constant_error),
+        "released_profile_two_probe_ols_mse": held_ols_error,
         "autodiff_used_for_learning": False,
     })
 
diff --git a/sdil/rain_ep_adapter.py b/sdil/rain_ep_adapter.py
index b16e96d..7dd47eb 100644
--- a/sdil/rain_ep_adapter.py
+++ b/sdil/rain_ep_adapter.py
@@ -18,6 +18,7 @@ import torch
 from sdil.two_state_debias import (
     BatchedLocalAffineDebiaser,
     LocalAffineDebiaser,
+    LocalLeastSquaresDebiaser,
 )
 
 
@@ -306,6 +307,7 @@ class DillavouUpdateCorrector:
         neutral_cadence: int = 1,
         drift_ratio: float = 0.0,
         empirical_profile: DillavouBiasProfile | None = None,
+        predictor_kind: str = "nlms",
         seed: int = 1729,
     ) -> None:
         if mode not in self.MODES:
@@ -321,6 +323,8 @@ class DillavouUpdateCorrector:
             raise ValueError("neutral cadence must be nonnegative")
         if calibration_steps < 0:
             raise ValueError("calibration steps must be nonnegative")
+        if predictor_kind not in {"nlms", "ols"}:
+            raise ValueError("predictor kind must be nlms or ols")
         self.mode = mode
         self.bias_ratio = bias_ratio
         self.predictor_rate = predictor_rate
@@ -328,9 +332,10 @@ class DillavouUpdateCorrector:
         self.neutral_cadence = neutral_cadence
         self.drift_ratio = drift_ratio
         self.empirical_profile = empirical_profile
+        self.predictor_kind = predictor_kind
         self.seed = seed
         self.steps = 0
-        self.debiaser: LocalAffineDebiaser | None = None
+        self.debiaser: LocalAffineDebiaser | LocalLeastSquaresDebiaser | None = None
         self._offsets: list[Tensor] | None = None
         self._slopes: list[Tensor] | None = None
         self._centers: list[Tensor] | None = None
@@ -402,12 +407,19 @@ class DillavouUpdateCorrector:
             self._centers.append(parameter.clone())
             self._scales.append(parameter_scale.clone())
             feature_scales.append(torch.ones_like(parameter_scale))
-        self.debiaser = LocalAffineDebiaser(
-            clean,
-            feature_centers=[0.0] * len(clean),
-            feature_scales=feature_scales,
-            affine=self.mode == "innovation",
-        )
+        if self.mode == "innovation" and self.predictor_kind == "ols":
+            self.debiaser = LocalLeastSquaresDebiaser(
+                clean,
+                feature_centers=[0.0] * len(clean),
+                feature_scales=feature_scales,
+            )
+        else:
+            self.debiaser = LocalAffineDebiaser(
+                clean,
+                feature_centers=[0.0] * len(clean),
+                feature_scales=feature_scales,
+                affine=self.mode == "innovation",
+            )
 
     @torch.no_grad()
     def _measure(
diff --git a/sdil/two_state_debias.py b/sdil/two_state_debias.py
index b1dfdde..6273b2f 100644
--- a/sdil/two_state_debias.py
+++ b/sdil/two_state_debias.py
@@ -180,6 +180,143 @@ class LocalAffineDebiaser:
         return clone
 
 
+@dataclass
+class LocalLeastSquaresState:
+    intercept: Tensor
+    slope: Tensor
+    feature_center: Tensor
+    feature_scale: Tensor
+    mean_feature: Tensor
+    mean_measurement: Tensor
+    feature_sum_squares: Tensor
+    cross_sum: Tensor
+
+
+class LocalLeastSquaresDebiaser:
+    """Per-element online affine regression from neutral observations.
+
+    Every element keeps its own scalar sufficient statistics. No observation,
+    coefficient, or update is shared across elements, and no autograd graph is
+    created. This is the direct online analogue of fitting a local baseline
+    relation before taking its innovation.
+    """
+
+    def __init__(
+        self,
+        templates: Iterable[Tensor],
+        *,
+        feature_centers: Iterable[Tensor | float],
+        feature_scales: Iterable[Tensor | float],
+        variance_floor: float = 1e-12,
+    ) -> None:
+        templates = _detached(templates)
+        centers = list(feature_centers)
+        scales = list(feature_scales)
+        if not (len(templates) == len(centers) == len(scales)):
+            raise ValueError("template and feature metadata lengths disagree")
+        if variance_floor <= 0.0:
+            raise ValueError("variance floor must be positive")
+        self.variance_floor = variance_floor
+        self.states = []
+        with torch.no_grad():
+            for template, center, scale in zip(templates, centers, scales):
+                center_tensor = torch.as_tensor(
+                    center, dtype=template.dtype, device=template.device)
+                scale_tensor = torch.as_tensor(
+                    scale, dtype=template.dtype, device=template.device)
+                if torch.any(scale_tensor <= 0):
+                    raise ValueError("feature scales must be positive")
+                self.states.append(LocalLeastSquaresState(
+                    intercept=torch.zeros_like(template),
+                    slope=torch.zeros_like(template),
+                    feature_center=center_tensor.clone(),
+                    feature_scale=scale_tensor.clone(),
+                    mean_feature=torch.zeros_like(template),
+                    mean_measurement=torch.zeros_like(template),
+                    feature_sum_squares=torch.zeros_like(template),
+                    cross_sum=torch.zeros_like(template),
+                ))
+        self.neutral_observations = 0
+
+    @staticmethod
+    def _feature(
+        state: LocalLeastSquaresState, local_feature: Tensor
+    ) -> Tensor:
+        return (local_feature - state.feature_center) / state.feature_scale
+
+    @torch.no_grad()
+    def predict(self, local_features: Iterable[Tensor]) -> list[Tensor]:
+        features = _detached(local_features)
+        if len(features) != len(self.states):
+            raise ValueError("feature collection length changed")
+        return [
+            (
+                state.intercept
+                + state.slope * self._feature(state, feature)
+            ).clone()
+            for state, feature in zip(self.states, features)
+        ]
+
+    @torch.no_grad()
+    def update_neutral(
+        self,
+        local_features: Iterable[Tensor],
+        neutral_measurements: Iterable[Tensor],
+        learning_rate: float,
+    ) -> list[Tensor]:
+        if learning_rate != 1.0:
+            raise ValueError(
+                "online least squares requires predictor_rate=1")
+        features = _detached(local_features)
+        measurements = _detached(neutral_measurements)
+        if not (len(features) == len(measurements) == len(self.states)):
+            raise ValueError("neutral tuple lengths disagree")
+        count = self.neutral_observations + 1
+        residuals = []
+        for state, feature, measurement in zip(
+            self.states, features, measurements
+        ):
+            normalized = self._feature(state, feature)
+            residuals.append(
+                measurement
+                - state.intercept
+                - state.slope * normalized)
+            delta_feature = normalized - state.mean_feature
+            delta_measurement = measurement - state.mean_measurement
+            state.mean_feature.add_(delta_feature / count)
+            state.mean_measurement.add_(delta_measurement / count)
+            state.feature_sum_squares.add_(
+                delta_feature * (normalized - state.mean_feature))
+            state.cross_sum.add_(
+                delta_feature * (measurement - state.mean_measurement))
+            identifiable = state.feature_sum_squares > self.variance_floor
+            state.slope.copy_(torch.where(
+                identifiable,
+                state.cross_sum
+                / state.feature_sum_squares.clamp_min(self.variance_floor),
+                torch.zeros_like(state.slope),
+            ))
+            state.intercept.copy_(
+                state.mean_measurement - state.slope * state.mean_feature)
+        self.neutral_observations = count
+        return residuals
+
+    @torch.no_grad()
+    def residual(
+        self,
+        local_features: Iterable[Tensor],
+        teaching_measurements: Iterable[Tensor],
+    ) -> list[Tensor]:
+        measurements = _detached(teaching_measurements)
+        predictions = self.predict(local_features)
+        if len(measurements) != len(predictions):
+            raise ValueError("teaching tuple lengths disagree")
+        return [
+            measurement - prediction
+            for measurement, prediction in zip(measurements, predictions)
+        ]
+
+
 class BatchedLocalAffineDebiaser:
     """Per-cell LMS whose coefficients are shared across observations.