From 2e2d166f5ee63526fa99bec3a82320c365f76f62 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Thu, 6 Aug 2026 16:53:27 -0500 Subject: feat: add per-cell batched local predictor --- experiments/two_state_debias_smoke.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) (limited to 'experiments') diff --git a/experiments/two_state_debias_smoke.py b/experiments/two_state_debias_smoke.py index 9b57819..fc659e6 100644 --- a/experiments/two_state_debias_smoke.py +++ b/experiments/two_state_debias_smoke.py @@ -12,6 +12,7 @@ ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from sdil.two_state_debias import ( # noqa: E402 + BatchedLocalAffineDebiaser, LocalAffineDebiaser, two_state_difference, ) @@ -105,6 +106,24 @@ def main() -> None: else: raise AssertionError("requires-grad input was not rejected") + feature = torch.linspace(-1.0, 1.0, 64).reshape(64, 1, 1) + feature = torch.cat((feature, feature.square()), dim=2) + target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * feature + batched_affine = BatchedLocalAffineDebiaser( + [feature], feature_centers=[0.0], feature_scales=[1.0], affine=True) + batched_constant = BatchedLocalAffineDebiaser( + [feature], feature_centers=[0.0], feature_scales=[1.0], affine=False) + for _ in range(20): + batched_affine.update_neutral([feature], [target], 0.2) + batched_constant.update_neutral([feature], [target], 0.2) + held_feature = torch.tensor([[[-1.3, 1.4]], [[1.3, 1.4]]]) + held_target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * held_feature + held_affine = batched_affine.residual([held_feature], [held_target])[0] + held_constant = batched_constant.residual([held_feature], [held_target])[0] + batched_affine_rmse = float(held_affine.square().mean().sqrt()) + batched_constant_rmse = float(held_constant.square().mean().sqrt()) + assert batched_affine_rmse < 0.05 * batched_constant_rmse + print({ "affine_heldout_rmse": float(affine_error), "constant_heldout_rmse": float(constant_error), @@ -112,6 +131,9 @@ def main() -> None: "common_mode_max_float_error": common_mode_error, "downstream_independence_exact": True, "requires_grad_rejected": True, + "batched_affine_heldout_rmse": batched_affine_rmse, + "batched_constant_heldout_rmse": batched_constant_rmse, + "batched_neutral_observations": batched_affine.neutral_observations, }) -- cgit v1.2.3