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.
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.
\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.

\subsection{IFD-Based Instruction Difficulty Analysis}

Instruction Fulfillment Difficulty (IFD)~\cite{li2023ifd} quantifies
how challenging an instruction is for a model by comparing conditional
and unconditional per-token losses over a response
$\mathbf{y} = (y_1,\dots,y_T)$:
\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}
An IFD $>1$ indicates the instruction increases the loss relative to
unconditional generation (the model struggles to follow it), while
IFD $<1$ means the instruction provides useful guidance.

We compute IFD for $N=3000$ SFT samples drawn from the
Alpaca-GPT4 dataset~\cite{alpaca} using both the pretrained
base model (after 15B tokens of pretraining) and a supervised
fine-tuned checkpoint (after 1K SFT steps).
Figure~\ref{fig:ifd} shows the distribution.

\begin{figure}[H]
\centering
\includegraphics[width=0.80\linewidth]{data/ifd_compare_clean.png}
\caption{IFD scatter: base model vs.\ trained checkpoint. The
diagonal line marks $\mathrm{IFD}_{\text{base}} = \mathrm{IFD}_{\text{ckpt}}$.}
\label{fig:ifd}
\end{figure}

\begin{figure}[H]
\centering
\includegraphics[width=0.80\linewidth]{data/ifd_density_dist.png}
\caption{IFD density distribution: base model and SFT checkpoint.}
\label{fig:ifd_density}
\end{figure}

Figure~\ref{fig:ifd_density} shows the corresponding density
estimates, confirming the systematic leftward shift after SFT.
The pretrained base model (15B tokens) has mean IFD $0.9625$;
$29.8\%$ of samples exceed $1.0$. After 1K SFT steps, mean IFD drops
to $0.7539$, with only $0.4\%$ of samples above $1.0$. The average
per-sample IFD reduction is $0.2086$. Conditional loss drops
$5.3\times$ more than unconditional loss, confirming that SFT teaches
instruction following rather than merely improving generic language
modeling. Detailed analysis is provided in Appendix~\ref{app:ifd}.

\subsubsection{IFD vs.\ Loss Ratio}

We further define the \emph{loss ratio}---the fraction of
conditional loss retained after SFT---as:
\begin{equation}
\text{Loss Ratio} = \frac{L_{\text{cond}}^{\text{ckpt}}}{L_{\text{cond}}^{\text{base}}}.
\end{equation}
Table~\ref{tab:ifd_lossratio_corr} reports the pairwise correlations.

\begin{table}[H]
\centering
\caption{Pairwise correlations among IFD and Loss Ratio.}
\label{tab:ifd_lossratio_corr}
\small
\begin{tabular}{@{}lcc@{}}
\toprule
\textbf{Pair} & \textbf{Pearson $r$} & \textbf{Spearman $\rho$} \\
\midrule
IFD\textsubscript{base} vs.\ Loss Ratio & $+0.10$ & $+0.05$ \\
IFD\textsubscript{ckpt} vs.\ Loss Ratio & $+0.90$ & $+0.91$ \\
IFD\textsubscript{base} vs.\ IFD\textsubscript{ckpt} & $+0.38$ & $+0.49$ \\
\bottomrule
\end{tabular}
\end{table}

The near-perfect correlation between IFD\textsubscript{ckpt} and
Loss Ratio ($r = 0.90$) reflects a mathematical near-identity:
both are dominated by $L_{\text{cond}}^{\text{ckpt}}$ in the
numerator. Consequently, IFD\textsubscript{ckpt} is
redundant---it essentially measures how much the conditional loss
has dropped after SFT, i.e., the learning speed of each sample. In contrast, IFD\textsubscript{base} and Loss
Ratio are nearly orthogonal ($r = 0.10$), forming a complementary
two-dimensional screening space: IFD\textsubscript{base} measures
``how hard does the base model find this,'' while Loss Ratio
measures ``how much did SFT improve it.'' Samples with high
IFD\textsubscript{base} \emph{and} low Loss Ratio are the most
informative for training.

\begin{figure}[H]
\centering
\includegraphics[width=0.80\linewidth]{data/ifd_both_vs_lossratio.png}
\caption{IFD\textsubscript{base} vs.\ Loss Ratio (left),
IFD\textsubscript{ckpt} vs.\ Loss Ratio (right).}
\label{fig:ifd_lossratio}
\end{figure}

% ======================================================================
\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{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 Examples}
\label{app:ifd}
% ======================================================================

Table~\ref{tab:ifd_examples} lists representative samples from the
IFD evaluation set, covering high, medium, and low IFD values
for the base model.

