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
|
"""Analog-hardware simulation layer for hardware-in-the-loop ZBP (paper part 3).
Models the fully-analog machine sketched in the plan: crossbar MVMs with programming quantization and
write noise, per-device saturating transfer functions, input DACs, a Walsh-dither + lock-in measurement
channel (gain error / DC offset / readout noise that trades against integration time), sample-and-hold
error transport with droop, and pulse-quantized weight updates. Everything is exposed to training only
through forward evaluations -- the simulation's autograd is used exclusively by the validation harness.
"""
import math
from dataclasses import dataclass, replace
import torch
import torch.nn as nn
@dataclass
class HWConfig:
w_bits: int = 7 # crossbar programming resolution (0 = ideal)
w_write_noise: float = 0.01 # relative programming noise per write
w_range: float = 2.0 # programmable weight range [-w_range, w_range]
dev_gain_std: float = 0.10 # per-unit transfer-function variation phi_i(z) = g tanh(a z + c) + d
dev_a_std: float = 0.10
dev_c_std: float = 0.05
dev_d_std: float = 0.02
dac_bits: int = 8 # input DAC (0 = ideal)
dac_range: float = 4.0
meas_gain_std: float = 0.02 # lock-in channel: per-unit gain error
meas_offset: float = 0.01 # per-unit DC offset (cancels in the +/- difference)
meas_sigma0: float = 0.01 # readout noise std at integration time 1
t_int: float = 1.0 # integration time (noise scales as sigma0 / sqrt(t_int))
sh_droop: float = 0.0 # fraction of the held error lost during a block's probe phase
sh_noise: float = 0.0 # additive noise on the held error
update_lsb: float = 0.0 # weight-update quantum (0 = continuous)
def replace(self, **kw):
return replace(self, **kw)
def _quantize(x, bits, rng):
if bits <= 0:
return x.clamp(-rng, rng)
step = 2 * rng / (2 ** bits - 1)
return (x.clamp(-rng, rng) / step).round() * step
class AnalogLinear(nn.Module):
"""Crossbar MVM: the ideal parameter W is 'programmed' into W_eff = quantize(W) + write noise.
program() is called after every optimizer step (one write per update)."""
def __init__(self, d_in, d_out, hw: HWConfig, gen=None):
super().__init__()
self.hw = hw
self.weight = nn.Parameter(torch.empty(d_out, d_in))
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
self.register_buffer("w_eff", torch.zeros_like(self.weight))
self.gen = gen
self.program()
@torch.no_grad()
def program(self):
w = _quantize(self.weight, self.hw.w_bits, self.hw.w_range)
if self.hw.w_write_noise > 0:
scale = self.weight.abs().mean().clamp_min(1e-12)
noise = torch.randn(w.shape, generator=self.gen) # CPU generator; move to the weight's device
w = w + self.hw.w_write_noise * scale * noise.to(w.device)
self.w_eff.copy_(w)
def forward(self, x):
# physics uses the programmed weights; the ideal parameter only exists inside the digital optimizer.
# (straight-through: gradients w.r.t. `weight` flow as if W_eff = weight, which is exactly the local
# in-block rule a physical outer-product update implements.)
return x @ (self.w_eff + (self.weight - self.weight.detach())).t()
class DeviceNonlinearity(nn.Module):
"""phi_i(z) = g_i tanh(a_i z + c_i) + d_i with fixed per-device variation (no analytic form assumed
by training: only this forward is ever called)."""
def __init__(self, d, hw: HWConfig, gen=None):
super().__init__()
g = 1 + hw.dev_gain_std * torch.randn(d, generator=gen)
a = 1 + hw.dev_a_std * torch.randn(d, generator=gen)
c = hw.dev_c_std * torch.randn(d, generator=gen)
dd = hw.dev_d_std * torch.randn(d, generator=gen)
for n, t in [("g", g), ("a", a), ("c", c), ("d", dd)]:
self.register_buffer(n, t)
def forward(self, z):
return self.g * torch.tanh(self.a * z + self.c) + self.d
class DAC(nn.Module):
def __init__(self, hw: HWConfig):
super().__init__()
self.hw = hw
def forward(self, x):
q = _quantize(x, self.hw.dac_bits, self.hw.dac_range)
return x + (q - x).detach() # straight-through for the validation autograd only
class AnalogBranch(nn.Module):
"""DAC -> crossbar -> device nonlinearity -> crossbar (one residual branch of the analog machine)."""
def __init__(self, d, hidden, hw: HWConfig, gen=None, out_scale=1.0):
super().__init__()
self.dac = DAC(hw)
self.a1 = AnalogLinear(d, hidden, hw, gen)
self.phi = DeviceNonlinearity(hidden, hw, gen)
self.a2 = AnalogLinear(hidden, d, hw, gen)
with torch.no_grad():
self.a2.weight.mul_(out_scale)
self.a2.program()
def forward(self, x):
return self.a2(self.phi(self.a1(self.dac(x))))
class Measure:
"""Lock-in measurement channel applied to the block output before the digital dot product with v:
y_meas = (1 + gamma) * y + offset + sigma0/sqrt(t_int) * xi (fresh xi per query)."""
def __init__(self, d, hw: HWConfig, gen=None, device="cpu"):
self.gamma = (hw.meas_gain_std * torch.randn(d, generator=gen)).to(device)
self.offset = (hw.meas_offset * torch.randn(d, generator=gen)).to(device)
self.sigma = hw.meas_sigma0 / math.sqrt(max(hw.t_int, 1e-12))
self.device = device
def __call__(self, y):
out = (1 + self.gamma) * y + self.offset
if self.sigma > 0:
out = out + self.sigma * torch.randn_like(y)
return out
class SampleHold:
"""Error-transport channel: the held error droops and picks up noise while the block is probed."""
def __init__(self, hw: HWConfig):
self.droop, self.noise = hw.sh_droop, hw.sh_noise
def __call__(self, v):
out = (1 - self.droop) * v
if self.noise > 0:
out = out + self.noise * v.std() * torch.randn_like(v)
return out
class PulseQuantizedSGD:
"""Wraps an optimizer: applied weight changes are rounded to multiples of update_lsb, and every
AnalogLinear is re-programmed (quantize + write noise) after the update."""
def __init__(self, opt, model, hw: HWConfig):
self.opt, self.model, self.hw = opt, model, hw
self._prev = None
def zero_grad(self, set_to_none=True):
self.opt.zero_grad(set_to_none=set_to_none)
@torch.no_grad()
def _snapshot(self):
return [p.detach().clone() for p in self.model.parameters()]
def step(self):
if self.hw.update_lsb > 0:
prev = self._snapshot()
self.opt.step()
with torch.no_grad():
for p, q in zip(self.model.parameters(), prev):
delta = p.detach() - q
p.copy_(q + (delta / self.hw.update_lsb).round() * self.hw.update_lsb)
else:
self.opt.step()
for m in self.model.modules():
if isinstance(m, AnalogLinear):
m.program()
|
