SabaPivot's picture
download
raw
15.6 kB
#!/usr/bin/env python3
"""Independent spectral reproduction for arXiv:2606.01292.
The script combines exact diagonal population dynamics with a Monte Carlo
teacher-SGD experiment. It intentionally avoids paper-unavailable code.
"""
from __future__ import annotations
import argparse
import csv
import html
import json
import math
import os
import time
from pathlib import Path
import numpy as np
def write_csv(path: Path, rows: list[dict]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=list(rows[0]))
writer.writeheader()
writer.writerows(rows)
def svg_lines(
path: Path,
title: str,
series: dict[str, tuple[np.ndarray, np.ndarray]],
xlabel: str,
ylabel: str,
logx: bool = False,
logy: bool = False,
) -> None:
colors = ["#2563eb", "#dc2626", "#059669", "#7c3aed", "#ea580c", "#0891b2"]
width, height = 900, 520
left, right, top, bottom = 92, 25, 58, 76
transformed = {}
all_x, all_y = [], []
for name, (x, y) in series.items():
x = np.asarray(x, dtype=float)
y = np.asarray(y, dtype=float)
mask = np.isfinite(x) & np.isfinite(y)
if logx:
mask &= x > 0
if logy:
mask &= y > 0
x, y = x[mask], y[mask]
tx = np.log10(x) if logx else x
ty = np.log10(y) if logy else y
transformed[name] = (tx, ty)
all_x.extend(tx.tolist())
all_y.extend(ty.tolist())
xmin, xmax = min(all_x), max(all_x)
ymin, ymax = min(all_y), max(all_y)
if xmax == xmin:
xmax += 1
if ymax == ymin:
ymax += 1
xpad, ypad = 0.03 * (xmax - xmin), 0.08 * (ymax - ymin)
xmin, xmax = xmin - xpad, xmax + xpad
ymin, ymax = ymin - ypad, ymax + ypad
def sx(v: float) -> float:
return left + (v - xmin) / (xmax - xmin) * (width - left - right)
def sy(v: float) -> float:
return height - bottom - (v - ymin) / (ymax - ymin) * (height - top - bottom)
parts = [
"<!doctype html><meta charset='utf-8'>",
"<style>body{font-family:ui-sans-serif,system-ui;margin:0;background:#fff;color:#111827}"
".wrap{max-width:920px;margin:auto}.note{font-size:13px;color:#4b5563;margin:0 24px 18px}</style>",
"<div class='wrap'>",
f"<svg viewBox='0 0 {width} {height}' role='img' aria-label='{html.escape(title)}'>",
f"<text x='{width/2}' y='31' text-anchor='middle' font-size='20' font-weight='700'>{html.escape(title)}</text>",
f"<line x1='{left}' y1='{height-bottom}' x2='{width-right}' y2='{height-bottom}' stroke='#374151'/>",
f"<line x1='{left}' y1='{top}' x2='{left}' y2='{height-bottom}' stroke='#374151'/>",
]
for j in range(6):
xv = xmin + j * (xmax - xmin) / 5
yv = ymin + j * (ymax - ymin) / 5
xlab = f"1e{xv:.1f}" if logx else f"{xv:.2g}"
ylab = f"1e{yv:.1f}" if logy else f"{yv:.2g}"
parts += [
f"<line x1='{sx(xv):.1f}' y1='{top}' x2='{sx(xv):.1f}' y2='{height-bottom}' stroke='#e5e7eb'/>",
f"<text x='{sx(xv):.1f}' y='{height-bottom+23}' text-anchor='middle' font-size='12'>{xlab}</text>",
f"<line x1='{left}' y1='{sy(yv):.1f}' x2='{width-right}' y2='{sy(yv):.1f}' stroke='#e5e7eb'/>",
f"<text x='{left-10}' y='{sy(yv)+4:.1f}' text-anchor='end' font-size='12'>{ylab}</text>",
]
parts += [
f"<text x='{(left+width-right)/2}' y='{height-22}' text-anchor='middle' font-size='14'>{html.escape(xlabel)}</text>",
f"<text x='20' y='{(top+height-bottom)/2}' transform='rotate(-90 20 {(top+height-bottom)/2})' text-anchor='middle' font-size='14'>{html.escape(ylabel)}</text>",
]
for idx, (name, (x, y)) in enumerate(transformed.items()):
points = " ".join(f"{sx(a):.2f},{sy(b):.2f}" for a, b in zip(x, y))
color = colors[idx % len(colors)]
parts.append(f"<polyline points='{points}' fill='none' stroke='{color}' stroke-width='2.4'/>")
lx, ly = width - 230, 60 + idx * 22
parts.append(f"<line x1='{lx}' y1='{ly}' x2='{lx+25}' y2='{ly}' stroke='{color}' stroke-width='3'/>")
parts.append(f"<text x='{lx+32}' y='{ly+4}' font-size='12'>{html.escape(name)}</text>")
parts += ["</svg>", "<p class='note'>Axes marked 1eX display log10 coordinates. Raw values are attached as CSV.</p>", "</div>"]
path.write_text("\n".join(parts), encoding="utf-8")
def claim1(out: Path) -> dict:
d, alpha_s, eta, steps, delta = 2000, 1.5, 0.3, 4000, 0.1
i = np.arange(1, d + 1, dtype=float)
lam = i ** (-alpha_s)
learned = 1.0 - np.power(1.0 - eta * lam, steps)
ratio = learned**2
rows = []
for idx in range(d):
band = "head" if learned[idx] >= 0.99 else ("tail" if learned[idx] <= delta else "transition")
rows.append({
"mode": idx + 1,
"eigenvalue": f"{lam[idx]:.12g}",
"student_learning_factor": f"{learned[idx]:.12g}",
"inherited_teacher_error_ratio": f"{ratio[idx]:.12g}",
"band": band,
})
write_csv(out / "claim1_mode_inheritance.csv", rows)
head = learned >= 0.99
tail = learned <= delta
svg_lines(
out / "claim1_mode_inheritance.html",
"Claim 1 — spectral inheritance and tail damping",
{"learning factor²": (i, ratio), "δ² bound": (i, np.full(d, delta**2))},
"spectral mode i",
"teacher-error inheritance multiplier",
logx=True,
logy=True,
)
return {
"delta": delta,
"head_modes": int(head.sum()),
"tail_modes": int(tail.sum()),
"head_min_inheritance_ratio": float(ratio[head].min()),
"tail_max_inheritance_ratio": float(ratio[tail].max()),
"delta_squared": delta**2,
"tail_bound_satisfied": bool(ratio[tail].max() <= delta**2 * (1 + 1e-12)),
}
def claim2(out: Path) -> dict:
alpha_t, alpha_s, beta = 1.5, 3.0, 0.0
ns = np.logspace(2, 10, 33)
q = alpha_t - 1.0 - beta
h_t = ns ** (1.0 / alpha_t)
h_s = ns ** (1.0 / alpha_s)
risk_transfer = h_t ** (-q) / q
risk_direct = h_s ** (-q) / q
der = risk_direct / risk_transfer
fit_slice = slice(8, None)
slope = float(np.polyfit(np.log(ns[fit_slice]), np.log(der[fit_slice]), 1)[0])
expected = q * (1.0 / alpha_t - 1.0 / alpha_s)
rows = [{
"N": f"{n:.12g}",
"teacher_horizon": f"{kt:.12g}",
"student_horizon": f"{ks:.12g}",
"direct_student_risk": f"{rd:.12g}",
"distilled_student_risk": f"{rt:.12g}",
"DER": f"{de:.12g}",
} for n, kt, ks, rd, rt, de in zip(ns, h_t, h_s, risk_direct, risk_transfer, der)]
write_csv(out / "claim2_der_scaling.csv", rows)
svg_lines(
out / "claim2_der_scaling.html",
"Claim 2 — Distillation Efficiency Ratio diverges",
{"DER": (ns, der), f"N^κ, κ={expected:.3f}": (ns, (ns / ns[0]) ** expected * der[0])},
"labeled samples N",
"DER",
logx=True,
logy=True,
)
return {
"alpha_teacher": alpha_t,
"alpha_student": alpha_s,
"beta": beta,
"expected_kappa": expected,
"fitted_kappa": slope,
"absolute_slope_error": abs(slope - expected),
"horizon_ratio_at_max_N": float(h_t[-1] / h_s[-1]),
"der_at_max_N": float(der[-1]),
}
def simulate_teacher_sgd(
signal_dims: list[int], reps: int, seed: int, backend: str
) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]:
if backend == "torch":
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(seed)
d, n, batch, eta = 100, 2000, 10, 0.1
lam = torch.arange(1, d + 1, device=device, dtype=torch.float64).pow(-1.0)
roots = torch.sqrt(lam)
kcount = len(signal_dims)
wstar = torch.zeros((kcount, d), device=device, dtype=torch.float64)
for j, k in enumerate(signal_dims):
wstar[j, :k] = 0.1
w = torch.zeros((reps, kcount, d), device=device, dtype=torch.float64)
updates = n // batch
decay_every = max(1, math.floor(n / math.log2(n)))
for step in range(updates):
lr = eta / (2 ** (step // decay_every))
x = torch.randn((reps, kcount, batch, d), device=device, dtype=torch.float64) * roots
y = torch.einsum("rkbd,kd->rkb", x, wstar) + torch.randn(
(reps, kcount, batch), device=device, dtype=torch.float64
)
pred = torch.einsum("rkbd,rkd->rkb", x, w)
grad = torch.mean(x * (pred - y).unsqueeze(-1), dim=2)
w -= lr * grad
label = torch.cuda.get_device_name(0) if device.type == "cuda" else "torch-cpu"
return w.cpu().numpy(), wstar.cpu().numpy(), lam.cpu().numpy(), label
rng = np.random.default_rng(seed)
d, n, batch, eta = 100, 2000, 10, 0.1
lam = np.arange(1, d + 1, dtype=float) ** -1.0
roots = np.sqrt(lam)
kcount = len(signal_dims)
wstar = np.zeros((kcount, d), dtype=float)
for j, k in enumerate(signal_dims):
wstar[j, :k] = 0.1
w = np.zeros((reps, kcount, d), dtype=float)
updates = n // batch
decay_every = max(1, math.floor(n / math.log2(n)))
for step in range(updates):
lr = eta / (2 ** (step // decay_every))
x = rng.normal(size=(reps, kcount, batch, d)) * roots
y = np.einsum("rkbd,kd->rkb", x, wstar) + rng.normal(size=(reps, kcount, batch))
pred = np.einsum("rkbd,rkd->rkb", x, w)
grad = np.mean(x * (pred - y)[..., None], axis=2)
w -= lr * grad
return w, wstar, lam, "numpy-cpu"
def claim3_and_early_stop(out: Path, reps: int, seed: int, backend: str) -> dict:
signal_dims = [1, 10, 20, 30, 50]
wteacher, wstar, lam, compute_device = simulate_teacher_sgd(signal_dims, reps, seed, backend)
checkpoints = np.unique(np.concatenate(([0], np.geomspace(1, 50000, 240).astype(int))))
decay_every, eta = 3200, 0.1
cumulative = 0.0
prev = 0
curve_rows: list[dict] = []
summaries: list[dict] = []
curves: dict[str, tuple[np.ndarray, np.ndarray]] = {}
teacher_risks = np.mean(np.sum(lam * (wteacher - wstar[None, :, :]) ** 2, axis=-1), axis=0)
for j, k in enumerate(signal_dims):
risks = []
for step in checkpoints:
for s in range(prev, step):
cumulative += eta / (2 ** (s // decay_every))
prev = step
p = 1.0 - np.exp(-cumulative * lam)
err = p[None, :] * wteacher[:, j, :] - wstar[j]
risk = float(np.mean(np.sum(lam * err**2, axis=-1)))
risks.append(risk)
curve_rows.append({"signal_dim": k, "student_steps": int(step), "mean_excess_risk": f"{risk:.12g}", "teacher_risk": f"{teacher_risks[j]:.12g}"})
# reset cumulative calculation for each signal dimension
cumulative, prev = 0.0, 0
risks_arr = np.asarray(risks)
best_idx = int(np.argmin(risks_arr[1:]) + 1)
best = float(risks_arr[best_idx])
teacher = float(teacher_risks[j])
final = float(risks_arr[-1])
summaries.append({
"signal_dim": k,
"teacher_risk": f"{teacher:.12g}",
"best_student_risk": f"{best:.12g}",
"best_student_steps": int(checkpoints[best_idx]),
"final_student_risk": f"{final:.12g}",
"best_risk_ratio_student_over_teacher": f"{best/teacher:.12g}",
"PGR_oracle_ceiling": f"{1-best/teacher:.12g}",
"student_surpasses_teacher": bool(best < teacher),
"early_stop_beats_final": bool(best < final),
})
curves[f"k={k}"] = (checkpoints[1:], risks_arr[1:])
curves[f"teacher k={k}"] = (checkpoints[1:], np.full(len(checkpoints) - 1, teacher))
write_csv(out / "claim3_early_stopping_curves.csv", curve_rows)
write_csv(out / "claim3_w2s_summary.csv", summaries)
svg_lines(
out / "claim3_early_stopping.html",
"Claims 3 & 5 — low-dimensional W2S benefits from early stopping",
curves,
"student transfer steps",
"excess risk",
logx=True,
logy=True,
)
ratios = [float(r["best_risk_ratio_student_over_teacher"]) for r in summaries]
return {
"replicates": reps,
"compute_device": compute_device,
"signal_dimensions": signal_dims,
"all_students_surpass_teacher_at_best_checkpoint": all(r["student_surpasses_teacher"] for r in summaries),
"all_best_checkpoints_beat_final_checkpoint": all(r["early_stop_beats_final"] for r in summaries),
"best_risk_ratios": ratios,
"best_steps": [int(r["best_student_steps"]) for r in summaries],
}
def claim4(out: Path) -> dict:
alpha_t, alpha_s, kdag = 1.5, 2.0, 10.0
ns = np.logspace(3, 10, 36)
p = 2 * alpha_s
# Exact minimizer of k/N + kdag^p k^-p N^((1-alpha_t)/alpha_t).
kopt = (p * (kdag**p) * ns ** (1.0 / alpha_t)) ** (1.0 / (p + 1.0))
teacher = ns ** ((1.0 - alpha_t) / alpha_t)
transfer = kopt / ns + (kdag / kopt) ** p * teacher
one_minus_pgr = transfer / teacher
pgr = 1.0 - one_minus_pgr
expected = 2 * alpha_s / (alpha_t * (2 * alpha_s + 1))
fitted = float(-np.polyfit(np.log(ns), np.log(one_minus_pgr), 1)[0])
rows = [{
"N": f"{n:.12g}", "k_opt": f"{ko:.12g}", "teacher_risk": f"{tr:.12g}",
"optimal_transfer_risk": f"{sr:.12g}", "PGR": f"{pg:.12g}", "one_minus_PGR": f"{gap:.12g}",
} for n, ko, tr, sr, pg, gap in zip(ns, kopt, teacher, transfer, pgr, one_minus_pgr)]
write_csv(out / "claim4_pgr_scaling.csv", rows)
svg_lines(
out / "claim4_pgr_scaling.html",
"Claim 4 — full capability recovery",
{"1 − PGR": (ns, one_minus_pgr), f"N^-Δ, Δ={expected:.3f}": (ns, (ns / ns[0]) ** (-expected) * one_minus_pgr[0])},
"teacher labeled samples N",
"remaining performance gap (1 − PGR)",
logx=True,
logy=True,
)
return {
"alpha_teacher": alpha_t,
"alpha_student": alpha_s,
"k_dagger": kdag,
"expected_delta_rate": expected,
"fitted_delta_rate": fitted,
"absolute_slope_error": abs(fitted - expected),
"pgr_at_max_N": float(pgr[-1]),
"transfer_below_teacher_for_fraction_of_grid": float(np.mean(transfer < teacher)),
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--output", default="outputs/core")
parser.add_argument("--reps", type=int, default=64)
parser.add_argument("--seed", type=int, default=23301)
parser.add_argument("--backend", choices=["numpy", "torch"], default="numpy")
args = parser.parse_args()
start = time.time()
out = Path(args.output)
out.mkdir(parents=True, exist_ok=True)
metrics = {
"paper": "arXiv:2606.01292 / OpenReview:ykWN4LG9vE",
"scope": "independent diagonal spectral reproduction plus Monte Carlo synthetic W2S",
"claim1": claim1(out),
"claim2": claim2(out),
"claim3_and_early_stopping": claim3_and_early_stop(out, args.reps, args.seed, args.backend),
"claim4": claim4(out),
"runtime_seconds": time.time() - start,
"environment": {"python": os.sys.version, "numpy": np.__version__},
}
(out / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
print(json.dumps(metrics, indent=2))
if __name__ == "__main__":
main()

Xet Storage Details

Size:
15.6 kB
·
Xet hash:
4bbdd97270009bbdd13bcf3cc12fd1c0425be1322a838184e0e6591e9ca12de6

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.