Fix training budget (19k steps / ~20B tokens), add Muon optimizer, correct variance scaling math, move figures to §Training Config, tighten abstract

main.tex

@@ -29,7 +29,7 @@
 
 \begin{abstract}
 We present {\sc AstrAI}, an open-source framework for end-to-end training
-of a 1.2B-parameter Transformer on 15B tokens. The pipeline covers
+of a 1.2B-parameter Transformer on $\sim$20B tokens. The pipeline covers
 JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap
 storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash
 deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA, SwiGLU,
@@ -191,7 +191,6 @@ non-linearity:
 \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})$.
 
@@ -201,7 +200,6 @@ Each decoder block $\ell$ then applies pre-norm residual connections:
 \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}
@@ -224,8 +222,8 @@ The model is trained on next-token cross-entropy loss:
 \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,
+Training uses a hybrid optimizer: Muon for 2D weight matrices and AdamW~\cite{loshchilov2019adamw} for 1D parameters (embeddings, biases, LayerNorm), with cosine learning rate
+scheduling (2\% 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.
 
@@ -237,18 +235,35 @@ Table~\ref{tab:train_params} lists the key hyperparameters.
 \toprule
 \textbf{Hyperparameter} & \textbf{Value} \\
 \midrule
-Precision & BF16 (weights + AdamW states) \\
+Precision & BF16 (weights + optimizer 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 \\
+Total steps & 19,000 \\
 \bottomrule
 \end{tabular}
 \end{table}
 
+\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 $\sim$20B tokens.}
+\label{fig:loss}
+\end{figure}
+
+\begin{figure}[H]
+\centering
+\includegraphics[width=0.95\linewidth]{data/ckpt_comparison.png}
+\caption{Optimizer and initialization comparison.}
+\label{fig:ckpt_comparison}
+\end{figure}
+
+Figure~\ref{fig:ckpt_comparison} compares four configurations. The left panel shows training loss for Muon (Embedding Adam + 1D Adam), Muon (Embedding Muon + 1D Adam), Kaiming init, and Normal init; the center panel zooms in on the two current Muon variants; and the right panel shows gradient norms over optimizer steps. The older Kaiming and Normal initializations converge more slowly and plateau at higher loss. Between the current variants, using Adam for the embedding layer yields lower loss and more stable gradients than using Muon embeddings.
+
 % ======================================================================
 \section{Numerical Stability via Residual Scaling}
 \label{sec:num-stability}
@@ -276,20 +291,20 @@ at each sub-stage are:
 \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$ \\
+Attention out & SDPA + $\mathbf{W}_o$ (scaled) & $0.378 / (2L)$ \\
 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$ \\
+Down proj & Linear(6912, 1536) (scaled) & $0.311 / (2L)$ \\
 \midrule
-Per-block residual & $\mathbf{R}_\ell = R_{\text{attn}} + R_{\text{ffn}}$ & $0.689 / L$ (scaled) \\
+Per-block residual & $\mathbf{R}_\ell = R_{\text{attn}} + R_{\text{ffn}}$ & $0.689 / (2L)$ (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:
+instead of $0.689/(2L) \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,\\
@@ -315,26 +330,9 @@ 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{figure}[H]
-\centering
-\includegraphics[width=0.95\linewidth]{data/ckpt_comparison.png}
-\caption{Optimizer and initialization comparison.}
-\label{fig:ckpt_comparison}
-\end{figure}
-
-Figure~\ref{fig:ckpt_comparison} compares four configurations. The left panel shows training loss for Muon (Embedding Adam + 1D Adam), Muon (Embedding Muon + 1D Adam), Kaiming init, and Normal init; the center panel zooms in on the two current Muon variants; and the right panel shows gradient norms over optimizer steps. The older Kaiming and Normal initializations converge more slowly and plateau at higher loss. Between the current variants, using Adam for the embedding layer yields lower loss and more stable gradients than using Muon embeddings.
-
 \begin{table}[H]
 \centering
 \caption{Loss at 0.125B-interval milestones, 0--1B tokens.}
@@ -370,8 +368,8 @@ identified in the theoretical analysis (Section~\ref{sec:num-stability}).
 
 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
+a 24-layer GQA-SwiGLU architecture, callback-based training with a hybrid
+Muon/AdamW optimizer under 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
@@ -519,17 +517,17 @@ evaluation samples. Instead, the fine-tuning distribution shift
 increases NLL on most held-out instructions.
 
 Table~\ref{tab:ifd_lr_corr} reports the pairwise correlations.
-Despite the theoretical expectation that IFD\textsubscript{ckpt} and
-Loss Ratio share $L_{\text{cond}}^{\text{1K}}$ in their numerators
-and should therefore correlate strongly~\cite{li2023ifd}, the
-observed correlation is near zero ($r = -0.02$, $\rho = 0.04$).
-This is because the {\em relative} ordering of
-$L_{\text{cond}}^{\text{base}}$ and
+Although IFD\textsubscript{ckpt} and Loss Ratio both depend on
+$L_{\text{cond}}^{\text{1K}}$, they need not correlate because their
+denominators vary independently across samples. The observed correlation
+is near zero ($r = -0.02$, $\rho = 0.04$), precisely because the
+{\em relative} ordering of $L_{\text{cond}}^{\text{base}}$ and
 $L_{\text{uncond}}^{\text{ckpt}}$ (which determine the slope
 $k_i = L_{\text{cond},i}^{\text{base}} / L_{\text{uncond},i}^{\text{ckpt}}$
 in the relationship $\text{IFD}_{\text{ckpt},i} = k_i \cdot
-\text{Loss Ratio}_i$) varies widely across samples, breaking the
-proportionality at the sample level.
+\text{Loss Ratio}_i$) varies widely, breaking the proportionality
+at the sample level. This invalidates the naive expectation that a shared
+numerator guarantees correlation~\cite{li2023ifd}.
 
 \begin{table}[H]
 \centering