Refactor experiment pipeline and add regime paper

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+paper/*.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -2,6 +2,12 @@
 __pycache__/
 *.py[cod]
 *.pdf
+*.aux
+*.blg
+*.fdb_latexmk
+*.fls
+*.log
+*.out
 *.npy
 .venv/
 *.egg-info/

paper/dropout_decay_regime_specific_pressure_law.bbl

@@ -0,0 +1,73 @@
+\begin{thebibliography}{8}
+\providecommand{\natexlab}[1]{#1}
+\providecommand{\url}[1]{\texttt{#1}}
+\expandafter\ifx\csname urlstyle\endcsname\relax
+  \providecommand{\doi}[1]{doi: #1}\else
+  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
+
+\bibitem[Eldan and Li(2023)]{eldan2023tinystories}
+Ronen Eldan and Yuanzhi Li.
+\newblock Tinystories: How small can language models be and still speak
+  coherent english?
+\newblock \emph{arXiv preprint arXiv:2305.07759}, 2023.
+\newblock URL \url{https://arxiv.org/abs/2305.07759}.
+
+\bibitem[Gokaslan et~al.(2019)Gokaslan, Cohen, Pavlick, and
+  Tellex]{gokaslan2019openwebtext}
+Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex.
+\newblock Openwebtext corpus.
+\newblock \url{https://zenodo.org/records/3834942}, 2019.
+\newblock Open reproduction of the WebText-style corpus.
+
+\bibitem[Hoffmann et~al.(2022)Hoffmann, Borgeaud, Mensch, Buchatskaya, Cai,
+  Rutherford, Casas, Hendricks, Welbl, Clark, Hennigan, Noland, Millican,
+  van~den Driessche, Damoc, Guy, Osindero, Simonyan, Elsen, Rae, Vinyals, and
+  Sifre]{hoffmann2022training}
+Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor
+  Cai, Eliza Rutherford, Diego de~Las Casas, Lisa~Anne Hendricks, Johannes
+  Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van~den
+  Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich
+  Elsen, Jack~W. Rae, Oriol Vinyals, and Laurent Sifre.
+\newblock Training compute-optimal large language models.
+\newblock \emph{arXiv preprint arXiv:2203.15556}, 2022.
+\newblock URL \url{https://arxiv.org/abs/2203.15556}.
+
+\bibitem[Kaplan et~al.(2020)Kaplan, McCandlish, Henighan, Brown, Chess, Child,
+  Gray, Radford, Wu, and Amodei]{kaplan2020scaling}
+Jared Kaplan, Sam McCandlish, Tom Henighan, Tom~B. Brown, Benjamin Chess, Rewon
+  Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.
+\newblock Scaling laws for neural language models.
+\newblock \emph{arXiv preprint arXiv:2001.08361}, 2020.
+\newblock URL \url{https://arxiv.org/abs/2001.08361}.
+
+\bibitem[Karpathy(2025)]{karpathy_nanochat}
+Andrej Karpathy.
+\newblock nanochat.
+\newblock \url{https://github.com/karpathy/nanochat}, 2025.
+\newblock MIT-licensed repository, accessed 2026-05-29.
+
+\bibitem[Merity et~al.(2017)Merity, Xiong, Bradbury, and
+  Socher]{merity2017pointer}
+Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.
+\newblock Pointer sentinel mixture models.
+\newblock In \emph{International Conference on Learning Representations}, 2017.
+\newblock URL \url{https://arxiv.org/abs/1609.07843}.
+
+\bibitem[Srivastava et~al.(2014)Srivastava, Hinton, Krizhevsky, Sutskever, and
+  Salakhutdinov]{srivastava2014dropout}
+Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan
+  Salakhutdinov.
+\newblock Dropout: A simple way to prevent neural networks from overfitting.
+\newblock \emph{Journal of Machine Learning Research}, 15\penalty0
+  (56):\penalty0 1929--1958, 2014.
+\newblock URL \url{https://www.jmlr.org/papers/v15/srivastava14a.html}.
+
+\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez,
+  Kaiser, and Polosukhin]{vaswani2017attention}
+Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
+  Aidan~N. Gomez, Lukasz Kaiser, and Illia Polosukhin.
+\newblock Attention is all you need.
+\newblock In \emph{Advances in Neural Information Processing Systems}, 2017.
+\newblock URL \url{https://arxiv.org/abs/1706.03762}.
+
+\end{thebibliography}

paper/dropout_decay_regime_specific_pressure_law.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46ba08475ba2566bc3733eed5403db061a3ce4145ed6bd79965796719a396bc8
+size 267426

paper/dropout_decay_regime_specific_pressure_law.tex

@@ -0,0 +1,427 @@
+\documentclass[11pt]{article}
+
+\usepackage[margin=1in]{geometry}
+\usepackage{amsmath}
+\usepackage{amssymb}
+\usepackage{booktabs}
+\usepackage{graphicx}
+\usepackage{microtype}
+\usepackage{tikz}
+\usepackage{pgfplots}
+\usepgfplotslibrary{groupplots}
+\pgfplotsset{compat=1.18}
+\usepackage[numbers,sort&compress]{natbib}
+\usepackage[colorlinks=true,linkcolor=blue,citecolor=blue,urlcolor=blue]{hyperref}
+
+\title{Regime-Specific Pressure Laws for Dropout Decay in Streaming Language Model Training}
+\author{Anonymous Authors}
+\date{June 2026}
+
+\newcommand{\clamp}{\operatorname{clamp}}
+\newcommand{\params}{P}
+\newcommand{\unique}{U_t}
+\newcommand{\sampled}{C_t}
+\newcommand{\dropout}{p_t}
+
+\begin{document}
+\maketitle
+
+\begin{abstract}
+Dropout is usually selected as a fixed hyperparameter before training, but
+streaming language-model training changes the regularization problem over time:
+early stages expose few unique tokens to a comparatively large model, while
+later stages reveal more data and reduce the need for strong stochastic
+regularization. We study this setting with nanochat-derived causal
+Transformers trained under an expanding-prefix protocol on three text regimes:
+OpenWebText10K, TinyStories, and WikiText-103. Static dropout sweeps show a
+repeatable pattern: for a fixed model, the validation-loss curve over dropout
+is typically unimodal or plateau-shaped, and the optimum falls as unique data
+increases; for a fixed prefix, the optimum rises with model capacity. We turn
+these per-cell optima into a regime-specific pressure law,
+\[
+\dropout=\clamp(p_{\min},p_{\max},
+ A\log_{10}(\params/\unique)
++B\log_{10}(\sampled/\unique)
++D\log_{10}(\params/\unique)\log_{10}(\sampled/\unique)+C_0),
+\]
+where $\params$ is parameter count, $\unique$ is currently available unique
+tokens, and $\sampled$ is cumulative sampled training tokens. The fitted
+interaction law explains static optimum cells with mean absolute errors of
+0.0158, 0.0180, and 0.0162 dropout on OpenWebText10K, TinyStories, and
+WikiText-103 respectively. In locked streaming validation, fitted or
+formula-derived decay schedules beat the best fixed-dropout baseline on final
+validation loss in all three regimes at 4M-token validation scale and again in
+larger 51.84M-parameter, 8M-token extrapolations. The large extrapolation tier
+wins all paired seeds across all three regimes (9/9 paired wins), but stage-wise
+analysis shows the schedules are final-loss optimized rather than best at every
+intermediate prefix. The results support a practical hypothesis: dropout decay
+should be calibrated as a regime-specific pressure law, not chosen as a
+universal constant or an arbitrary high-to-low schedule.
+\end{abstract}
+
+\section{Introduction}
+
+Dropout was introduced as a simple and effective regularizer for neural
+networks \citep{srivastava2014dropout}. In language model practice, however,
+dropout is often treated as a fixed scalar selected before a run. That choice
+is mismatched to streaming settings. When a model begins training on a small
+prefix of a stream, its parameter count can be large relative to the number of
+unique tokens currently available; later, the same model receives more unique
+data and can safely expose more functional capacity. A fixed dropout rate must
+therefore compromise between early overfitting risk and late underfitting risk.
+
+The motivating hypothesis of this project was intentionally ambitious: a
+high-capacity model could start with strong regularization on small initial
+data and then decay dropout as streaming data accumulates, allowing capacity to
+come online without catastrophic overfitting. The experiments refined that
+hypothesis. Extremely high initial dropout is not automatically useful; early
+tests with dropout near 0.8 underfit badly. The stronger and better-supported
+claim is that useful dropout is a function of training pressure: model capacity
+relative to the current data prefix, and cumulative optimizer sampling relative
+to that prefix.
+
+This paper makes four empirical contributions.
+
+\begin{enumerate}
+    \item We define calibration cells for streaming dropout: one model, one
+    unique-token prefix, one sampled-token budget, and a sweep over fixed
+    dropout rates.
+    \item We show that calibration cells often have a parabolic or plateau-like
+    validation-loss curve over dropout, making the local optimum a usable
+    supervised target.
+    \item We fit a regime-specific interaction pressure law that maps
+    $(\params,\unique,\sampled)$ to a dropout rate.
+    \item We validate the resulting schedules against broad fixed-dropout
+    baselines across three text regimes and a larger-model/larger-stream
+    extrapolation tier.
+\end{enumerate}
+
+The central result is not that one coefficient vector works everywhere. It does
+not. The central result is that the pressure-law structure is useful, while the
+coefficients are regime-specific.
+
+\section{Related Work}
+
+\paragraph{Dropout and regularization.}
+Dropout applies multiplicative Bernoulli noise during training and can be
+understood as an approximate ensemble over subnetworks
+\citep{srivastava2014dropout}. With inverted dropout, an activation $h_i$ is
+transformed as $h'_i=h_i m_i/(1-p)$, $m_i\sim\operatorname{Bernoulli}(1-p)$.
+The expectation is preserved, but the variance grows as $h_i^2p/(1-p)$. This
+captures the key tradeoff in our experiments: low dropout can overfit repeated
+small prefixes, but high dropout can remove too much usable capacity.
+
+\paragraph{Transformer language models and scaling.}
+Our model family is a decoder-only Transformer following the attention-based
+architecture of \citet{vaswani2017attention}. Scaling-law work shows that loss
+varies systematically with model size, data size, and compute
+\citep{kaplan2020scaling,hoffmann2022training}. Our study is narrower: it does
+not propose a global loss scaling law. Instead, it asks whether a
+regularization hyperparameter can be mapped to local pressure variables during
+streaming.
+
+\paragraph{Text regimes.}
+We evaluate on three text regimes. TinyStories is a synthetic simple-English
+story corpus designed for small language-model studies \citep{eldan2023tinystories}.
+WikiText-103 is a standard word-level language-modeling benchmark derived from
+Wikipedia articles \citep{merity2017pointer}. OpenWebText is an open
+WebText-style corpus \citep{gokaslan2019openwebtext}. We use cached/tokenized
+subsets and therefore report results as regime-level evidence rather than
+claims about full-corpus state of the art.
+
+\section{Experimental Setting}
+
+\subsection{Model and implementation}
+
+All experiments use a standalone research codebase derived from Andrej
+Karpathy's MIT-licensed nanochat repository \citep{karpathy_nanochat}. Chat
+interfaces and deployment scripts are removed; the retained components are the
+core BPE tokenizer and causal Transformer template. Derived source files retain
+the original copyright and permission notice.
+
+The architecture family keeps the standard Transformer feed-forward expansion
+ratio of four. The central model sizes are L8-H8-D256 (8.39M parameters),
+L12-H8-D320 (17.37M), L16-H8-D384 (31.46M or 32.71M depending on
+vocabulary/run metadata), and L20-H8-D448 (51.84M). The tokenizer vocabulary is
+4096 in the reported runs. Unless noted otherwise, training uses AdamW,
+learning rate $3\times 10^{-4}$, weight decay 0.1, gradient clipping 1.0, batch
+size 16, and block size 128. All Torch training runs are MPS-only; the runner
+refuses CUDA, CPU execution, and MPS fallback.
+
+\subsection{Streaming protocol}
+
+We simulate streaming by revealing larger token prefixes over time. A locked
+streaming run keeps the same model and optimizer state across stages. At stage
+$t$, minibatches are sampled only from the currently revealed prefix
+$\unique$. The two main protocols are:
+
+\begin{itemize}
+    \item 4M validation: OpenWebText10K and WikiText-103 use 250k, 500k, 1M,
+    2M, and 4M prefixes. TinyStories uses 500k, 1M, 2M, and 4M prefixes. These
+    runs use five random seeds.
+    \item 8M extrapolation: L20-H8-D448 with an 8M final prefix. These runs use
+    three seeds because they are substantially more expensive.
+\end{itemize}
+
+Each run compares a dropout schedule against a grid of static-dropout
+baselines. The best static baseline is selected after training from the same
+run family, so the decay schedule competes against a strong fixed-rate oracle
+within the tested grid.
+
+\section{From Calibration Cells to a Pressure Law}
+
+\subsection{Calibration cell definition}
+
+A calibration cell is one static dropout sweep at a fixed tuple:
+\[
+(\text{model architecture}, \params, \unique, \sampled, \text{validation split}).
+\]
+Within that cell, we train/evaluate several fixed dropout values and obtain a
+curve
+\[
+p \mapsto L_{\mathrm{val}}(p \mid \params,\unique,\sampled).
+\]
+The target dropout $p^\star$ is the best observed grid point or, when the best
+point is bracketed by worse neighbors, a local quadratic interpolation of the
+minimum. Boundary optima are retained but downweighted in coefficient fitting
+because the true optimum may lie outside the tested grid.
+
+Figure~\ref{fig:cell-curves} shows representative cells. The curves are not
+perfect mathematical parabolas, but they show the practical pattern needed for
+calibration: validation loss improves as dropout moves away from zero, reaches
+a local minimum or plateau, and then rises when dropout becomes too large.
+
+\begin{figure}[t]
+\centering
+\input{figures/calibration_cell_curves}
+\caption{Representative static dropout calibration cells. The y-axis is
+validation loss above the best tested dropout in that cell, so each curve is
+normalized to show shape rather than absolute corpus difficulty.}
+\label{fig:cell-curves}
+\end{figure}
+
+\subsection{Pressure variables}
+
+The formula uses two dimensionless pressure variables:
+\[
+x_t=\log_{10}(\params/\unique), \qquad
+y_t=\log_{10}(\sampled/\unique).
+\]
+The first term measures capacity pressure. Increasing parameters at fixed
+prefix increases memorization capacity and tends to increase useful dropout.
+Increasing unique tokens at fixed model reduces this pressure. The second term
+measures repeated-sampling pressure. If the optimizer consumes many sampled
+tokens while the unique prefix remains small, it repeatedly revisits the same
+data and can benefit from stronger regularization.
+
+The fitted interaction family is
+\begin{equation}
+\dropout=\clamp(p_{\min},p_{\max},
+ A x_t+B y_t+D x_t y_t+C_0).
+\label{eq:interaction-law}
+\end{equation}
+For schedule generation we use $p_{\min}=0.02$ and $p_{\max}=0.65$ unless a
+run explicitly specifies a different clamp. Static sweeps may still include
+0.0 as an ablation.
+
+\subsection{Regime-specific coefficients}
+
+We define a regime as the combination of architecture family, tokenizer, corpus
+family, optimizer protocol, dropout placement, stream protocol, and evaluation
+distribution. Model size and prefix size are not regime changes; they are
+inputs to the formula. Corpus family, tokenizer, or optimizer changes can shift
+the coefficients.
+
+Table~\ref{tab:coefficients} gives the fitted interaction coefficients. Each
+regime uses 15 or 16 calibration cells. The fitted mean absolute error is about
+0.016--0.018 dropout, meaning the coefficient law typically predicts the
+calibrated static optimum to within roughly two hundredths of dropout.
+
+\begin{table}[t]
+\centering
+\caption{Regime-specific interaction coefficient fits. L-model and L-prefix
+are leave-model-out and leave-prefix-out mean absolute errors.}
+\label{tab:coefficients}
+\resizebox{\textwidth}{!}{\input{tables/coefficient_table}}
+\end{table}
+
+Figure~\ref{fig:static-optima} shows the fitted target dropout values from the
+calibration cells. Across regimes, useful dropout generally decreases as the
+unique prefix grows. At a fixed prefix, larger models tend to require larger
+dropout. The absolute surface differs by corpus regime, which is why the paper
+does not propose universal coefficients.
+
+\begin{figure}[t]
+\centering
+\input{figures/static_optima_by_regime}
+\caption{Calibration targets by regime. At fixed model size, target dropout
+tends to fall as more unique tokens are available; at fixed prefix, larger
+models tend to have higher target dropout. WikiText-103 reaches the dropout
+floor earlier than the other regimes under its calibration protocol.}
+\label{fig:static-optima}
+\end{figure}
+
+\section{Streaming Results}
+
+\subsection{Main validation}
+
+Table~\ref{tab:streaming-validation} summarizes the main streaming results.
+The ``Gain'' column is best static final validation loss minus decay final
+validation loss, so positive values favor decay. The best static baseline is
+chosen from the tested fixed-dropout grid in the same run family.
+
+\begin{table}[t]
+\centering
+\caption{Streaming validation against best fixed-dropout baselines.}
+\label{tab:streaming-validation}
+\resizebox{\textwidth}{!}{\input{tables/streaming_validation_table}}
+\end{table}
+
+OpenWebText10K provides the cleanest five-seed coefficient-derived validation:
+the interaction schedule improves final validation loss by 0.0473 over the
+best static baseline and wins all five paired seeds. TinyStories shows a
+smaller but consistent final-loss margin: 0.0133 over the best static baseline
+with five paired wins. WikiText-103 is more nuanced. The earlier
+formula-derived L12 schedule wins the five-seed 4M run, but the later
+static-cell-fitted interaction schedule decays too quickly for L12 and is only
+near the best static baseline. This negative result is important: a good static
+cell fit does not by itself guarantee the best streaming schedule unless the
+schedule objective is aligned with final-loss optimization.
+
+\subsection{L20 8M extrapolation}
+
+The strongest comparable evidence is the L20 extrapolation tier. Here the
+model is fixed to L20-H8-D448 with 51.84M parameters, and each regime uses
+its fitted interaction coefficients to produce a frozen schedule before
+training. No coefficients are refit on the L20 results. Across OpenWebText10K,
+TinyStories, and WikiText-103, the decay schedule beats the best static
+baseline on final validation loss and wins all paired seeds (3/3 in each
+regime, 9/9 total).
+
+Figure~\ref{fig:final-improvements} visualizes the final-loss gains. The
+largest L20 gain is OpenWebText10K (0.0363), followed by WikiText-103 (0.0184)
+and TinyStories (0.0100). The TinyStories gain is smaller because a static
+dropout near 0.20 is already a strong fixed compromise for that regime.
+
+\begin{figure}[t]
+\centering
+\input{figures/final_improvements}
+\caption{Final validation-loss improvement over the best fixed-dropout
+baseline. Positive values favor the decay schedule.}
+\label{fig:final-improvements}
+\end{figure}
+
+\subsection{Stage-wise behavior}
+
+Final-loss improvement does not imply per-prefix dominance. Table~\ref{tab:l20-stage}
+and Figure~\ref{fig:l20-stage-deltas} compare the L20 schedule against the best
+static baseline at each stage. OpenWebText10K starts worse than static 0.30 at
+250k but becomes better from 1M onward. TinyStories is nearly tied early,
+slightly worse at 1M/2M, then better at 4M/8M. WikiText-103 is worse through
+the middle of the stream and wins only at the final 8M prefix.
+
+\begin{table}[t]
+\centering
+\caption{L20 stage-wise comparison. Delta is decay validation loss minus the
+best static validation loss at the same prefix; negative favors decay.}
+\label{tab:l20-stage}
+\resizebox{\textwidth}{!}{\input{tables/l20_stage_table}}
+\end{table}
+
+\begin{figure}[t]
+\centering
+\input{figures/l20_stage_deltas}
+\caption{L20 stage-wise deltas versus the best static dropout at each prefix.
+The schedules are final-loss schedules, not greedy per-prefix optima.}
+\label{fig:l20-stage-deltas}
+\end{figure}
+
+\section{Patterns and Interpretation}
+
+\paragraph{Moderate high dropout, not massive high dropout.}
+The earliest experiments tested very large dropout and failed. A static 0.8
+dropout stream underfit badly, and a high-dropout decay also lost to moderate
+static dropout. This falsified the naive version of the hypothesis. The
+successful schedules start high only when the pressure law predicts high
+dropout, and high means roughly 0.25--0.45 in the L20 extrapolations, not 0.8.
+
+\paragraph{Dropout optima move with pressure.}
+The calibration cells support the core mechanism. For the same model, target
+dropout declines as the unique prefix expands. For the same prefix, target
+dropout generally increases with parameter count. This is consistent with the
+view that dropout is compensating for capacity pressure and repeated sampling.
+
+\paragraph{The interaction term matters.}
+The interaction term $xy$ captures the fact that model/data pressure and
+sampled-token pressure are not independent. For example, repeated sampling is
+more dangerous when a large model sees a small prefix than when a smaller model
+sees the same prefix. In WikiText-103, the interaction form reduces calibration
+MAE from 0.0449 for the first-order base formula to 0.0162. Similar interaction
+fits are used for OpenWebText10K and TinyStories.
+
+\paragraph{Regime specificity is real.}
+The coefficients differ substantially across regimes. TinyStories needs
+different signs and offsets than OpenWebText10K, and WikiText-103 reaches low
+dropout optima earlier under its calibration protocol. This is not a weakness
+of the pressure-law structure; it is evidence that the corpus/tokenizer/eval
+distribution defines the coefficient regime.
+
+\paragraph{Final-loss schedules can sacrifice early local optimality.}
+The L20 stage-wise analysis shows that a schedule can lose to a fixed static
+dropout at early or middle prefixes and still win final loss. This is expected
+if the goal is to train one model continuously through a growing stream, not to
+restart and optimize each prefix independently. Static baselines can be locally
+optimal at one prefix but cannot adapt when the stream changes.
+
+\section{Limitations}
+
+The experiments are intentionally small compared with frontier language-model
+training. The largest model has 51.84M parameters and the largest stream prefix
+has 8M tokens. All training uses Apple MPS, short sequence length, and a 4096
+vocabulary. We therefore claim a reproducible small-scale phenomenon, not a
+confirmed law at billion-parameter scale.
+
+The static coefficient fit optimizes dropout targets extracted from static
+sweeps, while the streaming objective is final validation loss after continued
+training. WikiText-103 shows that these objectives can diverge: the fitted
+interaction schedule is excellent as a static optimum model but decays too
+quickly for L12 streaming. Future work should fit coefficients directly against
+streaming objective functions or include trajectory-aware loss terms.
+
+The L20 extrapolation tier uses three seeds for cost reasons. This is balanced
+across all three regimes and supported by five-seed smaller-stream validation,
+but larger seed counts would narrow confidence intervals. We also do not claim
+that the tested fixed-dropout grids contain every possible fixed optimum.
+
+\section{Reproducibility}
+
+The repository is standalone: cached corpora/tokenizers, saved metrics,
+coefficient files, scripts, and reports are stored in the project. The main
+training entry point is \texttt{scripts/run\_experiments.py}; coefficient
+fitting uses \texttt{scripts/fit\_dropout\_coefficients.py}; this manuscript's
+figures and tables are generated by
+\texttt{scripts/generate\_paper\_assets.py}. Existing result artifacts are read
+from \texttt{runs/}. The experiment runner enforces MPS-only Torch execution.
+
+All generated tables in this paper are derived from saved metrics JSONL files
+and coefficient JSON files. The paper-generation summary is written to
+\texttt{paper/paper\_results\_summary.json}.
+
+\section{Conclusion}
+
+The evidence supports a practical, falsifiable hypothesis: in streaming
+language-model training, dropout should decay according to a regime-specific
+pressure law over model size, currently available unique tokens, and cumulative
+sampled tokens. The hypothesis is stronger than ``larger dropout always helps''
+and weaker than ``one universal formula works everywhere.'' Static calibration
+cells reveal local dropout optima; interaction pressure laws fit those optima
+with low error; and the resulting schedules improve final validation loss over
+the best tested fixed-dropout baselines across three regimes, including a
+larger-model 8M-token extrapolation tier. The next scientific step is to test
+the same fitting procedure at larger scale and to optimize coefficients against
+streaming final-loss objectives directly.
+
+\bibliographystyle{plainnat}
+\bibliography{references}
+
+\end{document}

paper/figures/calibration_cell_curves.tex

@@ -0,0 +1,51 @@
+\begin{tikzpicture}
+\begin{groupplot}[
+    group style={group size=3 by 1, horizontal sep=1.0cm},
+    width=0.32\textwidth,
+    height=0.26\textwidth,
+    legend style={font=\scriptsize},
+    tick label style={font=\scriptsize},
+    label style={font=\scriptsize},
+    title style={font=\scriptsize},
+]
+
+\nextgroupplot[
+    title={OpenWebText10K, L12, 1M},
+    xlabel={Dropout},
+    ylabel={Validation loss above best},
+    ymin=0,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.64089) (0.02,0.41177) (0.05,0.21561) (0.08,0.10277) (0.1,0.06467) (0.14,0.02993) (0.2,0.00000) (0.3,0.03437) (0.4,0.09347) (0.5,0.21260) (0.6,0.37234) (0.7,0.57061) (0.8,0.88258) (0.9,2.06141)};
+\addlegendentry{static sweep}
+\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.2,0)};
+\addlegendentry{target $p^\star=0.20$}
+
+
+\nextgroupplot[
+    title={TinyStories, L12, 1M},
+    xlabel={Dropout},
+    ylabel={Validation loss above best},
+    ymin=0,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.40226) (0.04,0.13035) (0.1,0.02502) (0.2,0.00000) (0.35,0.09066) (0.5,0.22441)};
+\addlegendentry{static sweep}
+\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.2,0)};
+\addlegendentry{target $p^\star=0.20$}
+
+
+\nextgroupplot[
+    title={WikiText-103, L12, 250k},
+    xlabel={Dropout},
+    ylabel={Validation loss above best},
+    ymin=0,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.17165) (0.02,0.10888) (0.04,0.07189) (0.06,0.04641) (0.08,0.02649) (0.1,0.00897) (0.14,0.00363) (0.18,0.00000) (0.2,0.01021) (0.26,0.03365) (0.3,0.04794) (0.4,0.13084)};
+\addlegendentry{static sweep}
+\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.18,0)};
+\addlegendentry{target $p^\star=0.18$}
+
+\end{groupplot}
+\end{tikzpicture}

paper/figures/final_improvements.tex

@@ -0,0 +1,23 @@
+
+\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={OWT10K,TinyStories,WikiText},
+    xtick=data,
+    xtick={OWT10K,TinyStories,WikiText},
+    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 {(OWT10K,0.04732) (TinyStories,0.01326) (WikiText,0.02966)};
+\addlegendentry{4M validation}
+\addplot coordinates {(OWT10K,0.03631) (TinyStories,0.01000) (WikiText,0.01838)};
+\addlegendentry{L20 8M extrapolation}
+\end{axis}
+\end{tikzpicture}

paper/figures/l20_stage_deltas.tex

@@ -0,0 +1,47 @@
+\begin{tikzpicture}
+\begin{groupplot}[
+    group style={group size=3 by 1, horizontal sep=1.0cm},
+    width=0.32\textwidth,
+    height=0.27\textwidth,
+    tick label style={font=\scriptsize},
+    label style={font=\scriptsize},
+    title style={font=\scriptsize},
+]
+
+\nextgroupplot[
+    title={OpenWebText10K},
+    xlabel={Prefix},
+    ylabel={Decay $-$ best static},
+    xtick={0,1,2,3,4,5},
+    xticklabels={250k,500k,1M,2M,4M,8M},
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.06300) (1,0.00103) (2,-0.01477) (3,-0.02370) (4,-0.02720) (5,-0.03631)};
+\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (5,0.00000)};
+
+
+\nextgroupplot[
+    title={TinyStories},
+    xlabel={Prefix},
+    ylabel={Decay $-$ best static},
+    xtick={0,1,2,3,4},
+    xticklabels={500k,1M,2M,4M,8M},
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.00050) (1,0.00729) (2,0.00347) (3,-0.00570) (4,-0.01000)};
+\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (4,0.00000)};
+
+
+\nextgroupplot[
+    title={WikiText-103},
+    xlabel={Prefix},
+    ylabel={Decay $-$ best static},
+    xtick={0,1,2,3,4,5},
+    xticklabels={250k,500k,1M,2M,4M,8M},
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.00803) (1,0.01591) (2,0.03177) (3,0.03542) (4,0.00371) (5,-0.01838)};
+\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (5,0.00000)};
+
+\end{groupplot}
+\end{tikzpicture}

paper/figures/static_optima_by_regime.tex

@@ -0,0 +1,66 @@
+\begin{tikzpicture}
+\begin{groupplot}[
+    group style={group size=3 by 1, horizontal sep=1.0cm},
+    width=0.32\textwidth,
+    height=0.28\textwidth,
+    legend style={font=\tiny},
+    tick label style={font=\scriptsize},
+    label style={font=\scriptsize},
+    title style={font=\scriptsize},
+]
+
+\nextgroupplot[
+    title={OpenWebText10K},
+    xlabel={Unique prefix $U$},
+    ylabel={Fitted target dropout},
+    xtick={0,1,2,3,4},
+    xticklabels={250k,500k,1M,2M,4M},
+    ymin=0,
+    ymax=0.65,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.5062) (1,0.3559) (2,0.2174) (3,0.1364) (4,0.0224)};
+\addlegendentry{L12\_H8\_D320}
+\addplot+[thick, mark=*] coordinates {(0,0.6039) (1,0.4475) (2,0.2837) (3,0.1388) (4,0.0272)};
+\addlegendentry{L16\_H8\_D384}
+\addplot+[thick, mark=*] coordinates {(0,0.3837) (1,0.2401) (2,0.1463) (3,0.0691) (4,0.0000)};
+\addlegendentry{L8\_H8\_D256}
+
+
+\nextgroupplot[
+    title={TinyStories},
+    xlabel={Unique prefix $U$},
+    ylabel={Fitted target dropout},
+    xtick={0,1,2,3,4},
+    xticklabels={250k,500k,1M,2M,4M},
+    ymin=0,
+    ymax=0.65,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.4575) (1,0.2947) (2,0.1866) (3,0.0703) (4,0.0303)};
+\addlegendentry{L12\_H8\_D320}
+\addplot+[thick, mark=*] coordinates {(0,0.5699) (2,0.2450) (4,0.0468)};
+\addlegendentry{L16\_H8\_D384}
+\addplot+[thick, mark=*] coordinates {(0,0.3491) (2,0.0977) (4,0.0000)};
+\addlegendentry{L8\_H8\_D256}
+
+
+\nextgroupplot[
+    title={WikiText-103},
+    xlabel={Unique prefix $U$},
+    ylabel={Fitted target dropout},
+    xtick={0,1,2,3,4},
+    xticklabels={250k,500k,1M,2M,4M},
+    ymin=0,
+    ymax=0.65,
+    grid=both,
+]
+\addplot+[thick, mark=*] coordinates {(0,0.1645) (1,0.0299) (2,0.0000) (3,0.0000) (4,0.0000)};
+\addlegendentry{L12\_H8\_D320}
+\addplot+[thick, mark=*] coordinates {(0,0.2459) (1,0.0822) (2,0.0225) (3,0.0000) (4,0.0000)};
+\addlegendentry{L16\_H8\_D384}
+\addplot+[thick, mark=*] coordinates {(0,0.0965) (1,0.0114) (2,0.0000) (3,0.0000) (4,0.0000)};
+\addlegendentry{L8\_H8\_D256}
+
+\end{groupplot}
+\end{tikzpicture}

paper/paper_results_summary.json

@@ -0,0 +1,520 @@
+{
+  "openweb": {
+    "coefficients": {
+      "regime": "OpenWebText10K",
+      "cells": 15,
+      "A": 0.08189560629877818,
+      "B": -0.023948107975524876,
+      "D": 0.1536984646521253,
+      "C0": -0.05559686586455526,
+      "mae": 0.01579101988615387,
+      "rmse": 0.018845563252758567,
+      "leave_model_mae": 0.024771721924833498,
+      "leave_prefix_mae": 0.030983393290384315,
+      "formula": "p = +0.081896 * log10(P / U) -0.023948 * log10(C / U) +0.153698 * log10(P / U) * log10(C / U) -0.055597"
+    },
+    "small_stream": {
+      "n": 5,
+      "final_stage": 4,
+      "final_prefix": 4000000,
+      "condition_final_mean": 4.3981304407119755,
+      "condition_final_std": 0.009545784836147743,
+      "best_static_condition": "static_dropout_0.14",
+      "best_static_mean": 4.44545366615057,
+      "best_static_std": 0.012017216742245517,
+      "improvement": 0.04732322543859446,
+      "paired_deltas": [
+        -0.03942108899354935,
+        -0.05825117975473404,
+        -0.032753005623817444,
+        -0.052590467035770416,
+        -0.053600385785102844
+      ],
+      "paired_wins": 5,
+      "paired_mean_delta": -0.04732322543859482,
+      "paired_std_delta": 0.010738969078098962,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 250000,
+          "condition_dropout": 0.385,
+          "condition_mean": 5.4946602284908295,
+          "condition_std": 0.01093302726132647,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 5.448330155014991,
+          "best_static_std": 0.0138285316736341,
+          "delta": 0.04633007347583806
+        },
+        {
+          "stage": 1,
+          "prefix": 500000,
+          "condition_dropout": 0.319,
+          "condition_mean": 5.071460470557213,
+          "condition_std": 0.01179360076463939,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 5.066737298667431,
+          "best_static_std": 0.017273470277214743,
+          "delta": 0.004723171889781774
+        },
+        {
+          "stage": 2,
+          "prefix": 1000000,
+          "condition_dropout": 0.227,
+          "condition_mean": 4.781069016456604,
+          "condition_std": 0.008428247627355752,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 4.79825523942709,
+          "best_static_std": 0.01441949497608529,
+          "delta": -0.01718622297048622
+        },
+        {
+          "stage": 3,
+          "prefix": 2000000,
+          "condition_dropout": 0.139,
+          "condition_mean": 4.558990895748138,
+          "condition_std": 0.014177173687439483,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 4.603498187661171,
+          "best_static_std": 0.014129740963263297,
+          "delta": -0.04450729191303271
+        },
+        {
+          "stage": 4,
+          "prefix": 4000000,
+          "condition_dropout": 0.066,
+          "condition_mean": 4.3981304407119755,
+          "condition_std": 0.009545784836147743,
+          "best_static_condition": "static_dropout_0.14",
+          "best_static_mean": 4.44545366615057,
+          "best_static_std": 0.012017216742245517,
+          "delta": -0.04732322543859446
+        }
+      ]
+    },
+    "l20_stream": {
+      "n": 3,
+      "final_stage": 5,
+      "final_prefix": 8000000,
+      "condition_final_mean": 4.179318442940712,
+      "condition_final_std": 0.01344880214961143,
+      "best_static_condition": "static_dropout_0.2",
+      "best_static_mean": 4.215624410659075,
+      "best_static_std": 0.014242200766109402,
+      "improvement": 0.03630596771836281,
+      "paired_deltas": [
+        -0.03935118764638901,
+        -0.025513865053653717,
+        -0.03717709332704544
+      ],
+      "paired_wins": 3,
+      "paired_mean_delta": -0.03401404867569605,
+      "paired_std_delta": 0.007441203715372809,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 250000,
+          "condition_dropout": 0.433,
+          "condition_mean": 5.460076779127121,
+          "condition_std": 0.0012588964293079841,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 5.39707263559103,
+          "best_static_std": 0.014451385161445679,
+          "delta": 0.06300414353609085
+        },
+        {
+          "stage": 1,
+          "prefix": 500000,
+          "condition_dropout": 0.367,
+          "condition_mean": 5.017915904521942,
+          "condition_std": 0.005764488447180914,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 5.01688560595115,
+          "best_static_std": 0.012061853156578773,
+          "delta": 0.0010302985707921763
+        },
+        {
+          "stage": 2,
+          "prefix": 1000000,
+          "condition_dropout": 0.27,
+          "condition_mean": 4.7070309941967325,
+          "condition_std": 0.015739562982191573,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 4.72180004666249,
+          "best_static_std": 0.02152489745362506,
+          "delta": -0.014769052465757326
+        },
+        {
+          "stage": 3,
+          "prefix": 2000000,
+          "condition_dropout": 0.17600000000000002,
+          "condition_mean": 4.487937837839127,
+          "condition_std": 0.009026555295310244,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 4.511637444297473,
+          "best_static_std": 0.016469397224544248,
+          "delta": -0.023699606458346345
+        },
+        {
+          "stage": 4,
+          "prefix": 4000000,
+          "condition_dropout": 0.09500000000000001,
+          "condition_mean": 4.311577834188938,
+          "condition_std": 0.015007670903308238,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 4.338780450324218,
+          "best_static_std": 0.01698295581756926,
+          "delta": -0.027202616135279634
+        },
+        {
+          "stage": 5,
+          "prefix": 8000000,
+          "condition_dropout": 0.033,
+          "condition_mean": 4.179318442940712,
+          "condition_std": 0.01344880214961143,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 4.215624410659075,
+          "best_static_std": 0.014242200766109402,
+          "delta": -0.03630596771836281
+        }
+      ]
+    }
+  },
+  "tinystories": {
+    "coefficients": {
+      "regime": "TinyStories",
+      "cells": 16,
+      "A": -0.08926115772929914,
+      "B": -0.1297539607354037,
+      "D": 0.2550686141985067,
+      "C0": 0.08152490344304414,
+      "mae": 0.01799638717496229,
+      "rmse": 0.023120178816187012,
+      "leave_model_mae": 0.029523601005481504,
+      "leave_prefix_mae": 0.03253257579449062,
+      "formula": "p = -0.089261 * log10(P / U) -0.129754 * log10(C / U) +0.255069 * log10(P / U) * log10(C / U) +0.081525"
+    },
+    "small_stream": {
+      "n": 5,
+      "final_stage": 3,
+      "final_prefix": 4000000,
+      "condition_final_mean": 2.5311126589775084,
+      "condition_final_std": 0.021347338722582713,
+      "best_static_condition": "static_dropout_0.08",
+      "best_static_mean": 2.5443769969046115,
+      "best_static_std": 0.021138839821398296,
+      "improvement": 0.013264337927103131,
+      "paired_deltas": [
+        -0.0005062669515609741,
+        -0.021110694855451584,
+        -0.009244173765182495,
+        -0.016586389392614365,
+        -0.014097623527050018
+      ],
+      "paired_wins": 5,
+      "paired_mean_delta": -0.012309029698371887,
+      "paired_std_delta": 0.00786865194496459,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 500000,
+          "condition_dropout": 0.184,
+          "condition_mean": 3.232573334872723,
+          "condition_std": 0.012322249630223228,
+          "best_static_condition": "static_dropout_0.12",
+          "best_static_mean": 3.2225750528275965,
+          "best_static_std": 0.014316453554089874,
+          "delta": 0.009998282045126317
+        },
+        {
+          "stage": 1,
+          "prefix": 1000000,
+          "condition_dropout": 0.141,
+          "condition_mean": 2.8907706700265408,
+          "condition_std": 0.00267614743609897,
+          "best_static_condition": "static_dropout_0.12",
+          "best_static_mean": 2.893001724779606,
+          "best_static_std": 0.012108793042410648,
+          "delta": -0.002231054753065287
+        },
+        {
+          "stage": 2,
+          "prefix": 2000000,
+          "condition_dropout": 0.084,
+          "condition_mean": 2.6690264880657195,
+          "condition_std": 0.020694473868817478,
+          "best_static_condition": "static_dropout_0.12",
+          "best_static_mean": 2.6794611223042013,
+          "best_static_std": 0.016337678176899628,
+          "delta": -0.010434634238481788
+        },
+        {
+          "stage": 3,
+          "prefix": 4000000,
+          "condition_dropout": 0.045,
+          "condition_mean": 2.5311126589775084,
+          "condition_std": 0.021347338722582713,
+          "best_static_condition": "static_dropout_0.08",
+          "best_static_mean": 2.5443769969046115,
+          "best_static_std": 0.021138839821398296,
+          "delta": -0.013264337927103131
+        }
+      ]
+    },
+    "l20_stream": {
+      "n": 3,
+      "final_stage": 4,
+      "final_prefix": 8000000,
+      "condition_final_mean": 2.208572393283248,
+      "condition_final_std": 0.004634657584019386,
+      "best_static_condition": "static_dropout_0.2",
+      "best_static_mean": 2.218574500953158,
+      "best_static_std": 0.003241222983607362,
+      "improvement": 0.010002107669909943,
+      "paired_deltas": [
+        -0.00869341753423214,
+        -0.006095590069890022,
+        -0.013002954423427582
+      ],
+      "paired_wins": 3,
+      "paired_mean_delta": -0.00926398734251658,
+      "paired_std_delta": 0.0034888512447114515,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 500000,
+          "condition_dropout": 0.253,
+          "condition_mean": 3.030243801573912,
+          "condition_std": 0.006391037258560546,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 3.02974146977067,
+          "best_static_std": 0.006865445597812618,
+          "delta": 0.0005023318032422175
+        },
+        {
+          "stage": 1,
+          "prefix": 1000000,
+          "condition_dropout": 0.209,
+          "condition_mean": 2.7195058229068914,
+          "condition_std": 0.01655123058742941,
+          "best_static_condition": "static_dropout_0.26",
+          "best_static_mean": 2.712211973965168,
+          "best_static_std": 0.01060853933383755,
+          "delta": 0.007293848941723358
+        },
+        {
+          "stage": 2,
+          "prefix": 2000000,
+          "condition_dropout": 0.137,
+          "condition_mean": 2.476810533553362,
+          "condition_std": 0.006082638165661795,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 2.4733379955093064,
+          "best_static_std": 0.009827865290864221,
+          "delta": 0.003472538044055451
+        },
+        {
+          "stage": 3,
+          "prefix": 4000000,
+          "condition_dropout": 0.077,
+          "condition_mean": 2.3110335220893226,
+          "condition_std": 0.015477651713059048,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 2.3167371824383736,
+          "best_static_std": 0.008877543510132386,
+          "delta": -0.00570366034905101
+        },
+        {
+          "stage": 4,
+          "prefix": 8000000,
+          "condition_dropout": 0.041,
+          "condition_mean": 2.208572393283248,
+          "condition_std": 0.004634657584019386,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 2.218574500953158,
+          "best_static_std": 0.003241222983607362,
+          "delta": -0.010002107669909943
+        }
+      ]
+    }
+  },
+  "wikitext": {
+    "coefficients": {
+      "regime": "WikiText-103",
+      "cells": 15,
+      "A": -0.002674916387321836,
+      "B": -0.2836210480242947,
+      "D": 0.2686063066585002,
+      "C0": -0.01899454503425616,
+      "mae": 0.016244092648220347,
+      "rmse": 0.019131667586788335,
+      "leave_model_mae": 0.030438322968748356,
+      "leave_prefix_mae": 0.04435212839791333,
+      "formula": "p = -0.002675 * log10(P / U) -0.283621 * log10(C / U) +0.268606 * log10(P / U) * log10(C / U) -0.018995"
+    },
+    "small_stream": {
+      "n": 5,
+      "final_stage": 4,
+      "final_prefix": 4000000,
+      "condition_final_mean": 4.080797865241766,
+      "condition_final_std": 0.019458062276034267,
+      "best_static_condition": "static_dropout_0.1",
+      "best_static_mean": 4.1104619488120075,
+      "best_static_std": 0.018806640195540816,
+      "improvement": 0.02966408357024175,
+      "paired_deltas": [
+        -0.018418047577142715,
+        -0.018113188445568085,
+        -0.027215413749217987,
+        -0.02507050335407257,
+        -0.037036627531051636
+      ],
+      "paired_wins": 5,
+      "paired_mean_delta": -0.0251707561314106,
+      "paired_std_delta": 0.007752385341386418,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 250000,
+          "condition_dropout": 0.3,
+          "condition_mean": 5.21483291387558,
+          "condition_std": 0.018060268226475495,
+          "best_static_condition": "static_dropout_0.18",
+          "best_static_mean": 5.16162400841713,
+          "best_static_std": 0.014987630048664583,
+          "delta": 0.05320890545845014
+        },
+        {
+          "stage": 1,
+          "prefix": 500000,
+          "condition_dropout": 0.26,
+          "condition_mean": 4.80534844994545,
+          "condition_std": 0.027829930328897787,
+          "best_static_condition": "static_dropout_0.2",
+          "best_static_mean": 4.787301687896251,
+          "best_static_std": 0.023601504941398347,
+          "delta": 0.01804676204919886
+        },
+        {
+          "stage": 2,
+          "prefix": 1000000,
+          "condition_dropout": 0.18,
+          "condition_mean": 4.49377838075161,
+          "condition_std": 0.014674108647494532,
+          "best_static_condition": "static_dropout_0.14",
+          "best_static_mean": 4.500055414438248,
+          "best_static_std": 0.014819407090300617,
+          "delta": -0.006277033686638234
+        },
+        {
+          "stage": 3,
+          "prefix": 2000000,
+          "condition_dropout": 0.09,
+          "condition_mean": 4.260725708305836,
+          "condition_std": 0.018098739272622868,
+          "best_static_condition": "static_dropout_0.14",
+          "best_static_mean": 4.281563547253609,
+          "best_static_std": 0.02138858846966249,
+          "delta": -0.020837838947772802
+        },
+        {
+          "stage": 4,
+          "prefix": 4000000,
+          "condition_dropout": 0.02,
+          "condition_mean": 4.080797865241766,
+          "condition_std": 0.019458062276034267,
+          "best_static_condition": "static_dropout_0.1",
+          "best_static_mean": 4.1104619488120075,
+          "best_static_std": 0.018806640195540816,
+          "delta": -0.02966408357024175
+        }
+      ]
+    },
+    "l20_stream": {
+      "n": 3,
+      "final_stage": 5,
+      "final_prefix": 8000000,
+      "condition_final_mean": 3.957240684578816,
+      "condition_final_std": 0.008487626829204101,
+      "best_static_condition": "static_dropout_0.14",
+      "best_static_mean": 3.9756230898201466,
+      "best_static_std": 0.008723729300606414,
+      "improvement": 0.018382405241330613,
+      "paired_deltas": [
+        -0.016518089920282364,
+        -0.00015844032168388367,
+        -0.027258697897195816
+      ],
+      "paired_wins": 3,
+      "paired_mean_delta": -0.014645076046387354,
+      "paired_std_delta": 0.013646872385042126,
+      "stage_deltas": [
+        {
+          "stage": 0,
+          "prefix": 250000,
+          "condition_dropout": 0.284,
+          "condition_mean": 5.2075243567427,
+          "condition_std": 0.013993230772716539,
+          "best_static_condition": "static_dropout_0.26",
+          "best_static_mean": 5.1994956483443575,
+          "best_static_std": 0.013634536236337126,
+          "delta": 0.008028708398342133
+        },
+        {
+          "stage": 1,
+          "prefix": 500000,
+          "condition_dropout": 0.211,
+          "condition_mean": 4.816951217750709,
+          "condition_std": 0.0272929978517565,
+          "best_static_condition": "static_dropout_0.3",
+          "best_static_mean": 4.8010413100322085,
+          "best_static_std": 0.023497612891923463,
+          "delta": 0.015909907718500094
+        },
+        {
+          "stage": 2,
+          "prefix": 1000000,
+          "condition_dropout": 0.116,
+          "condition_mean": 4.538959090908368,
+          "condition_std": 0.018682984974041847,
+          "best_static_condition": "static_dropout_0.26",
+          "best_static_mean": 4.507193277279536,
+          "best_static_std": 0.022063926876841285,
+          "delta": 0.03176581362883191
+        },
+        {
+          "stage": 3,
+          "prefix": 2000000,
+          "condition_dropout": 0.03599999999999999,
+          "condition_mean": 4.310809387514989,
+          "condition_std": 0.02431455246903313,
+          "best_static_condition": "static_dropout_0.26",
+          "best_static_mean": 4.275388439496358,
+          "best_static_std": 0.021064935851257058,
+          "delta": 0.03542094801863094
+        },
+        {
+          "stage": 4,
+          "prefix": 4000000,
+          "condition_dropout": 0.02,
+          "condition_mean": 4.1124994705120725,
+          "condition_std": 0.024692308091551114,
+          "best_static_condition": "static_dropout_0.18",
+          "best_static_mean": 4.108785220732291,
+          "best_static_std": 0.016546231012667934,
+          "delta": 0.003714249779781298
+        },
+        {
+          "stage": 5,
+          "prefix": 8000000,
+          "condition_dropout": 0.02,
+          "condition_mean": 3.957240684578816,
+          "condition_std": 0.008487626829204101,
+          "best_static_condition": "static_dropout_0.14",
+          "best_static_mean": 3.9756230898201466,
+          "best_static_std": 0.008723729300606414,
+          "delta": -0.018382405241330613
+        }
+      ]
+    }
+  }
+}

paper/references.bib

@@ -33,6 +33,14 @@
   url = {https://arxiv.org/abs/2203.15556}
 }
 
+@misc{gokaslan2019openwebtext,
+  title = {OpenWebText Corpus},
+  author = {Gokaslan, Aaron and Cohen, Vanya and Pavlick, Ellie and Tellex, Stefanie},
+  year = {2019},
+  howpublished = {\url{https://zenodo.org/records/3834942}},
+  note = {Open reproduction of the WebText-style corpus}
+}
+
 @inproceedings{merity2017pointer,
   title = {Pointer Sentinel Mixture Models},
   author = {Merity, Stephen and Xiong, Caiming and Bradbury, James and Socher, Richard},

paper/tables/coefficient_table.tex

@@ -0,0 +1,9 @@
+\begin{tabular}{lrrrrrrrr}
+\toprule
+Regime & Cells & A & B & D & C0 & MAE & L-model MAE & L-prefix MAE \\
+\midrule
+OpenWebText10K & 15 & 0.0819 & -0.0239 & 0.1537 & -0.0556 & 0.0158 & 0.0248 & 0.0310 \\
+TinyStories & 16 & -0.0893 & -0.1298 & 0.2551 & 0.0815 & 0.0180 & 0.0295 & 0.0325 \\
+WikiText-103 & 15 & -0.0027 & -0.2836 & 0.2686 & -0.0190 & 0.0162 & 0.0304 & 0.0444 \\
+\bottomrule
+\end{tabular}

paper/tables/l20_stage_table.tex

@@ -0,0 +1,23 @@
+\begin{tabular}{lrrrrr}
+\toprule
+Regime & Prefix & Decay dropout & Decay val & Best static val & Delta \\
+\midrule
+OpenWebText10K & 250k & 0.433 & 5.4601 & 5.3971 & +0.0630 \\
+OpenWebText10K & 500k & 0.367 & 5.0179 & 5.0169 & +0.0010 \\
+OpenWebText10K & 1M & 0.270 & 4.7070 & 4.7218 & -0.0148 \\
+OpenWebText10K & 2M & 0.176 & 4.4879 & 4.5116 & -0.0237 \\
+OpenWebText10K & 4M & 0.095 & 4.3116 & 4.3388 & -0.0272 \\
+OpenWebText10K & 8M & 0.033 & 4.1793 & 4.2156 & -0.0363 \\
+TinyStories & 500k & 0.253 & 3.0302 & 3.0297 & +0.0005 \\
+TinyStories & 1M & 0.209 & 2.7195 & 2.7122 & +0.0073 \\
+TinyStories & 2M & 0.137 & 2.4768 & 2.4733 & +0.0035 \\
+TinyStories & 4M & 0.077 & 2.3110 & 2.3167 & -0.0057 \\
+TinyStories & 8M & 0.041 & 2.2086 & 2.2186 & -0.0100 \\
+WikiText-103 & 250k & 0.284 & 5.2075 & 5.1995 & +0.0080 \\
+WikiText-103 & 500k & 0.211 & 4.8170 & 4.8010 & +0.0159 \\
+WikiText-103 & 1M & 0.116 & 4.5390 & 4.5072 & +0.0318 \\
+WikiText-103 & 2M & 0.036 & 4.3108 & 4.2754 & +0.0354 \\
+WikiText-103 & 4M & 0.020 & 4.1125 & 4.1088 & +0.0037 \\
+WikiText-103 & 8M & 0.020 & 3.9572 & 3.9756 & -0.0184 \\
+\bottomrule
+\end{tabular}

paper/tables/streaming_validation_table.tex

@@ -0,0 +1,12 @@
+\begin{tabular}{llrrrrr}
+\toprule
+Regime & Protocol & Decay final & Best static final & Gain & Seeds & Paired wins \\
+\midrule
+OpenWebText10K & 4M / 5 seeds & 4.3981 $\pm$ 0.0095 & 4.4455 $\pm$ 0.0120 & 0.0473 & 5 & 5/5 \\
+OpenWebText10K & 8M / 3 seeds & 4.1793 $\pm$ 0.0134 & 4.2156 $\pm$ 0.0142 & 0.0363 & 3 & 3/3 \\
+TinyStories & 4M / 5 seeds & 2.5311 $\pm$ 0.0213 & 2.5444 $\pm$ 0.0211 & 0.0133 & 5 & 5/5 \\
+TinyStories & 8M / 3 seeds & 2.2086 $\pm$ 0.0046 & 2.2186 $\pm$ 0.0032 & 0.0100 & 3 & 3/3 \\
+WikiText-103 & 4M / 5 seeds & 4.0808 $\pm$ 0.0195 & 4.1105 $\pm$ 0.0188 & 0.0297 & 5 & 5/5 \\
+WikiText-103 & 8M / 3 seeds & 3.9572 $\pm$ 0.0085 & 3.9756 $\pm$ 0.0087 & 0.0184 & 3 & 3/3 \\
+\bottomrule
+\end{tabular}

runs/smoke_compat/locked_stream/20260605-224834/RESULT_SUMMARY.md

@@ -0,0 +1,29 @@
+# Locked Streaming Dropout Summary
+
+Run directory: `runs/smoke_compat/locked_stream/20260605-224834`
+
+Model: `smoke` causal Transformer, 134,144 parameters, 1 layers, 1 heads, 16 embedding dim.
+Training per stage: 2 steps. Sampled tokens are cumulative in each stage row. Seeds present: 123.
+
+## Condition Ranking
+
+| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |
+|---|---|---:|---:|---:|---:|---|
+| `static_dropout_0.05` | static | 0.05 | 8.3189 | 8.3194 | 0.0069 | 0.05 -> 0.05 |
+| `smoke_decay` | anchor_decay | 0.05 | 8.3189 | 8.3194 | 0.0014 | 0.10 -> 0.05 |
+
+## Stage Trajectory
+
+### Stage 0: 512 Prefix Tokens
+
+| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |
+|---|---:|---:|---:|---:|---:|
+| `static_dropout_0.05` | 0.05 | 8.3184 | 8.3184 | -0.0001 | 1 |
+| `smoke_decay` | 0.10 | 8.3185 | 8.3181 | 0.0003 | 1 |
+
+### Stage 1: 768 Prefix Tokens
+
+| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |
+|---|---:|---:|---:|---:|---:|
+| `smoke_decay` | 0.05 | 8.3194 | 8.3180 | 0.0014 | 1 |
+| `static_dropout_0.05` | 0.05 | 8.3194 | 8.3125 | 0.0069 | 1 |

runs/smoke_compat/locked_stream/20260605-224834/config.json

@@ -0,0 +1,103 @@
+{
+  "args": {
+    "mode": "locked_stream",
+    "corpus": null,
+    "corpus_glob": null,
+    "text_column": "text",
+    "use_cached_data": true,
+    "output_dir": "runs/smoke_compat",
+    "resume_from": null,
+    "cache_dir": ".cache/dropout_decay",
+    "models": [
+      "smoke=1x1x16"
+    ],
+    "seeds": [
+      123
+    ],
+    "token_limits": [
+      5000000
+    ],
+    "stream_token_caps": [
+      512,
+      768
+    ],
+    "val_tokens": 128,
+    "allow_short_corpus": true,
+    "force_retokenize": false,
+    "vocab_size": 4096,
+    "tokenizer_train_chars": 10000000,
+    "block_size": 16,
+    "batch_size": 2,
+    "steps": 2000,
+    "stage_steps": 2,
+    "dropout_rates": [
+      0.05
+    ],
+    "decays": [],
+    "anchor_decays": [
+      {
+        "name": "smoke_decay",
+        "kind": "anchor_decay",
+        "initial": 0.1,
+        "final": 0.05,
+        "schedule": "log_prefix_anchor",
+        "decay_tokens": null,
+        "anchors": [
+          [
+            512,
+            0.1
+          ],
+          [
+            768,
+            0.05
+          ]
+        ]
+      }
+    ],
+    "decay_tokens": null,
+    "eval_batches": 1,
+    "train_eval_batches": 1,
+    "trace_eval_batches": 1,
+    "eval_every": 0,
+    "log_every": 1,
+    "lr": 0.0003,
+    "weight_decay": 0.1,
+    "grad_clip": 1.0,
+    "plateau_delta": 0.01,
+    "target_min_dropout": 0.1,
+    "min_nonzero_margin": 0.01,
+    "min_high_dropout_margin": 0.03,
+    "screen_early_stop": false,
+    "screen_prune_patience": 3,
+    "screen_prune_min_delta": 0.01
+  },
+  "mode": "locked_stream",
+  "seeds": [
+    123
+  ],
+  "models": [
+    {
+      "model_name": "smoke",
+      "n_layer": 1,
+      "n_head": 1,
+      "n_embd": 16
+    }
+  ],
+  "device": "mps",
+  "torch": "2.12.0",
+  "python": "3.11.15 (main, Mar  3 2026, 00:52:57) [Clang 21.0.0 (clang-2100.0.123.102)]",
+  "mps_available": true,
+  "attribution": "Derived from Andrej Karpathy's nanochat project (https://github.com/karpathy/nanochat), MIT License, Copyright (c) 2025 Andrej Karpathy.",
+  "tokenizer_path": ".cache/dropout_decay/tokenizer-v4096.json",
+  "encoded_path": ".cache/dropout_decay/tokens-v4096-uint16.npy",
+  "train_tokens": 5500842,
+  "val_tokens": 128,
+  "effective_token_limits": [
+    5000000
+  ],
+  "effective_stream_token_caps": [
+    512,
+    768
+  ],
+  "resume_from": null
+}

runs/smoke_compat/locked_stream/20260605-224834/metrics.jsonl

@@ -0,0 +1,4 @@
+{"condition": "smoke_decay", "condition_kind": "anchor_decay", "dropout_active_final": 0.1, "dropout_final": 0.05, "dropout_initial": 0.1, "dropout_schedule": "log_prefix_anchor", "elapsed_sec": 0.08690309524536133, "eval_loss": 8.318451881408691, "generalization_gap": 0.00031757354736328125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 0, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318134307861328, "train_loss_last": 8.317201614379883, "val_eval_loss": 8.318451881408691}
+{"condition": "smoke_decay", "condition_kind": "anchor_decay", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.1, "dropout_schedule": "log_prefix_anchor", "elapsed_sec": 0.018578290939331055, "eval_loss": 8.319382667541504, "generalization_gap": 0.0013904571533203125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 1, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_eval_loss": 8.317992210388184, "train_loss_last": 8.3181734085083, "val_eval_loss": 8.319382667541504}
+{"condition": "static_dropout_0.05", "condition_kind": "static", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.05, "dropout_schedule": "constant", "elapsed_sec": 0.019870996475219727, "eval_loss": 8.318363189697266, "generalization_gap": -8.20159912109375e-05, "model_config": {"block_size": 16, "dropout": 0.05, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 0, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318445205688477, "train_loss_last": 8.317625999450684, "val_eval_loss": 8.318363189697266}
+{"condition": "static_dropout_0.05", "condition_kind": "static", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.05, "dropout_schedule": "constant", "elapsed_sec": 0.017475128173828125, "eval_loss": 8.319437026977539, "generalization_gap": 0.0068874359130859375, "model_config": {"block_size": 16, "dropout": 0.05, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 1, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_eval_loss": 8.312549591064453, "train_loss_last": 8.316361427307129, "val_eval_loss": 8.319437026977539}

runs/smoke_compat/locked_stream/20260605-224834/summary.csv

@@ -0,0 +1,5 @@
+run_mode,condition,condition_kind,stage,token_limit,model_name,n_layer,n_head,n_embd,parameters,dropout_initial,dropout_final,dropout_schedule,n,mean_train_eval_loss,std_train_eval_loss,mean_val_eval_loss,std_val_eval_loss,mean_generalization_gap,std_generalization_gap
+locked_stream,smoke_decay,anchor_decay,0,512,smoke,1,1,16,134144,0.1,0.05,log_prefix_anchor,1,8.318134307861328,0.0,8.318451881408691,0.0,0.00031757354736328125,0.0
+locked_stream,static_dropout_0.05,static,0,512,smoke,1,1,16,134144,0.05,0.05,constant,1,8.318445205688477,0.0,8.318363189697266,0.0,-8.20159912109375e-05,0.0
+locked_stream,smoke_decay,anchor_decay,1,768,smoke,1,1,16,134144,0.1,0.05,log_prefix_anchor,1,8.317992210388184,0.0,8.319382667541504,0.0,0.0013904571533203125,0.0
+locked_stream,static_dropout_0.05,static,1,768,smoke,1,1,16,134144,0.05,0.05,constant,1,8.312549591064453,0.0,8.319437026977539,0.0,0.0068874359130859375,0.0

runs/smoke_compat/locked_stream/20260605-224834/summary.json

@@ -0,0 +1,90 @@
+[
+  {
+    "run_mode": "locked_stream",
+    "condition": "smoke_decay",
+    "condition_kind": "anchor_decay",
+    "stage": 0,
+    "token_limit": 512,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.1,
+    "dropout_final": 0.05,
+    "dropout_schedule": "log_prefix_anchor",
+    "n": 1,
+    "mean_train_eval_loss": 8.318134307861328,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.318451881408691,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": 0.00031757354736328125,
+    "std_generalization_gap": 0.0
+  },
+  {
+    "run_mode": "locked_stream",
+    "condition": "static_dropout_0.05",
+    "condition_kind": "static",
+    "stage": 0,
+    "token_limit": 512,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.05,
+    "dropout_final": 0.05,
+    "dropout_schedule": "constant",
+    "n": 1,
+    "mean_train_eval_loss": 8.318445205688477,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.318363189697266,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": -8.20159912109375e-05,
+    "std_generalization_gap": 0.0
+  },
+  {
+    "run_mode": "locked_stream",
+    "condition": "smoke_decay",
+    "condition_kind": "anchor_decay",
+    "stage": 1,
+    "token_limit": 768,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.1,
+    "dropout_final": 0.05,
+    "dropout_schedule": "log_prefix_anchor",
+    "n": 1,
+    "mean_train_eval_loss": 8.317992210388184,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.319382667541504,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": 0.0013904571533203125,
+    "std_generalization_gap": 0.0
+  },
+  {
+    "run_mode": "locked_stream",
+    "condition": "static_dropout_0.05",
+    "condition_kind": "static",
+    "stage": 1,
+    "token_limit": 768,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.05,
+    "dropout_final": 0.05,
+    "dropout_schedule": "constant",
+    "n": 1,
+    "mean_train_eval_loss": 8.312549591064453,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.319437026977539,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": 0.0068874359130859375,
+    "std_generalization_gap": 0.0
+  }
+]

runs/smoke_compat/locked_stream/20260605-224834/trace.jsonl

@@ -0,0 +1,8 @@
+{"condition": "smoke_decay", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.31665325164795}
+{"condition": "smoke_decay", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317201614379883}
+{"condition": "smoke_decay", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 1, "steps": 2, "token_limit": 768, "tokens_seen": 96, "train_batch_loss": 8.316765785217285}
+{"condition": "smoke_decay", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 2, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_batch_loss": 8.3181734085083}
+{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.317569732666016}
+{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317625999450684}
+{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 1, "steps": 2, "token_limit": 768, "tokens_seen": 96, "train_batch_loss": 8.318181037902832}
+{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 2, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_batch_loss": 8.316361427307129}

runs/smoke_compat/screen_static/20260605-224821/RESULT_SUMMARY.md

@@ -0,0 +1,24 @@
+# Static Dropout Screen Summary
+
+Run directory: `runs/smoke_compat/screen_static/20260605-224821`
+
+## Models
+
+| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |
+|---|---:|---:|---:|---:|---:|---:|---|
+| `smoke` | 134,144 | 1 | 1 | 16 | 16 | 4096 | 123 |
+
+## Best Dropout By Model And Prefix
+
+| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |
+|---|---:|---:|---:|---:|---:|---:|---:|---|
+| `smoke` | 512 | 0.12 | 0.10 | 8.3185 | 0.0000 | 8.3181 | 0.0003 | not bracketed; best at top of tested grid |
+
+## Model `smoke`
+
+### Prefix 512 Tokens (0.12 Effective Epochs)
+
+| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |
+|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
+| 0.00 | 1 | 8.3186 | 0.0000 | 8.3134 | 0.0000 | 0.0052 | 0.0000 | 64 | 134,144 |
+| 0.10 | 1 | 8.3185 | 0.0000 | 8.3181 | 0.0000 | 0.0003 | 0.0000 | 64 | 134,144 |

runs/smoke_compat/screen_static/20260605-224821/config.json

@@ -0,0 +1,89 @@
+{
+  "args": {
+    "mode": "screen_static",
+    "corpus": null,
+    "corpus_glob": null,
+    "text_column": "text",
+    "use_cached_data": true,
+    "output_dir": "runs/smoke_compat",
+    "resume_from": null,
+    "cache_dir": ".cache/dropout_decay",
+    "models": [
+      "smoke=1x1x16"
+    ],
+    "seeds": [
+      123
+    ],
+    "token_limits": [
+      512
+    ],
+    "stream_token_caps": [
+      5000000,
+      10000000,
+      20000000,
+      40000000
+    ],
+    "val_tokens": 128,
+    "allow_short_corpus": true,
+    "force_retokenize": false,
+    "vocab_size": 4096,
+    "tokenizer_train_chars": 10000000,
+    "block_size": 16,
+    "batch_size": 2,
+    "steps": 2,
+    "stage_steps": 1000,
+    "dropout_rates": [
+      0.0,
+      0.1
+    ],
+    "decays": [],
+    "anchor_decays": [],
+    "decay_tokens": null,
+    "eval_batches": 1,
+    "train_eval_batches": 1,
+    "trace_eval_batches": 1,
+    "eval_every": 0,
+    "log_every": 1,
+    "lr": 0.0003,
+    "weight_decay": 0.1,
+    "grad_clip": 1.0,
+    "plateau_delta": 0.01,
+    "target_min_dropout": 0.1,
+    "min_nonzero_margin": 0.01,
+    "min_high_dropout_margin": 0.03,
+    "screen_early_stop": false,
+    "screen_prune_patience": 3,
+    "screen_prune_min_delta": 0.01
+  },
+  "mode": "screen_static",
+  "seeds": [
+    123
+  ],
+  "models": [
+    {
+      "model_name": "smoke",
+      "n_layer": 1,
+      "n_head": 1,
+      "n_embd": 16
+    }
+  ],
+  "device": "mps",
+  "torch": "2.12.0",
+  "python": "3.11.15 (main, Mar  3 2026, 00:52:57) [Clang 21.0.0 (clang-2100.0.123.102)]",
+  "mps_available": true,
+  "attribution": "Derived from Andrej Karpathy's nanochat project (https://github.com/karpathy/nanochat), MIT License, Copyright (c) 2025 Andrej Karpathy.",
+  "tokenizer_path": ".cache/dropout_decay/tokenizer-v4096.json",
+  "encoded_path": ".cache/dropout_decay/tokens-v4096-uint16.npy",
+  "train_tokens": 5500842,
+  "val_tokens": 128,
+  "effective_token_limits": [
+    512
+  ],
+  "effective_stream_token_caps": [
+    5000000,
+    5500842,
+    5500842,
+    5500842
+  ],
+  "resume_from": null
+}

runs/smoke_compat/screen_static/20260605-224821/metrics.jsonl

@@ -0,0 +1,2 @@
+{"condition": "static_dropout_0", "condition_kind": "static", "dropout_active_final": 0.0, "dropout_final": 0.0, "dropout_initial": 0.0, "dropout_schedule": "constant", "elapsed_sec": 0.5898959636688232, "eval_loss": 8.318593978881836, "generalization_gap": 0.005218505859375, "model_config": {"block_size": 16, "dropout": 0.0, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "screen_static", "seed": 123, "stage": null, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.313375473022461, "train_loss_last": 8.315808296203613, "val_eval_loss": 8.318593978881836}
+{"condition": "static_dropout_0.1", "condition_kind": "static", "dropout_active_final": 0.1, "dropout_final": 0.1, "dropout_initial": 0.1, "dropout_schedule": "constant", "elapsed_sec": 0.10111308097839355, "eval_loss": 8.318451881408691, "generalization_gap": 0.00031757354736328125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "screen_static", "seed": 123, "stage": null, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318134307861328, "train_loss_last": 8.317201614379883, "val_eval_loss": 8.318451881408691}

runs/smoke_compat/screen_static/20260605-224821/model_selection.csv

@@ -0,0 +1,2 @@
+run_mode,token_limit,model_name,n_layer,n_head,n_embd,parameters,n,best_dropout,best_val_loss,best_val_std,plateau_start_dropout,plateau_end_dropout,plateau_delta,zero_dropout_val_loss,zero_minus_best,best_nonzero_dropout,best_nonzero_val_loss,zero_minus_best_nonzero,max_dropout,max_dropout_val_loss,max_dropout_minus_best,has_nonzero_optimum,meets_target_dropout,curve_json
+screen_static,512,smoke,1,1,16,134144,1,0.1,8.318451881408691,0.0,0.0,0.1,0.01,8.318593978881836,0.00014209747314453125,0.1,8.318451881408691,0.00014209747314453125,0.1,8.318451881408691,0.0,False,True,"[{""dropout"": 0.0, ""mean_generalization_gap"": 0.005218505859375, ""mean_train_loss"": 8.313375473022461, ""mean_val_loss"": 8.318593978881836, ""n"": 1, ""std_val_loss"": 0.0}, {""dropout"": 0.1, ""mean_generalization_gap"": 0.00031757354736328125, ""mean_train_loss"": 8.318134307861328, ""mean_val_loss"": 8.318451881408691, ""n"": 1, ""std_val_loss"": 0.0}]"

runs/smoke_compat/screen_static/20260605-224821/model_selection.json

@@ -0,0 +1,29 @@
+[
+  {
+    "run_mode": "screen_static",
+    "token_limit": 512,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "n": 1,
+    "best_dropout": 0.1,
+    "best_val_loss": 8.318451881408691,
+    "best_val_std": 0.0,
+    "plateau_start_dropout": 0.0,
+    "plateau_end_dropout": 0.1,
+    "plateau_delta": 0.01,
+    "zero_dropout_val_loss": 8.318593978881836,
+    "zero_minus_best": 0.00014209747314453125,
+    "best_nonzero_dropout": 0.1,
+    "best_nonzero_val_loss": 8.318451881408691,
+    "zero_minus_best_nonzero": 0.00014209747314453125,
+    "max_dropout": 0.1,
+    "max_dropout_val_loss": 8.318451881408691,
+    "max_dropout_minus_best": 0.0,
+    "has_nonzero_optimum": false,
+    "meets_target_dropout": true,
+    "curve_json": "[{\"dropout\": 0.0, \"mean_generalization_gap\": 0.005218505859375, \"mean_train_loss\": 8.313375473022461, \"mean_val_loss\": 8.318593978881836, \"n\": 1, \"std_val_loss\": 0.0}, {\"dropout\": 0.1, \"mean_generalization_gap\": 0.00031757354736328125, \"mean_train_loss\": 8.318134307861328, \"mean_val_loss\": 8.318451881408691, \"n\": 1, \"std_val_loss\": 0.0}]"
+  }
+]

runs/smoke_compat/screen_static/20260605-224821/summary.csv

@@ -0,0 +1,3 @@
+run_mode,condition,condition_kind,stage,token_limit,model_name,n_layer,n_head,n_embd,parameters,dropout_initial,dropout_final,dropout_schedule,n,mean_train_eval_loss,std_train_eval_loss,mean_val_eval_loss,std_val_eval_loss,mean_generalization_gap,std_generalization_gap
+screen_static,static_dropout_0,static,,512,smoke,1,1,16,134144,0.0,0.0,constant,1,8.313375473022461,0.0,8.318593978881836,0.0,0.005218505859375,0.0
+screen_static,static_dropout_0.1,static,,512,smoke,1,1,16,134144,0.1,0.1,constant,1,8.318134307861328,0.0,8.318451881408691,0.0,0.00031757354736328125,0.0

runs/smoke_compat/screen_static/20260605-224821/summary.json

@@ -0,0 +1,46 @@
+[
+  {
+    "run_mode": "screen_static",
+    "condition": "static_dropout_0",
+    "condition_kind": "static",
+    "stage": null,
+    "token_limit": 512,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.0,
+    "dropout_final": 0.0,
+    "dropout_schedule": "constant",
+    "n": 1,
+    "mean_train_eval_loss": 8.313375473022461,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.318593978881836,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": 0.005218505859375,
+    "std_generalization_gap": 0.0
+  },
+  {
+    "run_mode": "screen_static",
+    "condition": "static_dropout_0.1",
+    "condition_kind": "static",
+    "stage": null,
+    "token_limit": 512,
+    "model_name": "smoke",
+    "n_layer": 1,
+    "n_head": 1,
+    "n_embd": 16,
+    "parameters": 134144,
+    "dropout_initial": 0.1,
+    "dropout_final": 0.1,
+    "dropout_schedule": "constant",
+    "n": 1,
+    "mean_train_eval_loss": 8.318134307861328,
+    "std_train_eval_loss": 0.0,
+    "mean_val_eval_loss": 8.318451881408691,
+    "std_val_eval_loss": 0.0,
+    "mean_generalization_gap": 0.00031757354736328125,
+    "std_generalization_gap": 0.0
+  }
+]

runs/smoke_compat/screen_static/20260605-224821/trace.jsonl

@@ -0,0 +1,4 @@
+{"condition": "static_dropout_0", "dropout": 0.0, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.318573951721191}
+{"condition": "static_dropout_0", "dropout": 0.0, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.315808296203613}
+{"condition": "static_dropout_0.1", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.31665325164795}
+{"condition": "static_dropout_0.1", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317201614379883}

scripts/generate_paper_assets.py

@@ -0,0 +1,579 @@
+#!/usr/bin/env python3
+"""
+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 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 / 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 / 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 / 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 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 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 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_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 = []
+    for row in condition_rows:
+        seed = int(row["seed"])
+        paired_deltas.append(float(row["val_eval_loss"]) - final_static_by_seed[seed][1])
+    paired_mean, paired_std = mean_std(paired_deltas)
+
+    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_wins": sum(1 for delta in paired_deltas if delta < 0),
+        "paired_mean_delta": paired_mean,
+        "paired_std_delta": paired_std,
+        "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 key, paths in 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])
+        axes.append(
+            rf"""
+\nextgroupplot[
+    title={{{paths["curve_title"]}}},
+    xlabel={{Dropout}},
+    ylabel={{Validation loss above best}},
+    ymin=0,
+    grid=both,
+]
+\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{{{tex_escape(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_tables(summaries: dict[str, dict], coeffs: dict[str, dict]) -> None:
+    coeff_lines = [
+        "\\begin{tabular}{lrrrrrrrr}",
+        "\\toprule",
+        "Regime & Cells & A & B & D & C0 & MAE & 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"
+    )
+
+    val_lines = [
+        "\\begin{tabular}{llrrrrr}",
+        "\\toprule",
+        "Regime & Protocol & Decay final & Best static final & Gain & Seeds & Paired wins \\\\",
+        "\\midrule",
+    ]
+    for key, paths in REGIMES.items():
+        for protocol_key, protocol_label in [("small", paths["small_protocol"]), ("l20", "8M / 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"
+    )
+
+    l20_lines = [
+        "\\begin{tabular}{lrrrrr}",
+        "\\toprule",
+        "Regime & Prefix & Decay dropout & Decay val & Best static val & Delta \\\\",
+        "\\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_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()

src/dropout_decay/data/__init__.py

@@ -0,0 +1,45 @@
+"""
+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 dropout_decay.datasets import (
+    CachedTokenizer,
+    TokenSplits,
+    encode_corpus,
+    iter_documents,
+    load_cached_splits,
+    preview_documents,
+    required_token_count,
+    resolve_paths,
+    train_or_load_tokenizer,
+)
+from dropout_decay.specs import DropoutCondition, ModelSpec, anchor_dropout
+
+__all__ = [
+    "CachedTokenizer",
+    "DropoutCondition",
+    "ModelSpec",
+    "TokenSplits",
+    "anchor_dropout",
+    "encode_corpus",
+    "iter_documents",
+    "load_cached_splits",
+    "preview_documents",
+    "required_token_count",
+    "resolve_paths",
+    "train_or_load_tokenizer",
+]

src/dropout_decay/datasets/__init__.py

@@ -0,0 +1,41 @@
+"""
+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 dropout_decay.datasets.corpus import (
+    CachedTokenizer,
+    TokenSplits,
+    encode_corpus,
+    iter_documents,
+    load_cached_splits,
+    preview_documents,
+    required_token_count,
+    resolve_paths,
+    train_or_load_tokenizer,
+)
+
+__all__ = [
+    "CachedTokenizer",
+    "TokenSplits",
+    "encode_corpus",
+    "iter_documents",
+    "load_cached_splits",
+    "preview_documents",
+    "required_token_count",
+    "resolve_paths",
+    "train_or_load_tokenizer",
+]

src/dropout_decay/{data.py → datasets/corpus.py}

@@ -143,7 +143,7 @@ def train_or_load_tokenizer(
     text_column: str,
     force_retrain: bool = False,
 ) -> object:
-    from dropout_decay.tokenizer import BpeTokenizer
+    from dropout_decay.tokenization import BpeTokenizer
 
     tokenizer_path = output_dir / f"tokenizer-v{vocab_size}.json"
     if tokenizer_path.exists() and not force_retrain:

src/dropout_decay/experiment.py

@@ -16,1694 +16,13 @@ The above copyright notice and this permission notice shall be included in all
 copies or substantial portions of the Software.
 """
 
-from __future__ import annotations
+from dropout_decay.experiments.runner import build_parser, main, run
 
-import argparse
-import csv
-from collections import defaultdict
-from dataclasses import dataclass
-from datetime import datetime
-import json
-import math
-import os
-from pathlib import Path
-import random
-import statistics
-import sys
-import time
-from typing import Callable
-
-import numpy as np
-import torch
-
-from dropout_decay.data import (
-    encode_corpus,
-    load_cached_splits,
-    resolve_paths,
-    train_or_load_tokenizer,
-)
-from dropout_decay.license import NANOCHAT_ATTRIBUTION
-from dropout_decay.model import DropoutGPT, GPTConfig
-from dropout_decay.scheduler import DropoutDecayConfig, DropoutDecayScheduler
-
-
-DEFAULT_DROPOUT_RATES = [0.0, 0.02, 0.05, 0.08, 0.10, 0.14, 0.20, 0.30, 0.50]
-SUMMARY_FIELDS = [
-    "run_mode",
-    "condition",
-    "condition_kind",
-    "stage",
-    "token_limit",
-    "model_name",
-    "n_layer",
-    "n_head",
-    "n_embd",
-    "parameters",
-    "dropout_initial",
-    "dropout_final",
-    "dropout_schedule",
-    "n",
-    "mean_train_eval_loss",
-    "std_train_eval_loss",
-    "mean_val_eval_loss",
-    "std_val_eval_loss",
-    "mean_generalization_gap",
-    "std_generalization_gap",
+__all__ = [
+    "build_parser",
+    "main",
+    "run",
 ]
-SELECTION_FIELDS = [
-    "run_mode",
-    "token_limit",
-    "model_name",
-    "n_layer",
-    "n_head",
-    "n_embd",
-    "parameters",
-    "n",
-    "best_dropout",
-    "best_val_loss",
-    "best_val_std",
-    "plateau_start_dropout",
-    "plateau_end_dropout",
-    "plateau_delta",
-    "zero_dropout_val_loss",
-    "zero_minus_best",
-    "best_nonzero_dropout",
-    "best_nonzero_val_loss",
-    "zero_minus_best_nonzero",
-    "max_dropout",
-    "max_dropout_val_loss",
-    "max_dropout_minus_best",
-    "has_nonzero_optimum",
-    "meets_target_dropout",
-    "curve_json",
-]
-
-
-@dataclass(frozen=True)
-class ModelSpec:
-    name: str
-    n_layer: int
-    n_head: int
-    n_embd: int
-
-    def config(self, vocab_size: int, block_size: int, dropout: float) -> GPTConfig:
-        return GPTConfig(
-            block_size=block_size,
-            vocab_size=vocab_size,
-            n_layer=self.n_layer,
-            n_head=self.n_head,
-            n_embd=self.n_embd,
-            dropout=dropout,
-        )
-
-    def to_dict(self) -> dict[str, int | str]:
-        return {
-            "model_name": self.name,
-            "n_layer": self.n_layer,
-            "n_head": self.n_head,
-            "n_embd": self.n_embd,
-        }
-
-
-@dataclass(frozen=True)
-class DropoutCondition:
-    name: str
-    kind: str
-    initial: float
-    final: float
-    schedule: str = "constant"
-    decay_tokens: int | None = None
-    anchors: tuple[tuple[int, float], ...] = ()
-
-    def to_dict(self) -> dict:
-        return {
-            "name": self.name,
-            "kind": self.kind,
-            "initial": self.initial,
-            "final": self.final,
-            "schedule": self.schedule,
-            "decay_tokens": self.decay_tokens,
-            "anchors": [list(anchor) for anchor in self.anchors],
-        }
-
-    def make_fn(
-        self,
-        fallback_decay_tokens: int,
-        unique_tokens: int | None = None,
-    ) -> Callable[[int], float]:
-        if self.kind == "static":
-            return lambda _tokens_seen, p=self.initial: p
-        if self.kind == "anchor_decay":
-            if unique_tokens is None:
-                raise ValueError("anchor_decay conditions require unique_tokens")
-            p = anchor_dropout(unique_tokens, self.anchors)
-            return lambda _tokens_seen, p=p: p
-        scheduler = DropoutDecayScheduler(
-            DropoutDecayConfig(
-                initial_dropout=self.initial,
-                final_dropout=self.final,
-                decay_tokens=self.decay_tokens or fallback_decay_tokens,
-                schedule=self.schedule,
-            )
-        )
-        return scheduler.value
-
-
-def anchor_dropout(unique_tokens: int, anchors: tuple[tuple[int, float], ...]) -> float:
-    if not anchors:
-        raise ValueError("anchor dropout schedule requires at least one anchor")
-    ordered = sorted(anchors)
-    if unique_tokens <= ordered[0][0]:
-        return ordered[0][1]
-    if unique_tokens >= ordered[-1][0]:
-        return ordered[-1][1]
-    log_unique = math.log(unique_tokens)
-    for (left_tokens, left_dropout), (right_tokens, right_dropout) in zip(
-        ordered, ordered[1:]
-    ):
-        if left_tokens <= unique_tokens <= right_tokens:
-            left_log = math.log(left_tokens)
-            right_log = math.log(right_tokens)
-            mix = (log_unique - left_log) / (right_log - left_log)
-            return left_dropout + mix * (right_dropout - left_dropout)
-    return ordered[-1][1]
-
-
-def assert_mps_only() -> torch.device:
-    if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1":
-        raise SystemExit(
-            "Refusing to run: PYTORCH_ENABLE_MPS_FALLBACK=1 could execute CPU fallbacks."
-        )
-    if not torch.backends.mps.is_built():
-        raise SystemExit("PyTorch was not built with MPS. Stopping.")
-    if not torch.backends.mps.is_available():
-        raise SystemExit("MPS is not available. Stopping before any Torch experiment work.")
-    device = torch.device("mps")
-    torch.set_default_device(device)
-    return device
-
-
-def clean_name(value: str) -> str:
-    cleaned = "".join(
-        c if c.isalnum() or c in {"-", "_"} else "_" for c in value.strip()
-    )
-    return cleaned.strip("_") or "model"
-
-
-def rate_label(rate: float) -> str:
-    label = f"{rate:.3f}".rstrip("0").rstrip(".")
-    return label if label else "0"
-
-
-def parse_model_spec(raw: str) -> ModelSpec:
-    if "=" in raw:
-        name, dims = raw.split("=", 1)
-        name = clean_name(name)
-    else:
-        dims = raw
-        name = ""
-    parts = dims.lower().replace(",", "x").split("x")
-    if len(parts) != 3:
-        raise argparse.ArgumentTypeError(
-            "model specs must look like 8x8x256 or name=8x8x256"
-        )
-    try:
-        n_layer, n_head, n_embd = [int(part) for part in parts]
-    except ValueError as exc:
-        raise argparse.ArgumentTypeError("model dimensions must be integers") from exc
-    if n_layer <= 0 or n_head <= 0 or n_embd <= 0:
-        raise argparse.ArgumentTypeError("model dimensions must be positive")
-    if n_embd % n_head != 0:
-        raise argparse.ArgumentTypeError("n_embd must be divisible by n_head")
-    if (n_embd // n_head) % 2 != 0:
-        raise argparse.ArgumentTypeError("n_embd / n_head must be even for rotary attention")
-    return ModelSpec(
-        name or f"L{n_layer}_H{n_head}_D{n_embd}",
-        n_layer,
-        n_head,
-        n_embd,
-    )
-
-
-def parse_decay_spec(raw: str) -> DropoutCondition:
-    parts = raw.split(":")
-    if len(parts) not in {3, 4, 5}:
-        raise argparse.ArgumentTypeError(
-            "decay specs must look like "
-            "name:initial:final[:cosine|smoothstep|linear[:decay_tokens]]"
-        )
-    name = clean_name(parts[0])
-    try:
-        initial = float(parts[1])
-        final = float(parts[2])
-        decay_tokens = int(parts[4]) if len(parts) == 5 else None
-    except ValueError as exc:
-        raise argparse.ArgumentTypeError(
-            "decay dropout values and decay_tokens must be numeric"
-        ) from exc
-    schedule = parts[3] if len(parts) >= 4 else "cosine"
-    if schedule not in {"cosine", "smoothstep", "linear"}:
-        raise argparse.ArgumentTypeError(
-            "decay schedule must be cosine, smoothstep, or linear"
-        )
-    return DropoutCondition(
-        name=name,
-        kind="decay",
-        initial=initial,
-        final=final,
-        schedule=schedule,
-        decay_tokens=decay_tokens,
-    )
-
-
-def parse_anchor_decay_spec(raw: str) -> DropoutCondition:
-    if ":" not in raw:
-        raise argparse.ArgumentTypeError(
-            "anchor decay specs must look like name:250000=0.60,500000=0.40"
-        )
-    name, raw_anchors = raw.split(":", 1)
-    anchors: list[tuple[int, float]] = []
-    for piece in raw_anchors.split(","):
-        if "=" not in piece:
-            raise argparse.ArgumentTypeError(
-                "anchor decay anchors must look like token_count=dropout"
-            )
-        raw_tokens, raw_dropout = piece.split("=", 1)
-        try:
-            token_count = int(raw_tokens)
-            dropout = float(raw_dropout)
-        except ValueError as exc:
-            raise argparse.ArgumentTypeError(
-                "anchor token counts must be integers and dropout values numeric"
-            ) from exc
-        if token_count <= 0:
-            raise argparse.ArgumentTypeError("anchor token counts must be positive")
-        if not 0.0 <= dropout < 1.0:
-            raise argparse.ArgumentTypeError("anchor dropout must satisfy 0 <= p < 1")
-        anchors.append((token_count, dropout))
-    anchors = sorted(set(anchors))
-    if len(anchors) < 2:
-        raise argparse.ArgumentTypeError("provide at least two dropout anchors")
-    for left, right in zip(anchors, anchors[1:]):
-        if right[1] > left[1]:
-            raise argparse.ArgumentTypeError(
-                "anchor dropout values must be non-increasing as token counts grow"
-            )
-    return DropoutCondition(
-        name=clean_name(name),
-        kind="anchor_decay",
-        initial=anchors[0][1],
-        final=anchors[-1][1],
-        schedule="log_prefix_anchor",
-        anchors=tuple(anchors),
-    )
-
-
-def default_seeds(mode: str, seeds: list[int] | None) -> list[int]:
-    if seeds:
-        return seeds
-    return [1] if mode == "screen_static" else [1, 2, 3]
-
-
-def set_seed(seed: int) -> None:
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-
-def batch_seed(seed: int, model: ModelSpec, dropout_code: int, stage: int | None) -> int:
-    return (
-        seed * 1_000_003
-        + model.n_layer * 100_003
-        + model.n_head * 10_007
-        + model.n_embd * 101
-        + dropout_code * 37
-        + (stage or 0) * 997
-    )
-
-
-def make_batch(
-    tokens: np.ndarray,
-    token_limit: int,
-    batch_size: int,
-    block_size: int,
-    rng: np.random.Generator,
-    device: torch.device,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    limit = min(token_limit, len(tokens))
-    max_start = limit - block_size - 1
-    if max_start <= 0:
-        raise ValueError("token_limit is too small for the requested block_size")
-    starts = rng.integers(0, max_start, size=batch_size)
-    x_np = np.stack([tokens[start : start + block_size] for start in starts]).astype(
-        np.int64
-    )
-    y_np = np.stack(
-        [tokens[start + 1 : start + 1 + block_size] for start in starts]
-    ).astype(np.int64)
-    return (
-        torch.tensor(x_np, dtype=torch.long, device=device),
-        torch.tensor(y_np, dtype=torch.long, device=device),
-    )
-
-
-@torch.no_grad()
-def estimate_loss(
-    model: DropoutGPT,
-    tokens: np.ndarray,
-    token_limit: int,
-    batches: int,
-    args: argparse.Namespace,
-    device: torch.device,
-    rng_seed: int,
-) -> float:
-    if batches <= 0:
-        return float("nan")
-    model.eval()
-    rng = np.random.default_rng(rng_seed)
-    losses: list[float] = []
-    for _ in range(batches):
-        x, y = make_batch(tokens, token_limit, args.batch_size, args.block_size, rng, device)
-        _, loss = model(x, y)
-        losses.append(float(loss.item()))
-    model.train()
-    return float(statistics.fmean(losses))
-
-
-def write_jsonl_row(handle, row: dict) -> None:
-    handle.write(json.dumps(row, sort_keys=True) + "\n")
-    handle.flush()
-
-
-def metric_key(row: dict) -> tuple:
-    return (
-        row["run_mode"],
-        row["condition"],
-        row["condition_kind"],
-        row.get("stage"),
-        int(row["token_limit"]),
-        row["model_name"],
-        int(row["n_layer"]),
-        int(row["n_head"]),
-        int(row["n_embd"]),
-        int(row["seed"]),
-        float(row["dropout_initial"]),
-        float(row["dropout_final"]),
-        row["dropout_schedule"],
-    )
-
-
-def planned_metric_key(
-    *,
-    mode: str,
-    condition: DropoutCondition,
-    model_spec: ModelSpec,
-    seed: int,
-    token_limit: int,
-    stage: int | None = None,
-) -> tuple:
-    return (
-        mode,
-        condition.name,
-        condition.kind,
-        stage,
-        int(token_limit),
-        model_spec.name,
-        int(model_spec.n_layer),
-        int(model_spec.n_head),
-        int(model_spec.n_embd),
-        int(seed),
-        float(condition.initial),
-        float(condition.final),
-        condition.schedule,
-    )
-
-
-def load_metrics(path: Path) -> list[dict]:
-    if not path.exists():
-        return []
-    rows = []
-    for line in path.read_text(encoding="utf-8").splitlines():
-        if line.strip():
-            rows.append(json.loads(line))
-    return rows
-
-
-def train_segment(
-    *,
-    run_mode: str,
-    condition: DropoutCondition,
-    model_spec: ModelSpec,
-    config: GPTConfig,
-    train_tokens: np.ndarray,
-    val_tokens: np.ndarray,
-    token_limit: int,
-    steps: int,
-    seed: int,
-    args: argparse.Namespace,
-    device: torch.device,
-    dropout_fn: Callable[[int], float],
-    metrics_file,
-    trace_file,
-    stage: int | None = None,
-    model: DropoutGPT | None = None,
-    optimizer: torch.optim.Optimizer | None = None,
-    tokens_seen_start: int = 0,
-) -> tuple[DropoutGPT, torch.optim.Optimizer, int, dict]:
-    if model is None:
-        set_seed(seed)
-        model = DropoutGPT(config).to(device)
-        optimizer = torch.optim.AdamW(
-            model.parameters(),
-            lr=args.lr,
-            betas=(0.9, 0.95),
-            weight_decay=args.weight_decay,
-        )
-    else:
-        torch.manual_seed(seed + 10_000 + (stage or 0))
-    if optimizer is None:
-        raise ValueError("optimizer is required when reusing a model")
-
-    model.train()
-    dropout_code = int(round(condition.initial * 10_000))
-    rng = np.random.default_rng(batch_seed(seed, model_spec, dropout_code, stage))
-    tokens_seen = tokens_seen_start
-    last_loss = float("nan")
-    active_dropout = condition.initial
-    t0 = time.time()
-
-    for step in range(1, steps + 1):
-        active_dropout = dropout_fn(tokens_seen)
-        model.set_dropout(active_dropout)
-        x, y = make_batch(train_tokens, token_limit, args.batch_size, args.block_size, rng, device)
-        _, loss = model(x, y)
-        optimizer.zero_grad(set_to_none=True)
-        loss.backward()
-        if args.grad_clip > 0:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
-        optimizer.step()
-
-        tokens_seen += args.batch_size * args.block_size
-        last_loss = float(loss.item())
-        if args.log_every > 0 and step % args.log_every == 0:
-            write_jsonl_row(
-                trace_file,
-                {
-                    "event": "train_step",
-                    "run_mode": run_mode,
-                    "condition": condition.name,
-                    "model_name": model_spec.name,
-                    "seed": seed,
-                    "stage": stage,
-                    "step": step,
-                    "steps": steps,
-                    "token_limit": int(token_limit),
-                    "tokens_seen": int(tokens_seen),
-                    "dropout": float(active_dropout),
-                    "train_batch_loss": last_loss,
-                },
-            )
-        if args.eval_every > 0 and step % args.eval_every == 0:
-            train_eval = estimate_loss(
-                model,
-                train_tokens,
-                token_limit,
-                args.trace_eval_batches,
-                args,
-                device,
-                rng_seed=seed + 20_000 + step,
-            )
-            val_eval = estimate_loss(
-                model,
-                val_tokens,
-                len(val_tokens),
-                args.trace_eval_batches,
-                args,
-                device,
-                rng_seed=seed + 30_000 + step,
-            )
-            write_jsonl_row(
-                trace_file,
-                {
-                    "event": "eval_step",
-                    "run_mode": run_mode,
-                    "condition": condition.name,
-                    "model_name": model_spec.name,
-                    "seed": seed,
-                    "stage": stage,
-                    "step": step,
-                    "steps": steps,
-                    "token_limit": int(token_limit),
-                    "tokens_seen": int(tokens_seen),
-                    "dropout": float(active_dropout),
-                    "train_eval_loss": train_eval,
-                    "val_eval_loss": val_eval,
-                    "generalization_gap": val_eval - train_eval,
-                },
-            )
-
-    train_eval = estimate_loss(
-        model,
-        train_tokens,
-        token_limit,
-        args.train_eval_batches,
-        args,
-        device,
-        rng_seed=seed + 40_000 + (stage or 0),
-    )
-    val_eval = estimate_loss(
-        model,
-        val_tokens,
-        len(val_tokens),
-        args.eval_batches,
-        args,
-        device,
-        rng_seed=seed + 50_000 + (stage or 0),
-    )
-    row = {
-        "run_mode": run_mode,
-        "condition": condition.name,
-        "condition_kind": condition.kind,
-        "seed": seed,
-        "stage": stage,
-        "token_limit": int(token_limit),
-        "steps": int(steps),
-        "tokens_seen": int(tokens_seen),
-        "dropout_initial": float(condition.initial),
-        "dropout_final": float(condition.final),
-        "dropout_schedule": condition.schedule,
-        "dropout_active_final": float(active_dropout),
-        "train_loss_last": last_loss,
-        "train_eval_loss": train_eval,
-        "val_eval_loss": val_eval,
-        "eval_loss": val_eval,
-        "generalization_gap": val_eval - train_eval,
-        "elapsed_sec": time.time() - t0,
-        "parameters": model.num_parameters(),
-        "model_config": config.to_dict(),
-        **model_spec.to_dict(),
-    }
-    write_jsonl_row(metrics_file, row)
-    return model, optimizer, tokens_seen, row
-
-
-def summarize(rows: list[dict]) -> list[dict]:
-    groups: dict[tuple, list[dict]] = defaultdict(list)
-    for row in rows:
-        key = (
-            row["run_mode"],
-            row["condition"],
-            row["condition_kind"],
-            row["stage"],
-            row["token_limit"],
-            row["model_name"],
-            row["n_layer"],
-            row["n_head"],
-            row["n_embd"],
-            row["parameters"],
-            row["dropout_initial"],
-            row["dropout_final"],
-            row["dropout_schedule"],
-        )
-        groups[key].append(row)
-
-    summary: list[dict] = []
-    for group_rows in groups.values():
-        first = group_rows[0]
-        item = {field: first[field] for field in SUMMARY_FIELDS if field in first}
-        item["n"] = len(group_rows)
-        for source, mean_key, std_key in [
-            ("train_eval_loss", "mean_train_eval_loss", "std_train_eval_loss"),
-            ("val_eval_loss", "mean_val_eval_loss", "std_val_eval_loss"),
-            ("generalization_gap", "mean_generalization_gap", "std_generalization_gap"),
-        ]:
-            values = [float(row[source]) for row in group_rows]
-            item[mean_key] = statistics.fmean(values)
-            item[std_key] = statistics.stdev(values) if len(values) > 1 else 0.0
-        summary.append(item)
-    return sorted(
-        summary,
-        key=lambda row: (
-            row["run_mode"],
-            row["model_name"],
-            row["token_limit"],
-            row["condition"],
-            row["stage"] or -1,
-        ),
-    )
-
-
-def write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None:
-    with path.open("w", newline="", encoding="utf-8") as handle:
-        writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
-        writer.writeheader()
-        writer.writerows(rows)
-
-
-def format_duration(seconds: float) -> str:
-    seconds = max(0, int(seconds))
-    hours, remainder = divmod(seconds, 3600)
-    minutes, secs = divmod(remainder, 60)
-    if hours:
-        return f"{hours}h{minutes:02d}m"
-    if minutes:
-        return f"{minutes}m{secs:02d}s"
-    return f"{secs}s"
-
-
-class ProgressMeter:
-    def __init__(self, total: int):
-        self.total = max(0, total)
-        self.done = 0
-        self.started_at = time.time()
-
-    def mark_done(self, row: dict) -> None:
-        self.done += 1
-        elapsed = time.time() - self.started_at
-        mean = elapsed / self.done if self.done else 0.0
-        remaining = max(0, self.total - self.done)
-        eta = remaining * mean
-        stage = row.get("stage")
-        stage_text = f"stage={stage} " if stage is not None else ""
-        print(
-            "progress "
-            f"{self.done}/{self.total} "
-            f"eta={format_duration(eta)} "
-            f"mode={row['run_mode']} "
-            f"model={row['model_name']} "
-            f"params={int(row['parameters']):,} "
-            f"prefix={int(row['token_limit']):,} "
-            f"{stage_text}"
-            f"seed={row['seed']} "
-            f"condition={row['condition']} "
-            f"dropout={float(row['dropout_active_final']):.3f} "
-            f"val={row['val_eval_loss']:.4f} "
-            f"train={row['train_eval_loss']:.4f} "
-            f"gap={row['generalization_gap']:.4f} "
-            f"elapsed={format_duration(float(row['elapsed_sec']))}",
-            flush=True,
-        )
-
-
-def build_model_selection(summary: list[dict], args: argparse.Namespace) -> list[dict]:
-    groups: dict[tuple, list[dict]] = defaultdict(list)
-    for row in summary:
-        if row["condition_kind"] == "static" and row["stage"] is None:
-            groups[(row["run_mode"], row["token_limit"], row["model_name"])].append(row)
-
-    selection: list[dict] = []
-    for (run_mode, token_limit, model_name), rows in groups.items():
-        curve = sorted(rows, key=lambda row: row["dropout_initial"])
-        best = min(curve, key=lambda row: row["mean_val_eval_loss"])
-        plateau = [
-            row for row in curve
-            if row["mean_val_eval_loss"] <= best["mean_val_eval_loss"] + args.plateau_delta
-        ]
-        zero = next((row for row in curve if row["dropout_initial"] == 0.0), None)
-        nonzero = [row for row in curve if row["dropout_initial"] > 0.0]
-        best_nonzero = min(nonzero, key=lambda row: row["mean_val_eval_loss"]) if nonzero else None
-        max_dropout = max(curve, key=lambda row: row["dropout_initial"])
-
-        zero_loss = zero["mean_val_eval_loss"] if zero else None
-        zero_minus_best = zero_loss - best["mean_val_eval_loss"] if zero_loss is not None else None
-        zero_minus_best_nonzero = (
-            zero_loss - best_nonzero["mean_val_eval_loss"]
-            if zero_loss is not None and best_nonzero is not None
-            else None
-        )
-        max_minus_best = max_dropout["mean_val_eval_loss"] - best["mean_val_eval_loss"]
-        has_nonzero_optimum = (
-            best["dropout_initial"] > 0.0
-            and zero_minus_best is not None
-            and zero_minus_best >= args.min_nonzero_margin
-            and max_minus_best >= args.min_high_dropout_margin
-        )
-        curve_json = json.dumps(
-            [
-                {
-                    "dropout": row["dropout_initial"],
-                    "mean_val_loss": row["mean_val_eval_loss"],
-                    "std_val_loss": row["std_val_eval_loss"],
-                    "mean_train_loss": row["mean_train_eval_loss"],
-                    "mean_generalization_gap": row["mean_generalization_gap"],
-                    "n": row["n"],
-                }
-                for row in curve
-            ],
-            sort_keys=True,
-        )
-        selection.append(
-            {
-                "run_mode": run_mode,
-                "token_limit": token_limit,
-                "model_name": model_name,
-                "n_layer": best["n_layer"],
-                "n_head": best["n_head"],
-                "n_embd": best["n_embd"],
-                "parameters": best["parameters"],
-                "n": best["n"],
-                "best_dropout": best["dropout_initial"],
-                "best_val_loss": best["mean_val_eval_loss"],
-                "best_val_std": best["std_val_eval_loss"],
-                "plateau_start_dropout": min(row["dropout_initial"] for row in plateau),
-                "plateau_end_dropout": max(row["dropout_initial"] for row in plateau),
-                "plateau_delta": args.plateau_delta,
-                "zero_dropout_val_loss": zero_loss,
-                "zero_minus_best": zero_minus_best,
-                "best_nonzero_dropout": (
-                    best_nonzero["dropout_initial"] if best_nonzero else None
-                ),
-                "best_nonzero_val_loss": (
-                    best_nonzero["mean_val_eval_loss"] if best_nonzero else None
-                ),
-                "zero_minus_best_nonzero": zero_minus_best_nonzero,
-                "max_dropout": max_dropout["dropout_initial"],
-                "max_dropout_val_loss": max_dropout["mean_val_eval_loss"],
-                "max_dropout_minus_best": max_minus_best,
-                "has_nonzero_optimum": has_nonzero_optimum,
-                "meets_target_dropout": best["dropout_initial"] >= args.target_min_dropout,
-                "curve_json": curve_json,
-            }
-        )
-    return sorted(
-        selection,
-        key=lambda row: (
-            not row["has_nonzero_optimum"],
-            not row["meets_target_dropout"],
-            row["best_val_loss"],
-        ),
-    )
-
-
-def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
-    if not rows:
-        return
-    static_rows = [
-        row
-        for row in rows
-        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
-    ]
-    if not static_rows:
-        return
-
-    by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list)
-    for row in rows:
-        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static":
-            by_model_prefix_rate[
-                (
-                    row["model_name"],
-                    int(row["token_limit"]),
-                    float(row["dropout_initial"]),
-                )
-            ].append(row)
-
-    aggregates: list[dict] = []
-    for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items():
-        first = group_rows[0]
-        val_losses = [float(row["val_eval_loss"]) for row in group_rows]
-        train_losses = [float(row["train_eval_loss"]) for row in group_rows]
-        gaps = [float(row["generalization_gap"]) for row in group_rows]
-        aggregates.append(
-            {
-                "model_name": model_name,
-                "token_limit": prefix,
-                "dropout_initial": dropout,
-                "n": len(group_rows),
-                "mean_val_eval_loss": statistics.fmean(val_losses),
-                "std_val_eval_loss": statistics.stdev(val_losses)
-                if len(val_losses) > 1
-                else 0.0,
-                "mean_train_eval_loss": statistics.fmean(train_losses),
-                "std_train_eval_loss": statistics.stdev(train_losses)
-                if len(train_losses) > 1
-                else 0.0,
-                "mean_generalization_gap": statistics.fmean(gaps),
-                "std_generalization_gap": statistics.stdev(gaps)
-                if len(gaps) > 1
-                else 0.0,
-                "parameters": int(first["parameters"]),
-                "n_layer": int(first["n_layer"]),
-                "n_head": int(first["n_head"]),
-                "n_embd": int(first["n_embd"]),
-                "block_size": int(first["model_config"]["block_size"]),
-                "vocab_size": int(first["model_config"]["vocab_size"]),
-                "tokens_seen": int(first["tokens_seen"]),
-                "seeds": sorted({int(row["seed"]) for row in group_rows}),
-            }
-        )
-
-    by_model: dict[str, list[dict]] = defaultdict(list)
-    for row in aggregates:
-        by_model[row["model_name"]].append(row)
-
-    model_rows = []
-    for model_name, model_group in by_model.items():
-        first = model_group[0]
-        seeds = sorted({seed for row in model_group for seed in row["seeds"]})
-        model_rows.append(
-            {
-                "model_name": model_name,
-                "parameters": first["parameters"],
-                "n_layer": first["n_layer"],
-                "n_head": first["n_head"],
-                "n_embd": first["n_embd"],
-                "block_size": first["block_size"],
-                "vocab_size": first["vocab_size"],
-                "seeds": seeds,
-            }
-        )
-
-    lines = [
-        "# Static Dropout Screen Summary",
-        "",
-        f"Run directory: `{output_dir}`",
-        "",
-        "## Models",
-        "",
-        "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |",
-        "|---|---:|---:|---:|---:|---:|---:|---|",
-    ]
-    for model in sorted(model_rows, key=lambda item: item["parameters"]):
-        lines.append(
-            "| "
-            f"`{model['model_name']}` | {model['parameters']:,} | "
-            f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | "
-            f"{model['block_size']} | {model['vocab_size']} | "
-            f"{', '.join(str(seed) for seed in model['seeds'])} |"
-        )
-    lines.extend(
-        [
-            "",
-            "## Best Dropout By Model And Prefix",
-            "",
-            "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |",
-            "|---|---:|---:|---:|---:|---:|---:|---:|---|",
-        ]
-    )
-    for model_name, model_group in sorted(by_model.items()):
-        by_prefix: dict[int, list[dict]] = defaultdict(list)
-        for row in model_group:
-            by_prefix[int(row["token_limit"])].append(row)
-        for prefix, prefix_rows in sorted(by_prefix.items()):
-            best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"])
-            rates = [float(row["dropout_initial"]) for row in prefix_rows]
-            eff_epochs = float(best["tokens_seen"]) / prefix
-            if best["dropout_initial"] == max(rates):
-                note = "not bracketed; best at top of tested grid"
-            elif best["dropout_initial"] == min(rates):
-                note = "not bracketed; best at bottom of tested grid"
-            else:
-                note = "bracketed by tested grid"
-            lines.append(
-                "| "
-                f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | "
-                f"{best['dropout_initial']:.2f} | "
-                f"{best['mean_val_eval_loss']:.4f} | "
-                f"{best['std_val_eval_loss']:.4f} | "
-                f"{best['mean_train_eval_loss']:.4f} | "
-                f"{best['mean_generalization_gap']:.4f} | {note} |"
-            )
-
-    for model_name, model_group in sorted(by_model.items()):
-        by_prefix = defaultdict(list)
-        for row in model_group:
-            by_prefix[int(row["token_limit"])].append(row)
-        lines.extend(
-            [
-                "",
-                f"## Model `{model_name}`",
-            ]
-        )
-        for prefix, prefix_rows in sorted(by_prefix.items()):
-            eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix
-            lines.extend(
-                [
-                    "",
-                    f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)",
-                    "",
-                    "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |",
-                    "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
-                ]
-            )
-            for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]):
-                lines.append(
-                    "| "
-                    f"{row['dropout_initial']:.2f} | {row['n']} | "
-                    f"{row['mean_val_eval_loss']:.4f} | "
-                    f"{row['std_val_eval_loss']:.4f} | "
-                    f"{row['mean_train_eval_loss']:.4f} | "
-                    f"{row['std_train_eval_loss']:.4f} | "
-                    f"{row['mean_generalization_gap']:.4f} | "
-                    f"{row['std_generalization_gap']:.4f} | "
-                    f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |"
-                )
-
-    output = "\n".join(lines) + "\n"
-    (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8")
-
-
-def svg_escape(value: object) -> str:
-    return (
-        str(value)
-        .replace("&", "&amp;")
-        .replace("<", "&lt;")
-        .replace(">", "&gt;")
-        .replace('"', "&quot;")
-    )
-
-
-def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None:
-    rows = [
-        row
-        for row in summary
-        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
-    ]
-    if not rows:
-        return
-
-    grouped: dict[tuple[str, int], list[dict]] = defaultdict(list)
-    model_params: dict[str, int] = {}
-    for row in rows:
-        model_name = row["model_name"]
-        grouped[(model_name, int(row["token_limit"]))].append(row)
-        model_params[model_name] = int(row["parameters"])
-
-    models = sorted(model_params, key=lambda name: model_params[name])
-    prefixes = sorted({int(row["token_limit"]) for row in rows})
-    panel_w, panel_h = 230, 170
-    margin_l, margin_t = 58, 34
-    plot_w, plot_h = 142, 94
-    gap_x, gap_y = 18, 38
-    width = margin_l + len(prefixes) * panel_w + gap_x
-    height = 70 + len(models) * (panel_h + gap_y)
-    colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"]
-
-    parts = [
-        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
-        "<style>",
-        "text{font-family:Arial,Helvetica,sans-serif;fill:#111827}",
-        ".small{font-size:10px}.label{font-size:11px}.title{font-size:15px;font-weight:700}",
-        ".axis{stroke:#374151;stroke-width:1}.grid{stroke:#e5e7eb;stroke-width:1}.line{fill:none;stroke-width:2}",
-        "</style>",
-        '<rect width="100%" height="100%" fill="#ffffff"/>',
-        '<text x="24" y="28" class="title">Static dropout law: validation loss vs dropout</text>',
-        '<text x="24" y="48" class="label">Each panel uses its own y-scale. Points are one-seed means unless N &gt; 1.</text>',
-    ]
-
-    for col, prefix in enumerate(prefixes):
-        x = margin_l + col * panel_w + plot_w / 2
-        parts.append(
-            f'<text x="{x:.1f}" y="70" text-anchor="middle" class="label">{prefix:,} prefix tokens</text>'
-        )
-
-    for row_idx, model_name in enumerate(models):
-        row_y = 92 + row_idx * (panel_h + gap_y)
-        parts.append(
-            f'<text x="24" y="{row_y + 48}" class="label" transform="rotate(-90 24 {row_y + 48})">'
-            f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)</text>'
-        )
-        for col, prefix in enumerate(prefixes):
-            panel_x = margin_l + col * panel_w
-            panel_y = row_y
-            curve = sorted(
-                grouped.get((model_name, prefix), []),
-                key=lambda item: float(item["dropout_initial"]),
-            )
-            if not curve:
-                continue
-            losses = [float(item["mean_val_eval_loss"]) for item in curve]
-            min_loss, max_loss = min(losses), max(losses)
-            pad = max(0.02, (max_loss - min_loss) * 0.08)
-            y_min, y_max = min_loss - pad, max_loss + pad
-            best = min(curve, key=lambda item: float(item["mean_val_eval_loss"]))
-
-            def px(dropout: float) -> float:
-                return panel_x + (dropout / 0.9) * plot_w
-
-            def py(loss: float) -> float:
-                scale = (loss - y_min) / (y_max - y_min)
-                return panel_y + plot_h - scale * plot_h
-
-            parts.extend(
-                [
-                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
-                    f'<line x1="{panel_x:.1f}" y1="{panel_y + plot_h:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
-                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y:.1f}" class="grid"/>',
-                    f'<text x="{panel_x:.1f}" y="{panel_y - 6:.1f}" class="small">{y_max:.2f}</text>',
-                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 13:.1f}" class="small">{y_min:.2f}</text>',
-                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 28:.1f}" class="small">0</text>',
-                    f'<text x="{panel_x + plot_w:.1f}" y="{panel_y + plot_h + 28:.1f}" text-anchor="end" class="small">0.9</text>',
-                ]
-            )
-            points = " ".join(
-                f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}"
-                for item in curve
-            )
-            color = colors[row_idx % len(colors)]
-            parts.append(f'<polyline points="{points}" class="line" stroke="{color}"/>')
-            for item in curve:
-                dropout = float(item["dropout_initial"])
-                loss = float(item["mean_val_eval_loss"])
-                radius = 4 if item is best else 2.7
-                fill = "#111827" if item is best else "#ffffff"
-                parts.append(
-                    f'<circle cx="{px(dropout):.1f}" cy="{py(loss):.1f}" r="{radius}" fill="{fill}" stroke="{color}" stroke-width="1.5"/>'
-                )
-            parts.append(
-                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 14:.1f}" class="small">'
-                f'best p={float(best["dropout_initial"]):.2f}</text>'
-            )
-            parts.append(
-                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 28:.1f}" class="small">'
-                f'loss={float(best["mean_val_eval_loss"]):.3f}</text>'
-            )
-
-    parts.append("</svg>")
-    (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8")
-
-
-def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
-    stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"]
-    if not stream_rows:
-        return
-
-    by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list)
-    by_condition: dict[str, list[dict]] = defaultdict(list)
-    for row in stream_rows:
-        condition = row["condition"]
-        by_condition_stage[(condition, int(row["stage"]))].append(row)
-        by_condition[condition].append(row)
-
-    first = stream_rows[0]
-    seeds = sorted({int(row["seed"]) for row in stream_rows})
-    conditions = sorted(
-        by_condition,
-        key=lambda name: (
-            by_condition[name][0]["condition_kind"] != "anchor_decay",
-            by_condition[name][0]["dropout_initial"],
-            name,
-        ),
-    )
-    stages = sorted({int(row["stage"]) for row in stream_rows})
-
-    lines = [
-        "# Locked Streaming Dropout Summary",
-        "",
-        f"Run directory: `{output_dir}`",
-        "",
-        (
-            f"Model: `{first['model_name']}` causal Transformer, "
-            f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, "
-            f"{first['n_head']} heads, {first['n_embd']} embedding dim."
-        ),
-        (
-            f"Training per stage: {first['steps']:,} steps. "
-            "Sampled tokens are cumulative in each stage row. "
-            f"Seeds present: {', '.join(str(seed) for seed in seeds)}."
-        ),
-        "",
-        "## Condition Ranking",
-        "",
-        "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |",
-        "|---|---|---:|---:|---:|---:|---|",
-    ]
-
-    ranking = []
-    for condition in conditions:
-        stage_items = []
-        for stage in stages:
-            group = by_condition_stage.get((condition, stage), [])
-            if not group:
-                continue
-            stage_items.append(
-                {
-                    "stage": stage,
-                    "token_limit": int(group[0]["token_limit"]),
-                    "mean_val": statistics.fmean(
-                        float(row["val_eval_loss"]) for row in group
-                    ),
-                    "mean_gap": statistics.fmean(
-                        float(row["generalization_gap"]) for row in group
-                    ),
-                    "mean_dropout": statistics.fmean(
-                        float(row["dropout_active_final"]) for row in group
-                    ),
-                    "kind": group[0]["condition_kind"],
-                }
-            )
-        if not stage_items:
-            continue
-        final = max(stage_items, key=lambda item: item["stage"])
-        ranking.append(
-            {
-                "condition": condition,
-                "kind": stage_items[0]["kind"],
-                "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items),
-                "final_val": final["mean_val"],
-                "final_gap": final["mean_gap"],
-                "final_dropout": final["mean_dropout"],
-                "dropout_path": " -> ".join(
-                    f"{item['mean_dropout']:.2f}" for item in stage_items
-                ),
-            }
-        )
-
-    for item in sorted(ranking, key=lambda row: row["trajectory_val"]):
-        lines.append(
-            "| "
-            f"`{item['condition']}` | {item['kind']} | "
-            f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | "
-            f"{item['final_val']:.4f} | {item['final_gap']:.4f} | "
-            f"{item['dropout_path']} |"
-        )
-
-    lines.extend(["", "## Stage Trajectory", ""])
-    for stage in stages:
-        stage_groups = {
-            condition: by_condition_stage[(condition, stage)]
-            for condition in conditions
-            if (condition, stage) in by_condition_stage
-        }
-        if not stage_groups:
-            continue
-        prefix = int(next(iter(stage_groups.values()))[0]["token_limit"])
-        lines.extend(
-            [
-                f"### Stage {stage}: {prefix:,} Prefix Tokens",
-                "",
-                "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |",
-                "|---|---:|---:|---:|---:|---:|",
-            ]
-        )
-        for condition, group in sorted(
-            stage_groups.items(),
-            key=lambda item: statistics.fmean(
-                float(row["val_eval_loss"]) for row in item[1]
-            ),
-        ):
-            val = statistics.fmean(float(row["val_eval_loss"]) for row in group)
-            train = statistics.fmean(float(row["train_eval_loss"]) for row in group)
-            gap = statistics.fmean(float(row["generalization_gap"]) for row in group)
-            dropout = statistics.fmean(
-                float(row["dropout_active_final"]) for row in group
-            )
-            lines.append(
-                "| "
-                f"`{condition}` | {dropout:.2f} | {val:.4f} | "
-                f"{train:.4f} | {gap:.4f} | {len(group)} |"
-            )
-        lines.append("")
-
-    (output_dir / "RESULT_SUMMARY.md").write_text(
-        "\n".join(lines).rstrip() + "\n",
-        encoding="utf-8",
-    )
-
-
-def static_conditions(dropout_rates: list[float]) -> list[DropoutCondition]:
-    return [
-        DropoutCondition(
-            name=f"static_dropout_{rate_label(rate)}",
-            kind="static",
-            initial=rate,
-            final=rate,
-        )
-        for rate in dropout_rates
-    ]
-
-
-def run_fixed_static_sweep(
-    *,
-    args: argparse.Namespace,
-    model_specs: list[ModelSpec],
-    seeds: list[int],
-    train_tokens: np.ndarray,
-    val_tokens: np.ndarray,
-    tokenizer_vocab_size: int,
-    token_limits: list[int],
-    device: torch.device,
-    metrics_file,
-    trace_file,
-    completed_keys: set[tuple] | None = None,
-) -> list[dict]:
-    rows: list[dict] = []
-    completed_keys = completed_keys or set()
-    conditions = static_conditions(sorted(set(args.dropout_rates)))
-    planned = 0
-    for token_limit in token_limits:
-        for model_spec in model_specs:
-            for condition in conditions:
-                for seed in seeds:
-                    key = planned_metric_key(
-                        mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        seed=seed,
-                        token_limit=token_limit,
-                    )
-                    if key not in completed_keys:
-                        planned += 1
-    progress = ProgressMeter(planned)
-    for token_limit in token_limits:
-        for model_spec in model_specs:
-            best_val_loss = float("inf")
-            worse_streak = 0
-            for condition in conditions:
-                condition_rows: list[dict] = []
-                for seed in seeds:
-                    key = planned_metric_key(
-                        mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        seed=seed,
-                        token_limit=token_limit,
-                    )
-                    if key in completed_keys:
-                        write_jsonl_row(
-                            trace_file,
-                            {
-                                "event": "skipped_completed_condition",
-                                "run_mode": args.mode,
-                                "condition": condition.name,
-                                "model_name": model_spec.name,
-                                "seed": seed,
-                                "stage": None,
-                                "token_limit": int(token_limit),
-                            },
-                        )
-                        continue
-                    config = model_spec.config(
-                        tokenizer_vocab_size,
-                        args.block_size,
-                        condition.initial,
-                    )
-                    model, optimizer, _, row = train_segment(
-                        run_mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        config=config,
-                        train_tokens=train_tokens,
-                        val_tokens=val_tokens,
-                        token_limit=token_limit,
-                        steps=args.steps,
-                        seed=seed,
-                        args=args,
-                        device=device,
-                        dropout_fn=condition.make_fn(
-                            args.steps * args.batch_size * args.block_size
-                        ),
-                        metrics_file=metrics_file,
-                        trace_file=trace_file,
-                    )
-                    rows.append(row)
-                    condition_rows.append(row)
-                    completed_keys.add(metric_key(row))
-                    progress.mark_done(row)
-                    del model, optimizer
-                    torch.mps.empty_cache()
-
-                if not condition_rows:
-                    continue
-                mean_val_loss = statistics.fmean(
-                    float(row["val_eval_loss"]) for row in condition_rows
-                )
-                if mean_val_loss < best_val_loss - args.screen_prune_min_delta:
-                    best_val_loss = mean_val_loss
-                    worse_streak = 0
-                elif mean_val_loss > best_val_loss + args.screen_prune_min_delta:
-                    worse_streak += 1
-
-                if (
-                    args.mode == "screen_static"
-                    and args.screen_early_stop
-                    and worse_streak >= args.screen_prune_patience
-                    and condition.initial >= args.target_min_dropout
-                ):
-                    write_jsonl_row(
-                        trace_file,
-                        {
-                            "event": "screen_pruned_model",
-                            "run_mode": args.mode,
-                            "model_name": model_spec.name,
-                            "token_limit": int(token_limit),
-                            "best_val_loss": best_val_loss,
-                            "pruned_after_dropout": condition.initial,
-                            "worse_streak": worse_streak,
-                            "remaining_dropouts": [
-                                rate
-                                for rate in args.dropout_rates
-                                if rate > condition.initial
-                            ],
-                        },
-                    )
-                    break
-    return rows
-
-
-def run_locked_stream(
-    *,
-    args: argparse.Namespace,
-    model_specs: list[ModelSpec],
-    seeds: list[int],
-    train_tokens: np.ndarray,
-    val_tokens: np.ndarray,
-    tokenizer_vocab_size: int,
-    stream_caps: list[int],
-    device: torch.device,
-    metrics_file,
-    trace_file,
-    completed_keys: set[tuple] | None = None,
-) -> list[dict]:
-    rows: list[dict] = []
-    completed_keys = completed_keys or set()
-    conditions = args.anchor_decays + static_conditions(args.dropout_rates) + args.decays
-    fallback_decay_tokens = (
-        args.decay_tokens
-        or args.stage_steps * args.batch_size * args.block_size * len(stream_caps)
-    )
-    planned = 0
-    for model_spec in model_specs:
-        for condition in conditions:
-            for seed in seeds:
-                for stage, token_limit in enumerate(stream_caps):
-                    key = planned_metric_key(
-                        mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        seed=seed,
-                        token_limit=token_limit,
-                        stage=stage,
-                    )
-                    if key not in completed_keys:
-                        planned += 1
-    progress = ProgressMeter(planned)
-    for model_spec in model_specs:
-        for condition in conditions:
-            for seed in seeds:
-                model = None
-                optimizer = None
-                tokens_seen = 0
-                for stage, token_limit in enumerate(stream_caps):
-                    key = planned_metric_key(
-                        mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        seed=seed,
-                        token_limit=token_limit,
-                        stage=stage,
-                    )
-                    if key in completed_keys:
-                        write_jsonl_row(
-                            trace_file,
-                            {
-                                "event": "skipped_completed_condition",
-                                "run_mode": args.mode,
-                                "condition": condition.name,
-                                "model_name": model_spec.name,
-                                "seed": seed,
-                                "stage": stage,
-                                "token_limit": int(token_limit),
-                            },
-                        )
-                        continue
-                    config = model_spec.config(
-                        tokenizer_vocab_size,
-                        args.block_size,
-                        condition.initial,
-                    )
-                    model, optimizer, tokens_seen, row = train_segment(
-                        run_mode=args.mode,
-                        condition=condition,
-                        model_spec=model_spec,
-                        config=config,
-                        train_tokens=train_tokens,
-                        val_tokens=val_tokens,
-                        token_limit=token_limit,
-                        steps=args.stage_steps,
-                        seed=seed,
-                        args=args,
-                        device=device,
-                        dropout_fn=condition.make_fn(
-                            fallback_decay_tokens,
-                            unique_tokens=token_limit,
-                        ),
-                        metrics_file=metrics_file,
-                        trace_file=trace_file,
-                        stage=stage,
-                        model=model,
-                        optimizer=optimizer,
-                        tokens_seen_start=tokens_seen,
-                    )
-                    rows.append(row)
-                    completed_keys.add(metric_key(row))
-                    progress.mark_done(row)
-                del model, optimizer
-                torch.mps.empty_cache()
-    return rows
-
-
-def prepare_data(args: argparse.Namespace, output_dir: Path, required_train_tokens: int):
-    cache_dir = Path(args.cache_dir) if args.cache_dir else output_dir / "cache"
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    if args.use_cached_data:
-        if args.force_retokenize:
-            raise ValueError("--use-cached-data cannot be combined with --force-retokenize")
-        return load_cached_splits(
-            cache_dir=cache_dir,
-            vocab_size=args.vocab_size,
-            max_required_train_tokens=required_train_tokens,
-            val_tokens=args.val_tokens,
-            allow_short_corpus=args.allow_short_corpus,
-        )
-
-    paths = resolve_paths(args.corpus, args.corpus_glob)
-    tokenizer = train_or_load_tokenizer(
-        paths=paths,
-        output_dir=cache_dir,
-        vocab_size=args.vocab_size,
-        tokenizer_train_chars=args.tokenizer_train_chars,
-        text_column=args.text_column,
-        force_retrain=args.force_retokenize,
-    )
-    splits = encode_corpus(
-        paths=paths,
-        tokenizer=tokenizer,
-        output_dir=cache_dir,
-        max_required_train_tokens=required_train_tokens,
-        val_tokens=args.val_tokens,
-        text_column=args.text_column,
-        allow_short_corpus=args.allow_short_corpus,
-        force_reencode=args.force_retokenize,
-    )
-    return tokenizer, splits
-
-
-def run(args: argparse.Namespace) -> Path:
-    device = assert_mps_only()
-    seeds = default_seeds(args.mode, args.seeds)
-    model_specs = [parse_model_spec(spec) for spec in args.models]
-    if args.mode != "locked_stream" and (args.decays or args.anchor_decays):
-        raise ValueError("--decays and --anchor-decays are only used with --mode locked_stream")
-    if args.resume_from and args.mode == "locked_stream":
-        raise ValueError("--resume-from currently supports fixed static sweeps only")
-
-    if args.resume_from:
-        output_dir = Path(args.resume_from)
-        if not output_dir.exists():
-            raise FileNotFoundError(f"resume directory does not exist: {output_dir}")
-    else:
-        run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
-        output_dir = Path(args.output_dir) / args.mode / run_id
-    output_dir.mkdir(parents=True, exist_ok=True)
-    required_train_tokens = max(
-        args.stream_token_caps if args.mode == "locked_stream" else args.token_limits
-    )
-    tokenizer, splits = prepare_data(args, output_dir, required_train_tokens)
-    token_limits = [min(limit, len(splits.train)) for limit in args.token_limits]
-    stream_caps = [min(limit, len(splits.train)) for limit in args.stream_token_caps]
-
-    args_payload = vars(args).copy()
-    args_payload["decays"] = [condition.to_dict() for condition in args.decays]
-    args_payload["anchor_decays"] = [
-        condition.to_dict() for condition in args.anchor_decays
-    ]
-    config_payload = {
-        "args": args_payload,
-        "mode": args.mode,
-        "seeds": seeds,
-        "models": [model.to_dict() for model in model_specs],
-        "device": str(device),
-        "torch": torch.__version__,
-        "python": sys.version,
-        "mps_available": torch.backends.mps.is_available(),
-        "attribution": NANOCHAT_ATTRIBUTION,
-        "tokenizer_path": str(splits.tokenizer_path),
-        "encoded_path": str(splits.encoded_path),
-        "train_tokens": int(len(splits.train)),
-        "val_tokens": int(len(splits.val)),
-        "effective_token_limits": [int(limit) for limit in token_limits],
-        "effective_stream_token_caps": [int(limit) for limit in stream_caps],
-        "resume_from": str(args.resume_from) if args.resume_from else None,
-    }
-    config_name = "config.resume.json" if args.resume_from else "config.json"
-    (output_dir / config_name).write_text(
-        json.dumps(config_payload, indent=2),
-        encoding="utf-8",
-    )
-
-    metrics_path = output_dir / "metrics.jsonl"
-    trace_path = output_dir / "trace.jsonl"
-    existing_rows = load_metrics(metrics_path) if args.resume_from else []
-    completed_keys = {metric_key(row) for row in existing_rows}
-    with (
-        metrics_path.open("a" if args.resume_from else "w", encoding="utf-8") as metrics_file,
-        trace_path.open("a" if args.resume_from else "w", encoding="utf-8") as trace_file,
-    ):
-        if args.mode in {"screen_static", "confirm_static"}:
-            new_rows = run_fixed_static_sweep(
-                args=args,
-                model_specs=model_specs,
-                seeds=seeds,
-                train_tokens=splits.train,
-                val_tokens=splits.val,
-                tokenizer_vocab_size=tokenizer.vocab_size,
-                token_limits=token_limits,
-                device=device,
-                metrics_file=metrics_file,
-                trace_file=trace_file,
-                completed_keys=completed_keys,
-            )
-        else:
-            new_rows = run_locked_stream(
-                args=args,
-                model_specs=model_specs,
-                seeds=seeds,
-                train_tokens=splits.train,
-                val_tokens=splits.val,
-                tokenizer_vocab_size=tokenizer.vocab_size,
-                stream_caps=stream_caps,
-                device=device,
-                metrics_file=metrics_file,
-                trace_file=trace_file,
-            )
-    rows = existing_rows + new_rows
-
-    summary = summarize(rows)
-    (output_dir / "summary.json").write_text(
-        json.dumps(summary, indent=2),
-        encoding="utf-8",
-    )
-    write_csv(output_dir / "summary.csv", summary, SUMMARY_FIELDS)
-    if args.mode in {"screen_static", "confirm_static"}:
-        selection = build_model_selection(summary, args)
-        (output_dir / "model_selection.json").write_text(
-            json.dumps(selection, indent=2),
-            encoding="utf-8",
-        )
-        write_csv(output_dir / "model_selection.csv", selection, SELECTION_FIELDS)
-        write_screen_markdown_summary(output_dir, rows)
-        write_dropout_curve_svg(output_dir, summary)
-    elif args.mode == "locked_stream":
-        write_stream_markdown_summary(output_dir, rows)
-
-    print(
-        json.dumps(
-            {
-                "output_dir": str(output_dir),
-                "new_rows": len(new_rows),
-                "total_metric_rows": len(rows),
-                "summary_rows": len(summary),
-            },
-            indent=2,
-        )
-    )
-    return output_dir
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(
-        description="MPS-only dropout/model selection experiments"
-    )
-    parser.add_argument(
-        "--mode",
-        choices=["screen_static", "confirm_static", "locked_stream"],
-        default="screen_static",
-    )
-    parser.add_argument("--corpus", default=None, help="Text or parquet corpus path")
-    parser.add_argument("--corpus-glob", default=None, help="Glob of text/parquet corpus paths")
-    parser.add_argument("--text-column", default="text", help="Parquet text column")
-    parser.add_argument(
-        "--use-cached-data",
-        action="store_true",
-        help=(
-            "Load tokenizer-v{vocab}.json and tokens-v{vocab}-*.npy from --cache-dir "
-            "instead of requiring the original text/parquet corpus."
-        ),
-    )
-    parser.add_argument("--output-dir", default="runs")
-    parser.add_argument(
-        "--resume-from",
-        default=None,
-        help="Existing fixed-static run directory; completed metric rows are skipped",
-    )
-    parser.add_argument("--cache-dir", default=".cache/dropout_decay")
-    parser.add_argument(
-        "--models",
-        nargs="+",
-        default=["8x8x256"],
-        help="Model specs like 8x8x256 or name=8x8x256",
-    )
-    parser.add_argument("--seeds", nargs="+", type=int, default=None)
-    parser.add_argument("--token-limits", nargs="+", type=int, default=[5_000_000])
-    parser.add_argument(
-        "--stream-token-caps",
-        nargs="+",
-        type=int,
-        default=[5_000_000, 10_000_000, 20_000_000, 40_000_000],
-    )
-    parser.add_argument("--val-tokens", type=int, default=500_000)
-    parser.add_argument("--allow-short-corpus", action="store_true")
-    parser.add_argument("--force-retokenize", action="store_true")
-    parser.add_argument("--vocab-size", type=int, default=4096)
-    parser.add_argument("--tokenizer-train-chars", type=int, default=10_000_000)
-    parser.add_argument("--block-size", type=int, default=128)
-    parser.add_argument("--batch-size", type=int, default=16)
-    parser.add_argument("--steps", type=int, default=2000)
-    parser.add_argument("--stage-steps", type=int, default=1000)
-    parser.add_argument("--dropout-rates", nargs="*", type=float, default=DEFAULT_DROPOUT_RATES)
-    parser.add_argument("--decays", nargs="*", type=parse_decay_spec, default=[])
-    parser.add_argument(
-        "--anchor-decays",
-        nargs="*",
-        type=parse_anchor_decay_spec,
-        default=[],
-        help=(
-            "Prefix-token anchor schedules like "
-            "fit:250000=0.60,500000=0.40,1000000=0.30"
-        ),
-    )
-    parser.add_argument("--decay-tokens", type=int, default=None)
-    parser.add_argument("--eval-batches", type=int, default=64)
-    parser.add_argument("--train-eval-batches", type=int, default=32)
-    parser.add_argument("--trace-eval-batches", type=int, default=8)
-    parser.add_argument("--eval-every", type=int, default=0)
-    parser.add_argument("--log-every", type=int, default=100)
-    parser.add_argument("--lr", type=float, default=3e-4)
-    parser.add_argument("--weight-decay", type=float, default=0.1)
-    parser.add_argument("--grad-clip", type=float, default=1.0)
-    parser.add_argument("--plateau-delta", type=float, default=0.01)
-    parser.add_argument("--target-min-dropout", type=float, default=0.10)
-    parser.add_argument("--min-nonzero-margin", type=float, default=0.01)
-    parser.add_argument("--min-high-dropout-margin", type=float, default=0.03)
-    parser.add_argument("--screen-early-stop", action="store_true")
-    parser.add_argument("--screen-prune-patience", type=int, default=3)
-    parser.add_argument("--screen-prune-min-delta", type=float, default=0.01)
-    return parser
-
-
-def main() -> None:
-    args = build_parser().parse_args()
-    run(args)
 
 
 if __name__ == "__main__":

src/dropout_decay/experiments/__init__.py

@@ -0,0 +1,25 @@
+"""
+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 dropout_decay.experiments.runner import build_parser, main, run
+
+__all__ = [
+    "build_parser",
+    "main",
+    "run",
+]

src/dropout_decay/experiments/artifacts.py

@@ -0,0 +1,281 @@
+"""
+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 argparse
+import csv
+from collections import defaultdict
+import json
+from pathlib import Path
+import statistics
+
+from dropout_decay.specs import DropoutCondition, ModelSpec
+
+
+SUMMARY_FIELDS = [
+    "run_mode",
+    "condition",
+    "condition_kind",
+    "stage",
+    "token_limit",
+    "model_name",
+    "n_layer",
+    "n_head",
+    "n_embd",
+    "parameters",
+    "dropout_initial",
+    "dropout_final",
+    "dropout_schedule",
+    "n",
+    "mean_train_eval_loss",
+    "std_train_eval_loss",
+    "mean_val_eval_loss",
+    "std_val_eval_loss",
+    "mean_generalization_gap",
+    "std_generalization_gap",
+]
+SELECTION_FIELDS = [
+    "run_mode",
+    "token_limit",
+    "model_name",
+    "n_layer",
+    "n_head",
+    "n_embd",
+    "parameters",
+    "n",
+    "best_dropout",
+    "best_val_loss",
+    "best_val_std",
+    "plateau_start_dropout",
+    "plateau_end_dropout",
+    "plateau_delta",
+    "zero_dropout_val_loss",
+    "zero_minus_best",
+    "best_nonzero_dropout",
+    "best_nonzero_val_loss",
+    "zero_minus_best_nonzero",
+    "max_dropout",
+    "max_dropout_val_loss",
+    "max_dropout_minus_best",
+    "has_nonzero_optimum",
+    "meets_target_dropout",
+    "curve_json",
+]
+
+
+def write_jsonl_row(handle, row: dict) -> None:
+    handle.write(json.dumps(row, sort_keys=True) + "\n")
+    handle.flush()
+
+
+def metric_key(row: dict) -> tuple:
+    return (
+        row["run_mode"],
+        row["condition"],
+        row["condition_kind"],
+        row.get("stage"),
+        int(row["token_limit"]),
+        row["model_name"],
+        int(row["n_layer"]),
+        int(row["n_head"]),
+        int(row["n_embd"]),
+        int(row["seed"]),
+        float(row["dropout_initial"]),
+        float(row["dropout_final"]),
+        row["dropout_schedule"],
+    )
+
+
+def planned_metric_key(
+    *,
+    mode: str,
+    condition: DropoutCondition,
+    model_spec: ModelSpec,
+    seed: int,
+    token_limit: int,
+    stage: int | None = None,
+) -> tuple:
+    return (
+        mode,
+        condition.name,
+        condition.kind,
+        stage,
+        int(token_limit),
+        model_spec.name,
+        int(model_spec.n_layer),
+        int(model_spec.n_head),
+        int(model_spec.n_embd),
+        int(seed),
+        float(condition.initial),
+        float(condition.final),
+        condition.schedule,
+    )
+
+
+def load_metrics(path: Path) -> list[dict]:
+    if not path.exists():
+        return []
+    rows = []
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if line.strip():
+            rows.append(json.loads(line))
+    return rows
+
+
+def summarize(rows: list[dict]) -> list[dict]:
+    groups: dict[tuple, list[dict]] = defaultdict(list)
+    for row in rows:
+        key = (
+            row["run_mode"],
+            row["condition"],
+            row["condition_kind"],
+            row["stage"],
+            row["token_limit"],
+            row["model_name"],
+            row["n_layer"],
+            row["n_head"],
+            row["n_embd"],
+            row["parameters"],
+            row["dropout_initial"],
+            row["dropout_final"],
+            row["dropout_schedule"],
+        )
+        groups[key].append(row)
+
+    summary: list[dict] = []
+    for group_rows in groups.values():
+        first = group_rows[0]
+        item = {field: first[field] for field in SUMMARY_FIELDS if field in first}
+        item["n"] = len(group_rows)
+        for source, mean_key, std_key in [
+            ("train_eval_loss", "mean_train_eval_loss", "std_train_eval_loss"),
+            ("val_eval_loss", "mean_val_eval_loss", "std_val_eval_loss"),
+            ("generalization_gap", "mean_generalization_gap", "std_generalization_gap"),
+        ]:
+            values = [float(row[source]) for row in group_rows]
+            item[mean_key] = statistics.fmean(values)
+            item[std_key] = statistics.stdev(values) if len(values) > 1 else 0.0
+        summary.append(item)
+    return sorted(
+        summary,
+        key=lambda row: (
+            row["run_mode"],
+            row["model_name"],
+            row["token_limit"],
+            row["condition"],
+            row["stage"] or -1,
+        ),
+    )
+
+
+def write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None:
+    with path.open("w", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore")
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def build_model_selection(summary: list[dict], args: argparse.Namespace) -> list[dict]:
+    groups: dict[tuple, list[dict]] = defaultdict(list)
+    for row in summary:
+        if row["condition_kind"] == "static" and row["stage"] is None:
+            groups[(row["run_mode"], row["token_limit"], row["model_name"])].append(row)
+
+    selection: list[dict] = []
+    for (run_mode, token_limit, model_name), rows in groups.items():
+        curve = sorted(rows, key=lambda row: row["dropout_initial"])
+        best = min(curve, key=lambda row: row["mean_val_eval_loss"])
+        plateau = [
+            row for row in curve
+            if row["mean_val_eval_loss"] <= best["mean_val_eval_loss"] + args.plateau_delta
+        ]
+        zero = next((row for row in curve if row["dropout_initial"] == 0.0), None)
+        nonzero = [row for row in curve if row["dropout_initial"] > 0.0]
+        best_nonzero = min(nonzero, key=lambda row: row["mean_val_eval_loss"]) if nonzero else None
+        max_dropout = max(curve, key=lambda row: row["dropout_initial"])
+
+        zero_loss = zero["mean_val_eval_loss"] if zero else None
+        zero_minus_best = zero_loss - best["mean_val_eval_loss"] if zero_loss is not None else None
+        zero_minus_best_nonzero = (
+            zero_loss - best_nonzero["mean_val_eval_loss"]
+            if zero_loss is not None and best_nonzero is not None
+            else None
+        )
+        max_minus_best = max_dropout["mean_val_eval_loss"] - best["mean_val_eval_loss"]
+        has_nonzero_optimum = (
+            best["dropout_initial"] > 0.0
+            and zero_minus_best is not None
+            and zero_minus_best >= args.min_nonzero_margin
+            and max_minus_best >= args.min_high_dropout_margin
+        )
+        curve_json = json.dumps(
+            [
+                {
+                    "dropout": row["dropout_initial"],
+                    "mean_val_loss": row["mean_val_eval_loss"],
+                    "std_val_loss": row["std_val_eval_loss"],
+                    "mean_train_loss": row["mean_train_eval_loss"],
+                    "mean_generalization_gap": row["mean_generalization_gap"],
+                    "n": row["n"],
+                }
+                for row in curve
+            ],
+            sort_keys=True,
+        )
+        selection.append(
+            {
+                "run_mode": run_mode,
+                "token_limit": token_limit,
+                "model_name": model_name,
+                "n_layer": best["n_layer"],
+                "n_head": best["n_head"],
+                "n_embd": best["n_embd"],
+                "parameters": best["parameters"],
+                "n": best["n"],
+                "best_dropout": best["dropout_initial"],
+                "best_val_loss": best["mean_val_eval_loss"],
+                "best_val_std": best["std_val_eval_loss"],
+                "plateau_start_dropout": min(row["dropout_initial"] for row in plateau),
+                "plateau_end_dropout": max(row["dropout_initial"] for row in plateau),
+                "plateau_delta": args.plateau_delta,
+                "zero_dropout_val_loss": zero_loss,
+                "zero_minus_best": zero_minus_best,
+                "best_nonzero_dropout": (
+                    best_nonzero["dropout_initial"] if best_nonzero else None
+                ),
+                "best_nonzero_val_loss": (
+                    best_nonzero["mean_val_eval_loss"] if best_nonzero else None
+                ),
+                "zero_minus_best_nonzero": zero_minus_best_nonzero,
+                "max_dropout": max_dropout["dropout_initial"],
+                "max_dropout_val_loss": max_dropout["mean_val_eval_loss"],
+                "max_dropout_minus_best": max_minus_best,
+                "has_nonzero_optimum": has_nonzero_optimum,
+                "meets_target_dropout": best["dropout_initial"] >= args.target_min_dropout,
+                "curve_json": curve_json,
+            }
+        )
+    return sorted(
+        selection,
+        key=lambda row: (
+            not row["has_nonzero_optimum"],
+            not row["meets_target_dropout"],
+            row["best_val_loss"],
+        ),
+    )

src/dropout_decay/experiments/device.py

@@ -0,0 +1,37 @@
+"""
+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 os
+
+import torch
+
+
+def assert_mps_only() -> torch.device:
+    if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1":
+        raise SystemExit(
+            "Refusing to run: PYTORCH_ENABLE_MPS_FALLBACK=1 could execute CPU fallbacks."
+        )
+    if not torch.backends.mps.is_built():
+        raise SystemExit("PyTorch was not built with MPS. Stopping.")
+    if not torch.backends.mps.is_available():
+        raise SystemExit("MPS is not available. Stopping before any Torch experiment work.")
+    device = torch.device("mps")
+    torch.set_default_device(device)
+    return device

src/dropout_decay/experiments/parsing.py

@@ -0,0 +1,148 @@
+"""
+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 argparse
+
+from dropout_decay.specs import DropoutCondition, ModelSpec
+
+
+DEFAULT_DROPOUT_RATES = [0.0, 0.02, 0.05, 0.08, 0.10, 0.14, 0.20, 0.30, 0.50]
+
+
+def clean_name(value: str) -> str:
+    cleaned = "".join(
+        c if c.isalnum() or c in {"-", "_"} else "_" for c in value.strip()
+    )
+    return cleaned.strip("_") or "model"
+
+
+def rate_label(rate: float) -> str:
+    label = f"{rate:.3f}".rstrip("0").rstrip(".")
+    return label if label else "0"
+
+
+def parse_model_spec(raw: str) -> ModelSpec:
+    if "=" in raw:
+        name, dims = raw.split("=", 1)
+        name = clean_name(name)
+    else:
+        dims = raw
+        name = ""
+    parts = dims.lower().replace(",", "x").split("x")
+    if len(parts) != 3:
+        raise argparse.ArgumentTypeError(
+            "model specs must look like 8x8x256 or name=8x8x256"
+        )
+    try:
+        n_layer, n_head, n_embd = [int(part) for part in parts]
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError("model dimensions must be integers") from exc
+    if n_layer <= 0 or n_head <= 0 or n_embd <= 0:
+        raise argparse.ArgumentTypeError("model dimensions must be positive")
+    if n_embd % n_head != 0:
+        raise argparse.ArgumentTypeError("n_embd must be divisible by n_head")
+    if (n_embd // n_head) % 2 != 0:
+        raise argparse.ArgumentTypeError("n_embd / n_head must be even for rotary attention")
+    return ModelSpec(
+        name or f"L{n_layer}_H{n_head}_D{n_embd}",
+        n_layer,
+        n_head,
+        n_embd,
+    )
+
+
+def parse_decay_spec(raw: str) -> DropoutCondition:
+    parts = raw.split(":")
+    if len(parts) not in {3, 4, 5}:
+        raise argparse.ArgumentTypeError(
+            "decay specs must look like "
+            "name:initial:final[:cosine|smoothstep|linear[:decay_tokens]]"
+        )
+    name = clean_name(parts[0])
+    try:
+        initial = float(parts[1])
+        final = float(parts[2])
+        decay_tokens = int(parts[4]) if len(parts) == 5 else None
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(
+            "decay dropout values and decay_tokens must be numeric"
+        ) from exc
+    schedule = parts[3] if len(parts) >= 4 else "cosine"
+    if schedule not in {"cosine", "smoothstep", "linear"}:
+        raise argparse.ArgumentTypeError(
+            "decay schedule must be cosine, smoothstep, or linear"
+        )
+    return DropoutCondition(
+        name=name,
+        kind="decay",
+        initial=initial,
+        final=final,
+        schedule=schedule,
+        decay_tokens=decay_tokens,
+    )
+
+
+def parse_anchor_decay_spec(raw: str) -> DropoutCondition:
+    if ":" not in raw:
+        raise argparse.ArgumentTypeError(
+            "anchor decay specs must look like name:250000=0.60,500000=0.40"
+        )
+    name, raw_anchors = raw.split(":", 1)
+    anchors: list[tuple[int, float]] = []
+    for piece in raw_anchors.split(","):
+        if "=" not in piece:
+            raise argparse.ArgumentTypeError(
+                "anchor decay anchors must look like token_count=dropout"
+            )
+        raw_tokens, raw_dropout = piece.split("=", 1)
+        try:
+            token_count = int(raw_tokens)
+            dropout = float(raw_dropout)
+        except ValueError as exc:
+            raise argparse.ArgumentTypeError(
+                "anchor token counts must be integers and dropout values numeric"
+            ) from exc
+        if token_count <= 0:
+            raise argparse.ArgumentTypeError("anchor token counts must be positive")
+        if not 0.0 <= dropout < 1.0:
+            raise argparse.ArgumentTypeError("anchor dropout must satisfy 0 <= p < 1")
+        anchors.append((token_count, dropout))
+    anchors = sorted(set(anchors))
+    if len(anchors) < 2:
+        raise argparse.ArgumentTypeError("provide at least two dropout anchors")
+    for left, right in zip(anchors, anchors[1:]):
+        if right[1] > left[1]:
+            raise argparse.ArgumentTypeError(
+                "anchor dropout values must be non-increasing as token counts grow"
+            )
+    return DropoutCondition(
+        name=clean_name(name),
+        kind="anchor_decay",
+        initial=anchors[0][1],
+        final=anchors[-1][1],
+        schedule="log_prefix_anchor",
+        anchors=tuple(anchors),
+    )
+
+
+def default_seeds(mode: str, seeds: list[int] | None) -> list[int]:
+    if seeds:
+        return seeds
+    return [1] if mode == "screen_static" else [1, 2, 3]

src/dropout_decay/experiments/progress.py

@@ -0,0 +1,66 @@
+"""
+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 time
+
+
+def format_duration(seconds: float) -> str:
+    seconds = max(0, int(seconds))
+    hours, remainder = divmod(seconds, 3600)
+    minutes, secs = divmod(remainder, 60)
+    if hours:
+        return f"{hours}h{minutes:02d}m"
+    if minutes:
+        return f"{minutes}m{secs:02d}s"
+    return f"{secs}s"
+
+
+class ProgressMeter:
+    def __init__(self, total: int):
+        self.total = max(0, total)
+        self.done = 0
+        self.started_at = time.time()
+
+    def mark_done(self, row: dict) -> None:
+        self.done += 1
+        elapsed = time.time() - self.started_at
+        mean = elapsed / self.done if self.done else 0.0
+        remaining = max(0, self.total - self.done)
+        eta = remaining * mean
+        stage = row.get("stage")
+        stage_text = f"stage={stage} " if stage is not None else ""
+        print(
+            "progress "
+            f"{self.done}/{self.total} "
+            f"eta={format_duration(eta)} "
+            f"mode={row['run_mode']} "
+            f"model={row['model_name']} "
+            f"params={int(row['parameters']):,} "
+            f"prefix={int(row['token_limit']):,} "
+            f"{stage_text}"
+            f"seed={row['seed']} "
+            f"condition={row['condition']} "
+            f"dropout={float(row['dropout_active_final']):.3f} "
+            f"val={row['val_eval_loss']:.4f} "
+            f"train={row['train_eval_loss']:.4f} "
+            f"gap={row['generalization_gap']:.4f} "
+            f"elapsed={format_duration(float(row['elapsed_sec']))}",
+            flush=True,
+        )

src/dropout_decay/experiments/reports.py

@@ -0,0 +1,448 @@
+"""
+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
+
+from collections import defaultdict
+from pathlib import Path
+import statistics
+
+
+def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
+    if not rows:
+        return
+    static_rows = [
+        row
+        for row in rows
+        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
+    ]
+    if not static_rows:
+        return
+
+    by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list)
+    for row in rows:
+        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static":
+            by_model_prefix_rate[
+                (
+                    row["model_name"],
+                    int(row["token_limit"]),
+                    float(row["dropout_initial"]),
+                )
+            ].append(row)
+
+    aggregates: list[dict] = []
+    for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items():
+        first = group_rows[0]
+        val_losses = [float(row["val_eval_loss"]) for row in group_rows]
+        train_losses = [float(row["train_eval_loss"]) for row in group_rows]
+        gaps = [float(row["generalization_gap"]) for row in group_rows]
+        aggregates.append(
+            {
+                "model_name": model_name,
+                "token_limit": prefix,
+                "dropout_initial": dropout,
+                "n": len(group_rows),
+                "mean_val_eval_loss": statistics.fmean(val_losses),
+                "std_val_eval_loss": statistics.stdev(val_losses)
+                if len(val_losses) > 1
+                else 0.0,
+                "mean_train_eval_loss": statistics.fmean(train_losses),
+                "std_train_eval_loss": statistics.stdev(train_losses)
+                if len(train_losses) > 1
+                else 0.0,
+                "mean_generalization_gap": statistics.fmean(gaps),
+                "std_generalization_gap": statistics.stdev(gaps)
+                if len(gaps) > 1
+                else 0.0,
+                "parameters": int(first["parameters"]),
+                "n_layer": int(first["n_layer"]),
+                "n_head": int(first["n_head"]),
+                "n_embd": int(first["n_embd"]),
+                "block_size": int(first["model_config"]["block_size"]),
+                "vocab_size": int(first["model_config"]["vocab_size"]),
+                "tokens_seen": int(first["tokens_seen"]),
+                "seeds": sorted({int(row["seed"]) for row in group_rows}),
+            }
+        )
+
+    by_model: dict[str, list[dict]] = defaultdict(list)
+    for row in aggregates:
+        by_model[row["model_name"]].append(row)
+
+    model_rows = []
+    for model_name, model_group in by_model.items():
+        first = model_group[0]
+        seeds = sorted({seed for row in model_group for seed in row["seeds"]})
+        model_rows.append(
+            {
+                "model_name": model_name,
+                "parameters": first["parameters"],
+                "n_layer": first["n_layer"],
+                "n_head": first["n_head"],
+                "n_embd": first["n_embd"],
+                "block_size": first["block_size"],
+                "vocab_size": first["vocab_size"],
+                "seeds": seeds,
+            }
+        )
+
+    lines = [
+        "# Static Dropout Screen Summary",
+        "",
+        f"Run directory: `{output_dir}`",
+        "",
+        "## Models",
+        "",
+        "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |",
+        "|---|---:|---:|---:|---:|---:|---:|---|",
+    ]
+    for model in sorted(model_rows, key=lambda item: item["parameters"]):
+        lines.append(
+            "| "
+            f"`{model['model_name']}` | {model['parameters']:,} | "
+            f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | "
+            f"{model['block_size']} | {model['vocab_size']} | "
+            f"{', '.join(str(seed) for seed in model['seeds'])} |"
+        )
+    lines.extend(
+        [
+            "",
+            "## Best Dropout By Model And Prefix",
+            "",
+            "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |",
+            "|---|---:|---:|---:|---:|---:|---:|---:|---|",
+        ]
+    )
+    for model_name, model_group in sorted(by_model.items()):
+        by_prefix: dict[int, list[dict]] = defaultdict(list)
+        for row in model_group:
+            by_prefix[int(row["token_limit"])].append(row)
+        for prefix, prefix_rows in sorted(by_prefix.items()):
+            best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"])
+            rates = [float(row["dropout_initial"]) for row in prefix_rows]
+            eff_epochs = float(best["tokens_seen"]) / prefix
+            if best["dropout_initial"] == max(rates):
+                note = "not bracketed; best at top of tested grid"
+            elif best["dropout_initial"] == min(rates):
+                note = "not bracketed; best at bottom of tested grid"
+            else:
+                note = "bracketed by tested grid"
+            lines.append(
+                "| "
+                f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | "
+                f"{best['dropout_initial']:.2f} | "
+                f"{best['mean_val_eval_loss']:.4f} | "
+                f"{best['std_val_eval_loss']:.4f} | "
+                f"{best['mean_train_eval_loss']:.4f} | "
+                f"{best['mean_generalization_gap']:.4f} | {note} |"
+            )
+
+    for model_name, model_group in sorted(by_model.items()):
+        by_prefix = defaultdict(list)
+        for row in model_group:
+            by_prefix[int(row["token_limit"])].append(row)
+        lines.extend(
+            [
+                "",
+                f"## Model `{model_name}`",
+            ]
+        )
+        for prefix, prefix_rows in sorted(by_prefix.items()):
+            eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix
+            lines.extend(
+                [
+                    "",
+                    f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)",
+                    "",
+                    "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |",
+                    "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
+                ]
+            )
+            for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]):
+                lines.append(
+                    "| "
+                    f"{row['dropout_initial']:.2f} | {row['n']} | "
+                    f"{row['mean_val_eval_loss']:.4f} | "
+                    f"{row['std_val_eval_loss']:.4f} | "
+                    f"{row['mean_train_eval_loss']:.4f} | "
+                    f"{row['std_train_eval_loss']:.4f} | "
+                    f"{row['mean_generalization_gap']:.4f} | "
+                    f"{row['std_generalization_gap']:.4f} | "
+                    f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |"
+                )
+
+    output = "\n".join(lines) + "\n"
+    (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8")
+
+
+def svg_escape(value: object) -> str:
+    return (
+        str(value)
+        .replace("&", "&")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+        .replace('"', "&quot;")
+    )
+
+
+def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None:
+    rows = [
+        row
+        for row in summary
+        if row["run_mode"] == "screen_static" and row["condition_kind"] == "static"
+    ]
+    if not rows:
+        return
+
+    grouped: dict[tuple[str, int], list[dict]] = defaultdict(list)
+    model_params: dict[str, int] = {}
+    for row in rows:
+        model_name = row["model_name"]
+        grouped[(model_name, int(row["token_limit"]))].append(row)
+        model_params[model_name] = int(row["parameters"])
+
+    models = sorted(model_params, key=lambda name: model_params[name])
+    prefixes = sorted({int(row["token_limit"]) for row in rows})
+    panel_w, panel_h = 230, 170
+    margin_l, margin_t = 58, 34
+    plot_w, plot_h = 142, 94
+    gap_x, gap_y = 18, 38
+    width = margin_l + len(prefixes) * panel_w + gap_x
+    height = 70 + len(models) * (panel_h + gap_y)
+    colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"]
+
+    parts = [
+        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
+        "<style>",
+        "text{font-family:Arial,Helvetica,sans-serif;fill:#111827}",
+        ".small{font-size:10px}.label{font-size:11px}.title{font-size:15px;font-weight:700}",
+        ".axis{stroke:#374151;stroke-width:1}.grid{stroke:#e5e7eb;stroke-width:1}.line{fill:none;stroke-width:2}",
+        "</style>",
+        '<rect width="100%" height="100%" fill="#ffffff"/>',
+        '<text x="24" y="28" class="title">Static dropout law: validation loss vs dropout</text>',
+        '<text x="24" y="48" class="label">Each panel uses its own y-scale. Points are one-seed means unless N &gt; 1.</text>',
+    ]
+
+    for col, prefix in enumerate(prefixes):
+        x = margin_l + col * panel_w + plot_w / 2
+        parts.append(
+            f'<text x="{x:.1f}" y="70" text-anchor="middle" class="label">{prefix:,} prefix tokens</text>'
+        )
+
+    for row_idx, model_name in enumerate(models):
+        row_y = 92 + row_idx * (panel_h + gap_y)
+        parts.append(
+            f'<text x="24" y="{row_y + 48}" class="label" transform="rotate(-90 24 {row_y + 48})">'
+            f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)</text>'
+        )
+        for col, prefix in enumerate(prefixes):
+            panel_x = margin_l + col * panel_w
+            panel_y = row_y
+            curve = sorted(
+                grouped.get((model_name, prefix), []),
+                key=lambda item: float(item["dropout_initial"]),
+            )
+            if not curve:
+                continue
+            losses = [float(item["mean_val_eval_loss"]) for item in curve]
+            min_loss, max_loss = min(losses), max(losses)
+            pad = max(0.02, (max_loss - min_loss) * 0.08)
+            y_min, y_max = min_loss - pad, max_loss + pad
+            best = min(curve, key=lambda item: float(item["mean_val_eval_loss"]))
+
+            def px(dropout: float) -> float:
+                return panel_x + (dropout / 0.9) * plot_w
+
+            def py(loss: float) -> float:
+                scale = (loss - y_min) / (y_max - y_min)
+                return panel_y + plot_h - scale * plot_h
+
+            parts.extend(
+                [
+                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
+                    f'<line x1="{panel_x:.1f}" y1="{panel_y + plot_h:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y + plot_h:.1f}" class="axis"/>',
+                    f'<line x1="{panel_x:.1f}" y1="{panel_y:.1f}" x2="{panel_x + plot_w:.1f}" y2="{panel_y:.1f}" class="grid"/>',
+                    f'<text x="{panel_x:.1f}" y="{panel_y - 6:.1f}" class="small">{y_max:.2f}</text>',
+                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 13:.1f}" class="small">{y_min:.2f}</text>',
+                    f'<text x="{panel_x:.1f}" y="{panel_y + plot_h + 28:.1f}" class="small">0</text>',
+                    f'<text x="{panel_x + plot_w:.1f}" y="{panel_y + plot_h + 28:.1f}" text-anchor="end" class="small">0.9</text>',
+                ]
+            )
+            points = " ".join(
+                f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}"
+                for item in curve
+            )
+            color = colors[row_idx % len(colors)]
+            parts.append(f'<polyline points="{points}" class="line" stroke="{color}"/>')
+            for item in curve:
+                dropout = float(item["dropout_initial"])
+                loss = float(item["mean_val_eval_loss"])
+                radius = 4 if item is best else 2.7
+                fill = "#111827" if item is best else "#ffffff"
+                parts.append(
+                    f'<circle cx="{px(dropout):.1f}" cy="{py(loss):.1f}" r="{radius}" fill="{fill}" stroke="{color}" stroke-width="1.5"/>'
+                )
+            parts.append(
+                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 14:.1f}" class="small">'
+                f'best p={float(best["dropout_initial"]):.2f}</text>'
+            )
+            parts.append(
+                f'<text x="{panel_x + plot_w + 8:.1f}" y="{panel_y + 28:.1f}" class="small">'
+                f'loss={float(best["mean_val_eval_loss"]):.3f}</text>'
+            )
+
+    parts.append("</svg>")
+    (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8")
+
+
+def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None:
+    stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"]
+    if not stream_rows:
+        return
+
+    by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list)
+    by_condition: dict[str, list[dict]] = defaultdict(list)
+    for row in stream_rows:
+        condition = row["condition"]
+        by_condition_stage[(condition, int(row["stage"]))].append(row)
+        by_condition[condition].append(row)
+
+    first = stream_rows[0]
+    seeds = sorted({int(row["seed"]) for row in stream_rows})
+    conditions = sorted(
+        by_condition,
+        key=lambda name: (
+            by_condition[name][0]["condition_kind"] != "anchor_decay",
+            by_condition[name][0]["dropout_initial"],
+            name,
+        ),
+    )
+    stages = sorted({int(row["stage"]) for row in stream_rows})
+
+    lines = [
+        "# Locked Streaming Dropout Summary",
+        "",
+        f"Run directory: `{output_dir}`",
+        "",
+        (
+            f"Model: `{first['model_name']}` causal Transformer, "
+            f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, "
+            f"{first['n_head']} heads, {first['n_embd']} embedding dim."
+        ),
+        (
+            f"Training per stage: {first['steps']:,} steps. "
+            "Sampled tokens are cumulative in each stage row. "
+            f"Seeds present: {', '.join(str(seed) for seed in seeds)}."
+        ),
+        "",
+        "## Condition Ranking",
+        "",
+        "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |",
+        "|---|---|---:|---:|---:|---:|---|",
+    ]
+
+    ranking = []
+    for condition in conditions:
+        stage_items = []
+        for stage in stages:
+            group = by_condition_stage.get((condition, stage), [])
+            if not group:
+                continue
+            stage_items.append(
+                {
+                    "stage": stage,
+                    "token_limit": int(group[0]["token_limit"]),
+                    "mean_val": statistics.fmean(
+                        float(row["val_eval_loss"]) for row in group
+                    ),
+                    "mean_gap": statistics.fmean(
+                        float(row["generalization_gap"]) for row in group
+                    ),
+                    "mean_dropout": statistics.fmean(
+                        float(row["dropout_active_final"]) for row in group
+                    ),
+                    "kind": group[0]["condition_kind"],
+                }
+            )
+        if not stage_items:
+            continue
+        final = max(stage_items, key=lambda item: item["stage"])
+        ranking.append(
+            {
+                "condition": condition,
+                "kind": stage_items[0]["kind"],
+                "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items),
+                "final_val": final["mean_val"],
+                "final_gap": final["mean_gap"],
+                "final_dropout": final["mean_dropout"],
+                "dropout_path": " -> ".join(
+                    f"{item['mean_dropout']:.2f}" for item in stage_items
+                ),
+            }
+        )
+
+    for item in sorted(ranking, key=lambda row: row["trajectory_val"]):
+        lines.append(
+            "| "
+            f"`{item['condition']}` | {item['kind']} | "
+            f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | "
+            f"{item['final_val']:.4f} | {item['final_gap']:.4f} | "
+            f"{item['dropout_path']} |"
+        )
+
+    lines.extend(["", "## Stage Trajectory", ""])
+    for stage in stages:
+        stage_groups = {
+            condition: by_condition_stage[(condition, stage)]
+            for condition in conditions
+            if (condition, stage) in by_condition_stage
+        }
+        if not stage_groups:
+            continue
+        prefix = int(next(iter(stage_groups.values()))[0]["token_limit"])
+        lines.extend(
+            [
+                f"### Stage {stage}: {prefix:,} Prefix Tokens",
+                "",
+                "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |",
+                "|---|---:|---:|---:|---:|---:|",
+            ]
+        )
+        for condition, group in sorted(
+            stage_groups.items(),
+            key=lambda item: statistics.fmean(
+                float(row["val_eval_loss"]) for row in item[1]
+            ),
+        ):
+            val = statistics.fmean(float(row["val_eval_loss"]) for row in group)
+            train = statistics.fmean(float(row["train_eval_loss"]) for row in group)
+            gap = statistics.fmean(float(row["generalization_gap"]) for row in group)
+            dropout = statistics.fmean(
+                float(row["dropout_active_final"]) for row in group
+            )
+            lines.append(
+                "| "
+                f"`{condition}` | {dropout:.2f} | {val:.4f} | "
+                f"{train:.4f} | {gap:.4f} | {len(group)} |"
+            )
+        lines.append("")
+
+    (output_dir / "RESULT_SUMMARY.md").write_text(
+        "\n".join(lines).rstrip() + "\n",
+        encoding="utf-8",
+    )

src/dropout_decay/experiments/runner.py

@@ -0,0 +1,555 @@
+"""
+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 argparse
+from datetime import datetime
+import json
+from pathlib import Path
+import statistics
+import sys
+
+import numpy as np
+import torch
+
+from dropout_decay.datasets import (
+    encode_corpus,
+    load_cached_splits,
+    resolve_paths,
+    train_or_load_tokenizer,
+)
+from dropout_decay.experiments.artifacts import (
+    SELECTION_FIELDS,
+    SUMMARY_FIELDS,
+    build_model_selection,
+    load_metrics,
+    metric_key,
+    planned_metric_key,
+    summarize,
+    write_csv,
+    write_jsonl_row,
+)
+from dropout_decay.experiments.device import assert_mps_only
+from dropout_decay.experiments.parsing import (
+    DEFAULT_DROPOUT_RATES,
+    default_seeds,
+    parse_anchor_decay_spec,
+    parse_decay_spec,
+    parse_model_spec,
+    rate_label,
+)
+from dropout_decay.experiments.progress import ProgressMeter
+from dropout_decay.experiments.reports import (
+    write_dropout_curve_svg,
+    write_screen_markdown_summary,
+    write_stream_markdown_summary,
+)
+from dropout_decay.experiments.training import train_segment
+from dropout_decay.license import NANOCHAT_ATTRIBUTION
+from dropout_decay.specs import DropoutCondition, ModelSpec
+
+
+def static_conditions(dropout_rates: list[float]) -> list[DropoutCondition]:
+    return [
+        DropoutCondition(
+            name=f"static_dropout_{rate_label(rate)}",
+            kind="static",
+            initial=rate,
+            final=rate,
+        )
+        for rate in dropout_rates
+    ]
+
+
+def run_fixed_static_sweep(
+    *,
+    args: argparse.Namespace,
+    model_specs: list[ModelSpec],
+    seeds: list[int],
+    train_tokens: np.ndarray,
+    val_tokens: np.ndarray,
+    tokenizer_vocab_size: int,
+    token_limits: list[int],
+    device: torch.device,
+    metrics_file,
+    trace_file,
+    completed_keys: set[tuple] | None = None,
+) -> list[dict]:
+    rows: list[dict] = []
+    completed_keys = completed_keys or set()
+    conditions = static_conditions(sorted(set(args.dropout_rates)))
+    planned = 0
+    for token_limit in token_limits:
+        for model_spec in model_specs:
+            for condition in conditions:
+                for seed in seeds:
+                    key = planned_metric_key(
+                        mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        seed=seed,
+                        token_limit=token_limit,
+                    )
+                    if key not in completed_keys:
+                        planned += 1
+    progress = ProgressMeter(planned)
+    for token_limit in token_limits:
+        for model_spec in model_specs:
+            best_val_loss = float("inf")
+            worse_streak = 0
+            for condition in conditions:
+                condition_rows: list[dict] = []
+                for seed in seeds:
+                    key = planned_metric_key(
+                        mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        seed=seed,
+                        token_limit=token_limit,
+                    )
+                    if key in completed_keys:
+                        write_jsonl_row(
+                            trace_file,
+                            {
+                                "event": "skipped_completed_condition",
+                                "run_mode": args.mode,
+                                "condition": condition.name,
+                                "model_name": model_spec.name,
+                                "seed": seed,
+                                "stage": None,
+                                "token_limit": int(token_limit),
+                            },
+                        )
+                        continue
+                    config = model_spec.config(
+                        tokenizer_vocab_size,
+                        args.block_size,
+                        condition.initial,
+                    )
+                    model, optimizer, _, row = train_segment(
+                        run_mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        config=config,
+                        train_tokens=train_tokens,
+                        val_tokens=val_tokens,
+                        token_limit=token_limit,
+                        steps=args.steps,
+                        seed=seed,
+                        args=args,
+                        device=device,
+                        dropout_fn=condition.make_fn(
+                            args.steps * args.batch_size * args.block_size
+                        ),
+                        metrics_file=metrics_file,
+                        trace_file=trace_file,
+                    )
+                    rows.append(row)
+                    condition_rows.append(row)
+                    completed_keys.add(metric_key(row))
+                    progress.mark_done(row)
+                    del model, optimizer
+                    torch.mps.empty_cache()
+
+                if not condition_rows:
+                    continue
+                mean_val_loss = statistics.fmean(
+                    float(row["val_eval_loss"]) for row in condition_rows
+                )
+                if mean_val_loss < best_val_loss - args.screen_prune_min_delta:
+                    best_val_loss = mean_val_loss
+                    worse_streak = 0
+                elif mean_val_loss > best_val_loss + args.screen_prune_min_delta:
+                    worse_streak += 1
+
+                if (
+                    args.mode == "screen_static"
+                    and args.screen_early_stop
+                    and worse_streak >= args.screen_prune_patience
+                    and condition.initial >= args.target_min_dropout
+                ):
+                    write_jsonl_row(
+                        trace_file,
+                        {
+                            "event": "screen_pruned_model",
+                            "run_mode": args.mode,
+                            "model_name": model_spec.name,
+                            "token_limit": int(token_limit),
+                            "best_val_loss": best_val_loss,
+                            "pruned_after_dropout": condition.initial,
+                            "worse_streak": worse_streak,
+                            "remaining_dropouts": [
+                                rate
+                                for rate in args.dropout_rates
+                                if rate > condition.initial
+                            ],
+                        },
+                    )
+                    break
+    return rows
+
+
+def run_locked_stream(
+    *,
+    args: argparse.Namespace,
+    model_specs: list[ModelSpec],
+    seeds: list[int],
+    train_tokens: np.ndarray,
+    val_tokens: np.ndarray,
+    tokenizer_vocab_size: int,
+    stream_caps: list[int],
+    device: torch.device,
+    metrics_file,
+    trace_file,
+    completed_keys: set[tuple] | None = None,
+) -> list[dict]:
+    rows: list[dict] = []
+    completed_keys = completed_keys or set()
+    conditions = args.anchor_decays + static_conditions(args.dropout_rates) + args.decays
+    fallback_decay_tokens = (
+        args.decay_tokens
+        or args.stage_steps * args.batch_size * args.block_size * len(stream_caps)
+    )
+    planned = 0
+    for model_spec in model_specs:
+        for condition in conditions:
+            for seed in seeds:
+                for stage, token_limit in enumerate(stream_caps):
+                    key = planned_metric_key(
+                        mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        seed=seed,
+                        token_limit=token_limit,
+                        stage=stage,
+                    )
+                    if key not in completed_keys:
+                        planned += 1
+    progress = ProgressMeter(planned)
+    for model_spec in model_specs:
+        for condition in conditions:
+            for seed in seeds:
+                model = None
+                optimizer = None
+                tokens_seen = 0
+                for stage, token_limit in enumerate(stream_caps):
+                    key = planned_metric_key(
+                        mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        seed=seed,
+                        token_limit=token_limit,
+                        stage=stage,
+                    )
+                    if key in completed_keys:
+                        write_jsonl_row(
+                            trace_file,
+                            {
+                                "event": "skipped_completed_condition",
+                                "run_mode": args.mode,
+                                "condition": condition.name,
+                                "model_name": model_spec.name,
+                                "seed": seed,
+                                "stage": stage,
+                                "token_limit": int(token_limit),
+                            },
+                        )
+                        continue
+                    config = model_spec.config(
+                        tokenizer_vocab_size,
+                        args.block_size,
+                        condition.initial,
+                    )
+                    model, optimizer, tokens_seen, row = train_segment(
+                        run_mode=args.mode,
+                        condition=condition,
+                        model_spec=model_spec,
+                        config=config,
+                        train_tokens=train_tokens,
+                        val_tokens=val_tokens,
+                        token_limit=token_limit,
+                        steps=args.stage_steps,
+                        seed=seed,
+                        args=args,
+                        device=device,
+                        dropout_fn=condition.make_fn(
+                            fallback_decay_tokens,
+                            unique_tokens=token_limit,
+                        ),
+                        metrics_file=metrics_file,
+                        trace_file=trace_file,
+                        stage=stage,
+                        model=model,
+                        optimizer=optimizer,
+                        tokens_seen_start=tokens_seen,
+                    )
+                    rows.append(row)
+                    completed_keys.add(metric_key(row))
+                    progress.mark_done(row)
+                del model, optimizer
+                torch.mps.empty_cache()
+    return rows
+
+
+def prepare_data(args: argparse.Namespace, output_dir: Path, required_train_tokens: int):
+    cache_dir = Path(args.cache_dir) if args.cache_dir else output_dir / "cache"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    if args.use_cached_data:
+        if args.force_retokenize:
+            raise ValueError("--use-cached-data cannot be combined with --force-retokenize")
+        return load_cached_splits(
+            cache_dir=cache_dir,
+            vocab_size=args.vocab_size,
+            max_required_train_tokens=required_train_tokens,
+            val_tokens=args.val_tokens,
+            allow_short_corpus=args.allow_short_corpus,
+        )
+
+    paths = resolve_paths(args.corpus, args.corpus_glob)
+    tokenizer = train_or_load_tokenizer(
+        paths=paths,
+        output_dir=cache_dir,
+        vocab_size=args.vocab_size,
+        tokenizer_train_chars=args.tokenizer_train_chars,
+        text_column=args.text_column,
+        force_retrain=args.force_retokenize,
+    )
+    splits = encode_corpus(
+        paths=paths,
+        tokenizer=tokenizer,
+        output_dir=cache_dir,
+        max_required_train_tokens=required_train_tokens,
+        val_tokens=args.val_tokens,
+        text_column=args.text_column,
+        allow_short_corpus=args.allow_short_corpus,
+        force_reencode=args.force_retokenize,
+    )
+    return tokenizer, splits
+
+
+def run(args: argparse.Namespace) -> Path:
+    device = assert_mps_only()
+    seeds = default_seeds(args.mode, args.seeds)
+    model_specs = [parse_model_spec(spec) for spec in args.models]
+    if args.mode != "locked_stream" and (args.decays or args.anchor_decays):
+        raise ValueError("--decays and --anchor-decays are only used with --mode locked_stream")
+    if args.resume_from and args.mode == "locked_stream":
+        raise ValueError("--resume-from currently supports fixed static sweeps only")
+
+    if args.resume_from:
+        output_dir = Path(args.resume_from)
+        if not output_dir.exists():
+            raise FileNotFoundError(f"resume directory does not exist: {output_dir}")
+    else:
+        run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
+        output_dir = Path(args.output_dir) / args.mode / run_id
+    output_dir.mkdir(parents=True, exist_ok=True)
+    required_train_tokens = max(
+        args.stream_token_caps if args.mode == "locked_stream" else args.token_limits
+    )
+    tokenizer, splits = prepare_data(args, output_dir, required_train_tokens)
+    token_limits = [min(limit, len(splits.train)) for limit in args.token_limits]
+    stream_caps = [min(limit, len(splits.train)) for limit in args.stream_token_caps]
+
+    args_payload = vars(args).copy()
+    args_payload["decays"] = [condition.to_dict() for condition in args.decays]
+    args_payload["anchor_decays"] = [
+        condition.to_dict() for condition in args.anchor_decays
+    ]
+    config_payload = {
+        "args": args_payload,
+        "mode": args.mode,
+        "seeds": seeds,
+        "models": [model.to_dict() for model in model_specs],
+        "device": str(device),
+        "torch": torch.__version__,
+        "python": sys.version,
+        "mps_available": torch.backends.mps.is_available(),
+        "attribution": NANOCHAT_ATTRIBUTION,
+        "tokenizer_path": str(splits.tokenizer_path),
+        "encoded_path": str(splits.encoded_path),
+        "train_tokens": int(len(splits.train)),
+        "val_tokens": int(len(splits.val)),
+        "effective_token_limits": [int(limit) for limit in token_limits],
+        "effective_stream_token_caps": [int(limit) for limit in stream_caps],
+        "resume_from": str(args.resume_from) if args.resume_from else None,
+    }
+    config_name = "config.resume.json" if args.resume_from else "config.json"
+    (output_dir / config_name).write_text(
+        json.dumps(config_payload, indent=2),
+        encoding="utf-8",
+    )
+
+    metrics_path = output_dir / "metrics.jsonl"
+    trace_path = output_dir / "trace.jsonl"
+    existing_rows = load_metrics(metrics_path) if args.resume_from else []
+    completed_keys = {metric_key(row) for row in existing_rows}
+    with (
+        metrics_path.open("a" if args.resume_from else "w", encoding="utf-8") as metrics_file,
+        trace_path.open("a" if args.resume_from else "w", encoding="utf-8") as trace_file,
+    ):
+        if args.mode in {"screen_static", "confirm_static"}:
+            new_rows = run_fixed_static_sweep(
+                args=args,
+                model_specs=model_specs,
+                seeds=seeds,
+                train_tokens=splits.train,
+                val_tokens=splits.val,
+                tokenizer_vocab_size=tokenizer.vocab_size,
+                token_limits=token_limits,
+                device=device,
+                metrics_file=metrics_file,
+                trace_file=trace_file,
+                completed_keys=completed_keys,
+            )
+        else:
+            new_rows = run_locked_stream(
+                args=args,
+                model_specs=model_specs,
+                seeds=seeds,
+                train_tokens=splits.train,
+                val_tokens=splits.val,
+                tokenizer_vocab_size=tokenizer.vocab_size,
+                stream_caps=stream_caps,
+                device=device,
+                metrics_file=metrics_file,
+                trace_file=trace_file,
+            )
+    rows = existing_rows + new_rows
+
+    summary = summarize(rows)
+    (output_dir / "summary.json").write_text(
+        json.dumps(summary, indent=2),
+        encoding="utf-8",
+    )
+    write_csv(output_dir / "summary.csv", summary, SUMMARY_FIELDS)
+    if args.mode in {"screen_static", "confirm_static"}:
+        selection = build_model_selection(summary, args)
+        (output_dir / "model_selection.json").write_text(
+            json.dumps(selection, indent=2),
+            encoding="utf-8",
+        )
+        write_csv(output_dir / "model_selection.csv", selection, SELECTION_FIELDS)
+        write_screen_markdown_summary(output_dir, rows)
+        write_dropout_curve_svg(output_dir, summary)
+    elif args.mode == "locked_stream":
+        write_stream_markdown_summary(output_dir, rows)
+
+    print(
+        json.dumps(
+            {
+                "output_dir": str(output_dir),
+                "new_rows": len(new_rows),
+                "total_metric_rows": len(rows),
+                "summary_rows": len(summary),
+            },
+            indent=2,
+        )
+    )
+    return output_dir
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="MPS-only dropout/model selection experiments"
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["screen_static", "confirm_static", "locked_stream"],
+        default="screen_static",
+    )
+    parser.add_argument("--corpus", default=None, help="Text or parquet corpus path")
+    parser.add_argument("--corpus-glob", default=None, help="Glob of text/parquet corpus paths")
+    parser.add_argument("--text-column", default="text", help="Parquet text column")
+    parser.add_argument(
+        "--use-cached-data",
+        action="store_true",
+        help=(
+            "Load tokenizer-v{vocab}.json and tokens-v{vocab}-*.npy from --cache-dir "
+            "instead of requiring the original text/parquet corpus."
+        ),
+    )
+    parser.add_argument("--output-dir", default="runs")
+    parser.add_argument(
+        "--resume-from",
+        default=None,
+        help="Existing fixed-static run directory; completed metric rows are skipped",
+    )
+    parser.add_argument("--cache-dir", default=".cache/dropout_decay")
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        default=["8x8x256"],
+        help="Model specs like 8x8x256 or name=8x8x256",
+    )
+    parser.add_argument("--seeds", nargs="+", type=int, default=None)
+    parser.add_argument("--token-limits", nargs="+", type=int, default=[5_000_000])
+    parser.add_argument(
+        "--stream-token-caps",
+        nargs="+",
+        type=int,
+        default=[5_000_000, 10_000_000, 20_000_000, 40_000_000],
+    )
+    parser.add_argument("--val-tokens", type=int, default=500_000)
+    parser.add_argument("--allow-short-corpus", action="store_true")
+    parser.add_argument("--force-retokenize", action="store_true")
+    parser.add_argument("--vocab-size", type=int, default=4096)
+    parser.add_argument("--tokenizer-train-chars", type=int, default=10_000_000)
+    parser.add_argument("--block-size", type=int, default=128)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--steps", type=int, default=2000)
+    parser.add_argument("--stage-steps", type=int, default=1000)
+    parser.add_argument("--dropout-rates", nargs="*", type=float, default=DEFAULT_DROPOUT_RATES)
+    parser.add_argument("--decays", nargs="*", type=parse_decay_spec, default=[])
+    parser.add_argument(
+        "--anchor-decays",
+        nargs="*",
+        type=parse_anchor_decay_spec,
+        default=[],
+        help=(
+            "Prefix-token anchor schedules like "
+            "fit:250000=0.60,500000=0.40,1000000=0.30"
+        ),
+    )
+    parser.add_argument("--decay-tokens", type=int, default=None)
+    parser.add_argument("--eval-batches", type=int, default=64)
+    parser.add_argument("--train-eval-batches", type=int, default=32)
+    parser.add_argument("--trace-eval-batches", type=int, default=8)
+    parser.add_argument("--eval-every", type=int, default=0)
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--weight-decay", type=float, default=0.1)
+    parser.add_argument("--grad-clip", type=float, default=1.0)
+    parser.add_argument("--plateau-delta", type=float, default=0.01)
+    parser.add_argument("--target-min-dropout", type=float, default=0.10)
+    parser.add_argument("--min-nonzero-margin", type=float, default=0.01)
+    parser.add_argument("--min-high-dropout-margin", type=float, default=0.03)
+    parser.add_argument("--screen-early-stop", action="store_true")
+    parser.add_argument("--screen-prune-patience", type=int, default=3)
+    parser.add_argument("--screen-prune-min-delta", type=float, default=0.01)
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    run(args)
+
+
+if __name__ == "__main__":
+    main()

src/dropout_decay/experiments/training.py

@@ -0,0 +1,261 @@
+"""
+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 argparse
+import random
+import statistics
+import time
+from typing import Callable
+
+import numpy as np
+import torch
+
+from dropout_decay.experiments.artifacts import write_jsonl_row
+from dropout_decay.models import DropoutGPT, GPTConfig
+from dropout_decay.specs import DropoutCondition, ModelSpec
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def batch_seed(
+    seed: int, model: ModelSpec, dropout_code: int, stage: int | None
+) -> int:
+    return (
+        seed * 1_000_003
+        + model.n_layer * 100_003
+        + model.n_head * 10_007
+        + model.n_embd * 101
+        + dropout_code * 37
+        + (stage or 0) * 997
+    )
+
+
+def make_batch(
+    tokens: np.ndarray,
+    token_limit: int,
+    batch_size: int,
+    block_size: int,
+    rng: np.random.Generator,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    limit = min(token_limit, len(tokens))
+    max_start = limit - block_size - 1
+    if max_start <= 0:
+        raise ValueError("token_limit is too small for the requested block_size")
+    starts = rng.integers(0, max_start, size=batch_size)
+    x_np = np.stack([tokens[start : start + block_size] for start in starts]).astype(
+        np.int64
+    )
+    y_np = np.stack(
+        [tokens[start + 1 : start + 1 + block_size] for start in starts]
+    ).astype(np.int64)
+    return (
+        torch.tensor(x_np, dtype=torch.long, device=device),
+        torch.tensor(y_np, dtype=torch.long, device=device),
+    )
+
+
+@torch.no_grad()
+def estimate_loss(
+    model: DropoutGPT,
+    tokens: np.ndarray,
+    token_limit: int,
+    batches: int,
+    args: argparse.Namespace,
+    device: torch.device,
+    rng_seed: int,
+) -> float:
+    if batches <= 0:
+        return float("nan")
+    model.eval()
+    rng = np.random.default_rng(rng_seed)
+    losses: list[float] = []
+    for _ in range(batches):
+        x, y = make_batch(
+            tokens, token_limit, args.batch_size, args.block_size, rng, device
+        )
+        _, loss = model(x, y)
+        losses.append(float(loss.item()))
+    model.train()
+    return float(statistics.fmean(losses))
+
+
+def train_segment(
+    *,
+    run_mode: str,
+    condition: DropoutCondition,
+    model_spec: ModelSpec,
+    config: GPTConfig,
+    train_tokens: np.ndarray,
+    val_tokens: np.ndarray,
+    token_limit: int,
+    steps: int,
+    seed: int,
+    args: argparse.Namespace,
+    device: torch.device,
+    dropout_fn: Callable[[int], float],
+    metrics_file,
+    trace_file,
+    stage: int | None = None,
+    model: DropoutGPT | None = None,
+    optimizer: torch.optim.Optimizer | None = None,
+    tokens_seen_start: int = 0,
+) -> tuple[DropoutGPT, torch.optim.Optimizer, int, dict]:
+    if model is None:
+        set_seed(seed)
+        model = DropoutGPT(config).to(device)
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=args.lr,
+            betas=(0.9, 0.95),
+            weight_decay=args.weight_decay,
+        )
+    else:
+        torch.manual_seed(seed + 10_000 + (stage or 0))
+    if optimizer is None:
+        raise ValueError("optimizer is required when reusing a model")
+
+    model.train()
+    dropout_code = int(round(condition.initial * 10_000))
+    rng = np.random.default_rng(batch_seed(seed, model_spec, dropout_code, stage))
+    tokens_seen = tokens_seen_start
+    last_loss = float("nan")
+    active_dropout = condition.initial
+    t0 = time.time()
+
+    for step in range(1, steps + 1):
+        active_dropout = dropout_fn(tokens_seen)
+        model.set_dropout(active_dropout)
+        x, y = make_batch(
+            train_tokens, token_limit, args.batch_size, args.block_size, rng, device
+        )
+        _, loss = model(x, y)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        if args.grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step()
+
+        tokens_seen += args.batch_size * args.block_size
+        last_loss = float(loss.item())
+        if args.log_every > 0 and step % args.log_every == 0:
+            write_jsonl_row(
+                trace_file,
+                {
+                    "event": "train_step",
+                    "run_mode": run_mode,
+                    "condition": condition.name,
+                    "model_name": model_spec.name,
+                    "seed": seed,
+                    "stage": stage,
+                    "step": step,
+                    "steps": steps,
+                    "token_limit": int(token_limit),
+                    "tokens_seen": int(tokens_seen),
+                    "dropout": float(active_dropout),
+                    "train_batch_loss": last_loss,
+                },
+            )
+        if args.eval_every > 0 and step % args.eval_every == 0:
+            train_eval = estimate_loss(
+                model,
+                train_tokens,
+                token_limit,
+                args.trace_eval_batches,
+                args,
+                device,
+                rng_seed=seed + 20_000 + step,
+            )
+            val_eval = estimate_loss(
+                model,
+                val_tokens,
+                len(val_tokens),
+                args.trace_eval_batches,
+                args,
+                device,
+                rng_seed=seed + 30_000 + step,
+            )
+            write_jsonl_row(
+                trace_file,
+                {
+                    "event": "eval_step",
+                    "run_mode": run_mode,
+                    "condition": condition.name,
+                    "model_name": model_spec.name,
+                    "seed": seed,
+                    "stage": stage,
+                    "step": step,
+                    "steps": steps,
+                    "token_limit": int(token_limit),
+                    "tokens_seen": int(tokens_seen),
+                    "dropout": float(active_dropout),
+                    "train_eval_loss": train_eval,
+                    "val_eval_loss": val_eval,
+                    "generalization_gap": val_eval - train_eval,
+                },
+            )
+
+    train_eval = estimate_loss(
+        model,
+        train_tokens,
+        token_limit,
+        args.train_eval_batches,
+        args,
+        device,
+        rng_seed=seed + 40_000 + (stage or 0),
+    )
+    val_eval = estimate_loss(
+        model,
+        val_tokens,
+        len(val_tokens),
+        args.eval_batches,
+        args,
+        device,
+        rng_seed=seed + 50_000 + (stage or 0),
+    )
+    row = {
+        "run_mode": run_mode,
+        "condition": condition.name,
+        "condition_kind": condition.kind,
+        "seed": seed,
+        "stage": stage,
+        "token_limit": int(token_limit),
+        "steps": int(steps),
+        "tokens_seen": int(tokens_seen),
+        "dropout_initial": float(condition.initial),
+        "dropout_final": float(condition.final),
+        "dropout_schedule": condition.schedule,
+        "dropout_active_final": float(active_dropout),
+        "train_loss_last": last_loss,
+        "train_eval_loss": train_eval,
+        "val_eval_loss": val_eval,
+        "eval_loss": val_eval,
+        "generalization_gap": val_eval - train_eval,
+        "elapsed_sec": time.time() - t0,
+        "parameters": model.num_parameters(),
+        "model_config": config.to_dict(),
+        **model_spec.to_dict(),
+    }
+    write_jsonl_row(metrics_file, row)
+    return model, optimizer, tokens_seen, row

src/dropout_decay/model.py

@@ -16,212 +16,24 @@ The above copyright notice and this permission notice shall be included in all
 copies or substantial portions of the Software.
 """
 
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-@dataclass(frozen=True)
-class GPTConfig:
-    block_size: int
-    vocab_size: int
-    n_layer: int
-    n_head: int
-    n_embd: int
-    dropout: float
-
-    def to_dict(self) -> dict[str, int | float]:
-        return asdict(self)
-
-
-def rms_norm(x: torch.Tensor) -> torch.Tensor:
-    return F.rms_norm(x, (x.size(-1),))
-
-
-class Linear(nn.Linear):
-    """Bias-free linear layer matching the nanochat architectural style."""
-
-    def __init__(self, in_features: int, out_features: int):
-        super().__init__(in_features, out_features, bias=False)
-
-
-def apply_rotary_emb(
-    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
-) -> torch.Tensor:
-    half = x.shape[-1] // 2
-    x1, x2 = x[..., :half], x[..., half:]
-    return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
-
-
-class CausalSelfAttention(nn.Module):
-    def __init__(self, config: GPTConfig):
-        super().__init__()
-        if config.n_embd % config.n_head != 0:
-            raise ValueError("n_embd must be divisible by n_head")
-        self.n_head = config.n_head
-        self.head_dim = config.n_embd // config.n_head
-        if self.head_dim % 2 != 0:
-            raise ValueError("rotary attention requires an even head dimension")
-        self.dropout_p = config.dropout
-        self.c_q = Linear(config.n_embd, config.n_embd)
-        self.c_k = Linear(config.n_embd, config.n_embd)
-        self.c_v = Linear(config.n_embd, config.n_embd)
-        self.c_proj = Linear(config.n_embd, config.n_embd)
-        self.resid_dropout = nn.Dropout(config.dropout)
-        mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril()
-        self.register_buffer(
-            "causal_mask",
-            mask.view(1, 1, config.block_size, config.block_size),
-            persistent=False,
-        )
-
-    def set_dropout(self, p: float) -> None:
-        self.dropout_p = p
-        self.resid_dropout.p = p
-
-    def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
-    ) -> torch.Tensor:
-        batch, seq_len, channels = x.shape
-        q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
-        k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
-        v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
-
-        q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2
-        k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2
-
-        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        mask = self.causal_mask[:, :, :seq_len, :seq_len]
-        att = att.masked_fill(~mask, float("-inf"))
-        att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype)
-        att = F.dropout(att, p=self.dropout_p, training=self.training)
-        y = att @ v
-        y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels)
-        return self.resid_dropout(self.c_proj(y))
-
-
-class MLP(nn.Module):
-    def __init__(self, config: GPTConfig):
-        super().__init__()
-        self.c_fc = Linear(config.n_embd, 4 * config.n_embd)
-        self.c_proj = Linear(4 * config.n_embd, config.n_embd)
-        self.dropout = nn.Dropout(config.dropout)
-
-    def set_dropout(self, p: float) -> None:
-        self.dropout.p = p
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.c_fc(x)
-        x = F.relu(x).square()
-        x = self.c_proj(x)
-        return self.dropout(x)
-
-
-class Block(nn.Module):
-    def __init__(self, config: GPTConfig):
-        super().__init__()
-        self.attn = CausalSelfAttention(config)
-        self.mlp = MLP(config)
-
-    def set_dropout(self, p: float) -> None:
-        self.attn.set_dropout(p)
-        self.mlp.set_dropout(p)
-
-    def forward(
-        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
-    ) -> torch.Tensor:
-        x = x + self.attn(rms_norm(x), cos, sin)
-        x = x + self.mlp(rms_norm(x))
-        return x
-
-
-class DropoutGPT(nn.Module):
-    """Minimal nanochat-style causal Transformer with dynamic dropout control."""
-
-    def __init__(self, config: GPTConfig):
-        super().__init__()
-        self.config = config
-        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
-        self.embedding_dropout = nn.Dropout(config.dropout)
-        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer))
-        self.lm_head = Linear(config.n_embd, config.vocab_size)
-        self.register_buffer(
-            "cos",
-            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
-            persistent=False,
-        )
-        self.register_buffer(
-            "sin",
-            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
-            persistent=False,
-        )
-        self._init_weights()
-        self._init_rotary()
-
-    @torch.no_grad()
-    def _init_weights(self) -> None:
-        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8)
-        nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
-        scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5)
-        for block in self.blocks:
-            nn.init.uniform_(block.attn.c_q.weight, -scale, scale)
-            nn.init.uniform_(block.attn.c_k.weight, -scale, scale)
-            nn.init.uniform_(block.attn.c_v.weight, -scale, scale)
-            nn.init.zeros_(block.attn.c_proj.weight)
-            nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale)
-            nn.init.zeros_(block.mlp.c_proj.weight)
-
-    @torch.no_grad()
-    def _init_rotary(self) -> None:
-        head_dim = self.config.n_embd // self.config.n_head
-        channel_range = torch.arange(
-            0,
-            head_dim,
-            2,
-            dtype=torch.float32,
-            device=self.cos.device,
-        )
-        inv_freq = 1.0 / (100000 ** (channel_range / head_dim))
-        positions = torch.arange(
-            self.config.block_size,
-            dtype=torch.float32,
-            device=self.cos.device,
-        )
-        freqs = torch.outer(positions, inv_freq)
-        self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2))
-        self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2))
-
-    def set_dropout(self, p: float) -> None:
-        if not 0.0 <= p < 1.0:
-            raise ValueError("dropout must satisfy 0 <= p < 1")
-        self.embedding_dropout.p = p
-        for block in self.blocks:
-            block.set_dropout(p)
-
-    def num_parameters(self) -> int:
-        return sum(p.numel() for p in self.parameters())
-
-    def forward(
-        self, idx: torch.Tensor, targets: torch.Tensor | None = None
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        _, seq_len = idx.shape
-        if seq_len > self.config.block_size:
-            raise ValueError("sequence length exceeds block_size")
-        cos = self.cos[:, :, :seq_len, :]
-        sin = self.sin[:, :, :seq_len, :]
-        x = self.embedding_dropout(rms_norm(self.token_embedding(idx)))
-        for block in self.blocks:
-            x = block(x, cos, sin)
-        logits = self.lm_head(rms_norm(x)).float()
-        if targets is None:
-            return logits
-        loss = F.cross_entropy(
-            logits.reshape(-1, logits.size(-1)),
-            targets.reshape(-1),
-        )
-        return logits, loss
+from dropout_decay.models.gpt import (
+    Block,
+    CausalSelfAttention,
+    DropoutGPT,
+    GPTConfig,
+    Linear,
+    MLP,
+    apply_rotary_emb,
+    rms_norm,
+)
+
+__all__ = [
+    "Block",
+    "CausalSelfAttention",
+    "DropoutGPT",
+    "GPTConfig",
+    "Linear",
+    "MLP",
+    "apply_rotary_emb",
+    "rms_norm",
+]

src/dropout_decay/models/__init__.py

@@ -0,0 +1,39 @@
+"""
+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 dropout_decay.models.gpt import (
+    Block,
+    CausalSelfAttention,
+    DropoutGPT,
+    GPTConfig,
+    Linear,
+    MLP,
+    apply_rotary_emb,
+    rms_norm,
+)
+
+__all__ = [
+    "Block",
+    "CausalSelfAttention",
+    "DropoutGPT",
+    "GPTConfig",
+    "Linear",
+    "MLP",
+    "apply_rotary_emb",
+    "rms_norm",
+]

src/dropout_decay/models/gpt.py

@@ -0,0 +1,227 @@
+"""
+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
+
+from dataclasses import asdict, dataclass
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass(frozen=True)
+class GPTConfig:
+    block_size: int
+    vocab_size: int
+    n_layer: int
+    n_head: int
+    n_embd: int
+    dropout: float
+
+    def to_dict(self) -> dict[str, int | float]:
+        return asdict(self)
+
+
+def rms_norm(x: torch.Tensor) -> torch.Tensor:
+    return F.rms_norm(x, (x.size(-1),))
+
+
+class Linear(nn.Linear):
+    """Bias-free linear layer matching the nanochat architectural style."""
+
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__(in_features, out_features, bias=False)
+
+
+def apply_rotary_emb(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        if config.n_embd % config.n_head != 0:
+            raise ValueError("n_embd must be divisible by n_head")
+        self.n_head = config.n_head
+        self.head_dim = config.n_embd // config.n_head
+        if self.head_dim % 2 != 0:
+            raise ValueError("rotary attention requires an even head dimension")
+        self.dropout_p = config.dropout
+        self.c_q = Linear(config.n_embd, config.n_embd)
+        self.c_k = Linear(config.n_embd, config.n_embd)
+        self.c_v = Linear(config.n_embd, config.n_embd)
+        self.c_proj = Linear(config.n_embd, config.n_embd)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril()
+        self.register_buffer(
+            "causal_mask",
+            mask.view(1, 1, config.block_size, config.block_size),
+            persistent=False,
+        )
+
+    def set_dropout(self, p: float) -> None:
+        self.dropout_p = p
+        self.resid_dropout.p = p
+
+    def forward(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        batch, seq_len, channels = x.shape
+        q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+
+        q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2
+        k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2
+
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        mask = self.causal_mask[:, :, :seq_len, :seq_len]
+        att = att.masked_fill(~mask, float("-inf"))
+        att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype)
+        att = F.dropout(att, p=self.dropout_p, training=self.training)
+        y = att @ v
+        y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels)
+        return self.resid_dropout(self.c_proj(y))
+
+
+class MLP(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        self.c_fc = Linear(config.n_embd, 4 * config.n_embd)
+        self.c_proj = Linear(4 * config.n_embd, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def set_dropout(self, p: float) -> None:
+        self.dropout.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return self.dropout(x)
+
+
+class Block(nn.Module):
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        self.attn = CausalSelfAttention(config)
+        self.mlp = MLP(config)
+
+    def set_dropout(self, p: float) -> None:
+        self.attn.set_dropout(p)
+        self.mlp.set_dropout(p)
+
+    def forward(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        x = x + self.attn(rms_norm(x), cos, sin)
+        x = x + self.mlp(rms_norm(x))
+        return x
+
+
+class DropoutGPT(nn.Module):
+    """Minimal nanochat-style causal Transformer with dynamic dropout control."""
+
+    def __init__(self, config: GPTConfig):
+        super().__init__()
+        self.config = config
+        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
+        self.embedding_dropout = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer))
+        self.lm_head = Linear(config.n_embd, config.vocab_size)
+        self.register_buffer(
+            "cos",
+            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
+            persistent=False,
+        )
+        self.register_buffer(
+            "sin",
+            torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2),
+            persistent=False,
+        )
+        self._init_weights()
+        self._init_rotary()
+
+    @torch.no_grad()
+    def _init_weights(self) -> None:
+        nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8)
+        nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
+        scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5)
+        for block in self.blocks:
+            nn.init.uniform_(block.attn.c_q.weight, -scale, scale)
+            nn.init.uniform_(block.attn.c_k.weight, -scale, scale)
+            nn.init.uniform_(block.attn.c_v.weight, -scale, scale)
+            nn.init.zeros_(block.attn.c_proj.weight)
+            nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale)
+            nn.init.zeros_(block.mlp.c_proj.weight)
+
+    @torch.no_grad()
+    def _init_rotary(self) -> None:
+        head_dim = self.config.n_embd // self.config.n_head
+        channel_range = torch.arange(
+            0,
+            head_dim,
+            2,
+            dtype=torch.float32,
+            device=self.cos.device,
+        )
+        inv_freq = 1.0 / (100000 ** (channel_range / head_dim))
+        positions = torch.arange(
+            self.config.block_size,
+            dtype=torch.float32,
+            device=self.cos.device,
+        )
+        freqs = torch.outer(positions, inv_freq)
+        self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2))
+        self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2))
+
+    def set_dropout(self, p: float) -> None:
+        if not 0.0 <= p < 1.0:
+            raise ValueError("dropout must satisfy 0 <= p < 1")
+        self.embedding_dropout.p = p
+        for block in self.blocks:
+            block.set_dropout(p)
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+    def forward(
+        self, idx: torch.Tensor, targets: torch.Tensor | None = None
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        _, seq_len = idx.shape
+        if seq_len > self.config.block_size:
+            raise ValueError("sequence length exceeds block_size")
+        cos = self.cos[:, :, :seq_len, :]
+        sin = self.sin[:, :, :seq_len, :]
+        x = self.embedding_dropout(rms_norm(self.token_embedding(idx)))
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        logits = self.lm_head(rms_norm(x)).float()
+        if targets is None:
+            return logits
+        loss = F.cross_entropy(
+            logits.reshape(-1, logits.size(-1)),
+            targets.reshape(-1),
+        )
+        return logits, loss

src/dropout_decay/scheduler.py

@@ -16,41 +16,9 @@ The above copyright notice and this permission notice shall be included in all
 copies or substantial portions of the Software.
 """
 
-from __future__ import annotations
+from dropout_decay.schedules.dropout import DropoutDecayConfig, DropoutDecayScheduler
 
-from dataclasses import dataclass
-import math
-
-
-@dataclass(frozen=True)
-class DropoutDecayConfig:
-    initial_dropout: float = 0.80
-    final_dropout: float = 0.10
-    decay_tokens: int = 40_000_000
-    schedule: str = "cosine"
-
-
-class DropoutDecayScheduler:
-    """Smoothly maps total processed/observed tokens to an active dropout rate."""
-
-    def __init__(self, config: DropoutDecayConfig):
-        if not 0.0 <= config.final_dropout <= config.initial_dropout < 1.0:
-            raise ValueError("dropout values must satisfy 0 <= final <= initial < 1")
-        if config.decay_tokens <= 0:
-            raise ValueError("decay_tokens must be positive")
-        if config.schedule not in {"cosine", "smoothstep", "linear"}:
-            raise ValueError("schedule must be cosine, smoothstep, or linear")
-        self.config = config
-
-    def value(self, tokens_seen: int) -> float:
-        progress = min(max(tokens_seen / self.config.decay_tokens, 0.0), 1.0)
-        if self.config.schedule == "cosine":
-            mix = 0.5 * (1.0 + math.cos(math.pi * progress))
-        elif self.config.schedule == "smoothstep":
-            smooth = progress * progress * (3.0 - 2.0 * progress)
-            mix = 1.0 - smooth
-        else:
-            mix = 1.0 - progress
-        return self.config.final_dropout + (
-            self.config.initial_dropout - self.config.final_dropout
-        ) * mix
+__all__ = [
+    "DropoutDecayConfig",
+    "DropoutDecayScheduler",
+]

src/dropout_decay/schedules/__init__.py

@@ -0,0 +1,24 @@
+"""
+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 dropout_decay.schedules.dropout import DropoutDecayConfig, DropoutDecayScheduler
+
+__all__ = [
+    "DropoutDecayConfig",
+    "DropoutDecayScheduler",
+]

src/dropout_decay/schedules/dropout.py

@@ -0,0 +1,56 @@
+"""
+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
+
+from dataclasses import dataclass
+import math
+
+
+@dataclass(frozen=True)
+class DropoutDecayConfig:
+    initial_dropout: float = 0.80
+    final_dropout: float = 0.10
+    decay_tokens: int = 40_000_000
+    schedule: str = "cosine"
+
+
+class DropoutDecayScheduler:
+    """Smoothly maps total processed/observed tokens to an active dropout rate."""
+
+    def __init__(self, config: DropoutDecayConfig):
+        if not 0.0 <= config.final_dropout <= config.initial_dropout < 1.0:
+            raise ValueError("dropout values must satisfy 0 <= final <= initial < 1")
+        if config.decay_tokens <= 0:
+            raise ValueError("decay_tokens must be positive")
+        if config.schedule not in {"cosine", "smoothstep", "linear"}:
+            raise ValueError("schedule must be cosine, smoothstep, or linear")
+        self.config = config
+
+    def value(self, tokens_seen: int) -> float:
+        progress = min(max(tokens_seen / self.config.decay_tokens, 0.0), 1.0)
+        if self.config.schedule == "cosine":
+            mix = 0.5 * (1.0 + math.cos(math.pi * progress))
+        elif self.config.schedule == "smoothstep":
+            smooth = progress * progress * (3.0 - 2.0 * progress)
+            mix = 1.0 - smooth
+        else:
+            mix = 1.0 - progress
+        return self.config.final_dropout + (
+            self.config.initial_dropout - self.config.final_dropout
+        ) * mix

src/dropout_decay/specs/__init__.py

@@ -0,0 +1,26 @@
+"""
+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 dropout_decay.specs.dropout_condition import DropoutCondition, anchor_dropout
+from dropout_decay.specs.model_spec import ModelSpec
+
+__all__ = [
+    "DropoutCondition",
+    "ModelSpec",
+    "anchor_dropout",
+]

src/dropout_decay/specs/dropout_condition.py

@@ -0,0 +1,89 @@
+"""
+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
+
+from dataclasses import dataclass
+import math
+from typing import Callable
+
+from dropout_decay.schedules import DropoutDecayConfig, DropoutDecayScheduler
+
+
+@dataclass(frozen=True)
+class DropoutCondition:
+    name: str
+    kind: str
+    initial: float
+    final: float
+    schedule: str = "constant"
+    decay_tokens: int | None = None
+    anchors: tuple[tuple[int, float], ...] = ()
+
+    def to_dict(self) -> dict:
+        return {
+            "name": self.name,
+            "kind": self.kind,
+            "initial": self.initial,
+            "final": self.final,
+            "schedule": self.schedule,
+            "decay_tokens": self.decay_tokens,
+            "anchors": [list(anchor) for anchor in self.anchors],
+        }
+
+    def make_fn(
+        self,
+        fallback_decay_tokens: int,
+        unique_tokens: int | None = None,
+    ) -> Callable[[int], float]:
+        if self.kind == "static":
+            return lambda _tokens_seen, p=self.initial: p
+        if self.kind == "anchor_decay":
+            if unique_tokens is None:
+                raise ValueError("anchor_decay conditions require unique_tokens")
+            p = anchor_dropout(unique_tokens, self.anchors)
+            return lambda _tokens_seen, p=p: p
+        scheduler = DropoutDecayScheduler(
+            DropoutDecayConfig(
+                initial_dropout=self.initial,
+                final_dropout=self.final,
+                decay_tokens=self.decay_tokens or fallback_decay_tokens,
+                schedule=self.schedule,
+            )
+        )
+        return scheduler.value
+
+
+def anchor_dropout(unique_tokens: int, anchors: tuple[tuple[int, float], ...]) -> float:
+    if not anchors:
+        raise ValueError("anchor dropout schedule requires at least one anchor")
+    ordered = sorted(anchors)
+    if unique_tokens <= ordered[0][0]:
+        return ordered[0][1]
+    if unique_tokens >= ordered[-1][0]:
+        return ordered[-1][1]
+    log_unique = math.log(unique_tokens)
+    for (left_tokens, left_dropout), (right_tokens, right_dropout) in zip(
+        ordered, ordered[1:]
+    ):
+        if left_tokens <= unique_tokens <= right_tokens:
+            left_log = math.log(left_tokens)
+            right_log = math.log(right_tokens)
+            mix = (log_unique - left_log) / (right_log - left_log)
+            return left_dropout + mix * (right_dropout - left_dropout)
+    return ordered[-1][1]