Fractional Mahalanobis radial head / experiment.py
Mechanism confirmed, baseline not beaten
1import json, math, random
2from pathlib import Path
3import numpy as np
4import torch
5from torch import nn
6from torch.utils.data import TensorDataset, DataLoader
7
8SEED = 464
9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
10torch.set_num_threads(4)
11DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
12
13# ---------- core fractional Mahalanobis radial head ----------
14class FractionalRadialHead(nn.Module):
15 def __init__(self, dim, exponents=(.25, .5, 1., 1.5, 2.), eps=1e-3):
16 super().__init__()
17 self.dim, self.exponents, self.eps = dim, tuple(exponents), eps
18 self.mu = nn.Parameter(torch.zeros(2, dim))
19 self.raw_diag = nn.Parameter(torch.zeros(2, dim))
20 self.offdiag = nn.Parameter(torch.zeros(2, dim, dim))
21 self.b = nn.Parameter(torch.zeros(len(exponents)))
22 self.bias = nn.Parameter(torch.tensor(0.0))
23
24 def factors(self):
25 L = torch.tril(self.offdiag, diagonal=-1)
26 diag = torch.nn.functional.softplus(self.raw_diag) + 1e-3
27 return L + torch.diag_embed(diag)
28
29 def covariances(self):
30 L = self.factors()
31 return L @ L.transpose(-1, -2) + self.eps * torch.eye(self.dim, device=L.device)
32
33 def radii(self, z):
34 # Cholesky(Sigma) followed by triangular solves avoids an explicit inverse.
35 S = self.covariances()
36 C = torch.linalg.cholesky(S)
37 out = []
38 for c in range(2):
39 delta = z - self.mu[c]
40 v = torch.linalg.solve_triangular(C[c], delta.T, upper=False).T
41 out.append((v*v).sum(1))
42 return torch.stack(out, 1)
43
44 def forward(self, z):
45 r = self.radii(z).clamp_min(1e-8)
46 powers = torch.stack([r ** a for a in self.exponents], dim=-1)
47 h = (powers * self.b).sum(-1)
48 # b0 cancels. The remaining constant is the determinant/prior term.
49 logdet = torch.linalg.slogdet(self.covariances())[1]
50 return h[:, 1] - h[:, 0] + self.bias + .5 * (logdet[0] - logdet[1])
51
52
53def init_radial(head, x, y):
54 with torch.no_grad():
55 for c in (0, 1):
56 q = x[y == c]
57 head.mu[c].copy_(q.mean(0))
58 cov = torch.cov(q.T) + .05 * torch.eye(x.shape[1], device=x.device)
59 L = torch.linalg.cholesky(cov)
60 head.raw_diag[c].copy_(torch.log(torch.expm1(torch.diag(L).clamp_min(.002))))
61 head.offdiag[c].copy_(torch.tril(L, diagonal=-1))
62 # Start close to a Gaussian/QDA radial coefficient and let fractional terms adapt.
63 head.b.zero_(); head.b[2] = -.5
64
65class LinearHead(nn.Module):
66 def __init__(self, d):
67 super().__init__(); self.w = nn.Linear(d, 1)
68 def forward(self, z): return self.w(z).squeeze(1)
69
70class MLPHead(nn.Module):
71 def __init__(self, d):
72 super().__init__(); self.net = nn.Sequential(nn.Linear(d, 8), nn.Tanh(), nn.Linear(8, 1))
73 def forward(self, z): return self.net(z).squeeze(1)
74
75class QDAHead(nn.Module):
76 def __init__(self, d):
77 super().__init__(); self.d = d
78 def fit(self, z, y):
79 with torch.no_grad():
80 self.mu0 = z[y < .5].mean(0); self.mu1 = z[y >= .5].mean(0)
81 self.s0 = torch.cov(z[y < .5].T) + .05*torch.eye(self.d, device=z.device)
82 self.s1 = torch.cov(z[y >= .5].T) + .05*torch.eye(self.d, device=z.device)
83 self.i0 = torch.linalg.inv(self.s0); self.i1 = torch.linalg.inv(self.s1)
84 self.const = .5*(torch.linalg.slogdet(self.s0)[1]-torch.linalg.slogdet(self.s1)[1])
85 def forward(self, z):
86 a=z-self.mu0; b=z-self.mu1
87 r0=(a @ self.i0 * a).sum(1); r1=(b @ self.i1 * b).sum(1)
88 return -.5*r1 + .5*r0 + self.const
89
90# ---------- data and training ----------
91def make_data(kind, n=5000):
92 y = np.random.randint(0, 2, n)
93 means = np.where(y[:, None] == 0, [-1.15, 0.0], [1.15, 0.0])
94 # Slightly distinct anisotropic class geometry makes the radius distinction visible.
95 scales = np.where(y[:, None] == 0, [1.0, .65], [1.0, .65])
96 if kind == 'student':
97 df = 2.5
98 raw = np.random.randn(n, 2) / np.sqrt(np.random.chisquare(df, n)[:, None] / df)
99 else:
100 raw = np.random.randn(n, 2)
101 x = means + raw * scales
102 # fixed stratified-ish split by random permutation
103 p = np.random.RandomState(SEED + (1 if kind == 'student' else 2)).permutation(n)
104 cut = int(.7*n)
105 return (torch.tensor(x[p[:cut]], dtype=torch.float32, device=DEVICE),
106 torch.tensor(y[p[:cut]], dtype=torch.float32, device=DEVICE),
107 torch.tensor(x[p[cut:]], dtype=torch.float32, device=DEVICE),
108 torch.tensor(y[p[cut:]], dtype=torch.float32, device=DEVICE))
109
110def fit(model, x, y, radial=False, epochs=180):
111 model.to(DEVICE); model.train()
112 if radial: init_radial(model, x, y.long())
113 opt = torch.optim.Adam(model.parameters(), lr=.025 if radial else .01, weight_decay=1e-4)
114 for _ in range(epochs):
115 opt.zero_grad(); logits = model(x)
116 loss = nn.functional.binary_cross_entropy_with_logits(logits, y)
117 if radial:
118 L = model.factors()
119 # modest covariance and coefficient regularization for stable learned geometry
120 loss = loss + 1e-5*(L*L).sum() + 1e-5*(model.b*model.b).sum()
121 loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
122 model.eval()
123 with torch.no_grad():
124 logits = model(x_test_global)
125 yy = y_test_global
126 nll = nn.functional.binary_cross_entropy_with_logits(logits, yy).item()
127 acc = ((logits > 0) == (yy > .5)).float().mean().item()
128 return nll, acc
129
130# Exact Student-t radial log generator and basis approximation sanity check.
131def math_check():
132 t = np.geomspace(1e-5, 1e3, 500)
133 df = 2.5
134 exact = -0.5*(df+2)*np.log1p(t/df)
135 exps = np.array([.25,.5,1.,1.5,2.])
136 A = np.column_stack([np.ones_like(t)] + [t**a for a in exps])
137 # fit on the numerically relevant log-radius range; report in-range error
138 grid = np.geomspace(1e-4, 100, 400)
139 Ag = np.column_stack([np.ones_like(grid)] + [grid**a for a in exps])
140 eg = -0.5*(df+2)*np.log1p(grid/df)
141 coef, *_ = np.linalg.lstsq(Ag, eg, rcond=None)
142 fit = A @ coef
143 ii = (t >= 1e-4) & (t <= 100)
144 # Cholesky solve and explicit inverse agree for an SPD test matrix.
145 S = np.array([[1.4,.25],[.25,.8]])
146 L = np.linalg.cholesky(S); delta = np.array([.7,-1.1])
147 r_solve = np.sum(np.linalg.solve(L, delta)**2)
148 r_inv = delta @ np.linalg.inv(S) @ delta
149 return {'basis_max_abs_error_t_1e-4_to_100': float(np.max(np.abs(fit[ii]-exact[ii]))),
150 'basis_rmse_t_1e-4_to_100': float(np.sqrt(np.mean((fit[ii]-exact[ii])**2))),
151 'cholesky_vs_inverse_abs_error': float(abs(r_solve-r_inv)),
152 'all_basis_radii_finite': bool(np.isfinite(A).all())}
153
154if __name__ == '__main__':
155 result = {'device': DEVICE, 'math_check': math_check(), 'datasets': {}}
156 for kind in ('gaussian', 'student'):
157 xtr, ytr, xte, yte = make_data(kind)
158 x_test_global, y_test_global = xte, yte
159 vals = {}
160 qda = QDAHead(2).to(DEVICE); qda.fit(xtr, ytr)
161 with torch.no_grad():
162 qlog = qda(xte); qloss = nn.functional.binary_cross_entropy_with_logits(qlog, yte).item()
163 vals['qda_frozen'] = {'test_logloss': qloss, 'test_accuracy': float(((qlog > 0)==(yte>.5)).float().mean())}
164 for name, cls, radial in [('linear', LinearHead, False), ('mlp', MLPHead, False), ('fractional_radial', lambda d: FractionalRadialHead(d), True)]:
165 torch.manual_seed(SEED + len(name) + (0 if kind=='gaussian' else 10))
166 model = cls(2)
167 vals[name] = dict(zip(('test_logloss','test_accuracy'), fit(model, xtr, ytr, radial)))
168 if radial:
169 with torch.no_grad():
170 vals[name]['max_radius_test'] = float(model.radii(xte).max().cpu())
171 result['datasets'][kind] = vals
172 Path('results.json').write_text(json.dumps(result, indent=2))
173 print(json.dumps(result, indent=2))