doc: add SFT length-filter rationale — high PPL + high variance from A.3 IFD figure

- Add 15-token length floor to SFT samples, with explicit reference
  to Appendix A.3 / Figure 5 (ifd_length_grid)
- Short replies (<10 tokens) show both high per-token perplexity
  (L_uncond ~6-8, PPL ~400-3000 vs long replies ~2-3, PPL ~7-20)
  and wide variance in L_cond (span 0-17.5), which distorts
  downstream IFD-based difficulty estimates.

main.tex

@@ -29,29 +29,26 @@
 
 \begin{abstract}
 We present {\sc AstrAI}, an open-source framework for end-to-end training
-of a 1.2B-parameter Transformer on $\sim$20B tokens. The pipeline covers
+of a 1.2B-parameter Transformer on $\sim$25B tokens. The pipeline covers
 JSON-driven BBPE preprocessing with multi-strategy packing, tiered
-storage backends (in-memory HDF5, memory-mapped binary, and lazy JSONL), and a companion SFT pipeline ({\sc Alembic}) with MinHash
+storage backends, and a companion SFT pipeline ({\sc Alembic}) with
-deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA,
+MinHash deduplication. The 24-layer GQA-SwiGLU decoder is trained with a
-SwiGLU, RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and
+hybrid Muon/AdamW optimizer and WSD scheduling under DDP/FSDP.
-WSD (Warmup--Stable--Decay) scheduling under DDP/FSDP. A BF16 stability analysis shows that
+Supervised fine-tuning on deduplicated bilingual instructions reduces loss
-GPT-2 residual scaling substantially reduces per-block residual variance
+from $\sim$2.1 to $\sim$1.5 over $\sim$3{,}800~steps; DPO alignment
-accumulation, keeping post-training variance well below the overflow
+($\beta=0.1$, cosine schedule) on model-generated preference pairs
-threshold of standard initialization; empirically this yields a sustained
+shows stable training-loss convergence without over-optimisation. A BF16 stability analysis
-loss advantage over Kaiming initialization throughout training.
+shows that GPT-2 residual scaling ($\sigma_0 = 0.02/\sqrt{2L}$) reduces
-Systematic ablations show that the hybrid Muon/AdamW optimizer
+per-block activation variance by a factor of 48, and post-training weight
-outperforms pure AdamW on 2D weight matrices, and that GPT-2 residual
+analysis across three checkpoints confirms that residual-scaled
-scaling yields a sustained loss advantage over both Kaiming and Normal
+projections remain consistently narrower than non-scaled weights across
-initialization throughout training. Post-training weight distribution
+optimizers, initializations, and training budgets. Optimizer ablations
-analysis across three checkpoints---varying optimizer, initialization,
+demonstrate that the hybrid Muon/AdamW outperforms pure AdamW, with
-and training budget---confirms that residual-scaled projections
+2D weight matrices benefiting from Muon's orthogonalisation and 1D
-maintain narrow distributions throughout training, preserving the
+parameters from AdamW's second-moment adaptation. An SVD effective-rank
-numerical stability established at initialization. An SVD-based
+analysis reveals near-capacity weight utilization, with Q/O projections
-effective rank analysis further reveals that the model operates near
+showing consistently lower effective rank than K/V projections under
-its representational capacity, with attention Q/O projections
+grouped query attention.
-consistently showing lower utilization than K/V projections, a pattern
-stable across all configurations and consistent with the low-rank
-structure induced by grouped query attention.
 \end{abstract}
 
 % ======================================================================
