AstrAI-paper/main.tex

\documentclass[11pt,a4paper]{article}

% ===== Packages =====
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{newtxtext,newtxmath}
\usepackage[margin=1in]{geometry}
\usepackage{amsmath}
\usepackage{booktabs}
\usepackage{graphicx}
\usepackage{hyperref}
\usepackage{float}
\usepackage{caption}
\usepackage{enumitem}
\usepackage{url}
\usepackage{microtype}

\DeclareMathOperator{\Var}{Var}

\title{End-to-End Training of a 1.2B Transformer with AstrAI \\
       \large Data Pipeline, Distributed Training, and BF16 Numerical Stability via Residual Scaling}

\author{AstrAI Contributors}
\date{}

\begin{document}

\maketitle

\begin{abstract}
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
over 15B tokens. 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. 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: variance propagation
analysis shows 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 confirm that
residual scaling consistently outperforms Kaiming initialization, with the
loss gap peaking at 0.79 in the mid-training regime.
\end{abstract}

% ======================================================================
\section{Introduction}
% ======================================================================

Training a billion-parameter language model end-to-end involves far more than
model architecture. Data must be preprocessed and stored efficiently, the
training loop must handle distributed parallelism, gradient accumulation,
checkpointing, and logging---and numerical pitfalls must be diagnosed and
fixed. This paper describes the complete workflow using {\sc AstrAI}~\cite{astrai}, an
open-source framework for Transformer training and inference, and highlights a
BF16 precision issue encountered along the way.

% ======================================================================
\section{Data Pipeline}
% ======================================================================

\subsection{Preprocessing}

Raw data arrives as JSONL files. The preprocessing pipeline is configured
via a JSON specification that defines:

\begin{itemize}[nosep]
  \item \textbf{Tokenization}: BBPE tokenizer (100K vocabulary) with standard
        special tokens.
  \item \textbf{Masking}: Declarative loss mask assignment per section
        (e.g.,~mask user input, compute loss on assistant response).
  \item \textbf{Packing}: Documents concatenated via \texttt{simple}
        (sequential), \texttt{bfd} (best-fit decreasing), or
        \texttt{bfd\_\allowbreak{}split} strategies.
  \item \textbf{Position IDs}: \texttt{none}, \texttt{doc\_reset} (per-document
        boundary), or \texttt{continuous}.
  \item \textbf{Output}: Tokenized sequences written to \texttt{.h5} or
        \texttt{.bin} shards, auto-split at 100M tokens per shard.
\end{itemize}

Samples shorter than 50~chars or longer than 2M~chars are filtered out.

\subsection{Storage Backends}

Two storage backends serve the DataLoader:

\begin{itemize}[nosep]
  \item \textbf{H5Store}: HDF5-based, memory-loaded with shared-memory support
        for multi-worker access.
  \item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files shared
        via OS page cache.
\end{itemize}

A resumable distributed sampler provides seed-based shuffle with
epoch/iteration resume.

\subsection{SFT Data Cleaning}

For supervised fine-tuning (SFT), raw data requires additional curation
beyond pretraining tokenization. {\sc Alembic}~\cite{alembic} is a companion
pipeline that handles SFT data generation, cleaning, and quality scoring:
three-generation strategies (topic-driven, seed-driven, self-instruct),
built-in cleaning (HTML/URL/markdown removal, char/word repetition filters),
and a MinHash-based near-duplicate detection system~\cite{broder1997syntactic}.
Given a set of $P$ hash functions ($P=128$) and a text $T$, the MinHash
pipeline proceeds as follows:

