feat: add shared no-grad two-state debiaser

2 files changed, 300 insertions, 0 deletions

diff --git a/experiments/two_state_debias_smoke.py b/experiments/two_state_debias_smoke.py
new file mode 100644
index 0000000..9b57819
--- /dev/null
+++ b/experiments/two_state_debias_smoke.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""Strict-locality smoke test for the shared two-state SDIL filter."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import sys
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from sdil.two_state_debias import (  # noqa: E402
+    LocalAffineDebiaser,
+    two_state_difference,
+)
+
+
+def main() -> None:
+    torch.manual_seed(20260806)
+    templates = [torch.zeros(7), torch.zeros(3, 4)]
+    centers = [0.2, -0.1]
+    scales = [1.3, 0.8]
+    affine = LocalAffineDebiaser(
+        templates,
+        feature_centers=centers,
+        feature_scales=scales,
+        affine=True,
+    )
+    constant = LocalAffineDebiaser(
+        templates,
+        feature_centers=centers,
+        feature_scales=scales,
+        affine=False,
+    )
+    intercepts = [torch.linspace(-0.4, 0.5, 7), torch.randn(3, 4) * 0.2]
+    slopes = [torch.linspace(0.3, 1.0, 7), torch.randn(3, 4) * 0.4]
+
+    train_features = []
+    for value in torch.linspace(-1.5, 1.1, 50):
+        train_features.append([
+            torch.full_like(templates[0], value),
+            torch.full_like(templates[1], -0.6 * value + 0.2),
+        ])
+    for _ in range(20):
+        for features in train_features:
+            measurements = [
+                intercept + slope * ((feature - center) / scale)
+                for intercept, slope, feature, center, scale in zip(
+                    intercepts, slopes, features, centers, scales)
+            ]
+            affine.update_neutral(features, measurements, 0.2)
+            constant.update_neutral(features, measurements, 0.2)
+
+    held_features = [
+        torch.full_like(templates[0], 1.35),
+        torch.full_like(templates[1], -0.75),
+    ]
+    held_bias = [
+        intercept + slope * ((feature - center) / scale)
+        for intercept, slope, feature, center, scale in zip(
+            intercepts, slopes, held_features, centers, scales)
+    ]
+    affine_error = torch.mean(torch.stack([
+        (prediction - target).square().mean()
+        for prediction, target in zip(affine.predict(held_features), held_bias)
+    ])).sqrt()
+    constant_error = torch.mean(torch.stack([
+        (prediction - target).square().mean()
+        for prediction, target in zip(constant.predict(held_features), held_bias)
+    ])).sqrt()
+    assert affine_error < 1e-5
+    assert constant_error > 0.1
+
+    first = [torch.randn_like(template) for template in templates]
+    common = [torch.randn_like(template) for template in templates]
+    second = [value + 0.03 * torch.ones_like(value) for value in first]
+    difference_a = two_state_difference(first, second, 0.2)
+    difference_b = two_state_difference(
+        [value + offset for value, offset in zip(first, common)],
+        [value + offset for value, offset in zip(second, common)],
+        0.2,
+    )
+    common_mode_error = max(
+        float((a - b).abs().max()) for a, b in zip(difference_a, difference_b)
+    )
+    assert common_mode_error < 2e-6
+
+    teaching = [torch.randn_like(template) for template in templates]
+    eligibility = [torch.randn_like(template) for template in templates]
+    update_a = affine.replay_updates(
+        held_features, teaching, eligibility, 0.07)
+    downstream = torch.randn(1024, 1024)
+    downstream.normal_()
+    update_b = affine.replay_updates(
+        held_features, teaching, eligibility, 0.07)
+    assert all(torch.equal(a, b) for a, b in zip(update_a, update_b))
+    assert all(not update.requires_grad for update in update_a)
+
+    try:
+        affine.predict([held_features[0].requires_grad_(), held_features[1]])
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("requires-grad input was not rejected")
+
+    print({
+        "affine_heldout_rmse": float(affine_error),
+        "constant_heldout_rmse": float(constant_error),
+        "neutral_observations_each": affine.neutral_observations,
+        "common_mode_max_float_error": common_mode_error,
+        "downstream_independence_exact": True,
+        "requires_grad_rejected": True,
+    })
+
+
+if __name__ == "__main__":
+    main()
diff --git a/sdil/two_state_debias.py b/sdil/two_state_debias.py
new file mode 100644
index 0000000..fad980a
--- /dev/null
+++ b/sdil/two_state_debias.py
@@ -0,0 +1,181 @@
+"""Shared no-autograd debiasing primitive for two-state local learners."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+import torch
+
+
+Tensor = torch.Tensor
+
+
+def _detached(tensors: Iterable[Tensor]) -> list[Tensor]:
+    values = list(tensors)
+    if any(value.requires_grad for value in values):
+        raise ValueError("paper-facing local tensors must not require gradients")
+    return values
+
+
+@torch.no_grad()
+def two_state_difference(
+    first: Iterable[Tensor],
+    second: Iterable[Tensor],
+    denominator: float,
+) -> list[Tensor]:
+    """Form a two-state teaching measurement with common-mode cancellation."""
+    first = _detached(first)
+    second = _detached(second)
+    if len(first) != len(second):
+        raise ValueError("state collections have different lengths")
+    if denominator == 0.0:
+        raise ValueError("state-difference denominator must be nonzero")
+    return [
+        (value_second - value_first) / denominator
+        for value_first, value_second in zip(first, second)
+    ]
+
+
+@dataclass
+class LocalFilterState:
+    intercept: Tensor
+    slope: Tensor | None
+    feature_center: Tensor
+    feature_scale: Tensor
+
+
+class LocalAffineDebiaser:
+    """Per-element affine filter trained only by normalized local LMS.
+
+    The class is intentionally not a ``torch.nn.Module``.  Coefficients are
+    ordinary detached tensors, and every operation executes under
+    ``torch.no_grad``.
+    """
+
+    def __init__(
+        self,
+        templates: Iterable[Tensor],
+        *,
+        feature_centers: Iterable[Tensor | float],
+        feature_scales: Iterable[Tensor | float],
+        affine: bool = True,
+    ) -> 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")
+        self.affine = affine
+        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(LocalFilterState(
+                    intercept=torch.zeros_like(template),
+                    slope=torch.zeros_like(template) if affine else None,
+                    feature_center=center_tensor.clone(),
+                    feature_scale=scale_tensor.clone(),
+                ))
+        self.neutral_observations = 0
+
+    @staticmethod
+    def _feature(state: LocalFilterState, 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")
+        predictions = []
+        for state, feature in zip(self.states, features):
+            normalized = self._feature(state, feature)
+            prediction = state.intercept
+            if self.affine:
+                prediction = prediction + state.slope * normalized
+            predictions.append(prediction.clone())
+        return predictions
+
+    @torch.no_grad()
+    def update_neutral(
+        self,
+        local_features: Iterable[Tensor],
+        neutral_measurements: Iterable[Tensor],
+        learning_rate: float,
+    ) -> list[Tensor]:
+        features = _detached(local_features)
+        measurements = _detached(neutral_measurements)
+        if not (len(features) == len(measurements) == len(self.states)):
+            raise ValueError("neutral tuple lengths disagree")
+        residuals = []
+        for state, feature, measurement in zip(
+            self.states, features, measurements
+        ):
+            normalized = self._feature(state, feature)
+            prediction = state.intercept
+            normalization = torch.ones_like(normalized)
+            if self.affine:
+                prediction = prediction + state.slope * normalized
+                normalization = normalization + normalized.square()
+            residual = measurement - prediction
+            state.intercept.add_(learning_rate * residual / normalization)
+            if self.affine:
+                state.slope.add_(
+                    learning_rate * residual * normalized / normalization)
+            residuals.append(residual.clone())
+        self.neutral_observations += 1
+        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)
+        ]
+
+    @torch.no_grad()
+    def replay_updates(
+        self,
+        local_features: Iterable[Tensor],
+        teaching_measurements: Iterable[Tensor],
+        eligibilities: Iterable[Tensor],
+        learning_rate: float,
+    ) -> list[Tensor]:
+        residuals = self.residual(local_features, teaching_measurements)
+        eligibilities = _detached(eligibilities)
+        if len(residuals) != len(eligibilities):
+            raise ValueError("eligibility tuple length changed")
+        return [
+            learning_rate * residual * eligibility
+            for residual, eligibility in zip(residuals, eligibilities)
+        ]
+
+    @torch.no_grad()
+    def clone(self) -> "LocalAffineDebiaser":
+        clone = LocalAffineDebiaser(
+            [state.intercept for state in self.states],
+            feature_centers=[state.feature_center for state in self.states],
+            feature_scales=[state.feature_scale for state in self.states],
+            affine=self.affine,
+        )
+        for source, target in zip(self.states, clone.states):
+            target.intercept.copy_(source.intercept)
+            if self.affine:
+                target.slope.copy_(source.slope)
+        clone.neutral_observations = self.neutral_observations
+        return clone
+