\begin{table}[H]
\centering
\caption{Representative IFD samples covering four patterns.}
\label{tab:ifd_examples}
\small
\begin{tabular}{@{}c c c c c c c p{4.5cm}@{}}
\toprule
\textbf{Idx} &
  \textbf{Base IFD} &
  \textbf{Ckpt IFD} &
  \textbf{$L_{\text{cond}}^{\text{base}}$} &
  \textbf{$L_{\text{uncond}}^{\text{base}}$} &
  \textbf{$L_{\text{cond}}^{\text{ckpt}}$} &
  \textbf{$L_{\text{uncond}}^{\text{ckpt}}$} &
  \textbf{Instruction} \\
\midrule
0 & 4.605 & 1.525 & 12.38 & 2.69 & 3.77 & 2.47 & Complete analogy: loud is to quiet as day is to \\
1 & 3.741 & 0.702 & 11.75 & 3.14 & 2.17 & 3.09 & Label news article as ``Political'' or ``Entertainment'' \\
2 & 1.044 & 0.089 & 3.50 & 3.35 & 0.28 & 3.10 & Find the capital of Spain \\
3 & 1.056 & 0.147 & 4.09 & 3.88 & 0.60 & 4.07 & Edit sentence for correct grammar: ``I were just going to'' \\
4 & 0.977 & 0.904 & 2.57 & 2.63 & 2.24 & 2.48 & Describe the role of a project manager \\
5 & 0.370 & 0.249 & 1.37 & 3.70 & 0.85 & 3.42 & Convert the given paragraph to a list \\
6 & 0.307 & 0.062 & 0.70 & 2.29 & 0.15 & 2.43 & Remove third-person words from sentence \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Quantitative Summary}

Over $N=3000$ SFT samples:
\begin{itemize}[nosep]
  \item \textbf{Pretrained base model (15B tokens)}: mean IFD $= 0.9625$,
        median $= 0.9773$, std $= 0.1925$; $29.8\%$ of samples have
        IFD $> 1.0$.
  \item \textbf{SFT checkpoint (1K steps)}: mean IFD $= 0.7539$,
        median $= 0.8547$, std $= 0.2352$; only $0.4\%$ of samples
        exceed $1.0$.
  \item \textbf{Average IFD reduction}: $0.2086$ per sample.
  \item \textbf{Loss decomposition}: conditional loss drops by $0.9657$
        ($3.2424 \rightarrow 2.2767$), while unconditional loss drops by
        only $0.1838$ ($3.4142 \rightarrow 3.2303$). The $5.3\times$
        larger conditional reduction confirms the model primarily learns
        instruction following.
  \item \textbf{Correlation}: Pearson $r = 0.38$ between base and
        checkpoint IFD, indicating a moderate tendency for relatively
        hard instructions to remain relatively hard after training.
\end{itemize}

\subsection{Observed Patterns}

\paragraph{High-IFD samples (base IFD $> 3$, e.g.,~rows~0,~1).}
These are tasks requiring task-intent comprehension: analogy completion
and article labeling. In the base model (15B pretraining), conditional
loss is extremely high ($L_{\text{cond}} \approx 12$), meaning the
instruction still acts as noise. After 1K SFT steps, IFD drops
sharply (e.g., $4.605 \rightarrow 1.525$), demonstrating
that SFT teaches the model to interpret and follow abstract task
descriptions.

\paragraph{Low-IFD samples (base IFD $< 0.4$, e.g.,~rows~5,~6).}
These are formatting or extraction tasks: ``Convert paragraph to list,''
``Remove third-person words.'' Unconditional
loss is much higher than conditional loss even in the base model,
because the instruction naturally constrains the output space. The
pattern persists after SFT but with lower absolute values.

\paragraph{Mid-range with large drop (e.g.,~rows~2,~3).}
These are factual QA or grammar correction tasks. Base IFD is
$\approx 1.05$ (instruction has little effect), but after SFT
IFD drops to $\approx 0.1$ as the model learns the precise answer
(e.g., ``Madrid'' for ``capital of Spain''), making conditional loss
near-zero while unconditional loss remains high.

\paragraph{Mid-range with small drop (e.g.,~row~4).}
These are open-ended generation tasks (``Describe the role of a
project manager''). Base IFD $\approx 0.98$; after SFT it drops
only modestly to $\approx 0.9$, since both conditional and
unconditional losses decrease proportionally without a memorized
target.

\paragraph{Cross-model correlation.}
The moderate Pearson correlation ($r = 0.38$) suggests that while
training reshapes the model's perception of instruction difficulty,
a residual signal persists: instructions that require complex reasoning
tend to remain non-trivially harder than simple rewrite or extraction
tasks even after SFT.

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

