1import json, os, random
  2import numpy as np
  3from scipy.optimize import linear_sum_assignment
  4import torch
  5from torch import nn
  6from sklearn.datasets import load_digits
  7from sklearn.model_selection import train_test_split
  8from sklearn.preprocessing import StandardScaler
  9
 10SEED = 17
 11np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
 12if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)
 13DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 14
 15
 16def positive_axis_match(z_a, z_b, eps=1e-8):
 17    # Rows are examples, columns are axes. Cost is normalized residual after
 18    # independently fitting the allowed positive scale for each candidate pair.
 19    va = np.var(z_a, axis=0) + eps
 20    za2 = np.sum(z_a * z_a, axis=0) + eps
 21    cross = z_a.T @ z_b
 22    alpha = np.maximum(eps, cross / za2[:, None])
 23    residual = z_a[:, :, None] * alpha[None, :, :] - z_b[:, None, :]
 24    cost = np.mean(residual * residual, axis=0) / va[:, None]
 25    rows, cols = linear_sum_assignment(cost)
 26    return rows, cols, alpha[rows, cols], cost[rows, cols]
 27
 28
 29def affine_fit_residual(z_a, z_b):
 30    # Fit E(z_b)=z_a, including an intercept, and return RMS residual.
 31    xb = np.concatenate([z_b, np.ones((len(z_b), 1))], axis=1)
 32    coef, *_ = np.linalg.lstsq(xb, z_a, rcond=None)
 33    pred = xb @ coef
 34    return float(np.sqrt(np.mean((pred - z_a) ** 2)))
 35
 36
 37def math_sanity():
 38    rng = np.random.default_rng(SEED)
 39    n, d = 1200, 8
 40    za = rng.normal(size=(n, d))
 41    perm = rng.permutation(d)
 42    scales = rng.uniform(.35, 2.4, size=d)
 43    values = []
 44    for noise in [0.0, .02, .05, .10, .20, .40]:
 45        zb = za[:, perm] * scales[None, :] + noise * rng.normal(size=(n, d))
 46        rows, cols, alpha, costs = positive_axis_match(za, zb)
 47        recovered = np.mean(cols == np.argsort(perm))
 48        axis_quality = float(np.mean(np.exp(-costs)))
 49        weak = affine_fit_residual(za, zb)
 50        values.append({"noise": noise, "affine_rms": weak,
 51                       "axis_quality": axis_quality,
 52                       "permutation_recovery": float(recovered),
 53                       "median_normalized_cost": float(np.median(costs))})
 54    # A control with unrelated coordinates tests that matching is not automatic.
 55    zb = rng.normal(size=(n, d))
 56    _, _, _, costs = positive_axis_match(za, zb)
 57    control = {"random_axis_quality": float(np.mean(np.exp(-costs))),
 58               "random_median_cost": float(np.median(costs))}
 59    return {"sweep": values, "control": control}
 60
 61
 62class MLP(nn.Module):
 63    def __init__(self, hidden=64):
 64        super().__init__()
 65        self.fc1 = nn.Linear(64, hidden)
 66        self.fc2 = nn.Linear(hidden, hidden)
 67        self.out = nn.Linear(hidden, 10)
 68    def forward(self, x, return_z=False):
 69        z1 = self.fc1(x); h1 = torch.relu(z1)
 70        z2 = self.fc2(h1); h2 = torch.relu(z2)
 71        y = self.out(h2)
 72        return (y, [z1, z2]) if return_z else y
 73
 74
 75def torch_axis_loss(za, zb):
 76    # Differentiable once the permutation is stop-gradient/fixed for this batch.
 77    with torch.no_grad():
 78        _, cols, alpha, _ = positive_axis_match(za.detach().cpu().numpy(), zb.detach().cpu().numpy())
 79    loss = 0.0
 80    count = 0
 81    for j, i, a in zip(range(len(cols)), cols, alpha):
 82        # Match teacher axis j to assigned student axis i.
 83        denom = za[:, j].var(unbiased=False).detach() + 1e-5
 84        loss = loss + ((zb[:, i] - float(a) * za[:, j]) ** 2 / denom).mean()
 85        count += 1
 86    return loss / max(count, 1)
 87
 88
 89def train_teacher_and_students():
 90    d = load_digits()
 91    x = StandardScaler().fit_transform(d.data).astype("float32")
 92    y = d.target.astype("int64")
 93    xt, xv, yt, yv = train_test_split(x, y, test_size=.25, random_state=SEED, stratify=y)
 94    Xtr, Ytr = torch.tensor(xt, device=DEVICE), torch.tensor(yt, device=DEVICE)
 95    Xva, Yva = torch.tensor(xv, device=DEVICE), torch.tensor(yv, device=DEVICE)
 96    teacher = MLP().to(DEVICE)
 97    opt = torch.optim.Adam(teacher.parameters(), lr=2e-3)
 98    for step in range(350):
 99        idx = torch.randint(len(Xtr), (96,), device=DEVICE)
