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
|
#!/usr/bin/env python3
"""Strict-locality smoke test for the Rain neuron-state adapter."""
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
RainLayerStateCorrector,
attach_layer_to_rain_estimator,
)
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 network, cost, augmented, minimizer, estimator
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--author-root", type=Path, required=True)
return parser.parse_args()
def main() -> None:
args = parse_args()
torch.manual_seed(20260806)
network, cost, augmented, minimizer, estimator = build_estimator(
args.author_root)
x = torch.randn(64, 4)
labels = torch.arange(64) % 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()
oracle_corrector = RainLayerStateCorrector(
mode="oracle", bias_ratio=4.0, seed=19)
attach_layer_to_rain_estimator(estimator, oracle_corrector)
oracle = estimator.compute_gradient()
oracle_relative_error = max(
float((actual - target).norm() / target.norm().clamp_min(1e-30))
for actual, target in zip(oracle, clean)
)
assert oracle_relative_error < 2e-5, oracle_relative_error
first = {
"input": torch.randn(64, 4),
"hidden": torch.randn(64, 7),
"output": torch.randn(64, 3),
}
clean_difference = {
name: 0.01 * torch.randn_like(value)
for name, value in first.items()
}
second = {
name: value + clean_difference[name]
for name, value in first.items()
}
innovation = RainLayerStateCorrector(
mode="innovation", bias_ratio=4.0, predictor_rate=0.2,
calibration_steps=1, seed=31)
constant = RainLayerStateCorrector(
mode="constant", bias_ratio=4.0, predictor_rate=0.2,
calibration_steps=1, seed=31)
layer_names = ["hidden", "output"]
innovation.apply(first, second, layer_names)
constant.apply(first, second, layer_names)
held_first = {
name: 1.25 * value + 0.1 for name, value in first.items()
}
held_clean = {
name: 0.01 * torch.randn_like(value)
for name, value in first.items()
}
held_second = {
name: value + held_clean[name]
for name, value in held_first.items()
}
innovation_used = innovation.apply(
held_first, held_second, layer_names)
constant_used = constant.apply(held_first, held_second, layer_names)
def residual_rms(used):
errors = [
used[name] - held_second[name] for name in layer_names
]
return (
sum(float(error.square().sum()) for error in errors)
/ sum(error.numel() for error in errors)
) ** 0.5
innovation_error = residual_rms(innovation_used)
constant_error = residual_rms(constant_used)
assert innovation_error < 0.1 * constant_error, (
innovation_error, constant_error)
assert innovation.debiaser.neutral_observations == 64
assert constant.debiaser.neutral_observations == 64
assert all(not value.requires_grad for value in innovation_used.values())
print({
"oracle_parameter_gradient_relative_error": oracle_relative_error,
"innovation_heldout_state_residual_rms": innovation_error,
"constant_heldout_state_residual_rms": constant_error,
"matched_neutral_observations": 64,
"extra_equilibrium_phases": 0,
"requires_grad": False,
})
if __name__ == "__main__":
main()
|