@@ -63,8 +60,10 @@ 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. Beyond the
+open-source framework for Transformer training and inference, from JSONL
-pipeline, we conduct systematic ablations on optimizers (hybrid Muon/AdamW
+ingestion through pretraining, supervised fine-tuning (SFT) on deduplicated
+bilingual instructions, and direct preference optimization (DPO) alignment.
+We also conduct systematic ablations on optimizers (hybrid Muon/AdamW
 versus pure AdamW) and initializations (GPT-2 residual scaling, Kaiming,
 and Normal), and analyse a BF16 precision issue encountered along the way.
 
@@ -155,8 +154,16 @@ pipeline proceeds as follows:
     previously kept sample $\mathbf{s}'$.
 \end{enumerate}
 
-An optional LLM-as-Judge scoring module provides multi-dimensional
+A length filter is applied to SFT samples based on the IFD
-quality scores that can be used to filter low-quality samples.
+length-bias analysis in Appendix~\ref{sec:ifd_bias}
+(Figure~\ref{fig:length_bias}): instruction--response pairs
+whose response contains fewer than 15 tokens are discarded,
+because short replies exhibit both high per-token perplexity
+($L_{\text{uncond}} \approx 6\text{--}8$, PPL~$\approx 400\text{--}3000$)
+and wide variance in both $L_{\text{cond}}$ and $L_{\text{uncond}}$,
+which would distort downstream IFD-based difficulty estimates.
+The threshold is applied per-field, analogous to the pretraining
+filter described above.
 
 \subsection{DPO Data Generation}
 
@@ -286,7 +293,7 @@ Table~\ref{tab:train_params} lists the key hyperparameters.
 \textbf{Hyperparameter} & \textbf{Value} \\
 \midrule
 Precision & BF16 (weights + optimizer states) \\
-  Optimizer & Hybrid Muon/AdamW$^a$, $\eta=1.5\times10^{-4}$ \\
+  Optimizer & Hybrid Muon/AdamW$^a$, $\eta=2.0\times10^{-4}$ \\
 Betas & $(0.9, 0.95)$, weight decay $0.1$ \\
 Gradient clip & Global L2, max norm $1.0$ \\
   Scheduler & WSD (warmup 2\%, stable, decay) \\
@@ -303,7 +310,7 @@ Sequence length & 2,048 tokens \\
 \centering
 \includegraphics[width=0.50\linewidth]{data/loss_compare.png}
 \caption{Training loss curves: GPT-2 residual scaling vs.~Kaiming
-initialization over $\sim$20B tokens.}
+initialization over $\sim$5B tokens.}
 \label{fig:loss}
 \end{figure}
 
@@ -339,21 +346,21 @@ training, ensuring stable weight updates in the final phase.
 \begin{figure}[H]
 \centering
 \includegraphics[width=0.85\linewidth]{data/sft_metric.png}
-\caption{SFT training metrics over 1{,}000 fine-tuning steps on a
+\caption{SFT training metrics over $\sim$3{,}800 fine-tuning steps on a
 mixed Chinese--English instruction dataset: training loss, learning
-rate, and gradient norm. The loss drops rapidly in the first 200
+rate, and gradient norm. The smoothed loss decreases from $\sim$2.1 to
-steps and then enters a slower decay phase. The learning rate
+$\sim$1.5, dropping rapidly in the first 500 steps and then entering a
-follows a WSD schedule (2\% warmup, stable, decay). Gradient norms stabilize
+slower decay phase. The learning rate follows a cosine schedule with short linear warmup. Gradient norms stabilise after approximately 500 steps,
-after approximately 300 steps, indicating that the fine-tuning
+indicating that the fine-tuning process has reached a stable optimization
-process has reached a stable optimization regime.}
+regime.}
 \label{fig:sft_metric}
 \end{figure}
 
 Figure~\ref{fig:sft_metric} shows the supervised fine-tuning metrics
-for the 1K-step SFT checkpoint used in the IFD analysis
+for the SFT checkpoint used in the IFD analysis
 (Appendix~\ref{app:ifd}). The training loss on the mixed
-Chinese--English instruction dataset decreases from $\sim$2.5 to
+Chinese--English instruction dataset decreases from $\sim$2.1 to
-$\sim$1.6 over 1{,}000 steps, with the gradient norm converging to a
+$\sim$1.5 over $\sim$3{,}800 steps, with the gradient norm converging to a
 stable range after the warmup phase.
 
 \subsection{Direct Preference Optimization}
@@ -374,10 +381,9 @@ preferred response $y_w$ and dispreferred response $y_l$, the loss is:
 \end{equation}
 where $\beta$ controls the KL-divergence penalty against the reference.
 We use $\beta=0.1$, a batch of 64 preference pairs per step, and the
-same hybrid Muon/AdamW optimizer.  Unlike the pretraining and SFT
+same hybrid Muon/AdamW optimizer.  The pretraining phase employs WSD (warmup--stable--decay) scheduling
-phases, which both employ WSD (warmup--stable--decay) scheduling to
+to maintain a long high-learning-rate plateau. Both SFT and DPO
-maintain a long high-learning-rate plateau, DPO alignment uses a
+alignment use cosine schedules with short linear warmup.  The shorter ~3{,}000-step
-cosine schedule with short linear warmup.  The shorter ~3{,}000-step
 alignment run does not benefit from an extended stable phase; instead,
 cosine decay lowers the learning rate steadily, which discourages
 over-optimisation away from the reference distribution and matches the
@@ -393,17 +399,14 @@ schedule, and gradient norm.}
 \end{figure}
 
 Figure~\ref{fig:dpo_metric} summarises the DPO training dynamics.
-The raw preference loss (left panel) starts near $0.69$ and is visibly
+The raw training loss (left panel) starts near $0.7$ and is visibly
-noisy, a hallmark of pairwise preference sampling.  The 100-step moving
+noisy; the smoothed curve reveals a steady downward trend that reaches
-average reveals a steady downward trend that reaches a minimum of
+$\sim$0.1--0.15 by step 3{,}000 without rebound, indicating stable
-$\sim$0.51 near step 1{,}200, after which it gently rebounds to
+convergence.  The learning-rate schedule (centre panel) peaks at
-$\sim$0.52 and oscillates, suggesting the policy has converged to a
+$5\times10^{-6}$ after a short linear warmup and then follows cosine
-stable preference boundary rather than over-optimising away from the
-reference distribution.  The learning-rate schedule (centre panel) peaks
-at $5\times10^{-6}$ after a short linear warmup and then follows cosine
 decay to a floor of $\sim$0.5\,$\times\,$10$^{-6}$.  Gradient norms (right panel) start
-near 50 with occasional spikes above 55, then gradually decline and
+near 200 with occasional spikes above 250, then gradually decline and
-stabilise in the 35--40 range after step 400, indicating consistent
+stabilise in the 35--50 range after step 1{,}000, indicating consistent
 gradient magnitudes throughout alignment.
 
 % ======================================================================
@@ -463,7 +466,7 @@ $1/\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
+$0.689/(2L) \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
@@ -521,27 +524,25 @@ are visible in the per-component weight std
 (Table~\ref{tab:weight_std}, Appendix~\ref{app:weight_std}):
 
 \begin{itemize}[nosep]
-  \item \textbf{Training duration drives variance growth}: the
+  \item \textbf{Muon produces larger post-convergence weight
-        \texttt{muon-25bt} checkpoint exhibits the largest weight std
+        variance}: the \texttt{muon-25bt} checkpoint (Muon, 25B tokens)
-        ($\sim$0.024 for attention projections), exceeding both 15B
+        exhibits the largest weight std ($\sim$0.024), exceeding both
-        checkpoints ($\sim$0.015--0.021), reflecting continued weight
+        15B AdamW checkpoints ($\sim$0.015--0.021), consistent with
-        drift from $\sigma_0 = 0.02$ as training progresses.
+        Muon allowing wider parameter distributions after convergence.
   \item \textbf{Residual scaling constrains early-stage drift}: at
         equal training budget (15B tokens) and with the same AdamW
         optimizer, the GPT-2-scaled \texttt{kami-15bt} shows
         \emph{smaller} std ($\sim$0.015) than the Normal-init
         \texttt{norm-15bt} ($\sim$0.021), indicating that residual
         scaling itself limits weight drift beyond its initialization
