diff --git a/.gitignore b/.gitignore index f49d932..46aff1d 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,4 @@ # Only track .tex files and .gitignore itself !.gitignore !*.tex -!data/** +!data/**.png diff --git a/data/ifd_both_vs_lossratio.png b/data/ifd_both_vs_lossratio.png deleted file mode 100644 index 267e863..0000000 Binary files a/data/ifd_both_vs_lossratio.png and /dev/null differ diff --git a/data/ifd_compare_clean.png b/data/ifd_compare_clean.png deleted file mode 100644 index 087f6d3..0000000 Binary files a/data/ifd_compare_clean.png and /dev/null differ diff --git a/data/ifd_density_dist.png b/data/ifd_density_dist.png deleted file mode 100644 index e1a0118..0000000 Binary files a/data/ifd_density_dist.png and /dev/null differ diff --git a/data/ifd_length_grid.png b/data/ifd_length_grid.png index 9b53634..091a781 100644 Binary files a/data/ifd_length_grid.png and b/data/ifd_length_grid.png differ diff --git a/data/ifd_loss_ratio_density.png b/data/ifd_loss_ratio_density.png deleted file mode 100644 index d398f89..0000000 Binary files a/data/ifd_loss_ratio_density.png and /dev/null differ diff --git a/main.tex b/main.tex index ab8da88..cf703f5 100644 --- a/main.tex +++ b/main.tex @@ -31,25 +31,19 @@ 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 +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 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. +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} % ====================================================================== @@ -141,157 +135,8 @@ 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 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_{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. +An IFD (Instruction Fulfillment Difficulty) analysis is provided in +Appendix~\ref{app:ifd}. % ====================================================================== \section{Model Architecture} @@ -532,181 +377,104 @@ weights are available at \url{https://github.com/ViperEkura/AstrAI}. % ====================================================================== % ====================================================================== -\section{IFD Data Examples} +\section{IFD Data Analysis} \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. +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_{ 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 covering four patterns.} +\caption{Representative IFD samples.} \label{tab:ifd_examples} \small -\begin{tabular}{@{}c c c c c c c p{4.5cm}@{}} +\begin{tabular}{@{}c c c c c p{4.2cm}@{}} \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{$L_{\text{cond}}^{\text{1K}}$} & + \textbf{$L_{\text{uncond}}^{\text{1K}}$} & \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 \\ +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} -\subsection{Quantitative Summary} +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. -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} +\subsection{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. +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.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.} +\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} -\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}. +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.} +\caption{Pearson $r$ and Spearman $\rho$ between sample dimensions and IFD components (base model).} \label{tab:corr_bias} \small \begin{tabular}{@{}lcccccc@{}} @@ -717,21 +485,12 @@ $1/T$ variance decay predicted in Section~\ref{sec:ifd_bias}. \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$ \\ +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} -\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}