Add IFD loss ratio density analysis, weight distribution appendix, and rewrite abstract

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}