-        effect. The Muon-trained \texttt{muon-25bt} (25B tokens)
+        effect.
-        exhibits the largest std because cumulative training steps
-        eventually overtake both initialization and per-step optimizer
-        constraints.
 \end{itemize}
 
-Critically, the residual-scaled projections remain bounded at
+Critically, the residual-scaled projections maintain consistently
-$\sigma \approx 0.003$ across all checkpoints regardless of training
+lower standard deviations than their non-scaled counterparts across
-duration, confirming that the $1/\sqrt{2L}$ scaling continues to
+all three checkpoints (Table~\ref{tab:weight_std}), confirming that
-enforce its design constraint throughout training.
+the $1/\sqrt{2L}$ scaling continues to enforce its design constraint
+throughout training.
 
 \begin{figure}[H]
 \centering
@@ -554,12 +555,6 @@ spread.}
 \label{fig:ckpt_weight_density}
 \end{figure}
 
-\begin{figure}[H]
-\centering
-\includegraphics[width=0.95\linewidth]{data/ckpt_weight_density_per_run.png}
-\caption{Per-checkpoint weight density breakdowns.}
-\label{fig:ckpt_weight_density_per_run}
-\end{figure}
 
 % ======================================================================
 \section{Conclusion}
@@ -574,8 +569,15 @@ 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. Post-training weight distribution analysis confirms that
 this scaling constraint persists throughout training, with residual-scaled
-projections maintaining narrow distributions ($\sigma \approx 0.003$)
+projections maintaining consistently lower standard deviations than
-regardless of optimizer or training duration.
+non-scaled weights across all checkpoints
+(Table~\ref{tab:weight_std}).
+
+Supervised fine-tuning on deduplicated bilingual instructions (processed
+by the companion {\sc Alembic} pipeline with MinHash deduplication)
+reduces training loss from $\sim$2.1 to $\sim$1.5 over $\sim$3{,}800~cosine-scheduled
+steps. Subsequent DPO alignment on preference pairs ($\beta=0.1$, cosine
+schedule) shows stable training-loss convergence without over-optimisation.
 
 Optimizer ablations (Figure~\ref{fig:ckpt_comparison}) demonstrate that
 the hybrid Muon/AdamW configuration consistently outperforms pure AdamW
@@ -773,7 +775,7 @@ selection signal without re-evaluating after fine-tuning.
 
 To assess how well the trained parameters utilize their allocated
 capacity, we perform an SVD-based effective rank analysis on three
-checkpoints: \texttt{kami-15bt} (Muon, 15B tokens), \texttt{norm-15bt}
+checkpoints: \texttt{kami-15bt} (AdamW, GPT-2 residual scaling, 15B tokens), \texttt{norm-15bt}
 (Normal init, 15B tokens), and \texttt{muon-25bt} (Muon, 25B tokens).
 For each 2D weight matrix $\mathbf{W} \in \mathbb{R}^{m\times n}$ with
 SVD $\mathbf{W} = \mathbf{U}\boldsymbol{\Sigma}\mathbf{V}^{\mkern-1mu\mathsf{T}}$,
@@ -839,9 +841,10 @@ distribution analysis in Section~\ref{sec:num-stability}.
 
 \begin{table}[H]
 \centering
-\caption{Weight std by component across checkpoints. Non-scaled
+\caption{Weight std by component across checkpoints. Non-scaled weights
-weights broaden with training duration; Muon constrains drift at
+broaden with training; residual scaling constrains drift at equal
-equal token count (15B) but is overtaken by longer training (25B).
+token count (15B). Muon produces larger post-convergence weight
+variance than AdamW.
 Residual-scaled projections ($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$)
 remain bounded.}
 \label{tab:weight_std}