diff options
Diffstat (limited to 'src/zbp_scaling/zbp/config.py')
| -rw-r--r-- | src/zbp_scaling/zbp/config.py | 78 |
1 files changed, 78 insertions, 0 deletions
diff --git a/src/zbp_scaling/zbp/config.py b/src/zbp_scaling/zbp/config.py new file mode 100644 index 0000000..e6269e2 --- /dev/null +++ b/src/zbp_scaling/zbp/config.py @@ -0,0 +1,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) |