| |
| |
| """Toy collapse benchmark (CPU, seconds). NOT evidence about large models — a reproducible |
| illustration of the failure mode and of what the spring does about it. |
| |
| Task: h = x + g * branch(x), target y = x + A x. The branch is needed, but its random init |
| is harmful, so early on the loss pushes a free gate towards zero; with plain SGD a closed |
| gate also starves the branch of gradient. Variants: |
| open : g = 1 (no gate) — reference |
| free : g = sigmoid(beta * (s - theta)), theta learned, no floor, no spring |
| tether : TetherGate (floor 1.0 for `ramp_start` steps, then decay; spring towards g_target) |
| s = mean |branch(x)| per token (detached), the same for both gated variants. |
| |
| python demo_collapse.py [--steps 3000] [--seeds 0 1 2] |
| """ |
| import argparse |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from tethergate import TetherGate, tethergate_step |
|
|
|
|
| def run(variant, seed, steps, d=32, lr=0.05, beta=6.0, init_scale=3.0, ramp_start=300, tau=300.0, |
| g_target=0.5, wd=0.5): |
| torch.manual_seed(seed) |
| A = torch.randn(d, d) * 0.5 / d ** 0.5 |
| branch = nn.Linear(d, d, bias=False) |
| with torch.no_grad(): |
| branch.weight.mul_(init_scale) |
| theta = nn.Parameter(torch.tensor(0.0)) |
| gate = TetherGate(beta=beta, g_target=g_target) if variant == 'tether' else None |
| params = list(branch.parameters()) + ([theta] if variant == 'free' else []) + \ |
| (list(gate.parameters()) if gate is not None else []) |
| opt = torch.optim.SGD(params, lr=lr) |
| gen = torch.Generator().manual_seed(1000 + seed) |
| log = [] |
| for step in range(steps): |
| x = torch.randn(256, d, generator=gen) |
| u = branch(x) |
| s = u.detach().abs().mean(-1, keepdim=True) |
| if variant == 'open': |
| g = torch.ones_like(s) |
| elif variant == 'free': |
| g = torch.sigmoid((beta * (s - theta)).clamp(-8.0, 8.0)) |
| else: |
| g = gate(s) |
| loss = ((x + g * u) - (x + x @ A.t())).pow(2).mean() |
| opt.zero_grad(); loss.backward(); opt.step() |
| if gate is not None: |
| tethergate_step([gate], step, lr=lr, wd=wd, kappa_sleep=0.05, ramp_start=ramp_start, tau=tau) |
| if step % max(1, steps // 6) == 0 or step == steps - 1: |
| log.append((step, loss.item(), g.mean().item())) |
| with torch.no_grad(): |
| x = torch.randn(4096, d, generator=torch.Generator().manual_seed(77)) |
| u = branch(x) |
| s = u.abs().mean(-1, keepdim=True) |
| if variant == 'open': |
| g = torch.ones_like(s) |
| elif variant == 'free': |
| g = torch.sigmoid((beta * (s - theta)).clamp(-8.0, 8.0)) |
| else: |
| gate.eval(); g = gate(s) |
| held = float(((x + g * u) - (x + x @ A.t())).pow(2).mean()) |
| return held, float(g.mean()), log |
|
|
|
|
| if __name__ == '__main__': |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--steps', type=int, default=3000) |
| ap.add_argument('--seeds', type=int, nargs='+', default=[0, 1, 2]) |
| ap.add_argument('--verbose', action='store_true') |
| a = ap.parse_args() |
| print(f'steps={a.steps} seeds={a.seeds} (held-out MSE; mean gate on held-out data)') |
| for variant in ('open', 'free', 'tether'): |
| for seed in a.seeds: |
| held, gm, log = run(variant, seed, a.steps) |
| print(f' {variant:6s} seed {seed}: held-out MSE {held:.5f} | mean gate {gm:.3f}') |
| if a.verbose: |
| for st, l, g in log: |
| print(f' step {st:5d} loss {l:.5f} gate {g:.3f}') |
|
|
