feat: add local online least-squares residual predictor

2 files changed, 29 insertions, 0 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,
     })