Add IFD loss ratio density analysis, weight distribution appendix, and rewrite abstract
This commit is contained in:
parent
459523f0e4
commit
f2c2b159ba
Binary file not shown.
|
After Width: | Height: | Size: 223 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 100 KiB |
118
main.tex
118
main.tex
|
|
@ -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}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue