Tighten abstract, add Muon optimizer mention, update ckpt_comparison figure

main.tex

@@ -28,22 +28,18 @@
 \maketitle
 
 \begin{abstract}
-Training billion-parameter language models requires careful co-design of
+We present {\sc AstrAI}, an open-source framework for end-to-end training
-data infrastructure, distributed execution, and numerical precision
+of a 1.2B-parameter Transformer on 15B tokens. The pipeline covers
-management. This paper presents {\sc AstrAI}, an open-source framework
+JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap
-for end-to-end training of a 1.2B-parameter autoregressive Transformer
+storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash
-over 15B tokens. We describe the full pipeline: JSON-driven preprocessing
+deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA, SwiGLU,
-with BBPE tokenization and multi-strategy packing, HDF5 and memory-mapped
+RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and cosine scheduling under DDP/FSDP.
-storage backends, and a companion SFT pipeline ({\sc Alembic}) with
+A focused BF16 stability analysis shows that GPT-2 residual scaling
-MinHash-based near-duplicate detection and LLM-as-Judge scoring. The
+($\sigma = 0.02/\sqrt{2L}$) reduces per-block residual variance by a factor
-model is a 24-layer decoder-only Transformer with Grouped Query Attention,
+of 48, containing post-24-layer variance at 1.34 versus 17.5 under standard
-SwiGLU, RoPE, and RMSNorm, trained with AdamW and cosine scheduling via
+initialization. Empirically, this scaling yields a sustained loss advantage
-DDP/FSDP. A central focus is BF16 numerical stability: variance propagation
+over Kaiming initialization, with the gap peaking at $\Delta = 0.79$ in the
-analysis shows that GPT-2 residual scaling reduces per-block residual
+mid-training regime.
-variance by a factor of 48, containing post-24-layer variance at 1.34
-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}
 
 % ======================================================================