1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
"""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
|
