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
|
from dataclasses import dataclass, asdict, replace
@dataclass
class ZBPConfig:
"""Configuration of the zeroth-order VJP estimator used inside every ZBPBlock.
mode:
'bp' exact reverse-mode autograd through the block (baseline)
'zero' control: no activation error is propagated (upstream layers get no credit)
'noise' control: random error with the norm an n-probe oracle estimate would have
'fa' feedback alignment: fixed random linear map of the incoming error (skip path exact)
'dfa' direct feedback alignment: fixed random projection of the output error (no chain)
'exact' same as 'bp' but routed through the ZBP backward (debug/consistency check)
'oracle' Stage A: exact g = J^T v computed internally, only its random
projection (1/n) sum_i u_i (u_i^T g) is returned (projection noise only)
'forward' one-sided finite difference (h(x+eps u) - h(x)) / eps
'cd' Stage B: central difference (h(x+eps u) - h(x-eps u)) / (2 eps)
'richardson' Richardson-corrected central difference (4 D_{eps/2} - D_eps) / 3
'ml' Stage C: randomized multilevel debiased central difference (unbiased)
'ml_richardson' randomized multilevel built on the Richardson sequence (unbiased)
n_probes: number of random directions per sample unit per backward call.
probe: 'rademacher' | 'gaussian' | 'orthogonal' | 'hadamard' | 'coordinate'.
All families are normalized so that E[u u^T] = I_d (entries are O(1)).
eps: finite-difference step. With eps_mode='coord' the perturbation is x + eps*u
(each coordinate moves by ~eps, Euclidean norm eps*sqrt(d)); with
eps_mode='norm' it is x + eps*u/sqrt(d) (fixed Euclidean norm eps).
alpha: survival exponent of the multilevel truncation level, P(N >= k) = 2^{-alpha k}.
max_level: hard cap on the multilevel index (P(N > max_level) = 2^{-alpha (max_level+1)}).
level_sampling: 'per_sample' (independent N for every (probe, sample) pair — the physical
cost model) or 'per_probe' (one N per probe shared by the batch — cheaper to simulate).
batch_probes: evaluate all probes of a backward call in one batched forward (simulation speed).
"""
mode: str = "oracle"
n_probes: int = 4
probe: str = "rademacher"
eps: float = 1e-2
eps_mode: str = "coord"
alpha: float = 2.0
max_level: int = 12
level_sampling: str = "per_sample"
batch_probes: bool = True
probe_chunk: int = 0 # >0: evaluate at most this many probes per batched forward (bounds memory)
param_mode: str = "local" # 'local': J_theta^T v via in-block autograd from the incoming error
readout_noise: float = 0.0 # std of additive Gaussian noise on every measured output coordinate (per query)
eps_min: float = 0.0 # multilevel ladder is truncated at eps_k >= eps_min (residual bias O(eps_min^2))
surrogate: bool = False # learned linear control variate J^T ~ A per block (fitted online from the ZO estimates)
surrogate_ema: float = 0.9 # EMA factor of the regression statistics E[g v^T], E[v v^T]
surrogate_ridge: float = 1e-3
def replace(self, **kw):
return replace(self, **kw)
def asdict(self):
return asdict(self)
@property
def is_zo(self):
return self.mode in ("forward", "cd", "richardson", "ml", "ml_richardson")
def expected_queries_per_probe(self):
"""Expected number of block evaluations (per sample unit) per probe direction."""
if self.mode == "oracle" or self.mode in ("bp", "exact"):
return 0.0
if self.mode == "forward":
return 1.0 # plus one shared base evaluation per backward call
if self.mode == "cd":
return 2.0
if self.mode == "richardson":
return 4.0
q = 2.0 ** (-self.alpha)
# E[N+1] = sum_{k>=0} Q_k = 1/(1-q) (ignoring the cap)
en1 = 1.0 / (1.0 - q)
if self.mode == "ml":
return 2.0 * en1
if self.mode == "ml_richardson":
return 2.0 * (en1 + 1.0) # levels 0..N+1
raise ValueError(self.mode)
|
