diff --git a/data/ifd_compare_clean.png b/data/ifd_compare_clean.png new file mode 100644 index 0000000..087f6d3 Binary files /dev/null and b/data/ifd_compare_clean.png differ diff --git a/main.tex b/main.tex index 0e10a82..94a1107 100644 --- a/main.tex +++ b/main.tex @@ -26,18 +26,24 @@ \maketitle \begin{abstract} -We present an end-to-end framework for training a 1.2B-parameter autoregressive -language model using the {\sc AstrAI} open-source toolkit. The pipeline -encompasses data preprocessing (JSONL tokenization to BBPE, HDF5 and -memory-mapped storage), a 24-layer decoder-only architecture with Grouped -Query Attention and SwiGLU feed-forward blocks, and distributed training -via DDP/FSDP with cosine scheduling. We further examine BF16 numerical -stability in deep Transformers, demonstrating that GPT-2 residual scaling -($\sigma_o = 0.02 / \sqrt{2L}$) on output projections reduces per-block -residual variance by a factor of 48, yielding a post-24-layer variance of -$1.34$ versus $17.5$ without scaling. Empirical results across 15B training -tokens confirm that residual scaling maintains superior loss reduction over -Kaiming initialization, with a widening gap in the mid-training regime. +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 ($\sigma_o = 0.02 / \sqrt{2L}$) on output +projections reduces per-block residual variance by a factor of $2L=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} % ====================================================================== @@ -129,6 +135,39 @@ pipeline proceeds as follows: 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 losses: +\begin{equation} +\mathrm{IFD} = \frac{L_{\text{cond}}}{L_{\text{uncond}}}. +\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 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 comparison: base model vs.\ trained checkpoint. The +diagonal line marks $\mathrm{IFD}_{\text{base}} = \mathrm{IFD}_{\text{ckpt}}$.} +\label{fig:ifd} +\end{figure} + +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}. + % ====================================================================== \section{Model Architecture} % ====================================================================== @@ -325,6 +364,101 @@ 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 (base model sorted by descending IFD).} +\label{tab:ifd_examples} +\small +\begin{tabular}{@{}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{Instruction (truncated)} \\ +\midrule +0 & 4.605 & 1.525 & 12.38 & 2.69 & 3.77 & Complete the following analogy \dots \\ +1 & 4.331 & 0.645 & 11.44 & 2.64 & 1.66 & Classify the following text \dots \\ +2 & 3.741 & 0.702 & 11.75 & 3.14 & 2.17 & Label the following news article \dots \\ +3 & 0.977 & 0.904 & 2.57 & 2.63 & 2.24 & Describe the role of a project manager \\ +4 & 0.977 & 0.915 & 2.19 & 2.25 & 1.98 & Select a historical figure \dots \\ +5 & 0.977 & 0.949 & 2.57 & 2.63 & 2.26 & Write the lyrics for an upbeat song \dots \\ +6 & 0.977 & 0.925 & 2.94 & 3.00 & 2.62 & Explain how neural networks \dots \\ +7 & 0.370 & 0.249 & 1.37 & 3.70 & 0.85 & Convert the given paragraph to a list \\ +8 & 0.338 & 0.197 & 0.98 & 2.91 & 0.55 & Insert a suitable greeting \dots \\ +9 & 0.307 & 0.062 & 0.70 & 2.29 & 0.15 & Remove third-person words \dots \\ +\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--2).} +These are tasks requiring task-intent comprehension: analogy completion, +text classification, 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~7--9).} +These are formatting or extraction tasks: ``Convert paragraph to list,'' +``Remove third-person words,'' ``Insert a greeting.'' 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 samples (base IFD $\approx 0.98$, e.g.,~rows~3--6).} +These cover factual Q\&A and generation tasks: ``Describe the role of +a project manager,'' ``Write lyrics for a song,'' ``Explain how neural +networks work.'' In the base model IFD $\approx 1$ (instruction has +little effect); after SFT IFD drops to $\approx 0.9$, driven by +a clear reduction in conditional loss. + +\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. + % ====================================================================== \begin{thebibliography}{99} @@ -347,6 +481,12 @@ framework for Transformer language models.} 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.