Add IFD analysis, appendix, and improve abstract

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.