1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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()
|
