ViperEkura
4e70e827ff
The per-run weight density figure contained misleading legacy iteration labels (500k/1M iter) that contradicted the paper's stated token budgets (15B). Removing it avoids confusion; the per-category density plot (ckpt_weight_density.png) and Table 7 remain as the primary evidence.
928 lines
40 KiB
TeX
928 lines
40 KiB
TeX
\documentclass[11pt,a4paper]{article}
|
|
|
|
% ===== Packages =====
|
|
\usepackage[utf8]{inputenc}
|
|
\usepackage[T1]{fontenc}
|
|
\usepackage{newtxtext,newtxmath}
|
|
\usepackage[margin=1in]{geometry}
|
|
\usepackage{amsmath}
|
|
\usepackage{booktabs}
|
|
\usepackage{graphicx}
|
|
\usepackage{hyperref}
|
|
\usepackage{float}
|
|
\usepackage{caption}
|
|
\usepackage{enumitem}
|
|
\usepackage{url}
|
|
\usepackage{microtype}
|
|
|
|
\DeclareMathOperator{\Var}{Var}
|
|
|
|
\title{End-to-End Training of a 1.2B Transformer with AstrAI \\
|
|
\large Data Pipeline, Distributed Training, and Ablations on Optimizer, Initialization, and BF16 Numerical Stability}
|
|
|
|
\author{AstrAI Contributors}
|
|
\date{}
|
|
|
|
\begin{document}
|
|
|
|
\maketitle
|
|
|
|
\begin{abstract}
|
|
We present {\sc AstrAI}, an open-source framework for end-to-end training
|
|
of a 1.2B-parameter Transformer on $\sim$25B tokens. The pipeline covers
|
|
JSON-driven BBPE preprocessing with multi-strategy packing, tiered
|
|
storage backends, and a companion SFT pipeline ({\sc Alembic}) with
|
|
MinHash deduplication. The 24-layer GQA-SwiGLU decoder is trained with a
|
|
hybrid Muon/AdamW optimizer and WSD scheduling under DDP/FSDP.
|
|
Supervised fine-tuning on deduplicated bilingual instructions reduces loss
|
|
from $\sim$2.1 to $\sim$1.5 over $\sim$3{,}800~steps; DPO alignment
|
|
($\beta=0.1$, cosine schedule) on model-generated preference pairs
|
|
shows stable training-loss convergence without over-optimisation. A BF16 stability analysis
|
|
shows that GPT-2 residual scaling ($\sigma_0 = 0.02/\sqrt{2L}$) reduces
|
|
per-block activation variance by a factor of 48, and post-training weight
|
|
analysis across three checkpoints confirms that residual-scaled
|
|
projections remain consistently narrower than non-scaled weights across
|
|
optimizers, initializations, and training budgets. Optimizer ablations
|
|
demonstrate that the hybrid Muon/AdamW outperforms pure AdamW, with
|
|
2D weight matrices benefiting from Muon's orthogonalisation and 1D
|
|
parameters from AdamW's second-moment adaptation. An SVD effective-rank
|
|
analysis reveals near-capacity weight utilization, with Q/O projections
|
|
showing consistently lower effective rank than K/V projections under
|
|
grouped query attention.
|
|
\end{abstract}
|
|
|
|
% ======================================================================
|
|
\section{Introduction}
|
|
% ======================================================================
|
|
|
|
Training a billion-parameter language model end-to-end involves far more than
|
|
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, from JSONL
|
|
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.
|
|
|
|
% ======================================================================
|
|
\section{Data Pipeline}
|
|
% ======================================================================
|
|
|
|
\subsection{Preprocessing}
|
|
|
|
Raw data arrives as JSONL files. The preprocessing pipeline is configured
|
|
via a JSON specification that defines:
|
|
|
|
\begin{itemize}[nosep]
|
|
\item \textbf{Tokenization}: BBPE tokenizer (100K vocabulary) with standard
|
|
special tokens.
|
|
\item \textbf{Masking}: Declarative loss mask assignment per section
|
|
(e.g.,~mask user input, compute loss on assistant response).
|
|
\item \textbf{Packing}: Documents concatenated via \texttt{simple}
|
|
(sequential), \texttt{bfd} (best-fit decreasing), or
|
|
\texttt{bfd\_\allowbreak{}split} strategies.
|
|
\item \textbf{Position IDs}: \texttt{none}, \texttt{doc\_reset} (per-document
|
|
boundary), or \texttt{continuous}.
|
|
\item \textbf{Output}: Tokenized sequences written to \texttt{.h5} or
|
|
\texttt{.bin} shards, auto-split at 100M tokens per shard.
|
|
\end{itemize}
|
|
|
|
In text-mode sections, individual fields shorter than 50~chars or
|
|
longer than 2M~chars are skipped during tokenization.
|
|
|
|
\subsection{Storage Backends}
|
|
|
|
Three storage backends serve the DataLoader, trading off memory
|
|
footprint against access speed for datasets ranging from fine-tuning
|
|
scale to TB-level pretraining:
|
|
|
|
\begin{itemize}[nosep]
|
|
\item \textbf{H5Store}: HDF5-based, fully loaded into RAM at init
|
|
with \texttt{share\_memory\_()} for cross-worker sharing.
|
|
Fastest random access; requires the dataset to fit in memory.
|
|
\item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files
|
|
via \texttt{np.memmap}. Data stays on disk, managed by the OS
|
|
page cache; multiple workers share physical pages without
|
|
duplication. Suited for large pretraining corpora that exceed
|
|
RAM.
|
|
\item \textbf{JsonlStore}: Reads raw \texttt{.jsonl} directly. Lazy
|
|
mode (\texttt{processor=fn}) keeps only raw text records in
|
|
memory and defers per-sample tokenization to
|
|
\texttt{fetch\_record()}, avoiding a pre-tokenized copy
|
|
entirely---used by DPO/GRPO for on-the-fly training from
|
|
source files.
|
|
\end{itemize}
|
|
|
|
A resumable distributed sampler provides seed-based shuffle with
|
|
epoch/iteration resume.
|
|
|
|
\subsection{SFT Data Cleaning}
|
|
|
|
For supervised fine-tuning (SFT), raw data requires additional curation
|
|
beyond pretraining tokenization. {\sc Alembic}~\cite{alembic} is a companion
|
|
pipeline that handles SFT data generation, cleaning, and quality scoring:
|
|
three-generation strategies (topic-driven, seed-driven, self-instruct),
|
|
built-in cleaning (HTML/URL/markdown removal, char/word repetition filters),
|
|
and a MinHash-based near-duplicate detection system~\cite{broder1997syntactic}.
|
|
Given a set of $P$ hash functions ($P=128$) and a text $T$, the MinHash
|
|
pipeline proceeds as follows:
|
|
|
|
\begin{enumerate}[nosep,leftmargin=*]
|
|
\item \textbf{Tokenization}: $T$ is split into character $n$-grams
|
|
($n=3$):
|
|
\begin{equation}
|
|
\Gamma(T) = \{\,c_i c_{i+1} c_{i+2} \mid i = 1,\dots,|T|-2 \,\}.
|
|
\end{equation}
|
|
\item \textbf{Signature}: For each hash function $h_k$, the minimum hash
|
|
value over all $n$-grams forms the $k$-th element of the fingerprint:
|
|
\begin{equation}
|
|
s_k = \min_{t \in \Gamma(T)} h_k(t), \qquad
|
|
h_k(t) = \operatorname{SHA256}(42 : k : t)_{[0:63]}.
|
|
\end{equation}
|
|
The full fingerprint is $\mathbf{s} = (s_1,\dots,s_P)$.
|
|
\item \textbf{Similarity}: The Jaccard similarity between two sets is
|
|
estimated by the fraction of agreeing fingerprint positions:
|
|
\begin{equation}
|
|
\widehat{J}(\mathbf{s}^{(a)},\mathbf{s}^{(b)}) =
|
|
\frac{|\{\,k \mid s^{(a)}_k = s^{(b)}_k \,\}|}{P}.
|
|
\end{equation}
|
|
\item \textbf{Filtering}: Samples are processed sequentially; a sample
|
|
is dropped if $\widehat{J}(\mathbf{s}, \mathbf{s}') \ge 0.7$ for any
|
|
previously kept sample $\mathbf{s}'$.
|
|
\end{enumerate}
|
|
|
|
An optional LLM-as-Judge scoring module provides multi-dimensional
|
|
quality scores that can be used to filter low-quality samples.
|
|
|
|
\subsection{DPO Data Generation}
|
|
|
|
To construct pairwise preference data for Direct Preference
|
|
Optimization, we start from a bilingual (Chinese--English) instruction
|
|
set curated by the same MinHash pipeline described above. For each
|
|
prompt $x$ in the instruction set, we generate two responses:
|
|
|
|
\begin{itemize}[nosep]
|
|
\item \textbf{Chosen} $y_w$: generated by the reference model
|
|
$\pi_{\text{ref}}$, i.e.~the SFT checkpoint after supervised
|
|
fine-tuning on the curated instruction data.
|
|
\item \textbf{Rejected} $y_l$: generated by the base model
|
|
$\pi_{\text{base}}$, i.e.~the model at the end of pretraining
|
|
before any instruction tuning.
|
|
\end{itemize}
|
|
|
|
Both generations use greedy decoding to eliminate sampling variance and
|
|
ensure that the preference signal reflects model capability rather than
|
|
decoding randomness. The resulting preference pairs
|
|
$(x, y_w, y_l)$ are stored in the same JSONL format used for SFT and
|
|
fed directly into the DPO training loop
|
|
(Section~\ref{sec:dpo}). Because the base model and the reference
|
|
model share the same architecture but differ in instruction-following
|
|
ability, the contrast between $y_w$ and $y_l$ is sharp and consistent,
|
|
which stabilises the DPO gradient updates. All instructions are
|
|
balanced across Chinese and English domains to preserve bilingual
|
|
alignment capability during preference optimisation.
|
|
\section{Model Architecture}
|
|
% ======================================================================
|
|
|
|
The model is a 24-layer decoder-only Transformer with Grouped Query Attention
|
|
(GQA)~\cite{ainslie2023gqa}, SwiGLU feed-forward blocks~\cite{shazeer2020glu},
|
|
and Rotary Position Embedding (RoPE)~\cite{su2024roformer}.
|
|
Table~\ref{tab:model_config} summarizes the configuration.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Model configuration. Total: $\sim$1.2B parameters.}
|
|
\label{tab:model_config}
|
|
\begin{tabular}{@{}lrlr@{}}
|
|
\toprule
|
|
\textbf{Parameter} & \textbf{Value} & \textbf{Parameter} & \textbf{Value} \\
|
|
\midrule
|
|
Vocabulary ($V$) & 100,000 & Hidden dim ($d$) & 1,536 \\
|
|
Layers ($L$) & 24 & FFN dim ($d_{\textit{ffn}}$) & 6,912 \\
|
|
Query heads & 24 & KV heads & 4 \\
|
|
Head dim & 64 & Max length & 2,048 \\
|
|
Norm & RMSNorm ($\epsilon=10^{-5}$) & RoPE $\theta$ & 10,000 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
With Grouped Query Attention~\cite{ainslie2023gqa} ($n_q = 24$ query
|
|
heads, $n_{kv} = 4$ key/value heads, group size $g = n_q / n_{kv} = 6$):
|
|
\begin{equation}
|
|
\begin{aligned}
|
|
\operatorname{GQA}(\mathbf{X}) &= \operatorname{Concat}\bigl(\operatorname{head}_1,\dots,\operatorname{head}_{n_q}\bigr)\mathbf{W}_O,\\[2mm]
|
|
\operatorname{head}_i &= \operatorname{Attn}\Bigl(
|
|
\mathbf{X}\mathbf{W}_Q^{(i)},\,
|
|
\mathbf{X}\mathbf{W}_K^{(\lfloor i / g \rfloor)},\,
|
|
\mathbf{X}\mathbf{W}_V^{(\lfloor i / g \rfloor)}
|
|
\Bigr),
|
|
\end{aligned}
|
|
\end{equation}
|
|
where $\operatorname{Attn}(\mathbf{Q},\mathbf{K},\mathbf{V}) =
|
|
\operatorname{Softmax}(\mathbf{Q}\mathbf{K}^{\mkern-1mu\mathsf{T}} / \sqrt{d_h})\mathbf{V}$.
|
|
Rotary Position Embedding (RoPE)~\cite{su2024roformer} encodes position
|
|
$m$ by rotating pairs of hidden dimensions:
|
|
\begin{equation}
|
|
\operatorname{RoPE}(\mathbf{x}_m)_i =
|
|
\begin{cases}
|
|
x_{m,i}\cos(m\theta_{j}) - x_{m,i+1}\sin(m\theta_{j}), & i = 2j,\\[2mm]
|
|
x_{m,i-1}\sin(m\theta_{j}) + x_{m,i}\cos(m\theta_{j}), & i = 2j+1,
|
|
\end{cases}
|
|
\end{equation}
|
|
with frequency $\theta_j = 10000^{-2j/d}$ for $j = 0,\dots,d/2-1$.
|
|
The SwiGLU~\cite{shazeer2020glu} feed-forward applies a gated Swish
|
|
non-linearity:
|
|
\begin{equation}
|
|
\operatorname{MLP}(\mathbf{x}) = \mathbf{W}_{\text{down}}\Bigl(
|
|
\mathbf{W}_{\text{up}}\mathbf{x} \odot
|
|
\operatorname{SiLU}\bigl(\mathbf{W}_{\text{gate}}\mathbf{x}\bigr)
|
|
\Bigr),
|
|
\end{equation}
|
|
where $\operatorname{SiLU}(z) = z / (1 + e^{-z})$.
|
|
|
|
Each decoder block $\ell$ then applies pre-norm residual connections:
|
|
\begin{equation}
|
|
\begin{aligned}
|
|
\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}
|
|
\end{equation}
|
|
|
|
\subsection{Initialization}
|
|
|
|
Linear weights follow $\mathcal{N}(0, 0.02)$; embeddings follow
|
|
$\mathcal{N}(0, 0.02)$. The output projection $\mathbf{W}_o$ and FFN
|
|
down-projection $\mathbf{W}_{\text{down}}$ use residual-scaled
|
|
initialization~\cite{radford2019gpt2}:
|
|
\begin{equation}
|
|
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
|
|
\end{equation}
|
|
This scaling is critical for BF16 stability (Section~\ref{sec:num-stability}).
|
|
|
|
% ======================================================================
|
|
\section{Training Configuration}
|
|
% ======================================================================
|
|
|
|
The model is trained on next-token cross-entropy loss:
|
|
\begin{equation}
|
|
\mathcal{L} = -\sum_{t=1}^{T} \log P(x_t \mid x_{<t}; \theta).
|
|
\end{equation}
|
|
|
|
Training uses a hybrid optimizer: Muon for 2D weight matrices and AdamW~\cite{loshchilov2019adamw} for 1D parameters (embeddings, biases, LayerNorm), with WSD (Warmup--Stable--Decay)
|
|
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.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Training hyperparameters for the 1.2B run.}
|
|
\label{tab:train_params}
|
|
\begin{tabular}{@{}lr@{}}
|
|
\toprule
|
|
\textbf{Hyperparameter} & \textbf{Value} \\
|
|
\midrule
|
|
Precision & BF16 (weights + optimizer states) \\
|
|
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) \\
|
|
Batch size & 4 per device $\times$ 4 GPUs $\times$ 32 accumulation \\
|
|
Sequence length & 2,048 tokens \\
|
|
Total tokens & $\sim$25B ($\approx$23k steps) \\
|
|
\bottomrule
|
|
\multicolumn{2}{@{}l@{}}{\footnotesize $^a$Muon applied to 2D weight matrices (attention projections and FFN layers);}\\
|
|
\multicolumn{2}{@{}l@{}}{\footnotesize \phantom{$^a$}AdamW applied to 1D parameters (embeddings, biases, LayerNorm scales).}\\
|
|
\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.
|
|
|
|
\begin{figure}[H]
|
|
\centering
|
|
\includegraphics[width=0.85\linewidth]{data/pt_metric.png}
|
|
\caption{Muon optimizer training dynamics: training loss and learning
|
|
rate schedule for the hybrid Muon/AdamW configuration across $\sim$25B
|
|
tokens. The Muon optimizer is applied to all 2D weight matrices
|
|
(attention projections and FFN layers), while AdamW handles 1D
|
|
parameters (embeddings, biases, and normalization scales). The WSD
|
|
schedule (warmup, stable phase, and final decay) is visible in the lower panel.}
|
|
\label{fig:pt_metric}
|
|
\end{figure}
|
|
|
|
Figure~\ref{fig:pt_metric} shows the extended training dynamics of the
|
|
Muon optimizer configuration over $\sim$25B tokens. The loss curve
|
|
exhibits the expected power-law decay in the early phase (0--5B tokens),
|
|
followed by a gradual plateau as the model approaches convergence on
|
|
the pretraining distribution. The WSD schedule holds the learning rate
|
|
constant during the long stable phase and then decays at the end of
|
|
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 $\sim$3{,}800 fine-tuning steps on a
|
|
mixed Chinese--English instruction dataset: training loss, learning
|
|
rate, and gradient norm. The smoothed loss decreases from $\sim$2.1 to
|
|
$\sim$1.5, dropping rapidly in the first 500 steps and then entering a
|
|
slower decay phase. The learning rate follows a cosine schedule with short linear warmup. Gradient norms stabilise after approximately 500 steps,
|
|
indicating that the fine-tuning process has reached a stable optimization
|
|
regime.}
|
|
\label{fig:sft_metric}
|
|
\end{figure}
|
|
|
|
Figure~\ref{fig:sft_metric} shows the supervised fine-tuning metrics
|
|
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.1 to
|
|
$\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}
|
|
\label{sec:dpo}
|
|
|
|
Following supervised fine-tuning, we align the model with pairwise
|
|
human preferences via Direct Preference Optimization
|
|
(DPO)~\cite{rafailov2023dpo}. DPO avoids explicit reward-model training
|
|
by optimizing the policy $\pi$ directly against a frozen reference
|
|
policy $\pi_{\text{ref}}$ (the SFT checkpoint). For a prompt $x$ with
|
|
preferred response $y_w$ and dispreferred response $y_l$, the loss is:
|
|
\begin{equation}
|
|
\mathcal{L}_{\text{DPO}} = -\log \sigma\!\left(
|
|
\beta \Bigl[
|
|
\log\tfrac{\pi(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)}
|
|
-\log\tfrac{\pi(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}
|
|
\Bigr]\right),
|
|
\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. The pretraining phase employs WSD (warmup--stable--decay) scheduling
|
|
to maintain a long high-learning-rate plateau. Both SFT and DPO
|
|
alignment use cosine schedules 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
|
|
observed convergence pattern in Figure~\ref{fig:dpo_metric}.
|
|
|
|
\begin{figure}[H]
|
|
\centering
|
|
\includegraphics[width=0.95\linewidth]{data/dpo_metric.png}
|
|
\caption{DPO training metrics over $\sim$3{,}000 alignment steps:
|
|
preference loss (raw and 100-step moving average), learning rate
|
|
schedule, and gradient norm.}
|
|
\label{fig:dpo_metric}
|
|
\end{figure}
|
|
|
|
Figure~\ref{fig:dpo_metric} summarises the DPO training dynamics.
|
|
The raw training loss (left panel) starts near $0.7$ and is visibly
|
|
noisy; the smoothed curve reveals a steady downward trend that reaches
|
|
$\sim$0.1--0.15 by step 3{,}000 without rebound, indicating stable
|
|
convergence. 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 200 with occasional spikes above 250, then gradually decline and
|
|
stabilise in the 35--50 range after step 1{,}000, indicating consistent
|
|
gradient magnitudes throughout alignment.
|
|
|
|
% ======================================================================
|
|
\section{Numerical Stability via Residual Scaling}
|
|
\label{sec:num-stability}
|
|
% ======================================================================
|
|
|
|
Deep Transformers trained in BF16 face numerical stability challenges from
|
|
residual variance accumulation across layers. We evaluate the GPT-2
|
|
residual-scaling initialization~\cite{radford2019gpt2} as a mitigation
|
|
strategy.
|
|
|
|
\subsection{Variance Analysis}
|
|
|
|
At initialization with $\mathcal{N}(0, 0.02)$, a linear projection output has:
|
|
\begin{equation}
|
|
\Var(\mathbf{W}\mathbf{x}) = d_{\text{in}} \cdot (0.02)^2 \cdot \Var(\mathbf{x})
|
|
= 0.6144 \cdot \Var(\mathbf{x}) \quad (\text{for } d=1536).
|
|
\end{equation}
|
|
|
|
Within one block, attention and FFN each add a residual term. The variances
|
|
at each sub-stage are:
|
|
|
|
\begin{center}
|
|
\begin{tabular}{@{}lcc@{}}
|
|
\toprule
|
|
\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 / (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 / (2L)$ \\
|
|
\midrule
|
|
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/(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,\\
|
|
\text{With scaling: } \Var(\mathbf{x}_{24}) &\approx 1 + 24 \times 0.014 = 1.34.
|
|
\end{aligned}
|
|
\end{equation}
|
|
|
|
\subsection{GPT-2 Residual Scaling}
|
|
|
|
The GPT-2 initialization~\cite{radford2019gpt2} scales output projections by
|
|
$1/\sqrt{2L}$:
|
|
\begin{equation}
|
|
\sigma_o = \sigma_{\text{down}} = 0.02 / \sqrt{2L}.
|
|
\end{equation}
|
|
|
|
This reduces per-block residual variance contribution from $0.689$ to
|
|
$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
|
|
recommend storing AdamW moments in FP32 and logging per-layer gradient
|
|
histograms during early training.
|
|
|
|
\subsection{Empirical Training Results}
|
|
|
|
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{table}[H]
|
|
\centering
|
|
\caption{Loss at 0.125B-interval milestones, 0--1B tokens.}
|
|
\label{tab:loss_milestones}
|
|
\begin{tabular}{@{}lccc@{}}
|
|
\toprule
|
|
\textbf{Tokens (B)} &
|
|
\textbf{GPT-2 scaling} &
|
|
\textbf{Kaiming init} &
|
|
\textbf{$\Delta$} \\
|
|
\midrule
|
|
0.125 & 7.37 & 7.66 & 0.29 \\
|
|
0.250 & 5.80 & 6.14 & 0.34 \\
|
|
0.375 & 4.82 & 5.38 & 0.56 \\
|
|
0.500 & 4.06 & 4.80 & 0.74 \\
|
|
0.625 & 3.50 & 4.29 & 0.79 \\
|
|
0.750 & 3.24 & 3.80 & 0.56 \\
|
|
0.875 & 3.21 & 3.43 & 0.22 \\
|
|
1.000 & 2.80 & 3.18 & 0.38 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
Table~\ref{tab:loss_milestones} quantifies the per-milestone gap. GPT-2
|
|
residual scaling leads at every interval, with $\Delta$ growing from 0.29
|
|
at 0.125B to a peak of 0.79 at 0.625B, then narrowing to 0.38 at 1B.
|
|
The widening mid-range gap aligns with the variance accumulation region
|
|
identified in the theoretical analysis (Section~\ref{sec:num-stability}).
|
|
|
|
\subsection{Post-Training Weight Distribution}
|
|
|
|
To verify that the residual-scaling constraint persists throughout
|
|
training---not just at initialization---we compare weight value
|
|
distributions across three checkpoints:
|
|
\texttt{kami-15bt} (AdamW, GPT-2 residual scaling, 15B tokens), \texttt{norm-15bt} (AdamW, Normal
|
|
init, 15B tokens), and \texttt{muon-25bt} (Muon, 25B tokens).
|
|
|
|
Figure~\ref{fig:ckpt_weight_density} shows the kernel density
|
|
estimates. Training preserves the bimodal structure: residual-scaled
|
|
projections ($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$) remain
|
|
narrowly concentrated near zero, while non-scaled weights broaden
|
|
from their initial $\mathcal{N}(0,0.02)$ distribution. Two trends
|
|
are visible in the per-component weight std
|
|
(Table~\ref{tab:weight_std}, Appendix~\ref{app:weight_std}):
|
|
|
|
\begin{itemize}[nosep]
|
|
\item \textbf{Muon produces larger post-convergence weight
|
|
variance}: the \texttt{muon-25bt} checkpoint (Muon, 25B tokens)
|
|
exhibits the largest weight std ($\sim$0.024), exceeding both
|
|
15B AdamW checkpoints ($\sim$0.015--0.021), consistent with
|
|
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.
|
|
\end{itemize}
|
|
|
|
Critically, the residual-scaled projections maintain consistently
|
|
lower standard deviations than their non-scaled counterparts across
|
|
all three checkpoints (Table~\ref{tab:weight_std}), confirming that
|
|
the $1/\sqrt{2L}$ scaling continues to enforce its design constraint
|
|
throughout training.
|
|
|
|
\begin{figure}[H]
|
|
\centering
|
|
\includegraphics[width=0.95\linewidth]{data/ckpt_weight_density.png}
|
|
\caption{Weight value density estimates across three checkpoints,
|
|
grouped by component. Residual-scaled projections
|
|
($\mathbf{W}_o$, $\mathbf{W}_{\text{down}}$) maintain narrow
|
|
distributions near zero, while non-scaled weights show broader
|
|
spread.}
|
|
\label{fig:ckpt_weight_density}
|
|
\end{figure}
|
|
|
|
|
|
% ======================================================================
|
|
\section{Conclusion}
|
|
% ======================================================================
|
|
|
|
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 a hybrid
|
|
Muon/AdamW optimizer under DDP/FSDP executors, and WSD 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. Post-training weight distribution analysis confirms that
|
|
this scaling constraint persists throughout training, with residual-scaled
|
|
projections maintaining consistently lower standard deviations than
|
|
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
|
|
with both Kaiming and Normal initializations, achieving lower training loss
|
|
and faster convergence. Within the Muon family, applying AdamW to the
|
|
embedding layer rather than Muon yields a further loss reduction and more
|
|
stable gradient norms, confirming that 1D parameters benefit from
|
|
second-moment adaptation while 2D weight matrices gain from Muon's
|
|
orthogonalisation step. An SVD-based effective rank analysis
|
|
(Appendix~\ref{app:eff_rank}) further reveals that the model
|
|
operates near its representational capacity. The complete framework and
|
|
model weights are available at \url{https://github.com/ViperEkura/AstrAI}.
|
|
|
|
% ======================================================================
|
|
\appendix
|
|
% ======================================================================
|
|
|
|
% ======================================================================
|
|
\section{IFD Data Analysis}
|
|
\label{app:ifd}
|
|
% ======================================================================
|
|
|
|
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_{<t}),\\[2mm]
|
|
L_{\text{uncond}} &= -\frac{1}{T}\sum_{t=1}^T \log P(y_t \mid y_{<t}).
|
|
\end{aligned}
|
|
\end{equation}
|
|
We compute IFD for $N=3000$ SFT samples (Alpaca-GPT4~\cite{alpaca})
|
|
using the base model (15B tokens) and the 1K-step SFT checkpoint.
|
|
After 1K SFT steps, both losses increase slightly; the mean IFD
|
|
changes from $0.8263$ (base) to $0.8485$ (1K SFT).
|
|
|
|
\subsection{Quantitative Summary}
|
|
|
|
Over $N=3000$ SFT samples from Alpaca-GPT4:
|
|
\begin{itemize}[nosep]
|
|
\item \textbf{Base model}: mean IFD $= 0.8263$,
|
|
median $= 0.8858$, std $= 0.1699$; $1.9\%$ of samples
|
|
have IFD $> 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.}
|
|
\label{tab:ifd_examples}
|
|
\small
|
|
\begin{tabular}{@{}c c c c c p{4.2cm}@{}}
|
|
\toprule
|
|
\textbf{Idx} &
|
|
\textbf{$L_{\text{cond}}^{\text{base}}$} &
|
|
\textbf{$L_{\text{uncond}}^{\text{base}}$} &
|
|
\textbf{$L_{\text{cond}}^{\text{1K}}$} &
|
|
\textbf{$L_{\text{uncond}}^{\text{1K}}$} &
|
|
\textbf{Instruction} \\
|
|
\midrule
|
|
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}
|
|
|
|
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.
|
|
|
|
\subsection{IFD Bias from Response Length}
|
|
\label{sec:ifd_bias}
|
|
|
|
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.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}
|
|
|
|
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 (base model).}
|
|
\label{tab:corr_bias}
|
|
\small
|
|
\begin{tabular}{@{}lcccccc@{}}
|
|
\toprule
|
|
& \multicolumn{2}{c}{vs.\ $L_{\text{cond}}$}
|
|
& \multicolumn{2}{c}{vs.\ $L_{\text{uncond}}$}
|
|
& \multicolumn{2}{c}{vs.\ IFD} \\
|
|
\cmidrule(lr){2-3} \cmidrule(lr){4-5} \cmidrule(lr){6-7}
|
|
\textbf{Dimension} & $r$ & $\rho$ & $r$ & $\rho$ & $r$ & $\rho$ \\
|
|
\midrule
|
|
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}
|
|
|
|
\subsection{Loss Ratio}
|
|
\label{sec:loss_ratio}
|
|
|
|
We further define the \textbf{Loss Ratio} as the fraction of conditional
|
|
loss retained after SFT:
|
|
\begin{equation}
|
|
\text{Loss Ratio} = \frac{L_{\text{cond}}^{\text{1K}}}{L_{\text{cond}}^{\text{base}}}.
|
|
\end{equation}
|
|
|
|
Over $N=3000$ samples:
|
|
\begin{itemize}[nosep]
|
|
\item Mean $= 1.106$, median $= 1.084$, std $= 0.110$;
|
|
$90.8\%$ of samples exceed $1.0$.
|
|
\item Range: $[0.768, 1.962]$; only $9.2\%$ of samples show
|
|
a decrease ($<1.0$) in conditional loss after SFT.
|
|
\end{itemize}
|
|
|
|
The predominance of loss ratio $>1$ confirms that the 1K-step SFT
|
|
checkpoint has not converged to a lower-loss region for the
|
|
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.
|
|
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, breaking the proportionality
|
|
at the sample level. This invalidates the naive expectation that a shared
|
|
numerator guarantees correlation~\cite{li2023ifd}.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Pairwise correlations between IFD variants and Loss Ratio.}
|
|
\label{tab:ifd_lr_corr}
|
|
\small
|
|
\begin{tabular}{@{}lcc@{}}
|
|
\toprule
|
|
\textbf{Pair} & Pearson $r$ & Spearman $\rho$ \\
|
|
\midrule
|
|
IFD\textsubscript{base} vs.\ IFD\textsubscript{ckpt} & $+0.97$ & $+0.96$ \\
|
|
IFD\textsubscript{base} vs.\ Loss Ratio & $-0.15$ & $-0.06$ \\
|
|
IFD\textsubscript{ckpt} vs.\ Loss Ratio & $-0.02$ & $+0.04$ \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
The near-perfect correlation between IFD\textsubscript{base} and
|
|
IFD\textsubscript{ckpt} ($r = 0.97$) reveals that the IFD ranking is
|
|
highly robust to the choice of evaluation model: samples that the
|
|
base model finds difficult remain difficult after 1K SFT steps.
|
|
This stability justifies using the base-model IFD as a data
|
|
selection signal without re-evaluating after fine-tuning.
|
|
|
|
% ======================================================================
|
|
\section{Effective Rank Analysis}
|
|
\label{app:eff_rank}
|
|
% ======================================================================
|
|
|
|
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} (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}}$,
|
|
we compute the effective rank at 99\% energy:
|
|
\begin{equation}
|
|
\text{ER@99\%} = \frac{1}{\min(m,n)} \min_k \left\{ k \;\middle|\; \frac{\sum_{i=1}^{k} \sigma_i^2}{\sum_{i=1}^{\min(m,n)} \sigma_i^2} \ge 0.99 \right\}.
|
|
\end{equation}
|
|
|
|
Table~\ref{tab:eff_rank} summarizes the results. All three checkpoints
|
|
exhibit a high overall ER@99\% ($\sim$90\%), indicating that the 1.2B
|
|
model operates close to its representational capacity. Key findings:
|
|
|
|
\begin{itemize}[nosep]
|
|
\item Q/O projections show lower ER@99\% ($\sim$0.73--0.77) and
|
|
high condition numbers ($\kappa > 10^4$), consistent with the
|
|
low-rank structure of GQA (24 query heads sharing 4 KV heads).
|
|
\item K/V projections, FFN layers, and embeddings maintain high
|
|
ER@99\% ($\sim$0.96--0.98) and low condition numbers
|
|
($\kappa < 10$).
|
|
\item The overall ER@99\% varies by less than 0.01 across
|
|
checkpoints, indicating these properties are determined
|
|
primarily by architecture rather than optimizer or training
|
|
duration.
|
|
\end{itemize}
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{SVD effective rank (ER@99\%) and mean condition number
|
|
($\kappa$) by component across three checkpoints.}
|
|
\label{tab:eff_rank}
|
|
\small
|
|
\begin{tabular}{@{}lcccccc@{}}
|
|
\toprule
|
|
& \multicolumn{3}{c}{\textbf{ER@99\%}}
|
|
& \multicolumn{3}{c}{\textbf{Cond.\ Number $\kappa$}} \\
|
|
\cmidrule(lr){2-4} \cmidrule(lr){5-7}
|
|
\textbf{Component} & \textbf{kami} & \textbf{norm} & \textbf{muon}
|
|
& \textbf{kami} & \textbf{norm} & \textbf{muon} \\
|
|
\midrule
|
|
attn.k\_proj & 0.967 & 0.971 & 0.960 & 8.4 & 5.4 & 5.7 \\
|
|
attn.o\_proj & 0.766 & 0.708 & 0.730 & 31{,}589 & 57{,}284 & 22{,}838 \\
|
|
attn.q\_proj & 0.754 & 0.764 & 0.756 & 46{,}644 & 32{,}072 & 14{,}472 \\
|
|
attn.v\_proj & 0.976 & 0.976 & 0.971 & 2.5 & 2.4 & 2.9 \\
|
|
embed\_tokens & 0.985 & 0.987 & 0.984 & 4.9 & 1.9 & 3.5 \\
|
|
lm\_head & 0.969 & 0.981 & 0.980 & 21.0 & 13.5 & 16.1 \\
|
|
mlp.down & 0.961 & 0.961 & 0.963 & 6.4 & 7.2 & 7.2 \\
|
|
mlp.gate & 0.966 & 0.967 & 0.965 & 5.7 & 5.3 & 5.7 \\
|
|
mlp.up & 0.967 & 0.968 & 0.965 & 4.9 & 5.9 & 4.9 \\
|
|
\midrule
|
|
\textbf{Overall ER@99\%} & \textbf{0.909} & \textbf{0.903} & \textbf{0.903} & & & \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
% ======================================================================
|
|
\section{Per-Component Weight Statistics}
|
|
\label{app:weight_std}
|
|
% ======================================================================
|
|
|
|
Table~\ref{tab:weight_std} reports the weight standard deviation by
|
|
component for each checkpoint, supplementing the post-training weight
|
|
distribution analysis in Section~\ref{sec:num-stability}.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Weight std by component across checkpoints. Non-scaled weights
|
|
broaden with training; residual scaling constrains drift at equal
|
|
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}
|
|
\small
|
|
\begin{tabular}{@{}lccc@{}}
|
|
\toprule
|
|
\textbf{Component} & \textbf{kami-15bt} & \textbf{norm-15bt} & \textbf{muon-25bt} \\
|
|
& (AdamW, 15B) & (AdamW, 15B) & (Muon, 25B) \\
|
|
\midrule
|
|
attn.q\_proj & 0.0154 & 0.0207 & 0.0230 \\
|
|
attn.k\_proj & 0.0153 & 0.0206 & 0.0238 \\
|
|
attn.v\_proj & 0.0146 & 0.0202 & 0.0244 \\
|
|
attn.o\_proj$^*$ & 0.0148 & 0.0084 & 0.0177 \\
|
|
mlp.up & 0.0153 & 0.0204 & 0.0237 \\
|
|
mlp.gate & 0.0155 & 0.0204 & 0.0235 \\
|
|
mlp.down$^*$ & 0.0100 & 0.0089 & 0.0180 \\
|
|
embed\_tokens & 0.0205 & 0.0205 & 0.0239 \\
|
|
lm\_head & 0.0224 & 0.0257 & 0.0298 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\\[2pt]
|
|
\footnotesize $^*$Residual-scaled projection ($\sigma_0 = 0.02/\sqrt{2L}$).
|
|
\end{table}
|
|
|
|
% ======================================================================
|
|
\begin{thebibliography}{99}
|
|
|
|
\bibitem{alpaca}
|
|
R.~Taori, I.~Gulrajani, T.~Zhang, Y.~Dubois, X.~Li, C.~Guestrin,
|
|
P.~Liang, T.~B.~Hashimoto.
|
|
Alpaca: A strong, replicable instruction-following model.
|
|
\textit{Stanford Center for Research on Foundation Models (CRFM)}, 2023.
|
|
|
|
\bibitem{ainslie2023gqa}
|
|
J.~Ainslie, J.~Lee-Thorp, M.~de Jong, Y.~Zemlyanskiy, F.~Lebr\'on, S.~Sanghai.
|
|
GQA: Training generalized multi-query transformer models from multi-head
|
|
checkpoints. \textit{EMNLP}, 2023.
|
|
|
|
\bibitem{alembic}
|
|
Alembic Contributors. \textit{Alembic: A lightweight LLM-driven SFT data
|
|
generation, cleaning, and scoring pipeline.}
|
|
\url{https://github.com/ViperEkura/Alembic}, 2026.
|
|
|
|
\bibitem{astrai}
|
|
AstrAI Contributors. \textit{AstrAI: An open-source training and inference
|
|
framework for Transformer language models.}
|
|
\url{https://github.com/ViperEkura/AstrAI}, 2026.
|
|
|
|
\bibitem{broder1997syntactic}
|
|
A.~Z.~Broder. On the resemblance and containment of documents.
|
|
\textit{SEQUENCES '97}, 1997.
|
|
|
|
\bibitem{li2023ifd}
|
|
M.~Li, Y.~Zhang, Z.~Li, J.~Chen, L.~Chen, N.~Cheng, J.~Wang, T.~Zhou, J.~Xiao.
|
|
From quantity to quality: Boosting LLM performance with self-guided data
|
|
selection for instruction tuning.
|
|
\textit{NAACL}, 2024.
|
|
|
|
\bibitem{ieee754}
|
|
IEEE Computer Society. \textit{IEEE Standard for Floating-Point Arithmetic},
|
|
IEEE Std 754-2019, 2019.
|
|
|
|
\bibitem{loshchilov2019adamw}
|
|
I.~Loshchilov, F.~Hutter. Decoupled weight decay regularization.
|
|
\textit{ICLR}, 2019.
|
|
|
|
\bibitem{radford2019gpt2}
|
|
A.~Radford, J.~Wu, R.~Child, D.~Luan, D.~Amodei, I.~Sutskever.
|
|
Language models are unsupervised multitask learners.
|
|
\textit{OpenAI Blog}, 2019.
|
|
|
|
\bibitem{rafailov2023dpo}
|
|
R.~Rafailov, A.~Sharma, E.~Mitchell, C.~D.~Manning, S.~Ermon, C.~Finn.
|
|
Direct Preference Optimization: Your language model is secretly a reward model.
|
|
\textit{NeurIPS}, 2023.
|
|
|
|
\bibitem{shazeer2020glu}
|
|
N.~Shazeer. GLU variants improve Transformer.
|
|
\textit{arXiv:2002.05202}, 2020.
|
|
|
|
\bibitem{su2024roformer}
|
|
J.~Su, A.~Murtadha, Y.~Lu, S.~Pan, B.~Wen, Y.~Liu.
|
|
Roformer: Enhanced transformer with rotary position embedding.
|
|
\textit{Neurocomputing}, 568:127063, 2024.
|
|
|
|
\end{thebibliography}
|
|
|
|
\end{document}
|