"""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()