diff options
| -rw-r--r-- | experiments/rain_ep_bias_train.py | 4 | ||||
| -rw-r--r-- | experiments/rain_ep_dillavou_smoke.py | 25 | ||||
| -rw-r--r-- | sdil/rain_ep_adapter.py | 26 | ||||
| -rw-r--r-- | sdil/two_state_debias.py | 137 |
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. |