\begin{enumerate}[nosep,leftmargin=*]
  \item \textbf{Tokenization}: $T$ is split into character $n$-grams
    ($n=3$):
    \begin{equation}
      \Gamma(T) = \{\,c_i c_{i+1} c_{i+2} \mid i = 1,\dots,|T|-2 \,\}.
    \end{equation}
  \item \textbf{Signature}: For each hash function $h_k$, the minimum hash
    value over all $n$-grams forms the $k$-th element of the fingerprint:
    \begin{equation}
      s_k = \min_{t \in \Gamma(T)} h_k(t), \qquad
      h_k(t) = \operatorname{SHA256}(42 : k : t)_{[0:63]}.
    \end{equation}
    The full fingerprint is $\mathbf{s} = (s_1,\dots,s_P)$.
  \item \textbf{Similarity}: The Jaccard similarity between two sets is
    estimated by the fraction of agreeing fingerprint positions:
    \begin{equation}
      \widehat{J}(\mathbf{s}^{(a)},\mathbf{s}^{(b)}) =
      \frac{|\{\,k \mid s^{(a)}_k = s^{(b)}_k \,\}|}{P}.
    \end{equation}
  \item \textbf{Filtering}: Samples are processed sequentially; a sample
    is dropped if $\widehat{J}(\mathbf{s}, \mathbf{s}') \ge 0.7$ for any
    previously kept sample $\mathbf{s}'$.
\end{enumerate}

An optional LLM-as-Judge scoring module provides multi-dimensional
quality scores that can be used to filter low-quality samples.

An IFD (Instruction Fulfillment Difficulty) analysis is provided in
Appendix~\ref{app:ifd}.

% ======================================================================
\section{Model Architecture}
% ======================================================================

The model is a 24-layer decoder-only Transformer with Grouped Query Attention
(GQA)~\cite{ainslie2023gqa}, SwiGLU feed-forward blocks~\cite{shazeer2020glu},
and Rotary Position Embedding (RoPE)~\cite{su2024roformer}.
Table~\ref{tab:model_config} summarizes the configuration.

\begin{table}[H]
\centering
\caption{Model configuration. Total: $\sim$1.2B parameters.}
\label{tab:model_config}
\begin{tabular}{@{}lrlr@{}}
\toprule
\textbf{Parameter} & \textbf{Value} & \textbf{Parameter} & \textbf{Value} \\
\midrule
Vocabulary ($V$) & 100,000 & Hidden dim ($d$) & 1,536 \\
Layers ($L$) & 24 & FFN dim ($d_{\textit{ffn}}$) & 6,912 \\
Query heads & 24 & KV heads & 4 \\
Head dim & 64 & Max length & 2,048 \\
Norm & RMSNorm ($\epsilon=10^{-5}$) & RoPE $\theta$ & 10,000 \\
\bottomrule
\end{tabular}
\end{table}

With Grouped Query Attention~\cite{ainslie2023gqa} ($n_q = 24$ query
heads, $n_{kv} = 4$ key/value heads, group size $g = n_q / n_{kv} = 6$):
\begin{equation}
\begin{aligned}
\operatorname{GQA}(\mathbf{X}) &= \operatorname{Concat}\bigl(\operatorname{head}_1,\dots,\operatorname{head}_{n_q}\bigr)\mathbf{W}_O,\\[2mm]
\operatorname{head}_i &= \operatorname{Attn}\Bigl(
\mathbf{X}\mathbf{W}_Q^{(i)},\,
\mathbf{X}\mathbf{W}_K^{(\lfloor i / g \rfloor)},\,
\mathbf{X}\mathbf{W}_V^{(\lfloor i / g \rfloor)}
\Bigr),
\end{aligned}
\end{equation}
where $\operatorname{Attn}(\mathbf{Q},\mathbf{K},\mathbf{V}) =
\operatorname{Softmax}(\mathbf{Q}\mathbf{K}^{\mkern-1mu\mathsf{T}} / \sqrt{d_h})\mathbf{V}$.
Rotary Position Embedding (RoPE)~\cite{su2024roformer} encodes position
$m$ by rotating pairs of hidden dimensions:
\begin{equation}
\operatorname{RoPE}(\mathbf{x}_m)_i =
\begin{cases}
x_{m,i}\cos(m\theta_{j}) - x_{m,i+1}\sin(m\theta_{j}), & i = 2j,\\[2mm]
x_{m,i-1}\sin(m\theta_{j}) + x_{m,i}\cos(m\theta_{j}), & i = 2j+1,
\end{cases}
\end{equation}
with frequency $\theta_j = 10000^{-2j/d}$ for $j = 0,\dots,d/2-1$.
The SwiGLU~\cite{shazeer2020glu} feed-forward applies a gated Swish
non-linearity:
\begin{equation}
\operatorname{MLP}(\mathbf{x}) = \mathbf{W}_{\text{down}}\Bigl(
\mathbf{W}_{\text{up}}\mathbf{x} \odot
\operatorname{SiLU}\bigl(\mathbf{W}_{\text{gate}}\mathbf{x}\bigr)
\Bigr),
\label{eq:swiglu}
\end{equation}
where $\operatorname{SiLU}(z) = z / (1 + e^{-z})$.

Each decoder block $\ell$ then applies pre-norm residual connections:
\begin{equation}
\begin{aligned}
\mathbf{h}_\ell &= \mathbf{x}_\ell + \operatorname{GQA}\bigl(\operatorname{RMSNorm}(\mathbf{x}_\ell)\bigr),\\[2mm]
\mathbf{x}_{\ell+1} &= \mathbf{h}_\ell + \operatorname{MLP}\bigl(\operatorname{RMSNorm}(\mathbf{h}_\ell)\bigr).
\end{aligned}
\label{eq:decoder_block}
\end{equation}

\subsection{Initialization}

Linear weights follow $\mathcal{N}(0, 0.02)$; embeddings follow
$\mathcal{N}(0, 0.02)$. The output projection $\mathbf{W}_o$ and FFN
down-projection $\mathbf{W}_{\text{down}}$ use residual-scaled
initialization~\cite{radford2019gpt2}:
\begin{equation}
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
\end{equation}
This scaling is critical for BF16 stability (Section~\ref{sec:num-stability}).

% ======================================================================
\section{Training Configuration}
% ======================================================================

The model is trained on next-token cross-entropy loss:
\begin{equation}
\mathcal{L} = -\sum_{t=1}^{T} \log P(x_t \mid x_{<t}; \theta).
\end{equation}

Training uses AdamW~\cite{loshchilov2019adamw} with cosine learning rate
scheduling (5\% warmup) and global L2 gradient clipping. The framework supports DDP and FSDP for multi-GPU distribution,
with gradient accumulation to manage memory.
Table~\ref{tab:train_params} lists the key hyperparameters.

\begin{table}[H]
\centering
\caption{Training hyperparameters for the 1.2B run.}
\label{tab:train_params}
\begin{tabular}{@{}lr@{}}
\toprule
\textbf{Hyperparameter} & \textbf{Value} \\
\midrule
Precision & BF16 (weights + AdamW states) \\
Optimizer & AdamW, $\eta=1.5\times10^{-4}$ \\
Betas & $(0.9, 0.95)$, weight decay $0.1$ \\
Gradient clip & Global L2, max norm $1.0$ \\
Scheduler & Cosine, warmup ratio $0.02$ \\
Batch size & 4 per device $\times$ 4 GPUs $\times$ 32 accumulation \\
Sequence length & 2,048 tokens \\
Total steps & 950,000 \\
\bottomrule
\end{tabular}
\end{table}

% ======================================================================
\section{Numerical Stability via Residual Scaling}
\label{sec:num-stability}
% ======================================================================

Deep Transformers trained in BF16 face numerical stability challenges from
residual variance accumulation across layers. We evaluate the GPT-2
residual-scaling initialization~\cite{radford2019gpt2} as a mitigation
strategy.

\subsection{Variance Analysis}

At initialization with $\mathcal{N}(0, 0.02)$, a linear projection output has:
\begin{equation}
\Var(\mathbf{W}\mathbf{x}) = d_{\text{in}} \cdot (0.02)^2 \cdot \Var(\mathbf{x})
= 0.6144 \cdot \Var(\mathbf{x}) \quad (\text{for } d=1536).
\end{equation}

Within one block, attention and FFN each add a residual term. The variances
at each sub-stage are:

\begin{center}
\begin{tabular}{@{}lcc@{}}
\toprule
\textbf{Component} & \textbf{Operation} & $\Var$ (scaled by $\Var(\mathbf{x})$) \\
\midrule
Q/K/V proj & Linear(1536, $n_{\text{heads}}\cdot64$) & 0.6144 \\
Attention out & SDPA + $\mathbf{W}_o$ (scaled) & $0.378 / L$ \\
Gate/Up proj & Linear(1536, 6912) & 0.6144 \\
SiLU gate & $\operatorname{SiLU}(z) \approx 0.5z$ & $0.6144 \times 0.298 = 0.1831$ \\
Gated product & element-wise $\odot$ & $\approx 0.6144 \times 0.1831 = 0.1125$ \\
Down proj & Linear(6912, 1536) (scaled) & $0.311 / L$ \\
\midrule
Per-block residual & $\mathbf{R}_\ell = R_{\text{attn}} + R_{\text{ffn}}$ & $0.689 / L$ (scaled) \\
\bottomrule
\end{tabular}
\end{center}

Without the $1/\sqrt{2L}$ factor on $\mathbf{W}_o$ and
$\mathbf{W}_{\text{down}}$, the per-block residual variance becomes $0.689$
instead of $0.689/L \approx 0.014$. After $L=24$ blocks:
\begin{equation}
\begin{aligned}
\text{Without scaling: } \Var(\mathbf{x}_{24}) &\approx 1 + 24 \times 0.689 = 17.5,\\
\text{With scaling: } \Var(\mathbf{x}_{24}) &\approx 1 + 24 \times 0.014 = 1.34.
\end{aligned}
\end{equation}

\subsection{GPT-2 Residual Scaling}

The GPT-2 initialization~\cite{radford2019gpt2} scales output projections by
$1/\sqrt{2L}$:
\begin{equation}
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
\end{equation}

This reduces per-block residual variance contribution from $0.689$ to
$0.689/L \approx 0.014$, a factor of $2L = 48$. The post-24-block variance
drops from $17.5$ to $1.34$, a $13.1\times$ improvement. In BF16
($7$-bit mantissa, ULP $= 0.0078$ at $w = 1.0$)~\cite{ieee754},
this keeps weight magnitudes within stable precision bounds. We further
recommend storing AdamW moments in FP32 and logging per-layer gradient
histograms during early training.

\subsection{Empirical Training Results}

\begin{figure}[H]
\centering
\includegraphics[width=0.50\linewidth]{data/loss_compare.png}
\caption{Training loss curves: GPT-2 residual scaling vs.~Kaiming
initialization over 15B tokens.}
\label{fig:loss}
\end{figure}

Figure~\ref{fig:loss} shows both loss curves; GPT-2 residual scaling (lower
curve) maintains a clear advantage, particularly in the 0.3--0.8B token region.

\begin{figure}[H]
\centering
\includegraphics[width=0.95\linewidth]{data/ckpt_comparison.png}
\caption{Optimizer and initialization comparison.}
\label{fig:ckpt_comparison}
\end{figure}

Figure~\ref{fig:ckpt_comparison} compares four configurations. The left panel shows training loss for Muon (Embedding Adam + 1D Adam), Muon (Embedding Muon + 1D Adam), Kaiming init, and Normal init; the center panel zooms in on the two current Muon variants; and the right panel shows gradient norms over optimizer steps. The older Kaiming and Normal initializations converge more slowly and plateau at higher loss. Between the current variants, using Adam for the embedding layer yields lower loss and more stable gradients than using Muon embeddings.

\begin{table}[H]
\centering
\caption{Loss at 0.125B-interval milestones, 0--1B tokens.}
\label{tab:loss_milestones}
\begin{tabular}{@{}lccc@{}}
\toprule
\textbf{Tokens (B)} &
  \textbf{GPT-2 scaling} &
  \textbf{Kaiming init} &
  \textbf{$\Delta$} \\
\midrule
0.125 & 7.37 & 7.66 & 0.29 \\
0.250 & 5.80 & 6.14 & 0.34 \\
0.375 & 4.82 & 5.38 & 0.56 \\
0.500 & 4.06 & 4.80 & 0.74 \\
0.625 & 3.50 & 4.29 & 0.79 \\
0.750 & 3.24 & 3.80 & 0.56 \\
0.875 & 3.21 & 3.43 & 0.22 \\
1.000 & 2.80 & 3.18 & 0.38 \\
\bottomrule
\end{tabular}
\end{table}

Table~\ref{tab:loss_milestones} quantifies the per-milestone gap. GPT-2
residual scaling leads at every interval, with $\Delta$ growing from 0.29
at 0.125B to a peak of 0.79 at 0.625B, then narrowing to 0.38 at 1B.
The widening mid-range gap aligns with the variance accumulation region
identified in the theoretical analysis (Section~\ref{sec:num-stability}).

% ======================================================================
\section{Conclusion}
% ======================================================================

We have described the end-to-end pipeline for training a 1.2B Transformer with
{\sc AstrAI}: data preprocessing with JSON-driven tokenization and packing,
a 24-layer GQA-SwiGLU architecture, callback-based training with DDP/FSDP
executors, and cosine scheduling. We further analyzed numerical stability
under BF16, showing that GPT-2 residual scaling ($\sigma_o = 0.02/\sqrt{2L}$)
reduces per-block residual variance by a factor of 48, keeping post-24-layer
variance at $1.34$ versus $17.5$ without scaling. The complete framework and model
weights are available at \url{https://github.com/ViperEkura/AstrAI}.

% ======================================================================
\appendix
% ======================================================================

% ======================================================================
\section{IFD Data Analysis}
\label{app:ifd}
% ======================================================================

Instruction Fulfillment Difficulty (IFD)~\cite{li2023ifd} compares
conditional and unconditional per-token losses:
\begin{equation}
\begin{aligned}
\mathrm{IFD} &= \frac{L_{\text{cond}}}{L_{\text{uncond}}},\\[2mm]
L_{\text{cond}} &= -\frac{1}{T}\sum_{t=1}^T \log P(y_t \mid \mathbf{x}, y_{<t}),\\[2mm]
L_{\text{uncond}} &= -\frac{1}{T}\sum_{t=1}^T \log P(y_t \mid y_{<t}).
\end{aligned}
\end{equation}
We compute IFD for $N=3000$ SFT samples (Alpaca-GPT4~\cite{alpaca})
using the base model (15B tokens) and the 1K-step SFT checkpoint.
After 1K SFT steps, both losses increase slightly; the mean IFD
changes from $0.8263$ (base) to $0.8485$ (1K SFT).

\subsection{Quantitative Summary}

Over $N=3000$ SFT samples from Alpaca-GPT4:
\begin{itemize}[nosep]
  \item \textbf{Base model}: mean IFD $= 0.8263$,
        median $= 0.8858$, std $= 0.1699$; $1.9\%$ of samples
        have IFD $> 1.0$.
  \item \textbf{1K SFT}: mean IFD $= 0.8485$,
        median $= 0.9083$, std $= 0.1588$; $3.1\%$ of samples
        exceed $1.0$.
  \item \textbf{Stability}: Pearson $r > 0.97$ between base and
        1K SFT IFD. The slight upward shift ($0.8263 \to 0.8485$)
        reflects both losses increasing after SFT, consistent with
        distribution shift during fine-tuning rather than uniform
        instruction-following improvement.
\end{itemize}

\subsection{Representative Samples}

Table~\ref{tab:ifd_examples} lists samples spanning the IFD range.

\begin{table}[H]
\centering
\caption{Representative IFD samples.}
\label{tab:ifd_examples}
\small
\begin{tabular}{@{}c c c c c p{4.2cm}@{}}
\toprule
\textbf{Idx} &
  \textbf{$L_{\text{cond}}^{\text{base}}$} &
  \textbf{$L_{\text{uncond}}^{\text{base}}$} &
  \textbf{$L_{\text{cond}}^{\text{1K}}$} &
  \textbf{$L_{\text{uncond}}^{\text{1K}}$} &
  \textbf{Instruction} \\
\midrule
81   & 13.38 & 5.84  & 13.25 & 5.69  & Classify incident as breach of protocol \\
906  & 13.12 & 9.75  & 13.06 & 9.75  & Convert numbers from words to digits \\
1076 & 2.53  & 2.46  & 2.53  & 2.53  & Pick best synonym \\
7    & 2.62  & 2.70  & 2.68  & 2.77  & Write a short story in third person \\
2427 & 2.59  & 2.84  & 2.69  & 2.90  & Find five most similar sentences \\
798  & 2.02  & 2.75  & 2.11  & 2.31  & List four social media platforms \\
223  & 1.34  & 3.16  & 1.36  & 3.27  & Classify text as Fiction or Non-fiction \\
\bottomrule
\end{tabular}
\end{table}

Samples with the highest conditional loss (rows~81,~906) are
short-answer classification tasks ($L_{\text{cond}} \approx 13$).
Lowest-IFD samples (row~223) are tasks where the instruction constrains
the output space so tightly that unconditional loss far exceeds
conditional loss. The four loss values remain nearly unchanged after
SFT across all samples.

\subsection{IFD Bias from Response Length}
\label{sec:ifd_bias}

Both losses are per-token averages. The variance of
$L_{\text{uncond}} = \frac{1}{T} \sum_{t=1}^T \log P(x_t)$
scales as $1/T$, so shorter responses produce noisier estimates.
Figure~\ref{fig:length_bias} plots the three metrics against response
length for the base model; samples with $<20$ tokens ($21.9\%$ of
the dataset) exhibit substantially higher scatter.

\begin{figure}[H]
\centering
\includegraphics[width=0.95\linewidth]{data/ifd_length_grid.png}
\caption{Response length vs.\ $L_{\text{cond}}$, $L_{\text{uncond}}$,
and IFD (base model, log scale on $x$-axis).}
\label{fig:length_bias}
\end{figure}

Table~\ref{tab:corr_bias} reports the correlations. Response length
is the dominant confound: $L_{\text{uncond}}$ shows a strong negative
monotonic trend ($\rho = -0.79$), while $L_{\text{cond}}$ is less
affected ($\rho = -0.48$). The net effect on IFD is a positive
correlation ($\rho = +0.72$).

\begin{table}[H]
\centering
\caption{Pearson $r$ and Spearman $\rho$ between sample dimensions and IFD components (base model).}
\label{tab:corr_bias}
\small
\begin{tabular}{@{}lcccccc@{}}
\toprule
 & \multicolumn{2}{c}{vs.\ $L_{\text{cond}}$}
 & \multicolumn{2}{c}{vs.\ $L_{\text{uncond}}$}
 & \multicolumn{2}{c}{vs.\ IFD} \\
\cmidrule(lr){2-3} \cmidrule(lr){4-5} \cmidrule(lr){6-7}
\textbf{Dimension} & $r$ & $\rho$ & $r$ & $\rho$ & $r$ & $\rho$ \\
\midrule
Instruction length & $+0.07$ & $+0.06$ & $+0.15$ & $+0.24$ & $-0.25$ & $-0.34$ \\
Response length    & $-0.36$ & $-0.48$ & $-0.56$ & $-0.79$ & $+0.58$ & $+0.72$ \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Loss Ratio}
\label{sec:loss_ratio}

We further define the \textbf{Loss Ratio} as the fraction of conditional
loss retained after SFT:
\begin{equation}
\text{Loss Ratio} = \frac{L_{\text{cond}}^{\text{1K}}}{L_{\text{cond}}^{\text{base}}}.
\end{equation}

Over $N=3000$ samples:
\begin{itemize}[nosep]
  \item Mean $= 1.106$, median $= 1.084$, std $= 0.110$;
        $90.8\%$ of samples exceed $1.0$.
  \item Range: $[0.768, 1.962]$; only $9.2\%$ of samples show
        a decrease ($<1.0$) in conditional loss after SFT.
\end{itemize}

The predominance of loss ratio $>1$ confirms that the 1K-step SFT
checkpoint has not converged to a lower-loss region for the
evaluation samples. Instead, the fine-tuning distribution shift
increases NLL on most held-out instructions.

Table~\ref{tab:ifd_lr_corr} reports the pairwise correlations.
Despite the theoretical expectation that IFD\textsubscript{ckpt} and
Loss Ratio share $L_{\text{cond}}^{\text{1K}}$ in their numerators
and should therefore correlate strongly~\cite{li2023ifd}, the
observed correlation is near zero ($r = -0.02$, $\rho = 0.04$).
This is because the {\em relative} ordering of
$L_{\text{cond}}^{\text{base}}$ and
$L_{\text{uncond}}^{\text{ckpt}}$ (which determine the slope
$k_i = L_{\text{cond},i}^{\text{base}} / L_{\text{uncond},i}^{\text{ckpt}}$
in the relationship $\text{IFD}_{\text{ckpt},i} = k_i \cdot
\text{Loss Ratio}_i$) varies widely across samples, breaking the
proportionality at the sample level.

\begin{table}[H]
\centering
\caption{Pairwise correlations between IFD variants and Loss Ratio.}
\label{tab:ifd_lr_corr}
\small
\begin{tabular}{@{}lcc@{}}
\toprule
\textbf{Pair} & Pearson $r$ & Spearman $\rho$ \\
\midrule
IFD\textsubscript{base} vs.\ IFD\textsubscript{ckpt}  & $+0.97$ & $+0.96$ \\
IFD\textsubscript{base} vs.\ Loss Ratio       & $-0.15$ & $-0.06$ \\
IFD\textsubscript{ckpt} vs.\ Loss Ratio       & $-0.02$ & $+0.04$ \\
\bottomrule
\end{tabular}
\end{table}

The near-perfect correlation between IFD\textsubscript{base} and
IFD\textsubscript{ckpt} ($r = 0.97$) reveals that the IFD ranking is
highly robust to the choice of evaluation model: samples that the
base model finds difficult remain difficult after 1K SFT steps.
This stability justifies using the base-model IFD as a data
selection signal without re-evaluating after fine-tuning.

% ======================================================================
\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}