Both $L_{\text{cond}}$ and $L_{\text{uncond}}$ are reported as per-token
average losses. For a response of length $T$, the unconditional loss is
$L_{\text{uncond}} = \frac{1}{T} \sum_{t=1}^T \log P(x_t)$.
Since the variance of this average scales as $1/T$, shorter responses
exhibit much larger fluctuations in $L_{\text{uncond}}$---a mathematical
necessity, not a signal of instruction difficulty. Consequently, IFD,
being a ratio of two such averages, inherits a systematic length bias:
short responses inflate IFD variance.

Figure~\ref{fig:length_bias} confirms this artifact across a 9-panel
grid. The top row shows conditional loss, middle row unconditional
loss, and bottom row IFD---each plotted against response length and
loss magnitude. Short responses ($<20$ tokens, e.g., ``Paris,'' ``42'')
produce wildly scattered $L_{\text{uncond}}$ values, which in turn
generate spurious high or low IFD scores in the bottom panels.
Longer responses ($>50$ tokens) converge toward the model's intrinsic
mean loss, yielding stable IFD estimates across both base and
checkpoint models.

\begin{figure}[H]
\centering
\includegraphics[width=0.80\linewidth]{data/ifd_length_grid.png}
\caption{Response length vs.\ conditional loss, unconditional loss,
and IFD. Short responses produce high-variance $L_{\text{uncond}}$
estimates, inflating IFD noise.}
\label{fig:length_bias}
\end{figure}

\paragraph{Distribution summary.}
Over the full 3000-sample set, the base model's conditional loss
has median $2.56$, unconditional loss median $2.80$, and IFD median
$0.95$, concentrated in the $0.6$--$1.1$ range with a slight left
skew (cond $<$ uncond for most samples).

\paragraph{Correlation analysis.}
Table~\ref{tab:corr_bias} reports Pearson $r$ and Spearman $\rho$
between key dimensions and the three IFD components.

Three patterns stand out:

\begin{enumerate}[nosep]
  \item \textbf{Instruction length is nearly independent}
        ($r \approx 0$ for all three targets). The length of the
        instruction text itself has no meaningful correlation with
        either loss or IFD. The slight negative IFD correlation
        ($r = -0.24$, $\rho = -0.35$) is an indirect artifact driven
        by response length (longer instructions tend to elicit shorter
        answers in our Alpaca distribution).

  \item \textbf{Response length is the dominant confound.}
        $L_{\text{uncond}}$ shows a strong negative monotonic trend
        ($\rho = -0.70$), a direct consequence of the per-token
        average variance scaling as $1/T$ (Section~\ref{sec:ifd_bias}).
        $L_{\text{cond}}$ has a weaker negative correlation
        ($r = -0.38$), because conditional generation already
        constrains the output distribution regardless of length.
        The net effect on IFD is a moderate positive bias
        ($r = +0.31$, $\rho = +0.47$): long responses produce
        higher IFD not because they are harder, but because
        $L_{\text{uncond}}$ drops faster with length than
        $L_{\text{cond}}$.

  \item \textbf{The ratio (resp/inst) is collinear with response
        length} and provides no independent information.
        All three columns mirror those of response length with
        slightly attenuated magnitudes. Filtering by response
        length alone suffices.
\end{enumerate}

The consistently larger $\rho$ than $r$ across all rows confirms
that the relationships are monotonic but nonlinear---steep at
the short end and flat for long sequences, consistent with the
$1/T$ variance decay predicted in Section~\ref{sec:ifd_bias}.

\begin{table}[H]
\centering
\caption{Pearson $r$ and Spearman $\rho$ between sample dimensions and IFD components.}
\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.01$ & $+0.04$ & $+0.11$ & $+0.22$ & $-0.24$ & $-0.35$ \\
Response length    & $-0.38$ & $-0.46$ & $-0.52$ & $-0.70$ & $+0.31$ & $+0.47$ \\
Ratio (resp/inst)  & $-0.32$ & $-0.41$ & $-0.46$ & $-0.67$ & $+0.30$ & $+0.52$ \\
\bottomrule
\end{tabular}
\end{table}

\paragraph{Practical recommendation.}
Filter samples with response length $<20$ or $>300$ tokens before
computing IFD. This retains the middle interval where per-token
loss averages are stable and IFD rankings are most reliable.
In our Alpaca-style dataset, this removes approximately
$5$--$8\%$ of samples and substantially reduces false positives
in the high-IFD tail.

% ======================================================================
\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, S.~Li, Z.~Li, Z.~Li, L.~Zhu.
From quantity to quality: Boosting LLM performance with self-guided data
selection for instruction tuning.
\textit{NeurIPS}, 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}