581 lines
23 KiB
TeX
581 lines
23 KiB
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{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}
|
|
|
|
% ======================================================================
|
|
\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}
|