summaryrefslogtreecommitdiff
path: root/tools/contrast_floor_check.py
blob: bec72adcf2313aea376d961746cbc82753c1d78e (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
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
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()