diff options
Diffstat (limited to 'tools/contrast_floor_check.py')
| -rw-r--r-- | tools/contrast_floor_check.py | 168 |
1 files changed, 168 insertions, 0 deletions
diff --git a/tools/contrast_floor_check.py b/tools/contrast_floor_check.py new file mode 100644 index 0000000..bec72ad --- /dev/null +++ b/tools/contrast_floor_check.py @@ -0,0 +1,168 @@ +"""Contrast readout floor: a standalone check for equilibrium-propagation-style code. + +WHAT THIS DETECTS +----------------- +Learning rules in the equilibrium-propagation / coupled-learning family read the weight +update from a *contrast*: a small perturbation d applied on top of a much larger state o. +A very common implementation forms the perturbed state and then recovers the perturbation +by subtracting the state back out: + + z = o + d # perturbed state, stored in float32 + contrast = 0.5 * (z - o) ** 2 # read the contrast, differentiate w.r.t. parameters + +Algebraically z - o == d. In floating point it does not. Forming z rounds away every +component of d below eps * |o| (absorption). The subtraction z - o is then exact, so it +faithfully returns the *already damaged* copy of d and raises no error. The damage grows as +the displacement-to-activation ratio falls, which happens as models get wider or deeper, so +it imitates an algorithmic scaling wall. + +The fix costs nothing: read the contrast from the stored perturbation itself. The linear +form -<d, o> has the same derivative with respect to the parameters and never touches the +rounded copy. + +HOW TO USE +---------- +1. Zero-setup demo of the mechanism and the fix: + + python contrast_floor_check.py + +2. Check your own tensors. Pass the state and the perturbation you actually use: + + from contrast_floor_check import check_pair + check_pair(o, d, name="block3") + +3. Scan a whole model. Collect (state, perturbation) pairs per layer and call: + + from contrast_floor_check import report + report({"block0": (o0, d0), "block1": (o1, d1), ...}) + +INTERPRETATION +-------------- +The headline quantity is the displacement ratio RMS(d) / RMS(o). Treat it as a screening +number rather than a verdict, for two reasons. + +Two thresholds, worth separating. The recovered perturbation first picks up error that is +close to orthogonal (`recovered_rel_err` rises while `recovered_projection` stays near one). +That part behaves like noise and averages over steps. Only further down does the surviving +perturbation shrink along its own direction (`recovered_projection` falls), and that part is +a consistent shortfall which accumulates across training. The projection is the number to +watch, and it turns later than the error does. + +Synthetic pairs underestimate. Running this on random tensors at a given RMS ratio gives a +lower bound on the damage a real implementation suffers at the same ratio, because a real +system re-forms the perturbed state several times per step and transmits the perturbation +down a chain that has already been rounded, so the loss compounds, and because what matters +is the share of the *update-carrying* mass sitting under the grid rather than the share of +elements. In our transformer language models the per-layer onset sat near 5e-7 in float32, +about four times epsilon, and layers an order of magnitude below that were missing most of +their contribution to the parameter update, far more than a matched synthetic pair loses. +Measure with your own tensors, and if you can, measure the parameter update itself against a +double precision or backprop reference. + +This file is standalone: only PyTorch is required, and it is released for anyone to run +against their own implementation. +""" + +import torch + + +def _rms(t): + return float(t.detach().float().pow(2).mean().sqrt()) + + +def check_pair(o, d, name="", dtype=None, verbose=True): + """Measure how much of the perturbation d survives being added to the state o. + + o state tensor the perturbation is applied on top of + d perturbation tensor (the quantity the contrast is supposed to carry) + dtype storage dtype to test; defaults to o's dtype + + Returns a dict with the displacement ratio, the fraction of elements below epsilon, + and the relative error of the recovered perturbation against the exact one. + """ + o = o.detach() + d = d.detach() + st = dtype or o.dtype + eps = torch.finfo(st).eps + + o64, d64 = o.double(), d.double() + ratio = _rms(d64) / max(_rms(o64), 1e-300) + + # elementwise exposure: how many components sit under the storage resolution + rel = (d64.abs() / o64.abs().clamp_min(1e-300)).flatten() + frac_below = float((rel < eps).float().mean()) + q10 = float(rel.kthvalue(max(1, int(0.10 * rel.numel()))).values) + + # what the subtract-back-out readout actually returns + z = (o.to(st) + d.to(st)) + d_recovered = (z - o.to(st)).double() + err = float((d_recovered - d64).norm() / max(float(d64.norm()), 1e-300)) + + # the damage that matters is the signed part along the true perturbation + proj = float((d_recovered * d64).sum() / d64.pow(2).sum().clamp_min(1e-300)) + + out = dict(name=name, dtype=str(st), eps=eps, disp_ratio=ratio, + frac_below_eps=frac_below, ratio_q10=q10, + recovered_rel_err=err, recovered_projection=proj) + if verbose: + flag = "OK " if err < 1e-3 else ("WARN" if err < 0.1 else "DAMAGED") + print(f"[{flag}] {name or 'pair':<16} ratio {ratio:.2e} " + f"below-eps {100*frac_below:5.1f}% q10 {q10:.2e} " + f"recovered err {err:.3f} projection {proj:.3f}") + return out + + +def report(pairs, dtype=None): + """Run check_pair over a dict of {name: (state, perturbation)} and summarize.""" + eps = torch.finfo(dtype or torch.float32).eps + print(f"contrast readout floor check storage eps = {eps:.2e}") + print(f"{'':18}{'ratio':>10}{'<eps':>10}{'err':>9}{'proj':>8}") + rows = [check_pair(o, d, name=k, dtype=dtype, verbose=False) for k, (o, d) in pairs.items()] + for r in rows: + flag = "ok" if r["recovered_rel_err"] < 1e-3 else ("warn" if r["recovered_rel_err"] < 0.1 else "DAMAGED") + print(f"{r['name']:<18}{r['disp_ratio']:>10.2e}{100*r['frac_below_eps']:>9.1f}%" + f"{r['recovered_rel_err']:>9.3f}{r['recovered_projection']:>8.3f} {flag}") + worst = max(rows, key=lambda r: r["recovered_rel_err"]) + if worst["recovered_rel_err"] >= 0.1: + print(f"\nThe readout loses a large part of the perturbation, worst at " + f"'{worst['name']}' with {100*(1-worst['recovered_projection']):.0f}% of it missing. " + f"Read the contrast from the stored perturbation instead of differencing the states.") + elif worst["recovered_rel_err"] >= 1e-3: + print(f"\nSome loss at '{worst['name']}'. Watch this as models get wider, since the " + f"displacement ratio falls with width.") + else: + print("\nNo material loss at these magnitudes.") + return rows + + +def _demo(): + torch.manual_seed(0) + print(__doc__.split("HOW TO USE")[0].strip()[:0] or "", end="") + print("Demonstration: one state, perturbations spanning eight decades.\n") + o = torch.randn(4096, 512) + pairs = {} + for k in range(4, 12): + scale = 10.0 ** (-k) + pairs[f"ratio 1e-{k}"] = (o, torch.randn_like(o) * scale) + report(pairs, dtype=torch.float32) + + print("\nSame perturbations, but the contrast is read from the stored tensor.") + print("This is the fix: the derivative is identical and nothing is rounded away.\n") + o_ = o.float() + for k in (4, 8, 11): + d = torch.randn_like(o) * (10.0 ** (-k)) + # damaged path: differentiate 0.5*||z-o||^2 -> cotangent is (z-o) + cot_bad = ((o_ + d.float()) - o_).double() + # fixed path: differentiate -<d,o> -> cotangent is d itself + cot_good = d.double() + ref = d.double() + e_bad = float((cot_bad - ref).norm() / ref.norm()) + e_good = float((cot_good - ref).norm() / ref.norm()) + print(f" ratio 1e-{k:<3} subtract-back-out error {e_bad:8.3f} stored-perturbation error {e_good:.3e}") + + print("\nIf your code differences two large states to obtain a small contrast, the left " + "column is what your update carries.") + + +if __name__ == "__main__": + _demo() |