\bibitem{alpaca}
R.~Taori, I.~Gulrajani, T.~Zhang, Y.~Dubois, X.~Li, C.~Guestrin,
P.~Liang, T.~B.~Hashimoto.
Alpaca: A strong, replicable instruction-following model.
\textit{Stanford Center for Research on Foundation Models (CRFM)}, 2023.

\bibitem{ainslie2023gqa}
J.~Ainslie, J.~Lee-Thorp, M.~de Jong, Y.~Zemlyanskiy, F.~Lebr\'on, S.~Sanghai.
GQA: Training generalized multi-query transformer models from multi-head
checkpoints. \textit{EMNLP}, 2023.

\bibitem{alembic}
Alembic Contributors. \textit{Alembic: A lightweight LLM-driven SFT data
generation, cleaning, and scoring pipeline.}
\url{https://github.com/ViperEkura/Alembic}, 2026.

\bibitem{astrai}
AstrAI Contributors. \textit{AstrAI: An open-source training and inference
framework for Transformer language models.}
\url{https://github.com/ViperEkura/AstrAI}, 2026.

\bibitem{broder1997syntactic}
A.~Z.~Broder. On the resemblance and containment of documents.
\textit{SEQUENCES '97}, 1997.

\bibitem{li2023ifd}
M.~Li, Y.~Zhang, Z.~Li, J.~Chen, L.~Chen, N.~Cheng, J.~Wang, T.~Zhou, J.~Xiao.
From quantity to quality: Boosting LLM performance with self-guided data
selection for instruction tuning.
\textit{NAACL}, 2024.

\bibitem{ieee754}
IEEE Computer Society. \textit{IEEE Standard for Floating-Point Arithmetic},
IEEE Std 754-2019, 2019.

\bibitem{loshchilov2019adamw}
I.~Loshchilov, F.~Hutter. Decoupled weight decay regularization.
\textit{ICLR}, 2019.

\bibitem{radford2019gpt2}
A.~Radford, J.~Wu, R.~Child, D.~Luan, D.~Amodei, I.~Sutskever.
Language models are unsupervised multitask learners.
\textit{OpenAI Blog}, 2019.

\bibitem{shazeer2020glu}
N.~Shazeer. GLU variants improve Transformer.
\textit{arXiv:2002.05202}, 2020.

\bibitem{su2024roformer}
J.~Su, A.~Murtadha, Y.~Lu, S.~Pan, B.~Wen, Y.~Liu.
Roformer: Enhanced transformer with rotary position embedding.
\textit{Neurocomputing}, 568:127063, 2024.

\end{thebibliography}

\end{document}