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
|
#!/usr/bin/env python3
"""Mechanics checks for the post-estimator Dillavou update bias."""
from __future__ import annotations
import argparse
from pathlib import Path
import sys
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from sdil.rain_ep_adapter import ( # noqa: E402
DillavouUpdateCorrector,
attach_dillavou_to_rain_estimator,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--author-root", type=Path, required=True)
return parser.parse_args()
def build_estimator(author_root: Path):
sys.path.insert(0, str(author_root))
from model.function.cost import SquaredError
from model.function.network import Network
from model.hopfield.minimizer import FixedPointMinimizer
from model.hopfield.network import DeepHopfieldEnergy
from training.sgd import AugmentedFunction, EquilibriumProp
energy = DeepHopfieldEnergy([(4,), (7,), (3,)], [0.5, 0.5])
energy.set_device("cpu")
network = Network(energy)
cost = SquaredError(energy.layers()[-1])
augmented = AugmentedFunction(energy, cost)
minimizer = FixedPointMinimizer(augmented, network.free_layers())
minimizer.mode = "asynchronous"
minimizer.num_iterations = 12
estimator = EquilibriumProp(
energy.params(), energy.layers(), augmented, cost, minimizer)
estimator.variant = "positive"
estimator.nudging = 0.25
return energy, network, cost, augmented, minimizer, estimator
def main() -> None:
args = parse_args()
torch.manual_seed(20260807)
# The exact paper model has a fixed B_i after the estimator. Changing the
# clean signal or parameter state must not change that field.
clean_a = [torch.randn(11, 7), torch.randn(7)]
clean_b = [torch.randn_like(value) for value in clean_a]
parameters_a = [torch.randn_like(value) for value in clean_a]
parameters_b = [value + 0.3 for value in parameters_a]
raw = DillavouUpdateCorrector(
mode="raw", bias_ratio=0.2, seed=41)
measured_a = raw.apply(clean_a, parameters_a)
measured_b = raw.apply(clean_b, parameters_b)
bias_a = [value - clean for value, clean in zip(measured_a, clean_a)]
bias_b = [value - clean for value, clean in zip(measured_b, clean_b)]
fixed_relative_error = max(
float((first - second).norm() / first.norm().clamp_min(1e-30))
for first, second in zip(bias_a, bias_b)
)
assert fixed_relative_error < 2e-6, fixed_relative_error
constant = DillavouUpdateCorrector(
mode="constant", bias_ratio=0.2, predictor_rate=1.0, seed=41)
innovation = DillavouUpdateCorrector(
mode="innovation", bias_ratio=0.2, predictor_rate=1.0, seed=41)
corrected_constant = constant.apply(clean_a, parameters_a)
corrected_innovation = innovation.apply(clean_a, parameters_a)
constant_error = max(
float((actual - target).norm() / target.norm().clamp_min(1e-30))
for actual, target in zip(corrected_constant, clean_a)
)
innovation_error = max(
float((actual - target).norm() / target.norm().clamp_min(1e-30))
for actual, target in zip(corrected_innovation, clean_a)
)
assert constant_error < 2e-7, constant_error
assert innovation_error < 2e-7, innovation_error
# 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(
args.author_root)
x = torch.randn(8, 4)
labels = torch.arange(8) % 3
network.set_input(x, reset=True)
cost.set_target(labels)
augmented.nudging = 0.0
minimizer.compute_equilibrium()
free = [layer.state.clone() for layer in minimizer._layers]
clean = [value.clone() for value in estimator.compute_gradient()]
for layer, state in zip(minimizer._layers, free):
layer.state = state.clone()
integrated = DillavouUpdateCorrector(
mode="raw", bias_ratio=0.2, seed=53)
attach_dillavou_to_rain_estimator(estimator, integrated)
measured = estimator.compute_gradient()
assert all(not value.requires_grad for value in measured)
assert integrated.last_diagnostics["bias_model"] == (
"dillavou_constant_update")
observed_ratio = integrated.last_diagnostics["bias_to_clean_update_rms"]
assert abs(observed_ratio - 0.2) < 2e-6, observed_ratio
print({
"fixed_bias_relative_error_after_state_change": fixed_relative_error,
"constant_calibration_relative_error": constant_error,
"innovation_relative_error": innovation_error,
"integrated_bias_to_clean_update_rms": observed_ratio,
"neutral_observations": constant.debiaser.neutral_observations,
"autodiff_used_for_learning": False,
})
if __name__ == "__main__":
main()
|
