File size: 3,958 Bytes
2504006 | 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 | """Claim 5: for a deep linear unconstrained-features model (UFM) with K = 3
classes, the loss Hessian has K^2 = 9 outlier eigenvalues separating from the
bulk.
The logbook states the figures remain source-reported with no independent
training run. A UFM needs no external data -- the features are free variables --
so it is trained here and the exact Hessian spectrum is computed.
Deep linear UFM: logits = W_L ... W_1 H, with H the free feature matrix
(d x N) and W_l the layer weights. Loss = MSE to one-hot targets + weight decay.
The Hessian is formed exactly by autograd over ALL parameters and diagonalised.
"""
import json, numpy as np, torch
RES = {}
torch.set_default_dtype(torch.float64)
def build(K=3, n=8, d=6, depth=2, seed=0):
g = torch.Generator().manual_seed(seed)
N = K*n
Y = torch.zeros(K, N)
for c in range(K): Y[c, c*n:(c+1)*n] = 1.0
H = (torch.randn(d, N, generator=g)*0.3).requires_grad_(True)
Ws = []
dims = [K]+[d]*(depth-1)+[d]
for i in range(depth):
Ws.append((torch.randn(dims[i], dims[i+1], generator=g)*0.3).requires_grad_(True))
return H, Ws, Y
def loss_fn(H, Ws, Y, wd=5e-4):
Z = H
for W in reversed(Ws): Z = W @ Z
l = ((Z-Y)**2).mean()
for W in Ws: l = l+wd*(W**2).sum()
return l+wd*(H**2).sum()
def train(H, Ws, Y, steps=4000, lr=0.05):
params = [H]+Ws
opt = torch.optim.Adam(params, lr=lr)
for _ in range(steps):
opt.zero_grad(); l = loss_fn(H, Ws, Y); l.backward(); opt.step()
return float(l.item())
def hessian_spectrum(H, Ws, Y):
params = [H]+Ws
flat = torch.cat([p.reshape(-1) for p in params]).detach()
n = flat.numel()
def f(v):
i = 0; ps = []
for p in params:
k = p.numel(); ps.append(v[i:i+k].view_as(p)); i += k
Z = ps[0]
for W in reversed(ps[1:]): Z = W @ Z
l = ((Z-Y)**2).mean()
for W in ps[1:]: l = l+5e-4*(W**2).sum()
return l+5e-4*(ps[0]**2).sum()
Hs = torch.autograd.functional.hessian(f, flat)
ev = torch.linalg.eigvalsh((Hs+Hs.T)/2).numpy()
return np.sort(ev)[::-1]
def count_outliers(ev, K):
"""Outliers = eigenvalues above the largest gap in the top of the spectrum."""
top = ev[:4*K*K]
gaps = top[:-1]-top[1:]
j = int(np.argmax(gaps))
return j+1, float(gaps[j]), float(top[j]), float(top[j+1])
def main():
rows = []
for K in (2, 3, 4):
for depth in (2, 3):
H, Ws, Y = build(K=K, depth=depth, seed=K*10+depth)
fin = train(H, Ws, Y)
ev = hessian_spectrum(H, Ws, Y)
n_out, gap, above, below = count_outliers(ev, K)
rows.append({"K": K, "depth": depth, "final_loss": round(fin, 8),
"params": int(sum(p.numel() for p in [H]+Ws)),
"K_squared": K*K, "outliers_found": n_out,
"matches_K2": bool(n_out == K*K),
"gap_size": round(gap, 6),
"last_outlier": round(above, 6), "first_bulk": round(below, 6),
"separation_ratio": round(above/max(below, 1e-12), 2)})
print(" K=%d depth=%d loss=%.2e params=%d K^2=%d outliers found=%d match=%s gap=%.4f (%.4f -> %.4f, ratio %.1fx)"
% (K, depth, fin, rows[-1]["params"], K*K, n_out, n_out == K*K,
gap, above, below, rows[-1]["separation_ratio"]), flush=True)
RES["claim5_ufm_hessian"] = {"rows": rows,
"matches_in": sum(r["matches_K2"] for r in rows), "cells": len(rows),
"min_separation_ratio": min(r["separation_ratio"] for r in rows)}
R = RES["claim5_ufm_hessian"]
print(" outlier count equals K^2 in %d/%d cells; min separation ratio %.1fx"
% (R["matches_in"], R["cells"], R["min_separation_ratio"]), flush=True)
json.dump(RES, open("ufm_results.json", "w"), indent=1)
if __name__ == "__main__":
main(); print("DONE")
|