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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
|
#!/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
BatchedLocalAffineDebiaser,
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")
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),
"neutral_observations_each": affine.neutral_observations,
"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,
})
if __name__ == "__main__":
main()
|
