atlas.py
#!/usr/bin/env python3
"""Checkpoint atlas: analysis of exported per-checkpoint tables.
Reads the TSV the GPU side writes (one row per model, step, measure, layer) and
answers plan.md's three pre-registered questions per series: (a) above null at
step 0, (b) correlation with loss beyond what step explains, (c) change-point on
the log-step axis, with the Pile batch range before it. Also holds the reference
linear probe and its shuffled-label null. No model, no torch, no network.
Usage: atlas.py TABLE.tsv --threshold T [--alpha 0.05]
atlas.py --selftest
"""
import argparse
import csv
import math
import random
import sys
import time
# Pythia publishes 154 checkpoints: step 0, 1, 2, 4, ..., 512, then every 1000.
PYTHIA_STEPS = [0] + [1 << i for i in range(10)] + list(range(1000, 143001, 1000))
# plan.md "Design": the coarse grid, 20 checkpoints.
GRID = [0] + [1 << i for i in range(10)] + [1000, 2000, 4000, 8000, 16000, 32000, 64000, 100000, 143000]
# XXX: Pythia paper value (1024 sequences of 2048 tokens per step), not in plan.md
PYTHIA_BATCH_SEQS = 1024
COLUMNS = ["model", "step", "measure", "layer", "value", "null_mean", "null_sd", "null_p", "n_null", "loss"]
TYPES = [str, int, str, int, float, float, float, float, int, float]
# --- checkpoint schedule ---
def is_pythia_step(s):
return s == 0 or (s <= 512 and s & (s - 1) == 0) or (s % 1000 == 0 and 1000 <= s <= 143000)
def log_step(s):
"""Analysis axis: log2(1 + step), so step 0 sits at 0 and the grid is
roughly evenly spaced."""
return math.log2(1 + s)
def batch_range(prev_step, step):
"""Batches consumed between two checkpoints: step t has seen batches
0..t-1, so the interval (prev_step, step] is batches [prev_step, step).
Returns ((batch_lo, batch_hi), (seq_lo, seq_hi)), half-open."""
return (prev_step, step), (prev_step * PYTHIA_BATCH_SEQS, step * PYTHIA_BATCH_SEQS)
# --- small linear algebra ---
def ols(cols, y):
"""Least squares of y on the given columns (no implicit intercept).
Returns (coefficients, residuals). Raises on a singular design."""
k = len(cols)
a = [[sum(ci * cj for ci, cj in zip(cols[i], cols[j])) for j in range(k)] +
[sum(ci * yi for ci, yi in zip(cols[i], y))] for i in range(k)]
for i in range(k):
p = max(range(i, k), key=lambda r: abs(a[r][i]))
if abs(a[p][i]) < 1e-12:
raise ValueError("singular design")
a[i], a[p] = a[p], a[i]
for r in range(k):
if r != i:
f = a[r][i] / a[i][i]
a[r] = [x - f * z for x, z in zip(a[r], a[i])]
beta = [a[i][k] / a[i][i] for i in range(k)]
res = [yi - sum(b * c[n] for b, c in zip(beta, cols)) for n, yi in enumerate(y)]
return beta, res
def pearson(x, y):
mx, my = sum(x) / len(x), sum(y) / len(y)
sxy = sum((a - mx) * (b - my) for a, b in zip(x, y))
sxx = sum((a - mx) ** 2 for a in x)
syy = sum((b - my) ** 2 for b in y)
return sxy / math.sqrt(sxx * syy) if sxx > 0 and syy > 0 else 0.0
def perm_p(obs, null):
"""One-sided permutation p-value, counting the observed value as one draw."""
return (1 + sum(1 for v in null if v >= obs)) / (1 + len(null))
# --- linear probe ---
def sigmoid(z):
if z >= 0:
return 1.0 / (1.0 + math.exp(-z))
e = math.exp(z)
return e / (1.0 + e)
def probe_fit(X, y, iters=200, lr=0.5, l2=1e-3):
"""L2 logistic regression by full-batch gradient descent from zero.
WARNING: X must be standardised. Returns weights, bias last."""
d, n = len(X[0]), len(X)
w = [0.0] * (d + 1)
for _ in range(iters):
g = [0.0] * (d + 1)
for x, t in zip(X, y):
e = sigmoid(sum(wi * xi for wi, xi in zip(w, x)) + w[d]) - t
for i in range(d):
g[i] += e * x[i]
g[d] += e
for i in range(d):
w[i] -= lr * (g[i] / n + l2 * w[i])
w[d] -= lr * g[d] / n
return w
def probe_acc(X, y, seed, test_frac=0.3):
"""Held-out accuracy of a probe trained on a seeded split; features are
standardised with train-set statistics only."""
if len(X) != len(y) or len(X) < 4 or set(y) - {0, 1}:
raise ValueError("probe needs >= 4 rows and 0/1 labels")
idx = list(range(len(X)))
random.Random(seed).shuffle(idx)
nt = max(1, int(len(idx) * test_frac))
te, tr = idx[:nt], idx[nt:]
d = len(X[0])
mu = [sum(X[i][j] for i in tr) / len(tr) for j in range(d)]
sd = [math.sqrt(sum((X[i][j] - mu[j]) ** 2 for i in tr) / len(tr)) or 1.0 for j in range(d)]
def z(i):
return [(X[i][j] - mu[j]) / sd[j] for j in range(d)]
w = probe_fit([z(i) for i in tr], [y[i] for i in tr])
hit = sum((sum(wi * xi for wi, xi in zip(w, z(i))) + w[d] > 0) == (y[i] == 1) for i in te)
return hit / nt
def probe_null(X, y, n_perm=99, seed=0):
"""Accuracy, shuffled-label null and p. Each null draw permutes all labels
and reruns the identical split and fit. Returns (acc, null, p)."""
acc = probe_acc(X, y, seed)
rng = random.Random(seed + 1)
null = []
for _ in range(n_perm):
ys = list(y)
rng.shuffle(ys)
null.append(probe_acc(X, ys, seed))
return acc, null, perm_p(acc, null)
# --- per-series analysis ---
def above_null_at_0(rows, alpha):
"""(a): the step-0 row's exported null_p below alpha; None without a step-0 row."""
r0 = [r for r in rows if r["step"] == 0]
return r0[0]["null_p"] < alpha if r0 else None
def loss_partial(steps, value, loss, n_perm=999, seed=0):
"""(b): partial Pearson r of value and loss given log2(1 + step).
plan.md says only "beyond what step explains". Chosen version: regress value
and loss each on [1, log2(1 + step)], correlate the residuals, and get a
two-sided p by permuting one residual vector. Linear in log step is the
simplest control and the axis of (c).
XXX: loss is convex in log step, so curvature left in both residuals can
read as shared signal; add a quadratic term if residual plots show it."""
x = [log_step(s) for s in steps]
one = [1.0] * len(x)
_, rv = ols([one, x], value)
_, rl = ols([one, x], loss)
r = pearson(rv, rl)
rng = random.Random(seed)
null = []
for _ in range(n_perm):
p = list(rv)
rng.shuffle(p)
null.append(abs(pearson(p, rl)))
return r, perm_p(abs(r), null)
def change_point(steps, value, threshold, min_side=3):
"""(c): best single hinge on the log-step axis.
Fits value = a + b*x + c*max(0, x - x_k) at each grid point k with at least
min_side points on either side, keeps the k with least squared error, and
reports it if |c| (slope change per doubling of step) exceeds threshold.
Returns None or a dict with step, slope change, and the batch range from
the previous grid point."""
pts = sorted(zip(steps, value))
s = [p[0] for p in pts]
v = [p[1] for p in pts]
x = [log_step(t) for t in s]
one = [1.0] * len(x)
best = None
for k in range(min_side - 1, len(x) - min_side):
h = [max(0.0, xi - x[k]) for xi in x]
beta, res = ols([one, x, h], v)
sse = sum(e * e for e in res)
if best is None or sse < best[0]:
best = (sse, k, beta[2])
if best is None or abs(best[2]) <= threshold:
return None
k = best[1]
batches, seqs = batch_range(s[k - 1], s[k])
return {"step": s[k], "slope_change": best[2], "batches": batches, "sequences": seqs}
def analyse(rows, threshold, alpha=0.05):
"""Group rows by (model, measure, layer) and answer (a), (b), (c) for each."""
series = {}
for r in rows:
series.setdefault((r["model"], r["measure"], r["layer"]), []).append(r)
out = []
for key in sorted(series):
rs = sorted(series[key], key=lambda r: r["step"])
steps = [r["step"] for r in rs]
vals = [r["value"] for r in rs]
pr = loss_partial(steps, vals, [r["loss"] for r in rs]) if len(rs) >= 5 else (None, None)
out.append({"model": key[0], "measure": key[1], "layer": key[2], "n": len(rs),
"above_null_0": above_null_at_0(rs, alpha), "partial_r": pr[0], "partial_p": pr[1],
"change_point": change_point(steps, vals, threshold)})
return out
# --- exported table ---
def write_table(path, rows):
with open(path, "w", newline="") as f:
w = csv.writer(f, delimiter="\t", lineterminator="\n")
w.writerow(COLUMNS)
for r in rows:
w.writerow([repr(r[c]) if isinstance(r[c], float) else r[c] for c in COLUMNS])
def read_table(path):
"""Parse and validate an exported table. Raises ValueError naming the line
on a wrong header, bad value, non-Pythia step, duplicate key, or a loss
that differs between rows of the same checkpoint."""
rows, seen, loss = [], set(), {}
with open(path, newline="") as f:
rd = csv.reader(f, delimiter="\t")
if next(rd, None) != COLUMNS:
raise ValueError(f"{path}:1: header must be {' '.join(COLUMNS)}")
for ln, raw in enumerate(rd, 2):
if len(raw) != len(COLUMNS):
raise ValueError(f"{path}:{ln}: {len(raw)} fields, want {len(COLUMNS)}")
try:
r = {c: t(v) for c, t, v in zip(COLUMNS, TYPES, raw)}
except ValueError as e:
raise ValueError(f"{path}:{ln}: {e}") from None
if any(isinstance(v, float) and not math.isfinite(v) for v in r.values()):
raise ValueError(f"{path}:{ln}: non-finite number")
if not is_pythia_step(r["step"]):
raise ValueError(f"{path}:{ln}: step {r['step']} is not a Pythia checkpoint")
if not 0.0 <= r["null_p"] <= 1.0:
raise ValueError(f"{path}:{ln}: null_p out of [0, 1]")
key = (r["model"], r["step"], r["measure"], r["layer"])
if key in seen:
raise ValueError(f"{path}:{ln}: duplicate {key}")
seen.add(key)
ck = (r["model"], r["step"])
if abs(loss.setdefault(ck, r["loss"]) - r["loss"]) > 1e-9:
raise ValueError(f"{path}:{ln}: loss for {ck} differs from an earlier row")
rows.append(r)
return rows
def report(res):
print("model\tmeasure\tlayer\tn\tabove_null_0\tpartial_r\tpartial_p\tcp_step\tslope_change\tbatches")
for a in res:
cp = a["change_point"]
pr = "-\t-" if a["partial_r"] is None else f"{a['partial_r']:.3f}\t{a['partial_p']:.3f}"
cps = f"{cp['step']}\t{cp['slope_change']:.4g}\t{cp['batches'][0]}-{cp['batches'][1]}" if cp else "-\t-\t-"
print(f"{a['model']}\t{a['measure']}\t{a['layer']}\t{a['n']}\t{a['above_null_0']}\t{pr}\t{cps}")
# --- self-test ---
def selftest():
import os
import tempfile
t0 = time.time()
assert len(PYTHIA_STEPS) == 154 and len(set(PYTHIA_STEPS)) == 154
assert all(is_pythia_step(s) for s in PYTHIA_STEPS)
assert sum(is_pythia_step(s) for s in range(150000)) == 154
assert len(GRID) == 20 and all(is_pythia_step(s) for s in GRID) and GRID == sorted(GRID)
assert batch_range(1000, 2000) == ((1000, 2000), (1000 * 1024, 2000 * 1024))
# probe: separable vs label noise, 80 rows x 4 dims, 99 permutations
rng = random.Random(1)
y = [i % 2 for i in range(80)]
sep = [[(1.5 if t else -1.5) + rng.gauss(0, 1)] + [rng.gauss(0, 1) for _ in range(3)] for t in y]
noise = [[rng.gauss(0, 1) for _ in range(4)] for _ in y]
acc, null, p = probe_null(sep, y, seed=2)
assert acc > 0.85 and p <= 0.02, (acc, p)
acc_n, _, p_n = probe_null(noise, y, seed=2)
assert p_n > 0.05, (acc_n, p_n)
assert probe_null(sep, y, seed=2)[0] == acc # deterministic
print(f"probe separable acc={acc:.3f} p={p:.3f}; label-noise acc={acc_n:.3f} p={p_n:.3f}")
# change-point: flat to step 2000, then +0.1 per doubling
rng = random.Random(3)
xk = log_step(2000)
kink = [0.5 + 0.1 * max(0.0, log_step(s) - xk) + rng.gauss(0, 0.005) for s in GRID]
flat = [0.5 + rng.gauss(0, 0.005) for _ in GRID]
smooth = [0.2 + 0.03 * log_step(s) + rng.gauss(0, 0.005) for s in GRID]
cp = change_point(GRID, kink, 0.05)
assert cp and cp["step"] == 2000 and abs(cp["slope_change"] - 0.1) < 0.02, cp
assert cp["batches"] == (1000, 2000)
assert change_point(GRID, flat, 0.05) is None
assert change_point(GRID, smooth, 0.05) is None
print(f"change-point kink -> step {cp['step']} slope_change={cp['slope_change']:.4f} batches {cp['batches']}; flat, smooth -> none")
# (b): loss shares a step-independent wiggle with one series and not the other
wig = [rng.gauss(0, 0.05) for _ in GRID]
loss = [10.0 - 0.4 * log_step(s) + w for s, w in zip(GRID, wig)]
tied = [0.1 * log_step(s) - 2 * w for s, w in zip(GRID, wig)]
r, p = loss_partial(GRID, tied, loss)
assert r < -0.9 and p < 0.01, (r, p)
# step-only series with independent noise: one draw can reject by chance,
# so check the false-positive rate over 40 draws instead
fp = 0
for k in range(40):
g = random.Random(100 + k)
only = [0.2 + 0.03 * log_step(s) + g.gauss(0, 0.005) for s in GRID]
fp += loss_partial(GRID, only, loss, n_perm=199, seed=k)[1] < 0.05
assert fp <= 6, fp
print(f"partial r tied={r:.3f} p={p:.3f}; step-only false positives {fp}/40 at 0.05")
# table round trip, then analysis on it
rows = []
for s, v, l in zip(GRID, kink, loss):
rows.append({"model": "pythia-160m", "step": s, "measure": "M3_self_other", "layer": 6, "value": v,
"null_mean": 0.5, "null_sd": 0.03, "null_p": 0.4 if s == 0 else 0.01, "n_null": 99, "loss": l})
with tempfile.TemporaryDirectory() as d:
path = os.path.join(d, "t.tsv")
write_table(path, rows)
assert read_table(path) == rows
res = analyse(read_table(path), 0.05)
assert len(res) == 1 and res[0]["above_null_0"] is False and res[0]["change_point"]["step"] == 2000
bad = dict(rows[0], step=3)
write_table(path, rows + [bad])
try:
read_table(path)
raise AssertionError("step 3 accepted")
except ValueError as e:
assert "not a Pythia checkpoint" in str(e)
write_table(path, rows + [dict(rows[1], layer=7, loss=0.0)])
try:
read_table(path)
raise AssertionError("inconsistent loss accepted")
except ValueError as e:
assert "loss" in str(e)
print(f"selftest ok ({time.time() - t0:.1f}s)")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("table", nargs="?")
ap.add_argument("--threshold", type=float, help="pre-registered slope change per doubling of step")
ap.add_argument("--alpha", type=float, default=0.05)
ap.add_argument("--selftest", action="store_true")
a = ap.parse_args()
if a.selftest:
selftest()
return
if not a.table or a.threshold is None:
ap.error("TABLE and --threshold are required")
try:
rows = read_table(a.table)
except (OSError, ValueError) as e:
sys.exit(f"atlas: {e}")
report(analyse(rows, a.threshold, a.alpha))
if __name__ == "__main__":
main()