blob: 4fa37479f2c4cc813d95ec0a739eb4a305030346 (
plain)
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
|
"""Branch-gain profile rho_k and the NSR constant c(scale) = L * mean(rho) (THEORY.md T2/corollary).
Validation-only: uses the simulation autograd for exact J_F^T v."""
import torch
import torch.nn as nn
from .zbp import zbp_blocks, ZBPConfig
@torch.enable_grad()
def rho_profile(model, x, y):
blocks = zbp_blocks(model)
cfgs = [b.cfg for b in blocks]
for b in blocks:
b.cfg = b.cfg.replace(mode="exact")
b.capture = True
model.zero_grad(set_to_none=True)
nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
out = {}
for b, c in zip(blocks, cfgs):
b.cfg = c
b.capture = False
b.captured = None
# rho via a recorded pass: |J_F^T v| / |v| per block from the Recorder
from .zbp import Recorder
for b in blocks:
b.cfg = b.cfg.replace(mode="zero")
with Recorder() as rec:
model.zero_grad(set_to_none=True)
nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
for b, c in zip(blocks, cfgs):
b.cfg = c
model.zero_grad(set_to_none=True)
rows = rec.summary()
for name, r in rows.items():
out[name] = (r["gnorm"] / max(r["vnorm"], 1e-30)) ** 2 if "vnorm" in r else None
return out
|
