scaffold: model (OLMo2-ish + ZBP partition), trainer (DDP/config), data shards, bench

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe

30 files changed, 2073 insertions, 0 deletions

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..c9f5821
--- /dev/null
+++ b/README.md
@@ -0,0 +1,23 @@
+# zbp-scaling
+
+Scaling study for **ZBP (zeroth-order backpropagation)** vs BP on decoder-only language models
+(part 2 of the ZBP paper). Every nonlinear block is treated as a physical black box trained from
+forward queries only (score-space attention core + per-token SwiGLU FFN); linear maps and the score
+product are digital. See the main zobp repo for the method, theory (NSR ≈ c·d/n) and part-1/3 results.
+
+## Layout
+- `src/zbp_scaling/zbp/` — vendored ZBP package (probes, estimators, ZBPBlock autograd)
+- `src/zbp_scaling/model.py` — OLMo2-style transformer (RMSNorm, SwiGLU, untied embeddings; learned
+  positions in v1) with the ZBP physical/digital partition
+- `src/zbp_scaling/data.py` — FineWeb-Edu → GPT-2-BPE uint16 shards (WikiText-103 fallback for smoke)
+- `src/zbp_scaling/trainer.py` — DDP (torchrun, 4 or 8 GPUs), bf16 autocast with fp32 measurement
+  accumulation, gradient accumulation to a fixed global batch, cosine + warmup, resume, JSONL logs
+- `src/zbp_scaling/diagnostics.py` — per-block branch-gain rho_k profile and the NSR constant c(scale)
+- `configs/` — model sizes (60m/124m/350m/1b) x training arms (bp / zbp_n16 / zbp_n4)
+
+## Run
+```
+python scripts/prepare_data.py --dataset fineweb-edu --tokens 3e9 --out data/fineweb
+torchrun --nproc_per_node=8 scripts/train.py --model configs/model/m124.yaml --train configs/train/zbp_n16.yaml
+```
+Global batch is fixed in the train config; per-rank micro-batch and accumulation adapt to world size.
diff --git a/configs/model/l350.yaml b/configs/model/l350.yaml
new file mode 100644
index 0000000..cd36f07
--- /dev/null
+++ b/configs/model/l350.yaml
@@ -0,0 +1,7 @@
+d: 1024
+layers: 24
+heads: 16
+vocab: 50257
+seq_len: 1024
+ffn_mult: 8/3
+pos: learned
diff --git a/configs/model/m124.yaml b/configs/model/m124.yaml
new file mode 100644
index 0000000..ad3b5fa
--- /dev/null
+++ b/configs/model/m124.yaml
@@ -0,0 +1,7 @@
+d: 768
+layers: 12
+heads: 12
+vocab: 50257
+seq_len: 1024
+ffn_mult: 8/3
+pos: learned
diff --git a/configs/model/s60.yaml b/configs/model/s60.yaml
new file mode 100644
index 0000000..ebccc0a
--- /dev/null
+++ b/configs/model/s60.yaml
@@ -0,0 +1,7 @@
+d: 512
+layers: 8
+heads: 8
+vocab: 50257
+seq_len: 1024
+ffn_mult: 8/3        # SwiGLU: hidden = mult*d rounded to 64
+pos: learned
diff --git a/configs/train/bp.yaml b/configs/train/bp.yaml
new file mode 100644
index 0000000..4573a30
--- /dev/null
+++ b/configs/train/bp.yaml
@@ -0,0 +1,11 @@
+mode: bp
+tokens: 2e9
+global_batch_tokens: 524288      # 512 seqs x 1024
+lr: 6e-4
+warmup: 2000
+optimizer: adamw
+wd: 0.1
+betas: [0.9, 0.95]
+grad_clip: 1.0
+eval_every: 500
+seed: 0
diff --git a/configs/train/zbp_n16.yaml b/configs/train/zbp_n16.yaml
new file mode 100644
index 0000000..6f97a80
--- /dev/null
+++ b/configs/train/zbp_n16.yaml
@@ -0,0 +1,16 @@
+mode: cd
+n_probes: 16
+probe: rademacher
+eps: 0.1
+probe_chunk: 8
+measure_fp32: true
+tokens: 2e9
+global_batch_tokens: 524288
+lr: 6e-4
+warmup: 2000
+optimizer: adamw
+wd: 0.1
+betas: [0.9, 0.95]
+grad_clip: 1.0
+eval_every: 500
+seed: 0
diff --git a/configs/train/zbp_n4.yaml b/configs/train/zbp_n4.yaml
new file mode 100644
index 0000000..cb26322
--- /dev/null
+++ b/configs/train/zbp_n4.yaml
@@ -0,0 +1,16 @@
+mode: cd
+n_probes: 4
+probe: rademacher
+eps: 0.1
+probe_chunk: 8
+measure_fp32: true
+tokens: 2e9
+global_batch_tokens: 524288
+lr: 6e-4
+warmup: 2000
+optimizer: adamw
+wd: 0.1
+betas: [0.9, 0.95]
+grad_clip: 1.0
+eval_every: 500
+seed: 0
diff --git a/data/wikitext103 b/data/wikitext103
new file mode 120000
index 0000000..fbd5f64
--- /dev/null
+++ b/data/wikitext103
@@ -0,0 +1 @@
+/home/yurenh2/zobp/data/wikitext103
\ No newline at end of file
diff --git a/scripts/bench_probes.py b/scripts/bench_probes.py
new file mode 100644
index 0000000..fc5faa0
--- /dev/null
+++ b/scripts/bench_probes.py
@@ -0,0 +1,38 @@
+"""Throughput vs probe_chunk for ZBP steps (informs the H200 config)."""
+import os, sys, time, argparse
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+import torch, torch.nn as nn
+from zbp_scaling.model import ScalingLM
+from zbp_scaling.zbp import ZBPConfig, zbp_blocks
+from zbp_scaling.zbp.autograd import seed_probes
+
+p = argparse.ArgumentParser()
+p.add_argument("--device", default="cuda:3")
+p.add_argument("--d", type=int, default=512); p.add_argument("--layers", type=int, default=8)
+p.add_argument("--vocab", type=int, default=8192); p.add_argument("--seq", type=int, default=1024)
+p.add_argument("--bs", type=int, default=8)
+a = p.parse_args()
+dev = torch.device(a.device); seed_probes(0, dev)
+x = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev); y = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev)
+
+def bench(mode, n, chunk, iters=4):
+    cfg = ZBPConfig(mode=mode, n_probes=n, eps=0.1, probe_chunk=chunk)
+    m = ScalingLM(a.vocab, a.d, a.layers, 8, a.seq, cfg=cfg).to(dev)
+    opt = torch.optim.AdamW(m.parameters(), lr=1e-4)
+    for _ in range(2):
+        loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
+    torch.cuda.synchronize(dev); t = time.time()
+    for _ in range(iters):
+        loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
+    torch.cuda.synchronize(dev)
+    dt = (time.time() - t) / iters
+    print(f"{mode:3s} n={n:3d} chunk={chunk:3d}: {dt*1000:7.0f} ms/step  {a.bs*a.seq/dt/1000:7.1f} ktok/s  peakmem {torch.cuda.max_memory_allocated(dev)/2**30:.1f} GB")
+    torch.cuda.reset_peak_memory_stats(dev)
+    return dt
+
+t_bp = bench("bp", 0, 0)
+for n in (16, 64):
+    for chunk in (8, 16, 32, 64):
+        if chunk > 2 * n: continue
+        dt = bench("cd", n, chunk)
+        print(f"    -> multiplier vs BP: {dt/t_bp:.1f}x")
diff --git a/scripts/train.py b/scripts/train.py
new file mode 100644
index 0000000..af43918
--- /dev/null
+++ b/scripts/train.py
@@ -0,0 +1,97 @@
+"""torchrun-able BP/ZBP trainer.  Usage:
+  torchrun --nproc_per_node=8 scripts/train.py --model configs/model/m124.yaml --train configs/train/zbp_n16.yaml \
+      --data data/fineweb --out runs/m124_zbp16 [--set key=value ...]
+Single-GPU: python scripts/train.py ... [--device cuda:0]"""
+import os, sys, json, math, time, argparse
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+import yaml
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+
+from zbp_scaling.model import ScalingLM
+from zbp_scaling.data import Shards
+from zbp_scaling.zbp import ZBPConfig, zbp_blocks
+from zbp_scaling.zbp.autograd import seed_probes
+
+p = argparse.ArgumentParser()
+p.add_argument("--model", required=True)
+p.add_argument("--train", required=True)
+p.add_argument("--data", required=True)
+p.add_argument("--out", default="runs/dev")
+p.add_argument("--device", default=None)
+p.add_argument("--steps", type=int, default=None, help="override: stop after this many steps (smoke)")
+p.add_argument("--micro_bs", type=int, default=8)
+p.add_argument("--set", nargs="*", default=[], help="key=value overrides for the yaml configs")
+a = p.parse_args()
+M = yaml.safe_load(open(a.model)); Tr = yaml.safe_load(open(a.train))
+for kv in a.set:
+    k, v = kv.split("=", 1)
+    tgt = M if k in M else Tr
+    tgt[k] = yaml.safe_load(v)
+
+ddp = "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1
+rank, world = (int(os.environ.get("RANK", 0)), int(os.environ.get("WORLD_SIZE", 1)))
+if ddp:
+    dist.init_process_group("nccl")
+    dev = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
+    torch.cuda.set_device(dev)
+else:
+    dev = torch.device(a.device or ("cuda:0" if torch.cuda.is_available() else "cpu"))
+torch.manual_seed(Tr.get("seed", 0) + rank)
+seed_probes(Tr.get("seed", 0) * 1000 + rank, dev)
+
+cfg = ZBPConfig(mode=Tr["mode"], n_probes=Tr.get("n_probes", 4), probe=Tr.get("probe", "rademacher"),
+                eps=Tr.get("eps", 0.1), probe_chunk=Tr.get("probe_chunk", 8))
+model = ScalingLM(M["vocab"], M["d"], M["layers"], M["heads"], M["seq_len"], M.get("ffn_mult", "8/3"), cfg).to(dev)
+raw = model
+if ddp:
+    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[dev.index])
+n_params = sum(q.numel() for q in raw.parameters())
+
+seq = M["seq_len"]
+accum = max(1, int(Tr["global_batch_tokens"]) // (seq * a.micro_bs * world))
+total_steps = a.steps or int(float(Tr["tokens"]) // int(Tr["global_batch_tokens"]))
+opt = torch.optim.AdamW(raw.parameters(), lr=float(Tr["lr"]), weight_decay=Tr.get("wd", 0.1), betas=tuple(Tr.get("betas", [0.9, 0.95])))
+sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / Tr.get("warmup", 2000)) * 0.5 * (1 + math.cos(math.pi * min(s, total_steps) / total_steps)))
+data = Shards(a.data, seq, dev)
+gen = torch.Generator(); gen.manual_seed(Tr.get("seed", 0) * 77 + rank)
+
+os.makedirs(a.out, exist_ok=True)
+logf = open(os.path.join(a.out, f"log_rank{rank}.jsonl"), "a")
+def log(**kw):
+    if rank == 0:
+        kw["time"] = time.time()
+        logf.write(json.dumps(kw) + "\n"); logf.flush(); print(json.dumps(kw), flush=True)
+
+@torch.no_grad()
+def evaluate(nb=20):
+    model.eval(); tot = 0.0
+    for _ in range(nb):
+        x, y = data.batch("val", a.micro_bs, gen)
+        tot += nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).item()
+    model.train(); return tot / nb
+
+if rank == 0:
+    log(kind="meta", params=n_params, accum=accum, world=world, total_steps=total_steps, cfg=cfg.asdict(), model=M, train={k: str(v) for k, v in Tr.items()})
+t0 = time.time()
+for step in range(total_steps + 1):
+    if step % Tr.get("eval_every", 500) == 0 or step == total_steps:
+        log(kind="eval", step=step, val_loss=evaluate(), elapsed=time.time() - t0,
+            tok_per_s=step * int(Tr["global_batch_tokens"]) / max(time.time() - t0, 1e-9))
+        if rank == 0 and step % (10 * Tr.get("eval_every", 500)) == 0:
+            torch.save({"model": raw.state_dict(), "opt": opt.state_dict(), "step": step}, os.path.join(a.out, "ckpt.pt"))
+    if step == total_steps:
+        break
+    opt.zero_grad(set_to_none=True)
+    for micro in range(accum):
+        x, y = data.batch("train", a.micro_bs, gen)
+        loss = nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()) / accum
+        loss.backward()
+    if Tr.get("grad_clip", 0):
+        torch.nn.utils.clip_grad_norm_(raw.parameters(), Tr["grad_clip"])
+    opt.step(); sched.step()
+    if step % 50 == 0:
+        log(kind="train", step=step, loss=loss.item() * accum, lr=sched.get_last_lr()[0])
+if ddp:
+    dist.destroy_process_group()
diff --git a/src/zbp_scaling/__pycache__/data.cpython-313.pyc b/src/zbp_scaling/__pycache__/data.cpython-313.pyc
new file mode 100644
index 0000000..333cf27
--- /dev/null
+++ b/src/zbp_scaling/__pycache__/data.cpython-313.pyc
diff --git a/src/zbp_scaling/__pycache__/model.cpython-313.pyc b/src/zbp_scaling/__pycache__/model.cpython-313.pyc
new file mode 100644
index 0000000..6b9985f
--- /dev/null
+++ b/src/zbp_scaling/__pycache__/model.cpython-313.pyc
diff --git a/src/zbp_scaling/data.py b/src/zbp_scaling/data.py
new file mode 100644
index 0000000..c6946a0
--- /dev/null
+++ b/src/zbp_scaling/data.py
@@ -0,0 +1,18 @@
+"""Token shards: uint16 memmap files train.bin / val.bin in --data dir."""
+import os
+import numpy as np
+import torch
+
+
+class Shards:
+    def __init__(self, path, seq_len, device):
+        self.train = np.memmap(os.path.join(path, "train.bin"), dtype=np.uint16, mode="r")
+        self.val = np.memmap(os.path.join(path, "val.bin"), dtype=np.uint16, mode="r")
+        self.T, self.device = seq_len, device
+
+    def batch(self, split, bs, gen):
+        src = self.train if split == "train" else self.val
+        ix = torch.randint(0, len(src) - self.T - 1, (bs,), generator=gen)
+        x = torch.stack([torch.from_numpy(src[i:i + self.T].astype(np.int64)) for i in ix])
+        y = torch.stack([torch.from_numpy(src[i + 1:i + 1 + self.T].astype(np.int64)) for i in ix])
+        return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
diff --git a/src/zbp_scaling/model.py b/src/zbp_scaling/model.py
new file mode 100644
index 0000000..61ea4de
--- /dev/null
+++ b/src/zbp_scaling/model.py
@@ -0,0 +1,96 @@
+"""OLMo2-style decoder LM with the ZBP physical/digital partition (part-1 'mixscoreqk,ffn' convention):
+physical = score-space attention core (per head x row) + whole SwiGLU FFN branch (per token);
+digital = embeddings, RMSNorms outside blocks, q/k/v/proj linears, score product, readout head."""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .zbp import ZBPBlock, ZBPConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class _Core(nn.Module):
+    """softmax(QK^T/sqrt(dh)) V from concatenated qkv [B,T,3d] -> [B,T,d]; explicit (forward-AD friendly)."""
+
+    def __init__(self, d, heads):
+        super().__init__()
+        self.d, self.h = d, heads
+
+    def forward(self, qkv):
+        B, T, _ = qkv.shape
+        q, k, v = qkv.split(self.d, dim=-1)
+        dh = self.d // self.h
+        q = q.view(B, T, self.h, dh).transpose(1, 2)
+        k = k.view(B, T, self.h, dh).transpose(1, 2)
+        v = v.view(B, T, self.h, dh).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
+        mask = torch.ones(T, T, dtype=torch.bool, device=qkv.device).tril()
+        att = att.masked_fill(~mask, float("-inf")).softmax(-1)
+        return (att @ v).transpose(1, 2).reshape(B, T, self.d)
+
+
+class _SwiGLU(nn.Module):
+    """RMS -> (W1 x) * SiLU(W3 x) -> W2, one per-token physical branch (identity skip outside)."""
+
+    def __init__(self, d, hidden, scale):
+        super().__init__()
+        self.norm = RMSNorm(d)
+        self.w1 = nn.Linear(d, hidden, bias=False)
+        self.w3 = nn.Linear(d, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, d, bias=False)
+        with torch.no_grad():
+            self.w2.weight.mul_(scale)
+
+    def forward(self, x):
+        h = self.norm(x)
+        return self.w2(self.w1(h) * F.silu(self.w3(h)))
+
+
+class Block(nn.Module):
+    def __init__(self, d, heads, hidden, cfg, bp, scale, li):
+        super().__init__()
+        self.norm_a = RMSNorm(d)
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        core = ZBPBlock(_Core(d, heads), cfg, batch_dims=1, name=f"L{li}.core")
+        core.score_probe = {"d": d, "heads": heads, "dk": d, "dv": d, "window": None, "qk": True}
+        self.core = core
+        self.proj = nn.Linear(d, d, bias=False)
+        with torch.no_grad():
+            self.proj.weight.mul_(scale)
+        self.ffn = ZBPBlock(_SwiGLU(d, hidden, scale), cfg, skip="identity", batch_dims=2, name=f"L{li}.ffn")
+
+    def forward(self, x):
+        x = x + self.proj(self.core(self.qkv(self.norm_a(x))))
+        return self.ffn(x)
+
+
+class ScalingLM(nn.Module):
+    def __init__(self, vocab, d, layers, heads, seq_len, ffn_mult="8/3", cfg=None):
+        super().__init__()
+        cfg = cfg or ZBPConfig(mode="bp")
+        bp = cfg.replace(mode="bp")
+        mult = eval(str(ffn_mult)) if isinstance(ffn_mult, str) else ffn_mult
+        hidden = int(round(mult * d / 64) * 64)
+        scale = 1 / math.sqrt(2 * layers)
+        self.tok = nn.Embedding(vocab, d)
+        self.pos = nn.Embedding(seq_len, d)
+        nn.init.normal_(self.tok.weight, std=0.02); nn.init.normal_(self.pos.weight, std=0.02)
+        self.blocks = nn.Sequential(*[Block(d, heads, hidden, cfg, bp, scale, i) for i in range(layers)])
+        self.norm_f = RMSNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.seq_len = seq_len
+
+    def forward(self, idx):
+        B, T = idx.shape
+        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))[None]
+        return self.head(self.norm_f(self.blocks(x)))
diff --git a/src/zbp_scaling/zbp/__init__.py b/src/zbp_scaling/zbp/__init__.py
new file mode 100644
index 0000000..96988ca
--- /dev/null
+++ b/src/zbp_scaling/zbp/__init__.py
@@ -0,0 +1,16 @@
+"""ZBP — Zeroth-Order Backpropagation.
+
+Per-block zeroth-order estimation of the vector-Jacobian product J^T v in
+activation space, used as a drop-in replacement for reverse-mode credit
+propagation through black-box (physical) blocks.  Parameter gradients are
+built locally from the (noisy, centered) incoming activation error.
+"""
+from .config import ZBPConfig
+from .autograd import ZBPBlock, zbp_blocks, set_mode, total_queries, reset_queries, Recorder, DFA, opzo_update
+from .probes import sample_probes
+from .estimators import oracle_projection, oracle_vjp, zo_vjp, Local
+
+__all__ = [
+    "ZBPConfig", "ZBPBlock", "zbp_blocks", "set_mode", "total_queries", "reset_queries",
+    "Recorder", "DFA", "sample_probes", "oracle_projection", "oracle_vjp", "zo_vjp", "Local",
+]
diff --git a/src/zbp_scaling/zbp/__pycache__/__init__.cpython-313.pyc b/src/zbp_scaling/zbp/__pycache__/__init__.cpython-313.pyc
new file mode 100644
index 0000000..f0f7c35
--- /dev/null
+++ b/src/zbp_scaling/zbp/__pycache__/__init__.cpython-313.pyc
diff --git a/src/zbp_scaling/zbp/__pycache__/autograd.cpython-313.pyc b/src/zbp_scaling/zbp/__pycache__/autograd.cpython-313.pyc
new file mode 100644
index 0000000..9aafe6d
--- /dev/null
+++ b/src/zbp_scaling/zbp/__pycache__/autograd.cpython-313.pyc
diff --git a/src/zbp_scaling/zbp/__pycache__/config.cpython-313.pyc b/src/zbp_scaling/zbp/__pycache__/config.cpython-313.pyc
new file mode 100644
index 0000000..060c7ea
--- /dev/null
+++ b/src/zbp_scaling/zbp/__pycache__/config.cpython-313.pyc
diff --git a/src/zbp_scaling/zbp/__pycache__/estimators.cpython-313.pyc b/src/zbp_scaling/zbp/__pycache__/estimators.cpython-313.pyc
new file mode 100644
index 0000000..9039ace
--- /dev/null
+++ b/src/zbp_scaling/zbp/__pycache__/estimators.cpython-313.pyc
diff --git a/src/zbp_scaling/zbp/__pycache__/probes.cpython-313.pyc b/src/zbp_scaling/zbp/__pycache__/probes.cpython-313.pyc
new file mode 100644
index 0000000..b52bf47
--- /dev/null
+++ b/src/zbp_scaling/zbp/__pycache__/probes.cpython-313.pyc
diff --git a/src/zbp_scaling/zbp/autograd.py b/src/zbp_scaling/zbp/autograd.py
new file mode 100644
index 0000000..e7ca6ec
--- /dev/null
+++ b/src/zbp_scaling/zbp/autograd.py
@@ -0,0 +1,449 @@
+"""ZBPBlock: an nn.Module wrapper whose backward pass never differentiates through
+the wrapped block for the *activation* error.  The returned input gradient is
+built purely from (i) function queries of the block (or, in 'oracle' mode, a random
+projection of the exact VJP), plus (ii) the exact VJP of an optional *known* skip
+path (residual control variate).  Parameter gradients are constructed locally
+from the incoming activation error v (J_theta^T v via in-block autograd).
+"""
+import contextlib
+import math
+import torch
+import torch.nn as nn
+
+from .config import ZBPConfig
+from .estimators import (oracle_vjp, zo_vjp, zo_vjp_cd_parts, make_h, control_vjp, Local,
+                         causal_score_coordinate_k_vjp, causal_score_coordinate_qk_vjp)
+
+_GENERATORS = {}
+
+
+def get_generator(device, seed=None):
+    key = str(device)
+    if key not in _GENERATORS or seed is not None:
+        g = torch.Generator(device=device)
+        g.manual_seed(1234 if seed is None else seed)
+        _GENERATORS[key] = g
+    return _GENERATORS[key]
+
+
+def seed_probes(seed, device="cpu"):
+    get_generator(torch.device(device), seed)
+
+
+class DFA:
+    """Holds the output error e = dL/dlogits of the current backward pass for direct feedback alignment.
+    Training scripts call DFA.attach(logits) after the forward; blocks in 'dfa' mode read DFA.e.
+    DRTP (Frenkel et al. 2021) uses the negated one-hot target instead: set DFA.t via DFA.attach_target(y, n)."""
+    e = None
+    t = None
+
+    @staticmethod
+    def attach_target(y, n_out):
+        DFA.t = -torch.nn.functional.one_hot(y, n_out).float()
+
+    @staticmethod
+    def attach(logits):
+        if logits.requires_grad:
+            logits.register_hook(DFA._set)
+
+    @staticmethod
+    def _set(g):
+        DFA.e = g.detach()
+
+
+def _feedback_matrix(block, key, shape_in, shape_out, device, dtype):
+    """Fixed random feedback map R: shape_out -> shape_in (per feature dims), created once per block."""
+    mats = block.__dict__.setdefault("_fb", {})
+    if key not in mats:
+        g = torch.Generator(device="cpu")
+        g.manual_seed(hash((block.name, key)) % (2 ** 31))
+        d_in, d_out = math.prod(shape_in), math.prod(shape_out)
+        R = torch.randn(d_in, d_out, generator=g) / math.sqrt(d_out)
+        mats[key] = R.to(device=device, dtype=dtype)
+    return mats[key]
+
+
+class Recorder:
+    """Collects per-block backward diagnostics (exact vs estimated activation error)."""
+    _active = None
+
+    def __init__(self):
+        self.rows = []
+
+    def __enter__(self):
+        Recorder._active = self
+        return self
+
+    def __exit__(self, *a):
+        Recorder._active = None
+
+    @staticmethod
+    def active():
+        return Recorder._active
+
+    def add(self, **row):
+        self.rows.append(row)
+
+    def summary(self):
+        import collections
+        by = collections.defaultdict(list)
+        for r in self.rows:
+            by[r["name"]].append(r)
+        out = {}
+        for k, rows in by.items():
+            out[k] = {
+                "cos_mean": sum(r["cos"] for r in rows) / len(rows),
+                "relerr_mean": sum(r["relerr"] for r in rows) / len(rows),
+                "gnorm": sum(r["gnorm"] for r in rows) / len(rows),
+                "ghat_norm": sum(r["ghat_norm"] for r in rows) / len(rows),
+                "n": len(rows),
+            }
+        return out
+
+
+class ZBPFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, block, *params):
+        # Explicit no_grad: no graph through the black box is ever built in the forward pass.
+        with torch.no_grad():
+            y = block.f(x)
+            if block.skip is not None:
+                y = y + (x if block.skip_is_identity else block.skip(x))
+        ctx.save_for_backward(x)
+        ctx.block = block
+        return y
+
+    @staticmethod
+    def backward(ctx, v):
+        (x,) = ctx.saved_tensors
+        block = ctx.block
+        cfg = block.cfg
+        f = block.f
+        nbd = block.batch_dims
+        v = v.contiguous()
+        if block.v_channel is not None:                      # hardware: the held error droops / picks up noise
+            v = block.v_channel(v)
+        f_params = block._f_params()
+        s_params = block._skip_params()
+        need_gx = ctx.needs_input_grad[0]
+        rec = Recorder.active()
+        want_exact = need_gx and (cfg.mode in ("exact", "zero", "noise") or rec is not None
+                                  or block.probe_slice is not None or block.score_probe is not None)
+        if cfg.mode == "replay":
+            want_exact = False
+        if cfg.mode in ("dfa", "opzo", "drtp") and block.skip is not None and not block.skip_is_identity and need_gx:
+            pass  # DFA ignores the skip path entirely (standard DFA); skip params still get their local gradient
+
+        # --- local parameter gradients (and exact g_x if needed) via in-block autograd ---
+        g_exact = None
+        g_params = [None] * len(f_params)
+        if want_exact or (f_params and cfg.param_mode == "local"):
+            with torch.enable_grad():
+                xd = x.detach().requires_grad_(want_exact)
+                y = f(xd)
+                inputs = ([xd] if want_exact else []) + list(f_params)
+                if inputs:
+                    grads = torch.autograd.grad(y, inputs, v, allow_unused=True)
+                    if want_exact:
+                        g_exact, grads = grads[0], grads[1:]
+                    g_params = list(grads)
+
+        # --- activation error through the black-box branch ---
+        g_hat = None
+        queries = 0.0
+        if need_gx:
+            gen = get_generator(x.device)
+            if cfg.mode == "exact":
+                g_hat = g_exact
+            elif cfg.mode == "replay":
+                g_hat = None                                  # full input error injected below
+            elif cfg.mode in ("zero", "noise"):
+                g_hat, queries = control_vjp(g_exact, cfg, gen, nbd)
+            elif cfg.mode == "fa":
+                # feedback alignment: J_f^T v replaced by a fixed random linear map of the incoming error.
+                # Sequence models ([B, T, d]): one d_in x d_out map shared across tokens (standard practice);
+                # otherwise a full map over the flattened feature dims.
+                if x.dim() == 3 and v.dim() == 3 and x.shape[1] == v.shape[1]:
+                    R = _feedback_matrix(block, "fa_tok", x.shape[2:], v.shape[2:], x.device, x.dtype)
+                    g_hat = (v @ R.t()) * block.fb_scale
+                else:
+                    feat_in, feat_out = x.shape[nbd:], v.shape[nbd:]
+                    R = _feedback_matrix(block, "fa", feat_in, feat_out, x.device, x.dtype)
+                    g_hat = (v.reshape(*v.shape[:nbd], -1) @ R.t()).reshape(x.shape) * block.fb_scale
+            elif cfg.mode == "opzo":
+                # OPZO-style baseline (Xiao et al. 2024): DFA topology with a *learned* feedback matrix — the EMA of a
+                # zeroth-order estimate of the data-averaged Jacobian^T from this block's input to the logits,
+                # obtained from perturbed forward passes of the whole network (see opzo_update).  Biased by design.
+                e = DFA.e
+                if e is None:
+                    raise RuntimeError("DFA.e is unset: call DFA.attach(logits) after the forward pass")
+                M = block.__dict__.get("_opzo", {}).get("M")
+                if M is None:
+                    g_hat = torch.zeros_like(x)
+                elif e.dim() == 3 and x.dim() >= 3 and e.shape[1] == x.shape[1]:      # per-token feedback
+                    g_hat = (e @ M.t()).reshape(x.shape)
+                else:                                                                  # per-sample feedback
+                    g_hat = (e.reshape(e.shape[0], -1) @ M.t()).reshape(x.shape)
+            elif cfg.mode in ("dfa", "drtp"):
+                # direct feedback alignment: the output error is projected straight to this block's input;
+                # DRTP: the (negated) one-hot target takes the place of the error
+                e = DFA.e if cfg.mode == "dfa" else DFA.t
+                if e is None:
+                    raise RuntimeError("DFA.e / DFA.t is unset: call DFA.attach(logits) / DFA.attach_target(y, n) first")
+                e = e.to(x.dtype)
+                if e.dim() == 3 and x.dim() >= 3 and e.shape[1] == x.shape[1]:  # per-token error [B, T, V]
+                    R = _feedback_matrix(block, "dfa_tok", x.shape[2:], e.shape[2:], x.device, x.dtype)
+                    g_hat = (e @ R.t()).reshape(x.shape) * block.fb_scale
+                else:                                                            # one error vector per sample
+                    R = _feedback_matrix(block, "dfa", x.shape[1:], e.shape[1:], x.device, x.dtype)
+                    g_hat = (e.reshape(e.shape[0], -1) @ R.t()).reshape(x.shape) * block.fb_scale
+            elif block.score_probe is not None:
+                # Probe the causal softmax score rows and map their VJP to K through QK^T.
+                # Other input slices retain their exact diagnostic VJPs (notably A^T delta for V).
+                spec = block.score_probe
+                g_hat = g_exact.clone()
+                a = spec["d"]
+                if spec.get("qk"):     # Q and K both from the score-VJP; V exact via the known attention weights
+                    gq, gk, queries = causal_score_coordinate_qk_vjp(
+                        x, v, cfg, gen, spec["d"], spec["heads"], spec["dk"], spec["dv"], spec["window"])
+                    g_hat[..., :a] = gq
+                else:
+                    gk, queries = causal_score_coordinate_k_vjp(
+                        x, v, cfg, gen, spec["d"], spec["heads"], spec["dk"], spec["dv"], spec["window"])
+                g_hat[..., a:a + spec["dk"]] = gk
+            elif block.probe_slice is not None:
+                # ZO estimate only for coordinates [a, b) of the last dim (with its own unit structure);
+                # the other coordinates keep the exact Jacobian (diagnostic: which input is hard to probe)
+                a, b = block.probe_slice
+                sbd = block.slice_batch_dims or nbd
+                xs = x[..., a:b].contiguous()
+                def f_slice(xsl):
+                    k = xsl.shape[0] // x.shape[0]
+                    xx = x.repeat(k, *([1] * (x.dim() - 1))) if k > 1 else x
+                    return f(torch.cat([xx[..., :a], xsl, xx[..., b:]], dim=-1))
+                if cfg.mode == "oracle":
+                    gs, queries = oracle_vjp(f_slice, xs, v, cfg, gen, sbd, None)
+                else:
+                    h = make_h(f_slice, v, sbd, None, cfg.readout_noise, gen)
+                    gs, queries = zo_vjp(h, xs, cfg, gen, sbd, None)
+                g_hat = g_exact.clone()
+                g_hat[..., a:b] = gs
+            elif cfg.mode == "oracle":
+                g_hat, queries = oracle_vjp(f, x, v, cfg, gen, nbd, block.local)
+            elif cfg.mode == "cd" and cfg.surrogate and block.local is None:
+                # learned linear control variate: g_hat = g_raw - s * (P_u g_sur - g_sur), g_sur = A v,
+                # with A and s fitted on *previous* steps' data only (unbiased conditional on the past)
+                h = make_h(f, v, nbd, None, cfg.readout_noise, gen)
+                cyc = block.stats["backward_calls"] if cfg.probe == "hadamard_cycle" else None
+                u, D, g_raw, queries = zo_vjp_cd_parts(h, x, cfg, gen, nbd, cyc)
+                g_hat = block._apply_surrogate(x, v, u, D, g_raw, nbd)
+            else:
+                fq = f if block.measure is None else (lambda xp: block.measure(f(xp)))
+                h = make_h(fq, v, nbd, block.local, cfg.readout_noise, gen)
+                cyc = block.stats["backward_calls"] if cfg.probe == "hadamard_cycle" else None
+                g_hat, queries = zo_vjp(h, x, cfg, gen, nbd, block.local, cyc)
+            block.stats["queries"] += queries
+            block.stats["backward_calls"] += 1
+            if rec is not None and g_hat is not None:
+                ge = g_exact.reshape(g_exact.shape[0], -1)
+                gh = g_hat.reshape(g_hat.shape[0], -1)
+                cos = torch.nn.functional.cosine_similarity(ge, gh, dim=1)
+                rel = (gh - ge).norm(dim=1) / ge.norm(dim=1).clamp_min(1e-30)
+                rec.add(name=block.name, cos=cos.mean().item(), relerr=rel.mean().item(),
+                        gnorm=ge.norm(dim=1).mean().item(), ghat_norm=gh.norm(dim=1).mean().item(),
+                        vnorm=v.reshape(v.shape[0], -1).norm(dim=1).mean().item(), queries=queries,
+                        # batch-level (mean over samples) cosine, what parameter updates "see"
+                        cos_batchmean=torch.nn.functional.cosine_similarity(
+                            ge.mean(0, keepdim=True), gh.mean(0, keepdim=True), dim=1).item())
+
+        # --- exact skip path (residual control variate) ---
+        g_skip_params = [None] * len(s_params)
+        grad_x = None
+        if need_gx and cfg.mode == "replay":
+            grad_x = block.replay
+            if s_params:
+                with torch.enable_grad():
+                    xd = x.detach()
+                    s = block.skip(xd)
+                    g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
+        elif need_gx and cfg.mode in ("dfa", "opzo", "drtp"):
+            grad_x = g_hat                                    # DFA/DRTP/OPZO: direct projection replaces the whole backward
+            if s_params:                                      # skip-path parameters still get their local gradient
+                with torch.enable_grad():
+                    xd = x.detach()
+                    s = block.skip(xd)
+                    g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
+        elif need_gx:
+            if block.skip is None:
+                grad_x = g_hat
+            elif block.skip_is_identity:
+                grad_x = v + g_hat
+            else:
+                with torch.enable_grad():
+                    xd = x.detach().requires_grad_(True)
+                    s = block.skip(xd)
+                    grads = torch.autograd.grad(s, [xd] + list(s_params), v, allow_unused=True)
+                grad_x = grads[0] + g_hat
+                g_skip_params = list(grads[1:])
+        elif s_params:
+            with torch.enable_grad():
+                xd = x.detach()
+                s = block.skip(xd)
+                g_skip_params = list(torch.autograd.grad(s, list(s_params), v, allow_unused=True))
+        if block.capture and grad_x is not None:
+            block.captured = grad_x.detach()
+        return (grad_x, None, *g_params, *g_skip_params)
+
+
+class ZBPBlock(nn.Module):
+    """y = f(x)  or  y = skip(x) + f(x).
+
+    The activation error through f is estimated with zeroth-order queries (or the oracle
+    projection); the skip path (identity or a cheap differentiable module) is propagated
+    exactly.  With cfg.mode == 'bp' the block is ordinary autograd.
+    """
+
+    def __init__(self, f, cfg=None, skip=None, batch_dims=1, name="", estimate_input_grad=True, local=None):
+        super().__init__()
+        self.f = f
+        # local=(radius, stride): conv-like block whose output position p depends only on input
+        # positions within `radius` of stride*p -> per-position measurements (lower variance)
+        self.local = Local(*local) if isinstance(local, (tuple, list)) else local
+        self.cfg = cfg if cfg is not None else ZBPConfig()
+        self.batch_dims = batch_dims
+        self.name = name
+        self.estimate_input_grad = estimate_input_grad
+        if skip is None:
+            self.skip = None
+            self.skip_is_identity = False
+        elif skip == "identity" or isinstance(skip, nn.Identity):
+            self.skip = nn.Identity()
+            self.skip_is_identity = True
+        else:
+            self.skip = skip
+            self.skip_is_identity = False
+        self.stats = {"queries": 0.0, "backward_calls": 0}
+        self.fb_scale = 1.0       # scale of the fixed random feedback map in 'fa' / 'dfa' modes
+        self.probe_slice = None   # (a, b): ZO-probe only last-dim coordinates a:b, exact elsewhere
+        self.slice_batch_dims = None
+        self.score_probe = None   # metadata for causal score-coordinate K probing
+        self.capture = False      # store the full input error of the next backward in .captured
+        self.captured = None
+        self.replay = None        # 'replay' mode returns this tensor as the input error
+        self.measure = None       # hardware measurement channel applied to y inside the query readout
+        self.v_channel = None     # hardware error-transport channel (sample/hold) applied to the incoming v
+
+    def _f_params(self):
+        return [p for p in self.f.parameters() if p.requires_grad]
+
+    # ---- learned linear control variate (prior-free): A_t = EMA[g_raw v^T] EMA[v v^T]^-1, s_t from EMA[D P]/EMA[P^2] ----
+    def _apply_surrogate(self, x, v, u, D, g_raw, nbd):
+        cfg = self.cfg
+        n = u.shape[0]
+        bs = x.shape[:nbd]
+        B = math.prod(bs) if len(bs) else 1
+        xf, vf, gf = x.reshape(B, -1), v.reshape(B, -1), g_raw.reshape(B, -1)
+        uf, Df = u.reshape(n, B, -1), D.reshape(n, B)
+        din, dout = xf.shape[1], vf.shape[1]
+        st = self.__dict__.setdefault("_sur", None)
+        if st is None:
+            st = {"S_gv": torch.zeros(din, dout, device=x.device, dtype=x.dtype),
+                  "S_vv": torch.zeros(dout, dout, device=x.device, dtype=x.dtype),
+                  "A": None, "s": 0.0, "ema_dp": 0.0, "ema_pp": 0.0, "cnt": 0}
+            self.__dict__["_sur"] = st
+        ema = cfg.surrogate_ema
+        g_hat = gf
+        if st["A"] is not None:
+            with torch.no_grad():
+                g_sur = vf @ st["A"].t()                              # [B, din]  surrogate  A v
+                P = (uf * g_sur[None]).sum(-1)                        # [n, B]    u_i^T g_sur (digital)
+                g_sur_raw = (uf * P[..., None]).mean(0)               # same-probe projection of the surrogate
+                if st["s"] > 0:
+                    g_hat = gf - st["s"] * (g_sur_raw - g_sur)        # E[g_hat] = g for any A, s fixed before this step
+                dp, pp = (Df * P).sum().item(), (P * P).sum().item()
+                st["ema_dp"] = ema * st["ema_dp"] + (1 - ema) * dp
+                st["ema_pp"] = ema * st["ema_pp"] + (1 - ema) * pp
+                st["s"] = float(min(1.0, max(0.0, st["ema_dp"] / max(st["ema_pp"], 1e-30))))
+        with torch.no_grad():                                          # regression statistics for the NEXT steps
+            st["S_gv"] = ema * st["S_gv"] + (1 - ema) * (gf.t() @ vf) / B
+            st["S_vv"] = ema * st["S_vv"] + (1 - ema) * (vf.t() @ vf) / B
+            st["cnt"] += 1
+            if st["cnt"] >= 2:
+                ridge = cfg.surrogate_ridge * st["S_vv"].diagonal().mean().clamp_min(1e-30)
+                M = st["S_vv"] + ridge * torch.eye(dout, device=x.device, dtype=x.dtype)
+                st["A"] = torch.linalg.solve(M, st["S_gv"].t()).t()   # [din, dout]: least squares of g_raw on v
+        self.stats["surrogate_s"] = st["s"]
+        return g_hat.reshape(g_raw.shape)
+
+    def _skip_params(self):
+        if self.skip is None or self.skip_is_identity:
+            return []
+        return [p for p in self.skip.parameters() if p.requires_grad]
+
+    def forward(self, x):
+        pert = self.__dict__.get("opzo_perturb")
+        if pert is not None:                                   # OPZO estimation pass: perturb this block's input
+            alpha, gen = pert
+            z = torch.randn(x.shape, generator=gen, device=x.device, dtype=x.dtype)
+            self.__dict__.setdefault("_opzo", {"M": None})["z"] = z
+            x = x + alpha * z
+        if self.cfg.mode == "bp" or not torch.is_grad_enabled():
+            y = self.f(x)
+            if self.skip is not None:
+                y = y + (x if self.skip_is_identity else self.skip(x))
+            return y
+        if not self.estimate_input_grad:
+            x = x.detach()
+        return ZBPFunction.apply(x, self, *self._f_params(), *self._skip_params())
+
+    def extra_repr(self):
+        loc = f", local=({self.local.radius},{self.local.stride})" if self.local is not None else ""
+        return f"name={self.name}, mode={self.cfg.mode}, n={self.cfg.n_probes}, skip={'id' if self.skip_is_identity else (self.skip is not None)}{loc}"
+
+
+def zbp_blocks(model):
+    return [m for m in model.modules() if isinstance(m, ZBPBlock)]
+
+
+def opzo_update(model, forward_fn, out_clean, alpha, lam, gen):
+    """One OPZO estimation step: perturb the inputs of all 'opzo' blocks at once (alpha * Gaussian), run one extra
+    forward of the whole network, and update each block's feedback matrix
+        M <- lam * M + (1 - lam) * mean_units z (delta_logits / alpha)^T
+    (cross-block responses are zero-mean noise, as in OPZO's single noisy pass).  Per-token when the logits are
+    [B, T, V] and the block input has a matching token dim, per-sample otherwise."""
+    blocks = [b for b in zbp_blocks(model) if b.cfg.mode == "opzo"]
+    if not blocks:
+        return
+    for b in blocks:
+        b.__dict__["opzo_perturb"] = (alpha, gen)
+    with torch.no_grad():
+        out_p = forward_fn()
+    for b in blocks:
+        b.__dict__["opzo_perturb"] = None
+    delta = (out_p - out_clean.detach()) / alpha
+    with torch.no_grad():
+        for b in blocks:
+            st = b.__dict__["_opzo"]
+            z = st.pop("z")
+            if delta.dim() == 3 and z.dim() >= 3 and z.shape[1] == delta.shape[1]:
+                zf = z.reshape(z.shape[0], z.shape[1], -1)
+                M_new = torch.einsum("btd,btv->dv", zf, delta) / (z.shape[0] * z.shape[1])
+            else:
+                M_new = z.reshape(z.shape[0], -1).t() @ delta.reshape(delta.shape[0], -1) / z.shape[0]
+            st["M"] = M_new if st["M"] is None else lam * st["M"] + (1 - lam) * M_new
+            b.stats["queries"] += 1.0
+
+
+def set_mode(model, cfg=None, **kw):
+    """Replace / update the estimator config of every ZBPBlock in a model."""
+    for b in zbp_blocks(model):
+        b.cfg = (cfg if cfg is not None else b.cfg).replace(**kw)
+
+
+def total_queries(model):
+    return sum(b.stats["queries"] for b in zbp_blocks(model))
+
+
+def reset_queries(model):
+    for b in zbp_blocks(model):
+        b.stats = {"queries": 0.0, "backward_calls": 0}
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)
diff --git a/src/zbp_scaling/zbp/estimators.py b/src/zbp_scaling/zbp/estimators.py
new file mode 100644
index 0000000..81f29a3
--- /dev/null
+++ b/src/zbp_scaling/zbp/estimators.py
@@ -0,0 +1,421 @@
+"""Zeroth-order estimators of g = J^T v for a black-box block.
+
+Conventions
+-----------
+A block maps x (shape [*batch_shape, *feat]) to y.  A *measurement* function
+`h(x_pert)` returns v^T f(x_pert) per measurement unit:
+  * global units:  h -> [*batch_shape]              (sum over all feature dims)
+  * local units:   h -> [*batch_shape, H', W']      (sum over channels only; conv blocks)
+`h` accepts any integer multiple of the batch stacked along dim 0, so several
+probes can be evaluated in one call.
+
+For local units the directional derivatives D_p (one per output position p) are
+box-summed over the receptive field before multiplying the probe, which keeps the
+estimator unbiased while removing the variance contributed by far-away positions.
+
+Every estimator returns (g_hat, queries) where queries is the number of block
+evaluations *per sample unit* that a physical system would have to perform.
+"""
+import math
+import torch
+import torch.nn.functional as F
+
+from .probes import sample_probes
+
+
+class Local:
+    """Receptive-field geometry of a conv-like block: output position p depends on input
+    positions q with |q - stride*p|_inf <= radius (input coordinates)."""
+
+    def __init__(self, radius, stride=1):
+        self.radius, self.stride = int(radius), int(stride)
+
+    def gather(self, D, in_hw):
+        """D: [n, *bs, H', W'] per-output-position derivatives -> S: [n, *bs, H, W] box sums."""
+        lead = D.shape[:-2]
+        Hp, Wp = D.shape[-2:]
+        H, W = in_hw
+        k = 2 * self.radius + 1
+        Dm = D.reshape(-1, 1, Hp, Wp)
+        ones = torch.ones(1, 1, k, k, dtype=D.dtype, device=D.device)
+        oph = H - ((Hp - 1) * self.stride + 1)
+        opw = W - ((Wp - 1) * self.stride + 1)
+        S = F.conv_transpose2d(Dm, ones, stride=self.stride, padding=self.radius, output_padding=(oph, opw))
+        return S.reshape(*lead, H, W)
+
+
+class Local1D:
+    """Causal token-window locality for [B, T, d] blocks: output token t depends on input tokens
+    s in [t-window+1, t].  Measurement units = output tokens; gather = causal window box-sum."""
+    kind = "1d"
+
+    def __init__(self, window):
+        self.window = int(window)
+
+    def gather(self, D, in_shape):
+        """D: [n, *bs, T] per-output-token derivatives -> S: [n, *bs, T], S[s] = sum_{t=s}^{s+W-1} D[t]."""
+        T = D.shape[-1]
+        c = torch.cumsum(D, dim=-1)
+        pad = torch.zeros_like(c[..., :1])
+        c0 = torch.cat([pad, c], dim=-1)                                # c0[t] = sum_{<t}
+        hi = torch.clamp(torch.arange(T, device=D.device) + self.window, max=T)
+        lo = torch.arange(T, device=D.device)
+        return c0[..., hi] - c0[..., lo]
+
+
+def _probes_for(x, n_probes, family, gen, n_batch_dims, cycle=None):
+    bs = x.shape[:n_batch_dims]
+    feat = x.shape[n_batch_dims:]
+    d = math.prod(feat)
+    u = sample_probes(n_probes, tuple(bs), d, family, gen, x.device, x.dtype, cycle=cycle)
+    return u.reshape(n_probes, *bs, *feat), d
+
+
+def zo_vjp_cd_parts(h, x, cfg, gen, n_batch_dims=1, cycle=None):
+    """Central-difference measurements together with their probes (global units only):
+    returns (u [n,*bs,*feat], D [n,*bs] = u_i^T g + O(eps^2), g_hat, queries)."""
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
+    eps = _eps_for(cfg, d)
+    D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+    return u, D, _combine(u, D, n_batch_dims, None, x.shape), float(2 * cfg.n_probes)
+
+
+def _combine(u, D, n_batch_dims, local, in_shape):
+    """g_hat = (1/n) sum_i u_i * S_i, where S_i are the (gathered) directional derivatives."""
+    n = u.shape[0]
+    if local is None:
+        S = D.reshape(*D.shape, *([1] * (u.dim() - 1 - n_batch_dims)))
+    elif getattr(local, "kind", "2d") == "1d":
+        S = local.gather(D, in_shape).unsqueeze(-1)           # [n, *bs, T, 1] broadcast over features
+    else:
+        S = local.gather(D, in_shape[-2:])                   # [n, *bs, H, W]
+        S = S.unsqueeze(n_batch_dims + 1)                     # broadcast over channels: [n, *bs, 1, H, W]
+    return (u * S).mean(0)
+
+
+def _eps_for(cfg, d):
+    return cfg.eps if cfg.eps_mode == "coord" else cfg.eps / math.sqrt(d)
+
+
+def _eval_pairs(h, x, u, eps_list, n_batch_dims, batch_probes=True, chunk=0):
+    """D(eps, u_i) = (h(x+eps u_i) - h(x-eps u_i)) / (2 eps) for all probes/eps -> [len(eps), n, *bs, *meas]."""
+    n = u.shape[0]
+    B0 = x.shape[0]
+    outs = []
+    if batch_probes:
+        step = n if chunk <= 0 else max(1, min(chunk, n))
+        for eps in eps_list:
+            parts = []
+            for i0 in range(0, n, step):
+                xs = []
+                for i in range(i0, min(n, i0 + step)):
+                    xs.append(x + eps * u[i])
+                    xs.append(x - eps * u[i])
+                hv = h(torch.cat(xs, 0))
+                hv = hv.reshape(len(xs) // 2, 2, B0, *hv.shape[1:])
+                parts.append((hv[:, 0] - hv[:, 1]) / (2 * eps))
+            outs.append(torch.cat(parts, 0))
+    else:
+        for eps in eps_list:
+            outs.append(torch.stack([(h(x + eps * u[i]) - h(x - eps * u[i])) / (2 * eps) for i in range(n)]))
+    return torch.stack(outs)
+
+
+def exact_directional(f, x, v, u, n_batch_dims=1, local=None, chunk=16):
+    """Oracle: exact directional derivatives D_i = <v, J u_i> per measurement unit via forward-mode AD
+    (vmapped over probes in chunks; falls back to a loop if the block is not vmap-compatible)."""
+    def one(t):
+        return torch.func.jvp(f, (x,), (t,))[1]
+    outs = []
+    n = u.shape[0]
+    try:
+        for i in range(0, n, chunk):
+            outs.append(torch.func.vmap(one)(u[i:i + chunk]))
+        Ju = torch.cat(outs, 0)
+    except Exception:
+        Ju = torch.stack([one(u[i]) for i in range(n)])
+    prod = Ju * v.unsqueeze(0)
+    if local is None:
+        return prod.reshape(n, *x.shape[:n_batch_dims], -1).sum(-1)
+    if getattr(local, "kind", "2d") == "1d":
+        return prod.sum(-1)
+    return prod.sum(n_batch_dims + 1)                         # sum over channel dim only
+
+
+def oracle_vjp(f, x, v, cfg, gen, n_batch_dims=1, local=None):
+    """Stage A: (1/n) sum_i u_i <v, J u_i>  (== (1/n) sum_i u_i (u_i^T g)).  Conditionally unbiased."""
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims)
+    with torch.no_grad():
+        D = exact_directional(f, x, v, u, n_batch_dims, local)
+    return _combine(u, D, n_batch_dims, local, x.shape), 0.0
+
+
+def oracle_projection(g, cfg, gen, n_batch_dims=1):
+    """Projection form of the oracle for a given exact VJP g (global units only)."""
+    u, d = _probes_for(g, cfg.n_probes, cfg.probe, gen, n_batch_dims)
+    n = u.shape[0]
+    bs = g.shape[:n_batch_dims]
+    D = (u.reshape(n, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1)
+    return _combine(u, D, n_batch_dims, None, g.shape), 0.0
+
+
+def zo_vjp(h, x, cfg, gen, n_batch_dims=1, local=None, cycle=None):
+    """Function-query estimate of J^T v.  Returns (g_hat, queries_per_sample_unit)."""
+    if cfg.mode == "cd" and cfg.probe == "coordinate_all" and local is None:
+        # analytic combine: u_i = sqrt(d) e_i  ->  g_hat[..., i] = D_i / sqrt(d); the probe tensor is a broadcast
+        # view of eye(d), so memory stays O(d^2 + measurements) instead of O(n * batch * d)
+        bs = x.shape[:n_batch_dims]
+        d = math.prod(x.shape[n_batch_dims:])
+        if cfg.n_probes != d:
+            raise ValueError(f"coordinate_all needs n == d (n={cfg.n_probes}, d={d})")
+        eps = _eps_for(cfg, d)
+        ident = math.sqrt(d) * torch.eye(d, device=x.device, dtype=x.dtype)
+        uv = ident.reshape(d, *([1] * len(bs)), d).expand(d, *bs, d)
+        D = _eval_pairs(h, x, uv, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+        return (D.movedim(0, -1) / math.sqrt(d)).reshape(x.shape), float(2 * d)
+    if cfg.mode == "cd" and cfg.probe_chunk and cfg.n_probes > 4 * cfg.probe_chunk and cfg.probe != "hadamard_cycle":
+        # probe-chunked sample -> evaluate -> accumulate: memory O(chunk * batch * d) for any n
+        bs = x.shape[:n_batch_dims]
+        d = math.prod(x.shape[n_batch_dims:])
+        eps = _eps_for(cfg, d)
+        g_acc = torch.zeros_like(x)
+        done = 0
+        while done < cfg.n_probes:
+            c = min(cfg.probe_chunk, cfg.n_probes - done)
+            uc, _ = _probes_for(x, c, cfg.probe, gen, n_batch_dims)
+            D = _eval_pairs(h, x, uc, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+            g_acc += _combine(uc, D, n_batch_dims, local, x.shape) * c
+            done += c
+        return g_acc / cfg.n_probes, float(2 * cfg.n_probes)
+    u, d = _probes_for(x, cfg.n_probes, cfg.probe, gen, n_batch_dims, cycle)
+    n = cfg.n_probes
+    eps = _eps_for(cfg, d)
+    bs = x.shape[:n_batch_dims]
+    B0 = x.shape[0]
+
+    if cfg.mode == "forward":
+        if cfg.batch_probes:
+            hv = h(torch.cat([x] + [x + eps * u[i] for i in range(n)], 0))
+            hv = hv.reshape(n + 1, B0, *hv.shape[1:])
+            D = (hv[1:] - hv[:1]) / eps
+        else:
+            h0 = h(x)
+            D = torch.stack([(h(x + eps * u[i]) - h0) / eps for i in range(n)])
+        return _combine(u, D, n_batch_dims, local, x.shape), float(n + 1)
+
+    if cfg.mode == "cd":
+        D = _eval_pairs(h, x, u, [eps], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)[0]
+        return _combine(u, D, n_batch_dims, local, x.shape), float(2 * n)
+
+    if cfg.mode == "richardson":
+        D = _eval_pairs(h, x, u, [eps, eps / 2], n_batch_dims, cfg.batch_probes, cfg.probe_chunk)
+        R = (4 * D[1] - D[0]) / 3
+        return _combine(u, R, n_batch_dims, local, x.shape), float(4 * n)
+
+    if cfg.mode in ("ml", "ml_richardson"):
+        rich = cfg.mode == "ml_richardson"
+        if cfg.level_sampling == "per_sample":
+            U = torch.rand((n, *bs), generator=gen, device=x.device, dtype=torch.float64)
+        else:
+            U = torch.rand((n,), generator=gen, device=x.device, dtype=torch.float64)
+            U = U.reshape(n, *([1] * len(bs))).expand(n, *bs)
+        N = torch.floor(-torch.log2(U) / cfg.alpha).clamp(max=cfg.max_level).to(torch.int64)
+        if cfg.eps_min > 0:   # truncate the ladder where readout noise would dominate (bias O(eps_min^2))
+            kmax = max(0, int(math.floor(math.log2(eps / cfg.eps_min)))) if eps > cfg.eps_min else 0
+            N = N.clamp(max=kmax)
+        Nmax = int(N.max().item())
+        K = Nmax + (2 if rich else 1)
+        eps_list = [eps * 2.0 ** (-k) for k in range(K)]
+        D = _eval_pairs(h, x, u, eps_list, n_batch_dims, cfg.batch_probes, cfg.probe_chunk)   # [K, n, *bs, *meas]
+        S = (4 * D[1:] - D[:-1]) / 3 if rich else D
+        est = S[0].clone()
+        extra = S.dim() - 1 - N.dim()
+        for k in range(1, Nmax + 1):
+            mask = (N >= k).to(x.dtype).reshape(*N.shape, *([1] * extra))
+            est = est + mask * (S[k] - S[k - 1]) / (2.0 ** (-cfg.alpha * k))
+        per = 2.0 * (N.to(torch.float64) + 1 + (1 if rich else 0))
+        queries = float(per.sum().item() / (math.prod(bs) if len(bs) else 1))
+        return _combine(u, est, n_batch_dims, local, x.shape), queries
+
+    raise ValueError(f"unknown ZBP mode {cfg.mode!r}")
+
+
+def causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
+    """Estimate attention's K-side VJP by probing causal score rows.
+
+    The black-box interface here is ``softmax(scores) @ value`` rather than K itself.
+    One paired query perturbs one valid score coordinate in every (sample, head,
+    output-token) row, and the output supplies a separate scalar measurement for
+    every such row.  For a row with ``m`` valid keys, the estimator visits all
+    coordinates when ``m <= n``; otherwise it samples ``n`` coordinates without
+    replacement and applies the Horvitz--Thompson weight ``m / n``.  Thus the
+    row-gradient relative MSE is exactly ``max(m / n - 1, 0)`` in the directional-
+    derivative limit, instead of ``(m - 1) / n`` for i.i.d. dense probes.
+
+    The estimated score VJP is mapped to K through the digitally known bilinear
+    score map QK^T/sqrt(dh).  Q and V are held fixed during the function queries.
+    Returns ``(g_k, queries_per_sequence)``; full-attention K is [B,T,d], while
+    multi-query K is [B,T,dh].
+    """
+    if cfg.mode not in ("cd", "oracle"):
+        raise ValueError(f"causal score-coordinate probing supports cd/oracle, got {cfg.mode!r}")
+    B, T, _ = qkv.shape
+    H = int(n_heads)
+    dh = int(d) // H
+    dk = int(d if dk is None else dk)
+    dv = int(d if dv is None else dv)
+    if int(d) % H:
+        raise ValueError(f"attention width {d} is not divisible by {H} heads")
+    if dk not in (int(d), dh) or dv not in (int(d), dh):
+        raise ValueError("score-coordinate probing supports full or multi-query K/V layouts")
+
+    q0, k0, value0 = qkv.split([int(d), dk, dv], dim=-1)
+    q = q0.view(B, T, H, dh).transpose(1, 2)
+    if dk == dh:
+        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        k = k0.view(B, T, H, dh).transpose(1, 2)
+    if dv == dh:
+        value = value0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        value = value0.view(B, T, H, dh).transpose(1, 2)
+    delta = v_out.view(B, T, H, dh).transpose(1, 2)
+
+    idx = torch.arange(T, device=qkv.device)
+    valid = idx[None, :] <= idx[:, None]
+    if window is not None:
+        valid = valid & (idx[None, :] > idx[:, None] - int(window))
+    counts = valid.sum(-1)                                      # [T], valid coordinates per row
+    scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
+    scores = scores.masked_fill(~valid, float("-inf"))
+
+    # Analytic score VJP is useful for projection-only audits.  The CD path below
+    # obtains the same quantities solely from paired softmax/value evaluations.
+    if cfg.mode == "oracle":
+        att = scores.softmax(-1)
+        g_att = delta @ value.transpose(-2, -1)
+        g_score_exact = att * (g_att - (att * g_att).sum(-1, keepdim=True))
+        g_score_exact = g_score_exact.masked_fill(~valid, 0)
+
+    n = int(cfg.n_probes)
+    if n <= 0:
+        raise ValueError("score-coordinate probing needs at least one probe")
+    slots = min(n, T)
+    # Independent random subsets for every sample/head/row.  Invalid positions
+    # sort last, so the first min(n,m) entries are a uniform subset of the m-prefix.
+    priorities = torch.rand((B, H, T, T), generator=gen, device=qkv.device, dtype=torch.float32)
+    priorities.masked_fill_(~valid.view(1, 1, T, T), 2.0)
+    chosen = priorities.argsort(dim=-1)[..., :slots]             # [B,H,T,slots]
+    chosen = chosen.permute(3, 0, 1, 2).contiguous()             # [slots,B,H,T]
+    slot_id = torch.arange(slots, device=qkv.device).view(slots, 1, 1, 1)
+    active = slot_id < counts.view(1, 1, 1, T)
+    active = active.expand(slots, B, H, T)
+    weights = torch.where(counts <= n, torch.ones_like(counts, dtype=qkv.dtype),
+                          counts.to(qkv.dtype) / n)               # [T]
+    g_score = torch.zeros_like(scores)
+
+    if cfg.mode == "oracle":
+        picked = torch.gather(g_score_exact.unsqueeze(0).expand(slots, -1, -1, -1, -1),
+                              -1, chosen.unsqueeze(-1)).squeeze(-1)
+        contrib = picked * active.to(qkv.dtype) * weights.view(1, 1, 1, T)
+        for j in range(slots):
+            g_score.scatter_add_(-1, chosen[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
+        queries = 0.0
+    else:
+        eps = float(cfg.eps)  # directions are unit coordinates, so coord/norm conventions coincide
+        if eps <= 0:
+            raise ValueError("score-coordinate probing needs eps > 0")
+        chunk = slots if cfg.probe_chunk <= 0 else max(1, min(int(cfg.probe_chunk), slots))
+        for i0 in range(0, slots, chunk):
+            ids = chosen[i0:i0 + chunk]
+            act = active[i0:i0 + chunk]
+            c = ids.shape[0]
+            plus = scores.unsqueeze(0).expand(c, -1, -1, -1, -1).clone()
+            minus = plus.clone()
+            shift = act.to(qkv.dtype).unsqueeze(-1) * eps
+            plus.scatter_add_(-1, ids.unsqueeze(-1), shift)
+            minus.scatter_add_(-1, ids.unsqueeze(-1), -shift)
+            y_plus = plus.softmax(-1) @ value.unsqueeze(0)
+            y_minus = minus.softmax(-1) @ value.unsqueeze(0)
+            if cfg.readout_noise > 0:
+                y_plus = y_plus + cfg.readout_noise * torch.randn(
+                    y_plus.shape, generator=gen, device=y_plus.device, dtype=y_plus.dtype)
+                y_minus = y_minus + cfg.readout_noise * torch.randn(
+                    y_minus.shape, generator=gen, device=y_minus.device, dtype=y_minus.dtype)
+            deriv = ((y_plus - y_minus) * delta.unsqueeze(0)).sum(-1) / (2 * eps)
+            contrib = deriv * act.to(qkv.dtype) * weights.view(1, 1, 1, T)
+            for j in range(c):
+                g_score.scatter_add_(-1, ids[j].unsqueeze(-1), contrib[j].unsqueeze(-1))
+        queries = float(2 * n)
+
+    _LAST_SCORE_VJP['g'] = g_score
+    gk_heads = g_score.transpose(-2, -1) @ q / math.sqrt(dh)
+    if dk == dh:
+        gk = gk_heads.sum(1).transpose(1, 2)                      # shared K accumulates all heads
+    else:
+        gk = gk_heads.transpose(1, 2).reshape(B, T, int(d))
+    return gk, queries
+
+
+def make_h(f, v, n_batch_dims=1, local=None, readout_noise=0.0, gen=None):
+    """Measurement function h(x_pert) = <v, f(x_pert)> per unit, tolerant to stacked batches.
+    readout_noise > 0 adds i.i.d. Gaussian noise to every measured output coordinate of every query
+    (a physical readout-noise model; v is applied digitally afterwards)."""
+    B0 = v.shape[0]
+
+    def h(xp):
+        with torch.no_grad():
+            y = f(xp)
+            if readout_noise > 0:
+                y = y + readout_noise * torch.randn(y.shape, generator=gen, device=y.device, dtype=y.dtype)
+            k = y.shape[0] // B0
+            vv = v.repeat(k, *([1] * (v.dim() - 1))) if k > 1 else v
+            prod = y * vv
+            if local is None:
+                return prod.reshape(*prod.shape[:n_batch_dims], -1).sum(-1)
+            if getattr(local, "kind", "2d") == "1d":
+                return prod.sum(-1)                           # per output token
+            return prod.sum(n_batch_dims)                     # keep spatial dims
+    return h
+
+
+def control_vjp(g, cfg, gen, n_batch_dims=1):
+    """Control estimators.  'zero': no credit is propagated (downstream layers frozen).
+    'noise': isotropic Gaussian noise with the per-sample norm that the oracle estimate with
+    the same n would have (sqrt(1 + (d-1)/n) ||g||): same magnitude, zero information."""
+    if cfg.mode == "zero":
+        return torch.zeros_like(g), 0.0
+    if cfg.mode == "noise":
+        bs = g.shape[:n_batch_dims]
+        gf = g.reshape(*bs, -1)
+        d = gf.shape[-1]
+        z = torch.randn(gf.shape, generator=gen, device=g.device, dtype=g.dtype)
+        z = z / z.norm(dim=-1, keepdim=True).clamp_min(1e-30)
+        target = gf.norm(dim=-1, keepdim=True) * math.sqrt(1 + (d - 1) / cfg.n_probes)
+        return (z * target).reshape(g.shape), 0.0
+    raise ValueError(cfg.mode)
+
+
+def causal_score_coordinate_qk_vjp(qkv, v_out, cfg, gen, d, n_heads, dk=None, dv=None, window=None):
+    """Same score-space probing as causal_score_coordinate_k_vjp, but returns the Q-side VJP as well
+    (both follow from the score-VJP through the known bilinear map S = Q K^T / sqrt(d_h))."""
+    B, T, _ = qkv.shape
+    H = int(n_heads); dh = int(d) // H
+    dk = int(d if dk is None else dk)
+    gk, queries = causal_score_coordinate_k_vjp(qkv, v_out, cfg, gen, d, n_heads, dk, dv, window)
+    # recover g_score from gk is not possible in general; recompute the score VJP estimate via the same
+    # probes by calling the K routine on a transposed view is not exact either, so instead re-derive:
+    # g_K[s] = sum_t g_S[t,s] q_t / sqrt(dh)  and  g_Q[t] = sum_s g_S[t,s] k_s / sqrt(dh).
+    # We obtain g_S by solving nothing: the K routine exposes it through `_last_score_vjp`.
+    g_score = _LAST_SCORE_VJP["g"]
+    k0 = qkv[..., int(d):int(d) + dk]
+    if dk == dh:
+        k = k0.view(B, T, 1, dh).transpose(1, 2).expand(B, H, T, dh)
+    else:
+        k = k0.view(B, T, H, dh).transpose(1, 2)
+    gq = (g_score @ k) / math.sqrt(dh)                              # [B,H,T,dh]
+    gq = gq.transpose(1, 2).reshape(B, T, int(d))
+    return gq, gk, queries
+
+
+_LAST_SCORE_VJP = {"g": None}
diff --git a/src/zbp_scaling/zbp/fa.py b/src/zbp_scaling/zbp/fa.py
new file mode 100644
index 0000000..cd8830b
--- /dev/null
+++ b/src/zbp_scaling/zbp/fa.py
@@ -0,0 +1,84 @@
+"""Standard (layer-wise) feedback alignment: Linear / Conv2d layers whose backward uses a fixed random
+weight B in place of W^T (Lillicrap et al. 2016).  Nonlinearity derivatives are the true ones, taken at
+the forward activations, as in the original algorithm.  Use `apply_fa(module)` to convert all Linear/Conv2d
+layers of a block in place."""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class _FALinearFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W, b, B):
+        ctx.save_for_backward(x, B)
+        ctx.has_bias = b is not None
+        y = x.matmul(W.t())
+        return y + b if b is not None else y
+
+    @staticmethod
+    def backward(ctx, gy):
+        x, B = ctx.saved_tensors
+        gx = gy.matmul(B)                                   # random feedback instead of W
+        gW = gy.reshape(-1, gy.shape[-1]).t().matmul(x.reshape(-1, x.shape[-1]))
+        gb = gy.reshape(-1, gy.shape[-1]).sum(0) if ctx.has_bias else None
+        return gx, gW, gb, None
+
+
+class _FAConvFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W, b, B, stride, padding):
+        ctx.save_for_backward(x, W, B)
+        ctx.stride, ctx.padding = stride, padding
+        ctx.has_bias = b is not None
+        return F.conv2d(x, W, b, stride, padding)
+
+    @staticmethod
+    def backward(ctx, gy):
+        x, W, B = ctx.saved_tensors
+        gx = torch.nn.grad.conv2d_input(x.shape, B, gy, ctx.stride, ctx.padding)     # random feedback kernels
+        gW = torch.nn.grad.conv2d_weight(x, W.shape, gy, ctx.stride, ctx.padding)
+        gb = gy.sum((0, 2, 3)) if ctx.has_bias else None
+        return gx, gW, gb, None, None, None
+
+
+class FALinear(nn.Linear):
+    def __init__(self, *a, **k):
+        super().__init__(*a, **k)
+        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(self.in_features))
+
+    def forward(self, x):
+        return _FALinearFn.apply(x, self.weight, self.bias, self.B)
+
+
+class FAConv2d(nn.Conv2d):
+    def __init__(self, *a, **k):
+        super().__init__(*a, **k)
+        fan_in = self.in_channels * self.kernel_size[0] * self.kernel_size[1]
+        self.register_buffer("B", torch.randn_like(self.weight) / math.sqrt(fan_in))
+
+    def forward(self, x):
+        return _FAConvFn.apply(x, self.weight, self.bias, self.B, self.stride, self.padding)
+
+
+def apply_fa(module):
+    """Replace every nn.Linear / nn.Conv2d inside `module` (recursively) by its FA variant, keeping weights."""
+    for name, child in list(module.named_children()):
+        if type(child) is nn.Linear:
+            new = FALinear(child.in_features, child.out_features, bias=child.bias is not None)
+            new.weight = child.weight
+            if child.bias is not None:
+                new.bias = child.bias
+            new.B = new.B.to(child.weight.device)
+            setattr(module, name, new)
+        elif type(child) is nn.Conv2d:
+            new = FAConv2d(child.in_channels, child.out_channels, child.kernel_size, child.stride, child.padding,
+                           bias=child.bias is not None)
+            new.weight = child.weight
+            if child.bias is not None:
+                new.bias = child.bias
+            new.B = new.B.to(child.weight.device)
+            setattr(module, name, new)
+        else:
+            apply_fa(child)
+    return module
diff --git a/src/zbp_scaling/zbp/hardware.py b/src/zbp_scaling/zbp/hardware.py
new file mode 100644
index 0000000..32ae482
--- /dev/null
+++ b/src/zbp_scaling/zbp/hardware.py
@@ -0,0 +1,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()
diff --git a/src/zbp_scaling/zbp/joint.py b/src/zbp_scaling/zbp/joint.py
new file mode 100644
index 0000000..f29c210
--- /dev/null
+++ b/src/zbp_scaling/zbp/joint.py
@@ -0,0 +1,81 @@
+"""Model-level (non-chained) estimators that bracket chained ZBP:
+
+'np'      classic node perturbation with block-boundary structure: all block inputs are perturbed
+          simultaneously with independent probes, ONE scalar D = sum_l <u_l, g_l> (= the directional
+          derivative of the loss) is measured, and every block gets g_hat_l = (1/n) sum_i u_{l,i} D_i.
+          No compounding; variance of block l ~ (d_l/n) sum_l' |g_l'|^2; O(n) network forwards.
+'direct'  per-block projection of the EXACT input error (no compounding, no cross-block terms);
+          physically this is INP-like probing through all downstream blocks: O(n L^2) block forwards.
+
+Both are simulated in oracle form (exact directional derivatives) with a two-pass replay:
+pass 1 runs the network with exact backward and captures every block's full input error; the
+projected errors are then injected in pass 2, whose backward builds the local parameter gradients
+from the injected errors exactly as chained ZBP would.
+"""
+import torch
+from .autograd import zbp_blocks, get_generator
+from .probes import sample_probes
+
+
+def joint_backward(model, loss_fn, x, y, cfg, mode):
+    # only the physical blocks (those in the joint mode) are perturbed / replayed; digital blocks stay exact
+    blocks = [b for b in zbp_blocks(model) if b.estimate_input_grad and b.cfg.mode == mode]
+    cfgs = [b.cfg for b in blocks]
+    # ---- pass 1: exact input errors
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="exact")
+        b.capture = True
+    model.zero_grad(set_to_none=True)
+    loss = loss_fn(model(x), y)
+    loss.backward()
+    gs = [b.captured for b in blocks]
+    for b in blocks:
+        b.capture = False
+        b.captured = None
+    model.zero_grad(set_to_none=True)
+    # ---- projections (probes drawn in chunks so that memory is O(chunk * sum_l |g_l|), not O(n * ...))
+    gen = get_generator(x.device)
+    n = cfg.n_probes
+    chunk = cfg.probe_chunk if cfg.probe_chunk and cfg.probe_chunk > 0 else 8
+    acc = [torch.zeros_like(g) for g in gs]
+    for i0 in range(0, n, chunk):
+        c = min(chunk, n - i0)
+        us, Ds = [], []
+        for b, g in zip(blocks, gs):
+            bs = g.shape[:b.batch_dims]
+            d = g[0].numel() // (int(torch.tensor(bs[1:]).prod()) if len(bs) > 1 else 1)
+            u = sample_probes(c, tuple(bs), d, cfg.probe, gen, g.device, g.dtype).reshape(c, *g.shape)
+            D = (u.reshape(c, *bs, -1) * g.reshape(1, *bs, -1)).sum(-1)    # [c, *bs]
+            us.append(u); Ds.append(D)
+        if mode == "np":
+            # one scalar per (probe, sample): sum over blocks (and over tokens/positions within a sample)
+            Dtot = sum(D.reshape(c, D.shape[1], -1).sum(-1) for D in Ds)    # [c, B]
+            for a, u, g in zip(acc, us, gs):
+                shape = (c, g.shape[0]) + (1,) * (g.dim() - 1)
+                a.add_((u * Dtot.reshape(shape)).sum(0))
+        elif mode == "direct":
+            for a, u, D, g in zip(acc, us, Ds, gs):
+                extra = g.dim() - len(D.shape[1:])
+                a.add_((u * D.reshape(*D.shape, *([1] * extra))).sum(0))
+        else:
+            raise ValueError(mode)
+        del us, Ds
+    for b, a in zip(blocks, acc):
+        b.replay = a / n
+    if mode == "np":
+        queries = 2.0 * n                                                   # full-network forwards per sample
+    else:
+        L = len(blocks)
+        queries = 2.0 * n * sum(range(1, L + 1))                            # block forwards per sample
+    # ---- pass 2: replay projected errors, build local parameter gradients
+    for b in blocks:
+        b.cfg = b.cfg.replace(mode="replay")
+    model.zero_grad(set_to_none=True)
+    loss = loss_fn(model(x), y)
+    loss.backward()
+    for b, c in zip(blocks, cfgs):
+        b.cfg = c
+        b.replay = None
+        b.stats["queries"] += queries / max(1, len(blocks))
+        b.stats["backward_calls"] += 1
+    return loss
diff --git a/src/zbp_scaling/zbp/metrics.py b/src/zbp_scaling/zbp/metrics.py
new file mode 100644
index 0000000..99c3841
--- /dev/null
+++ b/src/zbp_scaling/zbp/metrics.py
@@ -0,0 +1,61 @@
+"""Bias / variance diagnostics for gradient estimators."""
+import math
+import torch
+
+
+def cosine(a, b):
+    a = a.reshape(a.shape[0], -1) if a.dim() > 1 else a.reshape(1, -1)
+    b = b.reshape(b.shape[0], -1) if b.dim() > 1 else b.reshape(1, -1)
+    return torch.nn.functional.cosine_similarity(a, b, dim=1)
+
+
+def relerr(a, b):
+    a = a.reshape(a.shape[0], -1)
+    b = b.reshape(b.shape[0], -1)
+    return (a - b).norm(dim=1) / b.norm(dim=1).clamp_min(1e-30)
+
+
+def audit(samples, g):
+    """samples: [M, *shape] independent estimates; g: exact [*shape].
+
+    Returns dict with relative bias, variance ratio, a chi-square test of zero bias,
+    and the distribution of cosines.
+    """
+    M = samples.shape[0]
+    S = samples.reshape(M, -1).double()
+    g = g.reshape(-1).double()
+    mean = S.mean(0)
+    bias = mean - g
+    var = S.var(0, unbiased=True)                     # per-coordinate variance
+    gn2 = g.dot(g).clamp_min(1e-300)
+    rel_bias = (bias.norm() / gn2.sqrt()).item()
+    var_ratio = (var.sum() / gn2).item()               # E||g_hat - mean||^2 / ||g||^2
+    # chi-square statistic of H0: bias = 0, per coordinate z_j = bias_j / (sd_j / sqrt(M)),
+    # restricted to coordinates that are actually random (exact coordinates have zero variance)
+    active = var > 1e-12 * var.max().clamp_min(1e-300)
+    se = (var[active] / M).sqrt().clamp_min(1e-300)
+    z = bias[active] / se
+    d = max(int(z.numel()), 1)
+    chi2 = (z * z).sum().item()
+    # under H0 each z_j is Student-t with nu = M-1 dof: E[t^2] = nu/(nu-2), Var[t^2] = 2 nu^2 (nu-1) / ((nu-2)^2 (nu-4))
+    nu = M - 1
+    if nu > 4:
+        mu2 = nu / (nu - 2)
+        v2 = 2 * nu ** 2 * (nu - 1) / ((nu - 2) ** 2 * (nu - 4))
+    else:
+        mu2, v2 = 1.0, 2.0
+    zscore = (chi2 - d * mu2) / math.sqrt(d * v2)      # approx N(0,1) under H0 for large d
+    # relative bias expected under H0 (pure noise): sqrt(sum var / M) / ||g||
+    rel_bias_null = ((var.sum() / M).sqrt() / gn2.sqrt()).item()
+    cos = torch.nn.functional.cosine_similarity(S, g[None], dim=1)
+    return {
+        "M": M, "d": d, "d_total": g.numel(),
+        "rel_bias": rel_bias,
+        "rel_bias_null": rel_bias_null,
+        "bias_ratio": rel_bias / max(rel_bias_null, 1e-300),
+        "var_ratio": var_ratio,
+        "chi2": chi2, "chi2_z": zscore,
+        "cos_mean": cos.mean().item(), "cos_std": cos.std().item(),
+        "cos_q10": cos.quantile(0.1).item(), "cos_q90": cos.quantile(0.9).item(),
+        "cos_of_mean": torch.nn.functional.cosine_similarity(mean[None], g[None], dim=1).item(),
+    }
diff --git a/src/zbp_scaling/zbp/probes.py b/src/zbp_scaling/zbp/probes.py
new file mode 100644
index 0000000..a8e7fe6
--- /dev/null
+++ b/src/zbp_scaling/zbp/probes.py
@@ -0,0 +1,78 @@
+"""Random probe directions with E[u u^T] = I_d per sample unit.
+
+All families return entries of O(1) magnitude so that `x + eps * u` moves every
+coordinate by ~eps.  Divide by sqrt(d) for the fixed-Euclidean-norm variant.
+"""
+import math
+import torch
+
+_HADAMARD_CACHE = {}
+
+
+def _hadamard(n_pow2, device, dtype):
+    key = (n_pow2, str(device), dtype)
+    H = _HADAMARD_CACHE.get(key)
+    if H is None:
+        H = torch.ones(1, 1, dtype=dtype, device=device)
+        while H.shape[0] < n_pow2:
+            H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
+        _HADAMARD_CACHE[key] = H
+    return H
+
+
+def sample_probes(n, batch_shape, d, family, gen, device, dtype, cycle=None):
+    """Return u of shape [n, *batch_shape, d].
+
+    For every sample unit (index into batch_shape) the n probes satisfy
+    E[u_i u_i^T] = I_d.  'orthogonal' and 'hadamard' additionally make the n probes of
+    one sample unit mutually orthogonal (n <= d).
+    """
+    B = math.prod(batch_shape) if len(batch_shape) else 1
+    if family == "rademacher":
+        u = torch.randint(0, 2, (n, B, d), generator=gen, device=device).to(dtype) * 2 - 1
+    elif family == "gaussian":
+        u = torch.randn((n, B, d), generator=gen, device=device, dtype=dtype)
+    elif family == "orthogonal":
+        if n > d:
+            raise ValueError(f"orthogonal probes need n <= d (n={n}, d={d})")
+        A = torch.randn((B, d, n), generator=gen, device=device, dtype=dtype)
+        Q, R = torch.linalg.qr(A)                       # [B, d, n], orthonormal columns
+        s = torch.sign(torch.diagonal(R, dim1=-2, dim2=-1))  # make the distribution Haar
+        s[s == 0] = 1
+        Q = Q * s[:, None, :]
+        u = Q.permute(2, 0, 1) * math.sqrt(d)           # columns are unit vectors -> scale by sqrt(d)
+    elif family == "hadamard":
+        d2 = 1 << (d - 1).bit_length()
+        if n > d2:
+            raise ValueError(f"hadamard probes need n <= {d2}")
+        H = _hadamard(d2, device, dtype)                # +-1 entries, orthogonal rows
+        signs = torch.randint(0, 2, (B, d2), generator=gen, device=device).to(dtype) * 2 - 1
+        rows = torch.argsort(torch.rand((B, d2), generator=gen, device=device), dim=1)[:, :n]
+        u = H[rows] * signs[:, None, :]                 # [B, n, d2]
+        u = u[:, :, :d].permute(1, 0, 2).contiguous()
+    elif family == "hadamard_cycle":
+        # Temporal QMC: every unit of a step shares the same n Hadamard rows (advancing by n per step) and the same
+        # sign vector (fixed for a whole cycle of d2/n steps), so the projection noise cancels exactly over a cycle
+        # for a slowly varying gradient -- at the price of no cross-unit averaging within a step.
+        d2 = 1 << (d - 1).bit_length()
+        H = _hadamard(d2, device, dtype)
+        step = int(cycle or 0)
+        start, cyc = (step * n) % d2, (step * n) // d2
+        rows = (start + torch.arange(n, device=device)) % d2
+        g2 = torch.Generator(device=device)
+        g2.manual_seed(1000003 * cyc + 17)
+        signs = torch.randint(0, 2, (1, d2), generator=g2, device=device).to(dtype) * 2 - 1
+        u = (H[rows] * signs)[:, None, :d].expand(n, B, d).contiguous()
+    elif family == "coordinate_all":
+        # deterministic sweep of every coordinate (requires n == d): u_i = sqrt(d) e_i, so (1/n) sum u_i u_i^T g = g
+        # exactly -- the coordinate-finite-difference estimator of HZO / BOND at its full budget of 2d queries per unit
+        if n != d:
+            raise ValueError(f"coordinate_all needs n == d (n={n}, d={d})")
+        u = (math.sqrt(d) * torch.eye(d, device=device, dtype=dtype))[:, None, :].expand(n, B, d).contiguous()
+    elif family == "coordinate":
+        idx = torch.randint(0, d, (n, B), generator=gen, device=device)
+        u = torch.zeros((n, B, d), device=device, dtype=dtype)
+        u.scatter_(2, idx[..., None], math.sqrt(d))
+    else:
+        raise ValueError(f"unknown probe family {family!r}")
+    return u.reshape(n, *batch_shape, d)
diff --git a/src/zbp_scaling/zbp/rules.py b/src/zbp_scaling/zbp/rules.py
new file mode 100644
index 0000000..d556201
--- /dev/null
+++ b/src/zbp_scaling/zbp/rules.py
@@ -0,0 +1,264 @@
+"""Textbook BP-free learning rules on the physical/digital partition.
+
+Every rule is expressed in the same terms as ZBP: it decides which error signal reaches each physical block's
+OUTPUT and what (if anything) is propagated to its INPUT; block parameters always update locally from the
+incoming error (in-block autograd = the physical-equivalent local rule) and digital layers use autograd.
+
+  drtp    Direct Random Target Projection (Frenkel, Lefebvre, Bol 2021): DFA whose feedback signal is the
+          (negated) one-hot target instead of the output error.  Handled inside ZBPFunction (mode 'drtp').
+  pepita  PEPITA (Dellaferrera & Kreiman 2022): a second forward pass with the network input modulated by a
+          fixed random projection of the output error; every parameterised unit's error is its own
+          (clean - modulated) output, gradients are taken on the modulated pass, nothing propagates.
+  wm      weight mirror (Akrout et al. 2019): textbook FA whose feedback matrices are learned in a mirror
+          phase from input noise / output readout of each linear layer (B -> W, so FA -> BP as B converges).
+  dtp     difference target propagation (Lee et al. 2015), MLP / residual MLP: learned inverses g_l propagate
+          targets, each block minimises |f_l(h_{l-1}) - t_l|^2 locally.
+  ff      forward-forward (Hinton 2022), MLP / residual MLP: per-block goodness on positive / negative data
+          (label embedded in the input), inputs length-normalised between blocks, no propagation.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .autograd import ZBPBlock, zbp_blocks, DFA
+from .fa import FALinear, FAConv2d
+
+
+# ----------------------------------------------------------------------------------------------- units
+def _inside(module, containers):
+    return any(m is module for m in containers)
+
+
+def learning_units(model):
+    """PEPITA units: every ZBPBlock (a physical block is one unit: only its output is observable) plus every
+    parameterised leaf module that is not inside a ZBPBlock."""
+    blocks = zbp_blocks(model)
+    inside = set()
+    for b in blocks:
+        for m in b.modules():
+            if m is not b:
+                inside.add(id(m))
+    units = []
+    for m in model.modules():
+        if isinstance(m, ZBPBlock):
+            units.append(m)
+        elif id(m) not in inside and not isinstance(m, ZBPBlock) and any(True for _ in m.parameters(recurse=False)):
+            units.append(m)
+    return units
+
+
+# ----------------------------------------------------------------------------------------------- PEPITA
+class Pepita:
+    """Holds the fixed random projection F of the output error onto the network input (or the embedding)."""
+
+    def __init__(self, n_out, in_shape, device, scale=0.05, seed=0, embed=None):
+        g = torch.Generator(device="cpu"); g.manual_seed(seed)
+        n_in = math.prod(in_shape)
+        # Dellaferrera & Kreiman: F ~ U(-a, a), a = sqrt(6 / (n_in + n_out)) * scale-ish; we use their 0.05 default
+        a = math.sqrt(6.0 / (n_in + n_out))
+        self.F = ((torch.rand(n_out, n_in, generator=g) * 2 - 1) * a * scale).to(device)
+        self.in_shape = tuple(in_shape)
+        self.embed = embed          # nn.Embedding to modulate instead of the raw input (language models)
+        self._mod = None
+        if embed is not None:
+            embed.register_forward_hook(self._embed_hook)
+
+    def _embed_hook(self, m, inp, out):
+        return out if self._mod is None else out - self._mod
+
+    def step(self, model, x, y, n_out, blocks_mode="zero"):
+        """One PEPITA step: returns the clean loss.  Leaves .grad on every parameter (gradient-like quantities)."""
+        units = learning_units(model)
+        blocks = zbp_blocks(model)
+        modes = [b.cfg for b in blocks]
+        for b in blocks:
+            b.cfg = b.cfg.replace(mode=blocks_mode)          # no propagation through physical blocks
+        # clean pass
+        outs_clean = {}
+        hs = [u.register_forward_hook(lambda m, i, o, u=u: outs_clean.__setitem__(id(u), o.detach())) for u in units]
+        with torch.no_grad():
+            logits = model(x)
+        for h in hs:
+            h.remove()
+        if logits.dim() == 3:                                # LM: per-token error
+            e = logits.softmax(-1) - F.one_hot(y, n_out).to(logits.dtype)
+            e = e / (logits.shape[0] * logits.shape[1])
+        else:
+            e = (logits.softmax(-1) - F.one_hot(y, n_out).to(logits.dtype)) / logits.shape[0]
+        loss = F.cross_entropy(logits.reshape(-1, n_out), y.reshape(-1))
+        # modulated pass: input (or embedding) minus F e, each unit's input detached so nothing chains
+        if self.embed is None:
+            x_mod = x - (e.reshape(-1, n_out) @ self.F).reshape(x.shape) * x.shape[0]   # undo the 1/B in e
+        else:
+            self._mod = (e @ self.F).reshape(*e.shape[:-1], -1) * (e.shape[0] * e.shape[1])
+            x_mod = x
+        outs_mod = {}
+        def pre(m, inp):
+            return tuple(t.detach().requires_grad_(t.is_floating_point()) if torch.is_tensor(t) else t for t in inp)
+        hs = [u.register_forward_pre_hook(pre) for u in units]
+        hs += [u.register_forward_hook(lambda m, i, o, u=u: outs_mod.__setitem__(id(u), o)) for u in units]
+        logits_mod = model(x_mod)
+        for h in hs:
+            h.remove()
+        self._mod = None
+        # inject: last unit (logits) gets the true error, every other unit its activation difference
+        tensors, grads = [], []
+        last = next((u for u in units if outs_mod.get(id(u)) is logits_mod), units[-1])
+        n_units = logits.shape[0] * (logits.shape[1] if logits.dim() == 3 else 1)   # batch (x tokens) average, as for e
+        for u in units:
+            o = outs_mod.get(id(u))
+            if o is None or not torch.is_tensor(o) or not o.requires_grad:
+                continue
+            g = e if u is last else (outs_clean[id(u)] - o.detach()) / n_units
+            tensors.append(o); grads.append(g)
+        torch.autograd.backward(tensors, grads)
+        for b, c in zip(blocks, modes):
+            b.cfg = c
+        return loss.detach()
+
+
+# ----------------------------------------------------------------------------------------------- weight mirror
+@torch.no_grad()
+def mirror_update(model, rate=0.01, n_noise=256, gen=None):
+    """Akrout et al. 2019 mirror phase: for every FA layer, drive the linear part with zero-mean noise xi, read
+    y = W xi, and update B <- B + rate * (y^T xi / n) - rate * B, whose fixed point is B = W."""
+    for m in model.modules():
+        if isinstance(m, FALinear):
+            xi = torch.randn(n_noise, m.in_features, generator=gen, device=m.weight.device, dtype=m.weight.dtype)
+            y = xi @ m.weight.t()
+            m.B.mul_(1 - rate).add_(rate * (y.t() @ xi) / n_noise)
+        elif isinstance(m, FAConv2d):
+            k = m.in_channels * m.kernel_size[0] * m.kernel_size[1]
+            xi = torch.randn(n_noise, k, generator=gen, device=m.weight.device, dtype=m.weight.dtype)
+            y = xi @ m.weight.flatten(1).t()
+            m.B.mul_(1 - rate).add_(rate * (y.t() @ xi).reshape(m.B.shape) / n_noise)
+
+
+# ----------------------------------------------------------------------------------------------- DTP
+class DTP:
+    """Difference target propagation on a chain of ZBP blocks (MLP / residual MLP) with a digital readout.
+    Inverses g_l (one per block, residual form h + V2 tanh(V1 h)) are trained on noisy forward passes of the block."""
+
+    def __init__(self, model, device, sigma=0.1, eta_hat=0.5, lr_inv=1e-3):
+        self.blocks = [b for b in zbp_blocks(model) if b.estimate_input_grad]   # blocks that need a target below them
+        self.first = [b for b in zbp_blocks(model) if not b.estimate_input_grad]
+        self.model, self.sigma, self.eta_hat = model, sigma, eta_hat
+        self.inv = nn.ModuleDict()
+        self.dims = {}
+        self.opt_inv = None
+        self.device, self.lr_inv = device, lr_inv
+
+    def _inverse(self, b, d_out, d_in):
+        k = b.name
+        if k not in self.inv:
+            g = nn.Sequential(nn.Linear(d_out, d_in), nn.Tanh(), nn.Linear(d_in, d_in)).to(self.device)
+            nn.init.zeros_(g[2].weight); nn.init.zeros_(g[2].bias)
+            self.inv[k] = g
+            self.opt_inv = torch.optim.Adam(self.inv.parameters(), lr=self.lr_inv)
+        g = self.inv[k]
+        return lambda h: h[..., :d_in] + g(h) if d_out >= d_in else g(h)
+
+    def step(self, x, y):
+        model = self.model
+        caps = {}
+        hs = [b.register_forward_hook(lambda m, i, o, b=b: caps.__setitem__(b.name, (i[0].detach(), o.detach()))) for b in zbp_blocks(model)]
+        with torch.no_grad():
+            hs_all = model(x)
+        for h in hs:
+            h.remove()
+        # digital readout: exact gradient w.r.t. its parameters and w.r.t. the top block's output
+        top = zbp_blocks(model)[-1]
+        h_top = caps[top.name][1].clone().requires_grad_(True)
+        logits = model.readout(h_top)
+        loss = F.cross_entropy(logits, y)
+        loss.backward()                                       # readout params get .grad; h_top.grad = dL/dh_top
+        t = h_top.detach() - self.eta_hat * h_top.grad
+        # targets downward, local losses, inverse training
+        order = zbp_blocks(model)
+        inv_loss_total = 0.0
+        for b in reversed(order):
+            x_in, h_out = caps[b.name]
+            # local loss for the block: |f(x_in) - t|^2  (gradient w.r.t. block parameters only)
+            with torch.enable_grad():
+                out = b(x_in.detach())
+                l_loc = 0.5 * ((out - t) ** 2).sum() / x_in.shape[0]
+                grads = torch.autograd.grad(l_loc, [p for p in b.parameters() if p.requires_grad], allow_unused=True)
+            for p, g in zip([p for p in b.parameters() if p.requires_grad], grads):
+                if g is not None:
+                    p.grad = g if p.grad is None else p.grad + g
+            if not b.estimate_input_grad:
+                break                                         # first block: nothing below
+            g_inv = self._inverse(b, h_out.shape[-1], x_in.shape[-1])
+            # inverse training on a noisy forward pass of the block (one extra physical query)
+            with torch.no_grad():
+                xn = x_in + self.sigma * x_in.std() * torch.randn_like(x_in)
+                hn = b(xn)
+            with torch.enable_grad():
+                l_inv = ((g_inv(hn) - xn) ** 2).mean()
+                self.opt_inv.zero_grad(set_to_none=True)
+                l_inv.backward()
+                self.opt_inv.step()
+            inv_loss_total += l_inv.item()
+            with torch.no_grad():                             # difference target for the block below
+                t = x_in + g_inv(t) - g_inv(h_out)
+        self.last_inv_loss = inv_loss_total
+        return loss.detach()
+
+
+# ----------------------------------------------------------------------------------------------- Forward-Forward
+class FF:
+    """Hinton 2022 on a chain of ZBP blocks: positive = input with the true label embedded in its first n_out
+    coordinates, negative = a wrong label; per block loss softplus(-(G_pos - theta)) + softplus(G_neg - theta) with
+    goodness G = mean(h^2); block inputs are length-normalised; prediction = label with the largest summed goodness
+    of all blocks after the first."""
+
+    def __init__(self, model, n_out=10, theta=2.0, label_scale=1.0):
+        self.model, self.n_out, self.theta, self.label_scale = model, n_out, theta, label_scale
+
+    def embed(self, x, labels):
+        x = x.flatten(1).clone()
+        x[:, :self.n_out] = F.one_hot(labels, self.n_out).to(x.dtype) * self.label_scale
+        return x
+
+    @staticmethod
+    def norm(h):
+        return h / (h.norm(dim=-1, keepdim=True) + 1e-4) * math.sqrt(h.shape[-1])
+
+    def goodness_chain(self, x, train=False):
+        """Returns the list of per-block goodness values; with train=True also leaves gradients (FF losses)."""
+        h = x
+        goods = []
+        for b in zbp_blocks(self.model):
+            h_in = self.norm(h).detach()
+            out = b(h_in) if train else b(h_in).detach()
+            goods.append(out.pow(2).mean(-1))
+            h = out.detach()
+        return goods
+
+    def step(self, x, y):
+        blocks = zbp_blocks(self.model)
+        x_pos = self.embed(x, y)
+        wrong = (y + torch.randint(1, self.n_out, y.shape, device=y.device)) % self.n_out
+        x_neg = self.embed(x, wrong)
+        h_pos, h_neg = x_pos, x_neg
+        total = 0.0
+        for b in blocks:
+            hp, hn = self.norm(h_pos).detach(), self.norm(h_neg).detach()
+            with torch.enable_grad():
+                op, on = b(hp), b(hn)
+                gp, gn = op.pow(2).mean(-1), on.pow(2).mean(-1)
+                loss = (F.softplus(-(gp - self.theta)) + F.softplus(gn - self.theta)).mean()
+                loss.backward()
+            total += loss.item()
+            h_pos, h_neg = op.detach(), on.detach()
+        return torch.tensor(total / len(blocks))
+
+    @torch.no_grad()
+    def predict(self, x):
+        scores = []
+        for c in range(self.n_out):
+            xc = self.embed(x, torch.full((x.shape[0],), c, device=x.device, dtype=torch.long))
+            goods = self.goodness_chain(xc)
+            scores.append(torch.stack(goods[1:] if len(goods) > 1 else goods).sum(0))
+        return torch.stack(scores, 1).argmax(1)
diff --git a/src/zbp_scaling/zbp/wp.py b/src/zbp_scaling/zbp/wp.py
new file mode 100644
index 0000000..ff2511d
--- /dev/null
+++ b/src/zbp_scaling/zbp/wp.py
@@ -0,0 +1,28 @@
+"""Weight perturbation (parameter-space zeroth order) baseline: MeZO / CD-RGE / weight-space forward
+gradient.  All parameters are perturbed by eps*u (Rademacher), the loss is measured at +/- and the
+gradient estimate is (1/n) sum_i u_i (L(theta+eps u_i) - L(theta-eps u_i)) / (2 eps).
+2n forward passes of the whole network per step, no backward pass, variance ~ P/n with P = #params."""
+import torch
+
+
+def wp_step(model, loss_fn, x, y, n, eps, gen):
+    params = [p for p in model.parameters() if p.requires_grad]
+    grads = [torch.zeros_like(p) for p in params]
+    with torch.no_grad():
+        for i in range(n):
+            us = [(torch.randint(0, 2, p.shape, generator=gen, device=p.device).to(p.dtype) * 2 - 1) for p in params]
+            for p, u in zip(params, us):
+                p.add_(eps * u)
+            lp = loss_fn(model(x), y).item()
+            for p, u in zip(params, us):
+                p.sub_(2 * eps * u)
+            lm = loss_fn(model(x), y).item()
+            for p, u in zip(params, us):
+                p.add_(eps * u)
+            D = (lp - lm) / (2 * eps)
+            for g, u in zip(grads, us):
+                g.add_(u, alpha=D / n)
+        for p, g in zip(params, grads):
+            p.grad = g
+        loss = loss_fn(model(x), y)
+    return loss