diff --git a/data/ifd_loss_ratio_density.png b/data/ifd_loss_ratio_density.png new file mode 100644 index 0000000..b785688 Binary files /dev/null and b/data/ifd_loss_ratio_density.png differ diff --git a/data/weight_dist_by_component.png b/data/weight_dist_by_component.png new file mode 100644 index 0000000..fc9c117 Binary files /dev/null and b/data/weight_dist_by_component.png differ diff --git a/main.tex b/main.tex index 6f30f9d..ab8da88 100644 --- a/main.tex +++ b/main.tex @@ -32,20 +32,24 @@ Training billion-parameter language models requires careful co-design of data infrastructure, distributed execution, and numerical precision management. This paper presents {\sc AstrAI}, an open-source framework for end-to-end training of a 1.2B-parameter autoregressive Transformer. -The system integrates a JSON-driven preprocessing pipeline (BBPE -tokenization, multi-strategy packing, HDF5 and memory-mapped storage), -a 24-layer decoder-only architecture with Grouped Query Attention and -SwiGLU, and distributed training via DDP/FSDP with cosine scheduling. -A central focus is the numerical stability of BF16-precision training -in deep Transformers. Through variance propagation analysis, we show -that GPT-2 residual scaling on output projections reduces per-block -residual variance by a factor of 48, containing post-24-layer variance -at 1.34 compared to 17.5 without scaling. Empirical evaluations over -15B training tokens demonstrate that residual scaling consistently -outperforms Kaiming initialization, with the gap widening to 0.79 in -the mid-training regime before narrowing to 0.38 at convergence. These -results establish residual scaling as a practical necessity for BF16 -Transformer training at scale. +We describe the full pipeline: JSON-driven preprocessing with BBPE +tokenization and multi-strategy packing, HDF5 and memory-mapped storage +backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash-based +near-duplicate detection and LLM-as-Judge scoring. Using IFD (Instruction +Fulfillment Difficulty) analysis on 3000 SFT samples, we find that Base +IFD and Loss Ratio are nearly orthogonal ($r=0.10$), forming a +complementary two-dimensional screening space, while Instruct IFD is +redundant with Loss Ratio ($r=0.90$) due to a shared numerator---a +tautological artifact we identify and warn against. The model is a 24-layer +decoder-only Transformer with Grouped Query Attention, SwiGLU, RoPE, and +RMSNorm, trained with AdamW and cosine scheduling via DDP/FSDP. +A central focus is BF16 numerical stability: through variance propagation +analysis we show that GPT-2 residual scaling reduces per-block residual +variance by a factor of 48, containing post-24-layer variance at 1.34 +compared to 17.5 without scaling. Empirical evaluations over 15B training +tokens demonstrate that residual scaling consistently outperforms Kaiming +initialization, with the gap peaking at 0.79 in the mid-training regime. +The complete framework and model weights are open-source. \end{abstract} % ====================================================================== @@ -231,6 +235,64 @@ IFD\textsubscript{ckpt} vs.\ Loss Ratio (right).} \label{fig:ifd_lossratio} \end{figure} +\subsubsection{Loss Ratio Density by IFD Group} +\label{sec:ifd_loss_ratio_density} + +Figure~\ref{fig:ifd_loss_ratio_density} compares the Loss Ratio +density grouped by base IFD (left) and instruct IFD (right). + +\begin{figure}[H] +\centering +\includegraphics[width=0.90\linewidth]{data/ifd_loss_ratio_density.png} +\caption{Loss Ratio density grouped by base IFD (left) and instruct IFD +(right).} +\label{fig:ifd_loss_ratio_density} +\end{figure} + +\textbf{Left panel (Base IFD grouping).} +The four density curves overlap almost completely, all peaking at +Loss Ratio $0.75$--$0.85$. Whether a sample has base IFD $< 0.85$, +$0.85$--$0.95$, $0.95$--$1.05$, or $> 1.05$, its Loss Ratio +distribution is nearly identical. Base IFD cannot distinguish +which samples learn during SFT and which do not. This +near-orthogonality ($r = 0.10$, Table~\ref{tab:ifd_lossratio_corr}) +implies that how \emph{hard} an instruction appears to the base +model carries almost no information about how much the model will +improve on it. The signal is either dominated by data quality +variation, or the current training budget is insufficient for +high-IFD samples to realize their potential. + +\textbf{Right panel (Instruct IFD grouping).} +The four curves separate into near-perfectly stratified layers: + +\medskip +\begin{minipage}{\linewidth} +\begin{tabular}{@{}lcc@{}} +\toprule +\textbf{Instruct IFD} & \textbf{\#Samples} & \textbf{Loss Ratio peak} \\ +\midrule +$< 0.50$ & 356 & $\sim 0.25$ (75\% drop) \\ +$0.50$--$0.70$ & 702 & $\sim 0.55$ (45\% drop) \\ +$0.70$--$0.85$ & 1056 & $\sim 0.78$ (22\% drop) \\ +$> 0.85$ & 886 & $\sim 0.95$ (5\% drop) \\ +\bottomrule +\end{tabular} +\end{minipage} +\medskip + +This separation, however, is a mathematical artifact. Instruct IFD +and Loss Ratio share the numerator $L_{\text{cond}}^{\text{ckpt}}$, +producing a tautological correlation ($r = 0.90$, $p \ll 0.001$). +Grouping by instruct IFD is equivalent to grouping by Loss Ratio +itself---explaining the outcome with the outcome, not predicting it +from input features. + +The contrast between the two panels is the central finding: +base IFD and Loss Ratio carry independent information +($r = 0.10$), forming a two-dimensional screening space. +Instruct IFD, despite its apparent predictive power, is redundant +with Loss Ratio and should not be used for data selection. + % ====================================================================== \section{Model Architecture} % ====================================================================== @@ -670,6 +732,34 @@ In our Alpaca-style dataset, this removes approximately $5$--$8\%$ of samples and substantially reduces false positives in the high-IFD tail. +% ====================================================================== +\section{Weight Distribution by Component} +\label{app:weight_dist} + +Figure~\ref{fig:weight_dist} shows the distribution of weight +magnitudes at initialization, grouped by component type. Embeddings +and non-residual-scaled projections (QKV, attention output, FFN +gate/up) follow $\mathcal{N}(0, 0.02)$, producing near-identical +bell curves centered at zero. The residual-scaled projections +(output projection $\mathbf{W}_o$ and FFN down-projection +$\mathbf{W}_{\text{down}}$) use $\sigma = 0.02 / \sqrt{2L} \approx 0.0029$, +visible as the narrow, sharply peaked distribution concentrated +near zero. This factor-48 variance reduction is the mechanism by +which GPT-2 residual scaling prevents BF16 underflow in deep +Transformers (Section~\ref{sec:num-stability}). + +\begin{figure}[H] +\centering +\includegraphics[width=0.85\linewidth]{data/weight_dist_by_component.png} +\caption{Weight distribution by component at initialization. +Each panel shows the histogram of weight values for a specific +module group (embedding, attention projections, FFN projections, +output projections). The narrow peaks correspond to the +residual-scaled $\mathbf{W}_o$ and $\mathbf{W}_{\text{down}}$ +projections.} +\label{fig:weight_dist} +\end{figure} + % ====================================================================== \begin{thebibliography}{99}