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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
|
"""No-autograd structured-bias adapter for the Rain EP implementation.
The adapter monkey-patches only the estimator's two-state parameter-gradient
measurement. Rain's interaction objects already compute dense, convolutional
and bias energy derivatives by explicit local tensor operations, so no reverse
mode is introduced here.
"""
from __future__ import annotations
from types import MethodType
from typing import Iterable
import torch
from sdil.two_state_debias import LocalAffineDebiaser
Tensor = torch.Tensor
class LocalStructuredBias:
"""Fixed per-element bias affine in a local first-state measurement."""
def __init__(self, ratio: float, seed: int = 1729) -> None:
if ratio < 0.0:
raise ValueError("bias ratio must be nonnegative")
self.ratio = ratio
self.seed = seed
self._metadata: list[tuple[Tensor, Tensor, Tensor]] | None = None
@torch.no_grad()
def _initialize(self, local_states: list[Tensor]) -> None:
self._metadata = []
for tensor_index, state in enumerate(local_states):
scale = state.square().mean().sqrt().clamp_min(1e-6)
flat_index = torch.arange(
state.numel(), dtype=state.dtype, device=state.device
).reshape(state.shape)
phase = flat_index + float(self.seed + 97 * tensor_index)
offset = torch.where(
torch.remainder(phase, 2.0) < 1.0,
torch.full_like(state, -0.5),
torch.full_like(state, 0.5),
)
slope = 0.5 + torch.remainder(phase * 0.61803398875, 1.0)
amplitude = self.ratio * scale
self._metadata.append((scale, offset, amplitude * slope))
@torch.no_grad()
def measure(self, local_states: Iterable[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
local_states = list(local_states)
if any(state.requires_grad for state in local_states):
raise ValueError("local bias state must be detached")
if self._metadata is None:
self._initialize(local_states)
if len(local_states) != len(self._metadata):
raise ValueError("parameter collection changed after bias initialization")
bases = []
biases = []
for state, (scale, offset, scaled_slope) in zip(
local_states, self._metadata
):
basis = torch.tanh(state / scale)
amplitude = self.ratio * scale
bias = amplitude * offset + scaled_slope * basis
bases.append(basis)
biases.append(bias)
return bases, biases
class RainGradientCorrector:
"""Apply clean/raw/constant/SDIL/oracle/noise measurement policies."""
MODES = {
"clean", "raw", "constant", "innovation", "oracle", "same_rms_noise"
}
def __init__(
self,
*,
mode: str,
bias_ratio: float,
predictor_rate: float = 0.05,
neutral_cadence: int = 1,
seed: int = 1729,
) -> None:
if mode not in self.MODES:
raise ValueError(f"unrecognized correction mode {mode}")
if neutral_cadence < 1:
raise ValueError("neutral cadence must be positive")
self.mode = mode
self.predictor_rate = predictor_rate
self.neutral_cadence = neutral_cadence
self.bias = LocalStructuredBias(bias_ratio, seed)
self.debiaser: LocalAffineDebiaser | None = None
self.steps = 0
self.last_diagnostics: dict[str, float | int] = {}
self._noise_generators: list[torch.Generator] | None = None
self.seed = seed
@torch.no_grad()
def _initialize_debiaser(self, templates: list[Tensor]) -> None:
self.debiaser = LocalAffineDebiaser(
templates,
feature_centers=[0.0] * len(templates),
feature_scales=[1.0] * len(templates),
affine=self.mode == "innovation",
)
@torch.no_grad()
def apply(self, clean: Iterable[Tensor], local_states: Iterable[Tensor]) -> list[Tensor]:
clean = list(clean)
local_states = list(local_states)
if any(value.requires_grad for value in clean + local_states):
raise ValueError("Rain adapter received a requires-grad tensor")
if self.mode == "clean":
return [value.clone() for value in clean]
bases, bias = self.bias.measure(local_states)
measured = [value + corruption for value, corruption in zip(clean, bias)]
if self.mode == "raw":
corrected = measured
elif self.mode == "oracle":
corrected = [value.clone() for value in clean]
elif self.mode == "same_rms_noise":
if self._noise_generators is None:
self._noise_generators = []
for index, template in enumerate(clean):
generator = torch.Generator(device=template.device)
generator.manual_seed(self.seed + 1009 * index)
self._noise_generators.append(generator)
corrected = []
for value, corruption, generator in zip(
clean, bias, self._noise_generators
):
noise = torch.randn(
value.shape, dtype=value.dtype, device=value.device,
generator=generator)
noise.mul_(corruption.square().mean().sqrt())
corrected.append(value + noise)
else:
if self.debiaser is None:
self._initialize_debiaser(clean)
if self.steps % self.neutral_cadence == 0:
self.debiaser.update_neutral(
bases, bias, self.predictor_rate)
corrected = self.debiaser.residual(bases, measured)
residual_bias = [
value - target for value, target in zip(corrected, clean)
]
total_elements = sum(value.numel() for value in bias)
bias_square = sum(float(value.square().sum()) for value in bias)
residual_square = sum(
float(value.square().sum()) for value in residual_bias)
self.last_diagnostics = {
"step": self.steps,
"bias_rms": (bias_square / total_elements) ** 0.5,
"residual_bias_rms": (residual_square / total_elements) ** 0.5,
"neutral_observations": (
0 if self.debiaser is None else self.debiaser.neutral_observations
),
}
self.steps += 1
return corrected
def attach_to_rain_estimator(estimator, corrector: RainGradientCorrector):
"""Replace Rain's standard two-state measurement with a corrected one."""
@torch.no_grad()
def corrected_standard_param_grads(self, layers_first, layers_second):
for layer in self._layers:
layer.state = layers_first[layer.name]
grads_first = [updater.grad() for updater in self._param_updaters]
for layer in self._layers:
layer.state = layers_second[layer.name]
grads_second = [updater.grad() for updater in self._param_updaters]
denominator = self._second_nudging - self._first_nudging
clean = [
(second - first) / denominator
for first, second in zip(grads_first, grads_second)
]
return corrector.apply(clean, grads_first)
estimator._standard_param_grads = MethodType(
corrected_standard_param_grads, estimator)
estimator.sdil_corrector = corrector
return estimator
|
