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

This commit is contained in:
ViperEkura 2026-07-05 12:31:40 +08:00
parent 459523f0e4
commit f2c2b159ba
3 changed files with 104 additions and 14 deletions

Binary file not shown.

After

Width:  |  Height:  |  Size: 223 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 100 KiB

118
main.tex
View File

@ -32,20 +32,24 @@ Training billion-parameter language models requires careful co-design of
data infrastructure, distributed execution, and numerical precision data infrastructure, distributed execution, and numerical precision
management. This paper presents {\sc AstrAI}, an open-source framework management. This paper presents {\sc AstrAI}, an open-source framework
for end-to-end training of a 1.2B-parameter autoregressive Transformer. for end-to-end training of a 1.2B-parameter autoregressive Transformer.
The system integrates a JSON-driven preprocessing pipeline (BBPE We describe the full pipeline: JSON-driven preprocessing with BBPE
tokenization, multi-strategy packing, HDF5 and memory-mapped storage), tokenization and multi-strategy packing, HDF5 and memory-mapped storage
a 24-layer decoder-only architecture with Grouped Query Attention and backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash-based
SwiGLU, and distributed training via DDP/FSDP with cosine scheduling. near-duplicate detection and LLM-as-Judge scoring. Using IFD (Instruction
A central focus is the numerical stability of BF16-precision training Fulfillment Difficulty) analysis on 3000 SFT samples, we find that Base
in deep Transformers. Through variance propagation analysis, we show IFD and Loss Ratio are nearly orthogonal ($r=0.10$), forming a
that GPT-2 residual scaling on output projections reduces per-block complementary two-dimensional screening space, while Instruct IFD is
residual variance by a factor of 48, containing post-24-layer variance redundant with Loss Ratio ($r=0.90$) due to a shared numerator---a
at 1.34 compared to 17.5 without scaling. Empirical evaluations over tautological artifact we identify and warn against. The model is a 24-layer
15B training tokens demonstrate that residual scaling consistently decoder-only Transformer with Grouped Query Attention, SwiGLU, RoPE, and
outperforms Kaiming initialization, with the gap widening to 0.79 in RMSNorm, trained with AdamW and cosine scheduling via DDP/FSDP.
the mid-training regime before narrowing to 0.38 at convergence. These A central focus is BF16 numerical stability: through variance propagation
results establish residual scaling as a practical necessity for BF16 analysis we show that GPT-2 residual scaling reduces per-block residual
Transformer training at scale. 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} \end{abstract}
% ====================================================================== % ======================================================================
@ -231,6 +235,64 @@ IFD\textsubscript{ckpt} vs.\ Loss Ratio (right).}
\label{fig:ifd_lossratio} \label{fig:ifd_lossratio}
\end{figure} \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} \section{Model Architecture}
% ====================================================================== % ======================================================================
@ -670,6 +732,34 @@ In our Alpaca-style dataset, this removes approximately
$5$--$8\%$ of samples and substantially reduces false positives $5$--$8\%$ of samples and substantially reduces false positives
in the high-IFD tail. 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} \begin{thebibliography}{99}