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

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

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

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.

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