100        loss = nn.functional.cross_entropy(teacher(Xtr[idx]), Ytr[idx])
101        opt.zero_grad(); loss.backward(); opt.step()
102    teacher.eval()
103    with torch.no_grad():
104        ty, tz = teacher(Xtr, True); tv, _ = teacher(Xva, True)
105        tacc = float((tv.argmax(1) == Yva).float().mean())
106    results = {}
107    for kind in ["baseline", "axis"]:
108        torch.manual_seed(SEED + (0 if kind == "baseline" else 1))
109        student = MLP().to(DEVICE)
110        opt = torch.optim.Adam(student.parameters(), lr=2e-3)
111        last_axis = None
112        for step in range(350):
113            idx = torch.randint(len(Xtr), (96,), device=DEVICE)
114            sx, sz = student(Xtr[idx], True)
115            with torch.no_grad():
116                target = ty[idx]
117            task = nn.functional.cross_entropy(sx, Ytr[idx])
118            kd = nn.functional.mse_loss(sx, target)
119            total = task + .15 * kd
120            if kind == "axis":
121                total = total + .03 * (torch_axis_loss(tz[0][idx], sz[0]) + torch_axis_loss(tz[1][idx], sz[1]))
122            opt.zero_grad(); total.backward(); opt.step()
123            if step % 100 == 0:
124                with torch.no_grad():
125                    _, szfull = student(Xtr, True)
126                    _, cols, _, costs = positive_axis_match(tz[0].cpu().numpy(), szfull[0].cpu().numpy())
127                    last_axis = {"quality": float(np.mean(np.exp(-costs))),
128                                 "matched_fraction_cost_lt_0.25": float(np.mean(costs < .25)),
129                                 "permutation": cols.tolist()}
130        student.eval()
131        with torch.no_grad():
132            sy, _ = student(Xva, True)
133            val_loss = float(nn.functional.cross_entropy(sy, Yva))
134            acc = float((sy.argmax(1) == Yva).float().mean())
135            train_loss = float(nn.functional.cross_entropy(student(Xtr), Ytr))
136        results[kind] = {"val_accuracy": acc, "val_loss": val_loss,
137                         "train_ce": train_loss, "final_assignment": last_axis}
138    return {"device": DEVICE, "teacher_accuracy": tacc, "students": results}
139
140
141def main():
142    out = {"seed": SEED, "math_sanity": math_sanity(), "distillation": train_teacher_and_students()}
143    with open("results.json", "w") as f: json.dump(out, f, indent=2)
144    print(json.dumps(out, indent=2))
145
146if __name__ == "__main__":
147    try:
148        main()
149    except (RuntimeError, AssertionError) as e:
150        if torch.cuda.is_available() and DEVICE == "cuda":
151            print("CUDA failed; rerun with CPU:", repr(e))
152            os.environ["CUDA_VISIBLE_DEVICES"] = ""
153            # Explicit CPU fallback in a fresh process avoids partially allocated CUDA state.
154            os.execv("/home/maxwelhelp/main/bin/python3", ["python3", __file__])
155        raise