| |
| """ |
| Derived from Andrej Karpathy's nanochat project. |
| |
| MIT License |
| |
| Copyright (c) 2025 Andrej Karpathy |
| |
| Permission is hereby granted, free of charge, to any person obtaining a copy |
| of this software and associated documentation files (the "Software"), to deal |
| in the Software without restriction, including without limitation the rights |
| to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| copies of the Software, and to permit persons to whom the Software is |
| furnished to do so, subject to the following conditions: |
| |
| The above copyright notice and this permission notice shall be included in all |
| copies or substantial portions of the Software. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import csv |
| import json |
| import math |
| from pathlib import Path |
| import random |
| import statistics |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| PAPER_DIR = ROOT / "paper" |
| FIG_DIR = PAPER_DIR / "figures" |
| TABLE_DIR = PAPER_DIR / "tables" |
|
|
|
|
| REGIMES = { |
| "openweb": { |
| "label": "OpenWebText10K", |
| "short": "OWT10K", |
| "coeff": ROOT |
| / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/coefficients.json", |
| "cells": ROOT |
| / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/calibration_cells.csv", |
| "curve_run": ROOT / "runs/screen_static/20260525-133008", |
| "curve_model": "L12_H8_D320", |
| "curve_prefix": 1_000_000, |
| "curve_title": "OpenWebText10K, L12, 1M", |
| "small_stream": ROOT |
| / "runs/openwebtext10k_l16_updated_formula_clean_5seed/locked_stream/20260530-174525", |
| "small_condition": "openwebtext10k_interaction", |
| "small_model": "L16_H8_D384", |
| "small_protocol": "4M final prefix / 5 seeds", |
| "l20_stream": ROOT |
| / "runs/openwebtext10k_l20_8m_extrapolation_3seed/locked_stream/20260601-191150", |
| "l20_condition": "openwebtext10k_l20_interaction_8m", |
| }, |
| "tinystories": { |
| "label": "TinyStories", |
| "short": "TinyStories", |
| "coeff": ROOT |
| / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/coefficients.json", |
| "cells": ROOT |
| / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/calibration_cells.csv", |
| "curve_run": ROOT |
| / "runs/regime_calibration_tinystories_coarse/screen_static/20260529-161726", |
| "curve_model": "L12_H8_D320", |
| "curve_prefix": 1_000_000, |
| "curve_title": "TinyStories, L12, 1M", |
| "small_stream": [ |
| ROOT |
| / "runs/streaming_tinystories_interaction_schedule_l12/locked_stream/20260530-053831", |
| ROOT |
| / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-111523", |
| ROOT |
| / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-141335", |
| ], |
| "small_condition": "interaction", |
| "small_model": "L12_H8_D320", |
| "small_protocol": "4M final prefix / 5 seeds", |
| "l20_stream": ROOT |
| / "runs/tinystories_l20_8m_extrapolation_3seed/locked_stream/20260602-112249", |
| "l20_condition": "tinystories_l20_interaction_8m", |
| }, |
| "wikitext": { |
| "label": "WikiText-103", |
| "short": "WikiText", |
| "coeff": ROOT / "runs/coefficient_calibration/wikitext103_interaction/coefficients.json", |
| "cells": ROOT / "runs/coefficient_calibration/wikitext103_interaction/calibration_cells.csv", |
| "curve_run": ROOT |
| / "runs/wikitext103_interaction_static_screen/screen_static/20260603-153314", |
| "curve_model": "L12_H8_D320", |
| "curve_prefix": 250_000, |
| "curve_title": "WikiText-103, L12, 250k", |
| "small_stream": ROOT |
| / "runs/wikitext103_l12_streaming_validation_5seed/locked_stream/20260531-093525", |
| "small_condition": "wikitext103_formula_l12", |
| "small_model": "L12_H8_D320", |
| "small_protocol": "4M final prefix / 5 seeds", |
| "l20_stream": ROOT |
| / "runs/wikitext103_l20_8m_extrapolation_3seed/locked_stream/20260603-231952", |
| "l20_condition": "wikitext103_l20_interaction_8m", |
| }, |
| } |
|
|
|
|
| def read_json(path: Path) -> dict | list: |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def read_jsonl(path: Path) -> list[dict]: |
| rows = [] |
| for line in path.read_text(encoding="utf-8").splitlines(): |
| if line.strip(): |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def read_csv(path: Path) -> list[dict]: |
| with path.open(newline="", encoding="utf-8") as handle: |
| return list(csv.DictReader(handle)) |
|
|
|
|
| def mean_std(values: list[float]) -> tuple[float, float]: |
| if not values: |
| return float("nan"), float("nan") |
| mean = statistics.fmean(values) |
| std = statistics.stdev(values) if len(values) > 1 else 0.0 |
| return mean, std |
|
|
|
|
| def bootstrap_ci( |
| values: list[float], *, confidence: float = 0.95, iterations: int = 20_000, seed: int = 1337 |
| ) -> tuple[float, float]: |
| if not values: |
| return float("nan"), float("nan") |
| if len(values) == 1: |
| return values[0], values[0] |
| rng = random.Random(seed) |
| n = len(values) |
| means = [] |
| for _ in range(iterations): |
| sample = [values[rng.randrange(n)] for _ in range(n)] |
| means.append(statistics.fmean(sample)) |
| means.sort() |
| lower_q = (1.0 - confidence) / 2.0 |
| upper_q = 1.0 - lower_q |
| lower_idx = min(iterations - 1, max(0, int(lower_q * iterations))) |
| upper_idx = min(iterations - 1, max(0, int(upper_q * iterations))) |
| return means[lower_idx], means[upper_idx] |
|
|
|
|
| def fmt(x: float, digits: int = 4) -> str: |
| return f"{x:.{digits}f}" |
|
|
|
|
| def tex_escape(value: object) -> str: |
| text = str(value) |
| return ( |
| text.replace("\\", "\\textbackslash{}") |
| .replace("_", "\\_") |
| .replace("&", "\\&") |
| .replace("%", "\\%") |
| .replace("#", "\\#") |
| ) |
|
|
|
|
| def model_label(model_name: str) -> str: |
| return tex_escape(model_name.replace("_", "-")) |
|
|
|
|
| def token_label(value: int) -> str: |
| if value >= 1_000_000: |
| whole = value / 1_000_000 |
| return f"{whole:g}M" |
| return f"{value // 1000}k" |
|
|
|
|
| def protocol_budget(condition_rows: list[dict], stages: list[int]) -> dict: |
| rows_by_seed: dict[int, list[dict]] = {} |
| for row in condition_rows: |
| rows_by_seed.setdefault(int(row["seed"]), []).append(row) |
| seed = sorted(rows_by_seed)[0] |
| seed_rows = sorted(rows_by_seed[seed], key=lambda item: int(item["stage"])) |
| sampled_deltas = [] |
| previous_tokens_seen = 0 |
| for row in seed_rows: |
| tokens_seen = int(row["tokens_seen"]) |
| sampled_deltas.append(tokens_seen - previous_tokens_seen) |
| previous_tokens_seen = tokens_seen |
| first = seed_rows[0] |
| steps_per_stage = sorted({int(row["steps"]) for row in seed_rows}) |
| block_size = int(first["model_config"]["block_size"]) |
| sampled_per_stage = sorted(set(sampled_deltas)) |
| batch_sizes = sorted( |
| { |
| round(delta / (int(row["steps"]) * block_size)) |
| for delta, row in zip(sampled_deltas, seed_rows, strict=True) |
| } |
| ) |
| return { |
| "steps_per_stage": steps_per_stage, |
| "batch_sizes": batch_sizes, |
| "block_size": block_size, |
| "sampled_per_stage": sampled_per_stage, |
| "final_sampled_tokens": int(seed_rows[-1]["tokens_seen"]), |
| "prefix_path": [int(row["token_limit"]) for row in seed_rows], |
| "stage_count": len(stages), |
| } |
|
|
|
|
| def metric_stage_summary(run_dir: Path | list[Path], condition_name: str) -> dict: |
| run_dirs = run_dir if isinstance(run_dir, list) else [run_dir] |
| rows = [] |
| for path in run_dirs: |
| rows.extend(read_jsonl(path / "metrics.jsonl")) |
| stages = sorted({int(row["stage"]) for row in rows if row.get("stage") is not None}) |
| final_stage = max(stages) |
| by_condition_stage: dict[tuple[str, int], list[dict]] = {} |
| for row in rows: |
| if row.get("stage") is None: |
| continue |
| key = (row["condition"], int(row["stage"])) |
| by_condition_stage.setdefault(key, []).append(row) |
|
|
| condition_rows_by_stage = [ |
| row |
| for (condition, _stage), group in by_condition_stage.items() |
| if condition == condition_name |
| for row in group |
| ] |
| condition_rows = by_condition_stage[(condition_name, final_stage)] |
| final_vals = [float(row["val_eval_loss"]) for row in condition_rows] |
| condition_final_mean, condition_final_std = mean_std(final_vals) |
|
|
| final_static_by_condition: dict[str, list[dict]] = {} |
| for row in rows: |
| if int(row.get("stage", -1)) != final_stage: |
| continue |
| if row["condition_kind"] == "static": |
| final_static_by_condition.setdefault(row["condition"], []).append(row) |
| static_means = { |
| condition: statistics.fmean(float(row["val_eval_loss"]) for row in group) |
| for condition, group in final_static_by_condition.items() |
| } |
| best_static_condition = min(static_means, key=static_means.get) |
| best_static_rows = final_static_by_condition[best_static_condition] |
| best_static_vals = [float(row["val_eval_loss"]) for row in best_static_rows] |
| best_static_mean, best_static_std = mean_std(best_static_vals) |
|
|
| final_static_by_seed: dict[int, tuple[str, float]] = {} |
| for row in rows: |
| if int(row.get("stage", -1)) != final_stage or row["condition_kind"] != "static": |
| continue |
| seed = int(row["seed"]) |
| value = float(row["val_eval_loss"]) |
| if seed not in final_static_by_seed or value < final_static_by_seed[seed][1]: |
| final_static_by_seed[seed] = (row["condition"], value) |
|
|
| paired_deltas = [] |
| paired_records = [] |
| for row in sorted(condition_rows, key=lambda item: int(item["seed"])): |
| seed = int(row["seed"]) |
| decay_loss = float(row["val_eval_loss"]) |
| static_condition, static_loss = final_static_by_seed[seed] |
| delta = decay_loss - static_loss |
| paired_deltas.append(delta) |
| paired_records.append( |
| { |
| "seed": seed, |
| "decay_loss": decay_loss, |
| "best_static_condition": static_condition, |
| "best_static_loss": static_loss, |
| "delta": delta, |
| "gain": -delta, |
| } |
| ) |
| paired_mean, paired_std = mean_std(paired_deltas) |
| paired_gains = [-delta for delta in paired_deltas] |
| paired_gain_mean, paired_gain_std = mean_std(paired_gains) |
| paired_gain_ci_low, paired_gain_ci_high = bootstrap_ci(paired_gains) |
|
|
| stage_deltas = [] |
| for stage in stages: |
| condition_group = by_condition_stage[(condition_name, stage)] |
| condition_mean = statistics.fmean( |
| float(row["val_eval_loss"]) for row in condition_group |
| ) |
| static_candidates = [] |
| for (condition, candidate_stage), group in by_condition_stage.items(): |
| if candidate_stage != stage: |
| continue |
| if group[0]["condition_kind"] != "static": |
| continue |
| static_candidates.append( |
| ( |
| condition, |
| statistics.fmean(float(row["val_eval_loss"]) for row in group), |
| statistics.stdev(float(row["val_eval_loss"]) for row in group) |
| if len(group) > 1 |
| else 0.0, |
| ) |
| ) |
| best_static_stage = min(static_candidates, key=lambda item: item[1]) |
| first = condition_group[0] |
| stage_deltas.append( |
| { |
| "stage": stage, |
| "prefix": int(first["token_limit"]), |
| "condition_dropout": statistics.fmean( |
| float(row["dropout_active_final"]) for row in condition_group |
| ), |
| "condition_mean": condition_mean, |
| "condition_std": statistics.stdev( |
| float(row["val_eval_loss"]) for row in condition_group |
| ) |
| if len(condition_group) > 1 |
| else 0.0, |
| "best_static_condition": best_static_stage[0], |
| "best_static_mean": best_static_stage[1], |
| "best_static_std": best_static_stage[2], |
| "delta": condition_mean - best_static_stage[1], |
| } |
| ) |
|
|
| return { |
| "n": len(condition_rows), |
| "final_stage": final_stage, |
| "final_prefix": int(condition_rows[0]["token_limit"]), |
| "condition_final_mean": condition_final_mean, |
| "condition_final_std": condition_final_std, |
| "best_static_condition": best_static_condition, |
| "best_static_mean": best_static_mean, |
| "best_static_std": best_static_std, |
| "improvement": best_static_mean - condition_final_mean, |
| "paired_deltas": paired_deltas, |
| "paired_records": paired_records, |
| "paired_wins": sum(1 for delta in paired_deltas if delta < 0), |
| "paired_mean_delta": paired_mean, |
| "paired_std_delta": paired_std, |
| "paired_gain_mean": paired_gain_mean, |
| "paired_gain_std": paired_gain_std, |
| "paired_gain_ci_low": paired_gain_ci_low, |
| "paired_gain_ci_high": paired_gain_ci_high, |
| "budget": protocol_budget(condition_rows_by_stage, stages), |
| "stage_deltas": stage_deltas, |
| } |
|
|
|
|
| def static_curve_points(run_dir: Path, model: str, prefix: int) -> list[tuple[float, float]]: |
| rows = read_json(run_dir / "summary.json") |
| curve = [] |
| for row in rows: |
| if row["run_mode"] != "screen_static": |
| continue |
| if row["condition_kind"] != "static": |
| continue |
| if row["model_name"] != model: |
| continue |
| if int(row["token_limit"]) != prefix: |
| continue |
| curve.append((float(row["dropout_initial"]), float(row["mean_val_eval_loss"]))) |
| return sorted(curve) |
|
|
|
|
| def calibration_points(cells_path: Path) -> dict[str, list[tuple[int, float]]]: |
| rows = read_csv(cells_path) |
| grouped: dict[str, dict[int, float]] = {} |
| for row in rows: |
| sampled = int(float(row["sampled_tokens"])) |
| unique = int(float(row["unique_tokens"])) |
| source = row["source"] |
| if "sample_pressure_low" in source: |
| continue |
| if "heldout_model" in source: |
| continue |
| model = row["model_name"] |
| target = max(0.0, float(row["target_dropout"])) |
| grouped.setdefault(model, {})[unique] = target |
| return { |
| model: sorted(points.items()) |
| for model, points in sorted(grouped.items()) |
| if len(points) >= 2 |
| } |
|
|
|
|
| def coeff_summary(regime_key: str, paths: dict) -> dict: |
| coeff = read_json(paths["coeff"]) |
| return { |
| "regime": paths["label"], |
| "cells": int(coeff["metrics"]["n"]), |
| "A": float(coeff["coefficients"]["A"]), |
| "B": float(coeff["coefficients"]["B"]), |
| "D": float(coeff["coefficients"]["D"]), |
| "C0": float(coeff["coefficients"]["C0"]), |
| "mae": float(coeff["metrics"]["mae"]), |
| "rmse": float(coeff["metrics"]["rmse"]), |
| "leave_model_mae": float(coeff["cv"]["leave_model"]["mae"]), |
| "leave_prefix_mae": float(coeff["cv"]["leave_prefix"]["mae"]), |
| "formula": coeff["formula"], |
| } |
|
|
|
|
| def tikz_coordinates(points: list[tuple[float, float]], digits: int = 4) -> str: |
| return " ".join(f"({x:.4g},{y:.{digits}f})" for x, y in points) |
|
|
|
|
| def write_calibration_curve_figure() -> None: |
| axes = [] |
| for idx, (key, paths) in enumerate(REGIMES.items()): |
| raw_curve = static_curve_points( |
| paths["curve_run"], paths["curve_model"], paths["curve_prefix"] |
| ) |
| minimum = min(loss for _, loss in raw_curve) |
| points = [(dropout, loss - minimum) for dropout, loss in raw_curve] |
| best = min(raw_curve, key=lambda item: item[1]) |
| axis_options = [ |
| f" title={{{paths['curve_title']}}},", |
| " xlabel={Dropout},", |
| ] |
| if idx == 0: |
| axis_options.append(" ylabel={Validation loss above best},") |
| axis_options.extend( |
| [ |
| " ymin=0,", |
| " grid=both,", |
| ] |
| ) |
| axes.append( |
| rf""" |
| \nextgroupplot[ |
| {chr(10).join(axis_options)} |
| ] |
| \addplot+[thick, mark=*] coordinates {{{tikz_coordinates(points, digits=5)}}}; |
| \addlegendentry{{static sweep}} |
| \addplot+[only marks, mark=star, mark size=3pt] coordinates {{({best[0]:.4g},0)}}; |
| \addlegendentry{{target $p^\star={best[0]:.2f}$}} |
| """ |
| ) |
| content = ( |
| "\\begin{tikzpicture}\n" |
| "\\begin{groupplot}[\n" |
| " group style={group size=3 by 1, horizontal sep=1.0cm},\n" |
| " width=0.32\\textwidth,\n" |
| " height=0.26\\textwidth,\n" |
| " legend style={font=\\scriptsize},\n" |
| " tick label style={font=\\scriptsize},\n" |
| " label style={font=\\scriptsize},\n" |
| " title style={font=\\scriptsize},\n" |
| "]\n" |
| + "\n".join(axes) |
| + "\n\\end{groupplot}\n\\end{tikzpicture}\n" |
| ) |
| (FIG_DIR / "calibration_cell_curves.tex").write_text(content, encoding="utf-8") |
|
|
|
|
| def write_static_optima_figure() -> None: |
| axes = [] |
| for key, paths in REGIMES.items(): |
| grouped = calibration_points(paths["cells"]) |
| labels = sorted({prefix for points in grouped.values() for prefix, _ in points}) |
| label_to_x = {prefix: idx for idx, prefix in enumerate(labels)} |
| plots = [] |
| for model, points in grouped.items(): |
| coords = [ |
| (label_to_x[prefix], target) |
| for prefix, target in points |
| if prefix in label_to_x |
| ] |
| plots.append( |
| "\\addplot+[thick, mark=*] coordinates {" |
| + tikz_coordinates(coords) |
| + "};\n" |
| + f"\\addlegendentry{{{model_label(model)}}}" |
| ) |
| xticks = ",".join(str(i) for i in range(len(labels))) |
| xticklabels = ",".join(token_label(prefix) for prefix in labels) |
| axes.append( |
| rf""" |
| \nextgroupplot[ |
| title={{{paths["label"]}}}, |
| xlabel={{Unique prefix $U$}}, |
| ylabel={{Fitted target dropout}}, |
| xtick={{{xticks}}}, |
| xticklabels={{{xticklabels}}}, |
| ymin=0, |
| ymax=0.65, |
| grid=both, |
| ] |
| {chr(10).join(plots)} |
| """ |
| ) |
| content = ( |
| "\\begin{tikzpicture}\n" |
| "\\begin{groupplot}[\n" |
| " group style={group size=3 by 1, horizontal sep=1.0cm},\n" |
| " width=0.32\\textwidth,\n" |
| " height=0.28\\textwidth,\n" |
| " legend style={font=\\tiny},\n" |
| " tick label style={font=\\scriptsize},\n" |
| " label style={font=\\scriptsize},\n" |
| " title style={font=\\scriptsize},\n" |
| "]\n" |
| + "\n".join(axes) |
| + "\n\\end{groupplot}\n\\end{tikzpicture}\n" |
| ) |
| (FIG_DIR / "static_optima_by_regime.tex").write_text(content, encoding="utf-8") |
|
|
|
|
| def write_l20_stage_delta_figure(summaries: dict[str, dict]) -> None: |
| axes = [] |
| for key, paths in REGIMES.items(): |
| stage_rows = summaries[key]["l20"]["stage_deltas"] |
| labels = [token_label(row["prefix"]) for row in stage_rows] |
| coords = [(idx, row["delta"]) for idx, row in enumerate(stage_rows)] |
| zero_coords = [(0, 0.0), (len(stage_rows) - 1, 0.0)] |
| xticks = ",".join(str(i) for i in range(len(labels))) |
| xticklabels = ",".join(labels) |
| axes.append( |
| rf""" |
| \nextgroupplot[ |
| title={{{paths["label"]}}}, |
| xlabel={{Prefix}}, |
| ylabel={{Decay $-$ best static}}, |
| xtick={{{xticks}}}, |
| xticklabels={{{xticklabels}}}, |
| grid=both, |
| ] |
| \addplot+[thick, mark=*] coordinates {{{tikz_coordinates(coords, digits=5)}}}; |
| \addplot+[black, dashed, mark=none] coordinates {{{tikz_coordinates(zero_coords, digits=5)}}}; |
| """ |
| ) |
| content = ( |
| "\\begin{tikzpicture}\n" |
| "\\begin{groupplot}[\n" |
| " group style={group size=3 by 1, horizontal sep=1.0cm},\n" |
| " width=0.32\\textwidth,\n" |
| " height=0.27\\textwidth,\n" |
| " tick label style={font=\\scriptsize},\n" |
| " label style={font=\\scriptsize},\n" |
| " title style={font=\\scriptsize},\n" |
| "]\n" |
| + "\n".join(axes) |
| + "\n\\end{groupplot}\n\\end{tikzpicture}\n" |
| ) |
| (FIG_DIR / "l20_stage_deltas.tex").write_text(content, encoding="utf-8") |
|
|
|
|
| def write_final_improvement_figure(summaries: dict[str, dict]) -> None: |
| labels = [] |
| coords_small = [] |
| coords_l20 = [] |
| for idx, (key, paths) in enumerate(REGIMES.items()): |
| labels.append(paths["short"]) |
| coords_small.append((idx, summaries[key]["small"]["improvement"])) |
| coords_l20.append((idx, summaries[key]["l20"]["improvement"])) |
| xticks = ",".join(str(i) for i in range(len(labels))) |
| xticklabels = ",".join(labels) |
| content = rf""" |
| \begin{{tikzpicture}} |
| \begin{{axis}}[ |
| ybar, |
| bar width=8pt, |
| width=0.72\textwidth, |
| height=0.34\textwidth, |
| ymin=0, |
| ylabel={{Best static final loss $-$ decay final loss}}, |
| symbolic x coords={{{xticklabels}}}, |
| xtick=data, |
| xtick={{{xticklabels}}}, |
| grid=both, |
| legend style={{font=\scriptsize, at={{(0.5,1.04)}}, anchor=south, legend columns=2}}, |
| tick label style={{font=\scriptsize}}, |
| label style={{font=\scriptsize}}, |
| ] |
| \addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_small))}}}; |
| \addlegendentry{{4M validation}} |
| \addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_l20))}}}; |
| \addlegendentry{{L20 8M extrapolation}} |
| \end{{axis}} |
| \end{{tikzpicture}} |
| """ |
| (FIG_DIR / "final_improvements.tex").write_text(content, encoding="utf-8") |
|
|
|
|
| def write_protocol_schematic() -> None: |
| content = r""" |
| \begin{tikzpicture}[font=\scriptsize] |
| \tikzstyle{box}=[ |
| draw, |
| rounded corners=2pt, |
| thick, |
| align=center, |
| text width=0.19\textwidth, |
| minimum height=1.5cm, |
| inner sep=5pt |
| ] |
| \node[box, fill=blue!7] at (0,0) (sweep) {Static calibration cells\\ |
| \footnotesize model $P$, prefix $U$, sampled tokens $C$, dropout grid}; |
| \node[box, fill=green!8] at (3.7,0) (target) {Extract $p^\star$\\ |
| \footnotesize grid best or bracketed quadratic minimum}; |
| \node[box, fill=yellow!12] at (7.4,0) (fit) {Fit pressure law\\ |
| \footnotesize $p^\star \approx A x + B y + Dxy + C_0$}; |
| \node[box, fill=red!7] at (11.1,0) (stream) {Frozen streaming test\\ |
| \footnotesize compare decay path to best static dropout grid}; |
| \draw[->, thick] (sweep.east) -- (target.west); |
| \draw[->, thick] (target.east) -- (fit.west); |
| \draw[->, thick] (fit.east) -- (stream.west); |
| \node[align=center, text width=0.9\textwidth] at (5.55,-1.45) { |
| Inputs to the fitted law are absolute quantities available before or during a run: |
| $x=\log_{10}(P/U)$ measures capacity pressure and |
| $y=\log_{10}(C/U)$ measures repeated-sampling pressure. |
| }; |
| \end{tikzpicture} |
| """ |
| (FIG_DIR / "calibration_to_streaming_protocol.tex").write_text( |
| content.strip() + "\n", encoding="utf-8" |
| ) |
|
|
|
|
| def write_tables(summaries: dict[str, dict], coeffs: dict[str, dict]) -> None: |
| symbol_lines = [ |
| "\\begin{tabular}{ll}", |
| "\\toprule", |
| "Symbol & Meaning \\\\", |
| "\\midrule", |
| "$P$ & Trainable parameter count \\\\", |
| "$U_t$ & Unique-token prefix currently revealed by stream stage $t$ \\\\", |
| "$C_t$ & Cumulative sampled training tokens consumed by the optimizer \\\\", |
| "$x_t$ & $\\log_{10}(P/U_t)$, capacity pressure \\\\", |
| "$y_t$ & $\\log_{10}(C_t/U_t)$, repeated-sampling pressure \\\\", |
| "$p_t$ & Dropout rate used at stream stage $t$ \\\\", |
| "$p^\\star$ & Static-sweep target dropout for one calibration cell \\\\", |
| "$A,B,D,C_0$ & Regime-specific fitted pressure-law coefficients \\\\", |
| "$p_{\\min},p_{\\max}$ & Clamp bounds applied to generated dropout schedules \\\\", |
| "\\bottomrule", |
| "\\end{tabular}", |
| ] |
| (TABLE_DIR / "symbol_table.tex").write_text( |
| "\n".join(symbol_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| regime_lines = [ |
| "\\begin{tabular}{llll}", |
| "\\toprule", |
| "Regime & Corpus source & Tokenizer / vocab & Streaming prefixes \\\\", |
| "\\midrule", |
| ( |
| "OpenWebText10K & Cached OpenWebText-style subset & BPE / 4096 & " |
| "250k, 500k, 1M, 2M, 4M; plus 8M L20 tier \\\\" |
| ), |
| ( |
| "TinyStories & Cached TinyStories subset & BPE / 4096 & " |
| "500k, 1M, 2M, 4M; plus 8M L20 tier \\\\" |
| ), |
| ( |
| "WikiText-103 & Cached WikiText-103 subset & BPE / 4096 & " |
| "250k, 500k, 1M, 2M, 4M; plus 8M L20 tier \\\\" |
| ), |
| "\\bottomrule", |
| "\\end{tabular}", |
| ] |
| (TABLE_DIR / "regime_setup_table.tex").write_text( |
| "\n".join(regime_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| repro_lines = [ |
| "\\begin{tabular}{ll}", |
| "\\toprule", |
| "Item & Value \\\\", |
| "\\midrule", |
| "Repository & \\url{https://huggingface.co/cuber12/dropout-decay} \\\\", |
| "Archived snapshot DOI & \\url{https://doi.org/10.5281/zenodo.20616633} \\\\", |
| "Repository snapshot & Git tag \\texttt{v1.1-preprint} \\\\", |
| "License & MIT-derived source with nanochat attribution retained \\\\", |
| "Python & 3.11.15 \\\\", |
| "PyTorch & 2.12.0 \\\\", |
| "Operating system & macOS 26.4 arm64 \\\\", |
| "Training backend & Apple Metal Performance Shaders (MPS); runner refuses CPU/CUDA \\\\", |
| "\\bottomrule", |
| "\\end{tabular}", |
| ] |
| (TABLE_DIR / "reproducibility_table.tex").write_text( |
| "\n".join(repro_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| weight_lines = [ |
| "\\begin{tabular}{lll}", |
| "\\toprule", |
| "Cell property & Weight multiplier & Reason \\\\", |
| "\\midrule", |
| "Best grid point on dropout boundary & $0.30$ & Optimum may lie outside tested grid \\\\", |
| "Quadratic minimum not bracketed & $0.50$ & Local parabola is less reliable \\\\", |
| "Loss span below $0.02$ nats & $0.50$ & Very flat curve weakly identifies $p^\\star$ \\\\", |
| ( |
| "Best-loss std. $\\sigma > 0$ & $1/(1+20\\sigma)$ & " |
| "Noisier minima receive less influence \\\\" |
| ), |
| "\\bottomrule", |
| "\\end{tabular}", |
| ] |
| (TABLE_DIR / "calibration_weight_table.tex").write_text( |
| "\n".join(weight_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| protocol_lines = [ |
| "\\begin{tabular}{llll}", |
| "\\toprule", |
| "Experiment layer & What is run & Question answered & Statistical unit \\\\", |
| "\\midrule", |
| ( |
| "Static calibration & Per-regime dropout grid at fixed $(P,U,C)$ cells & " |
| "Where does validation loss plateau over dropout? & 15--16 cells per regime \\\\" |
| ), |
| ( |
| "Coefficient fit & Per-regime regression from cell targets $p^\\star$ & " |
| "Can pressure variables predict useful dropout? & MAE / leave-out MAE \\\\" |
| ), |
| ( |
| "4M final-prefix locked stream & Per-regime continuing model over growing prefixes & " |
| "Does decay beat the best tested fixed dropout? & 5 paired seeds per regime \\\\" |
| ), |
| ( |
| "L20 8M extrapolation & Per-regime 51.84M-param model with frozen coefficients & " |
| "Does the law extend to larger model/data scale? & 3 paired seeds per regime \\\\" |
| ), |
| ( |
| "Stage-wise audit & Per-regime decay vs best static at every prefix & " |
| "Is the schedule final-loss or per-prefix optimal? & Prefix-level deltas \\\\" |
| ), |
| "\\bottomrule", |
| "\\end{tabular}", |
| ] |
| (TABLE_DIR / "experiment_protocol_table.tex").write_text( |
| "\n".join(protocol_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| budget_lines = [ |
| "\\begin{tabular}{llllrr}", |
| "\\toprule", |
| "Regime & Protocol & Prefix path & Sampled/stage & Batch & Block \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| for protocol_key, protocol_label in [ |
| ("small", paths["small_protocol"]), |
| ("l20", "8M final prefix / 3 seeds"), |
| ]: |
| budget = summaries[key][protocol_key]["budget"] |
| prefix_path = " $\\rightarrow$ ".join(token_label(prefix) for prefix in budget["prefix_path"]) |
| sampled = ", ".join(token_label(value) for value in budget["sampled_per_stage"]) |
| batch = "/".join(str(value) for value in budget["batch_sizes"]) |
| block = str(budget["block_size"]) |
| budget_lines.append( |
| f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & " |
| f"{prefix_path} & {sampled} & {batch} & {block} \\\\" |
| ) |
| budget_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "protocol_budget_table.tex").write_text( |
| "\n".join(budget_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| coeff_lines = [ |
| "\\begin{tabular}{lrrrrrrrr}", |
| "\\toprule", |
| "Regime & Cells & A & B & D & C0 & Dropout MAE (abs. $p$) & L-model MAE & L-prefix MAE \\\\", |
| "\\midrule", |
| ] |
| for key in REGIMES: |
| row = coeffs[key] |
| coeff_lines.append( |
| f"{tex_escape(row['regime'])} & {row['cells']} & " |
| f"{row['A']:.4f} & {row['B']:.4f} & {row['D']:.4f} & {row['C0']:.4f} & " |
| f"{row['mae']:.4f} & {row['leave_model_mae']:.4f} & {row['leave_prefix_mae']:.4f} \\\\" |
| ) |
| coeff_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "coefficient_table.tex").write_text( |
| "\n".join(coeff_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| evidence_lines = [ |
| "\\begin{tabular}{lrrrrr}", |
| "\\toprule", |
| "Regime & Calib. dropout MAE & 4M gain (nats) & 4M wins & 8M gain (nats) & 8M wins \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| small = summaries[key]["small"] |
| l20 = summaries[key]["l20"] |
| evidence_lines.append( |
| f"{tex_escape(paths['label'])} & {coeffs[key]['mae']:.4f} & " |
| f"{small['improvement']:.4f} & {small['paired_wins']}/{small['n']} & " |
| f"{l20['improvement']:.4f} & {l20['paired_wins']}/{l20['n']} \\\\" |
| ) |
| evidence_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "evidence_ladder_table.tex").write_text( |
| "\n".join(evidence_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| val_lines = [ |
| "\\begin{tabular}{llrrrrr}", |
| "\\toprule", |
| "Regime & Protocol & Decay final (nats) & Best static final (nats) & Gain (nats) & Seeds & Paired wins \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| for protocol_key, protocol_label in [ |
| ("small", paths["small_protocol"]), |
| ("l20", "8M final prefix / 3 seeds"), |
| ]: |
| row = summaries[key][protocol_key] |
| val_lines.append( |
| f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & " |
| f"{row['condition_final_mean']:.4f} $\\pm$ {row['condition_final_std']:.4f} & " |
| f"{row['best_static_mean']:.4f} $\\pm$ {row['best_static_std']:.4f} & " |
| f"{row['improvement']:.4f} & {row['n']} & {row['paired_wins']}/{row['n']} \\\\" |
| ) |
| val_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "streaming_validation_table.tex").write_text( |
| "\n".join(val_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| paired_lines = [ |
| "\\begin{tabular}{llrrrr}", |
| "\\toprule", |
| "Regime & Protocol & Paired gain mean & 95\\% bootstrap CI & Seeds & Wins \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| for protocol_key, protocol_label in [ |
| ("small", paths["small_protocol"]), |
| ("l20", "8M final prefix / 3 seeds"), |
| ]: |
| row = summaries[key][protocol_key] |
| paired_lines.append( |
| f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & " |
| f"{row['paired_gain_mean']:.4f} & " |
| f"[{row['paired_gain_ci_low']:.4f}, {row['paired_gain_ci_high']:.4f}] & " |
| f"{row['n']} & {row['paired_wins']}/{row['n']} \\\\" |
| ) |
| paired_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "paired_gain_ci_table.tex").write_text( |
| "\n".join(paired_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| path_lines = [ |
| "\\begin{tabular}{llll}", |
| "\\toprule", |
| "Regime & Locked-stream prefix path & Frozen decay path & Final best static \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| stage_rows = summaries[key]["l20"]["stage_deltas"] |
| prefix_path = " $\\rightarrow$ ".join(token_label(row["prefix"]) for row in stage_rows) |
| dropout_path = " $\\rightarrow$ ".join( |
| f"{row['condition_dropout']:.3f}" for row in stage_rows |
| ) |
| final_static = summaries[key]["l20"]["best_static_condition"].replace( |
| "static_dropout_", "p=" |
| ) |
| path_lines.append( |
| f"{tex_escape(paths['label'])} & {prefix_path} & {dropout_path} & " |
| f"{tex_escape(final_static)} \\\\" |
| ) |
| path_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "l20_schedule_path_table.tex").write_text( |
| "\n".join(path_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| l20_lines = [ |
| "\\begin{tabular}{lrrrrr}", |
| "\\toprule", |
| "Regime & Prefix & Decay dropout & Decay val (nats) & Best static val (nats) & Delta (nats) \\\\", |
| "\\midrule", |
| ] |
| for key, paths in REGIMES.items(): |
| for row in summaries[key]["l20"]["stage_deltas"]: |
| l20_lines.append( |
| f"{tex_escape(paths['label'])} & {token_label(row['prefix'])} & " |
| f"{row['condition_dropout']:.3f} & {row['condition_mean']:.4f} & " |
| f"{row['best_static_mean']:.4f} & {row['delta']:+.4f} \\\\" |
| ) |
| l20_lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (TABLE_DIR / "l20_stage_table.tex").write_text( |
| "\n".join(l20_lines) + "\n", encoding="utf-8" |
| ) |
|
|
| summary = { |
| key: { |
| "coefficients": coeffs[key], |
| "small_stream": summaries[key]["small"], |
| "l20_stream": summaries[key]["l20"], |
| } |
| for key in REGIMES |
| } |
| (PAPER_DIR / "paper_results_summary.json").write_text( |
| json.dumps(summary, indent=2), encoding="utf-8" |
| ) |
|
|
|
|
| def main() -> None: |
| FIG_DIR.mkdir(parents=True, exist_ok=True) |
| TABLE_DIR.mkdir(parents=True, exist_ok=True) |
| summaries = {} |
| coeffs = {} |
| for key, paths in REGIMES.items(): |
| coeffs[key] = coeff_summary(key, paths) |
| summaries[key] = { |
| "small": metric_stage_summary(paths["small_stream"], paths["small_condition"]), |
| "l20": metric_stage_summary(paths["l20_stream"], paths["l20_condition"]), |
| } |
| write_calibration_curve_figure() |
| write_static_optima_figure() |
| write_protocol_schematic() |
| write_l20_stage_delta_figure(summaries) |
| write_final_improvement_figure(summaries) |
| write_tables(summaries, coeffs) |
| print(json.dumps({"figures": str(FIG_DIR), "tables": str(TABLE_DIR)}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
