Add effective rank analysis, weight distribution evolution, Muon/SFT training figures
- Add SVD effective rank (ER@99%) and condition number analysis across three checkpoints (kami-15bt, norm-15bt, muon-25bt) in appendix - Add post-training weight distribution analysis as §5.4, verifying residual scaling persists throughout training - Add per-component weight std table in appendix showing Muon constrains early-stage drift vs Normal init - Add Muon optimizer training dynamics figure (muon_pt.png) - Add SFT training metrics figure (sft_metrics.png) on mixed CN-EN data - Restructure: move effective rank to appendix, weight distribution to §5.4, delete standalone §6, tighten abstract to prose-only - Update conclusion to reference new analyses
This commit is contained in:
Binary file not shown.
|
After Width: | Height: | Size: 940 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 315 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 128 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 176 KiB |
@@ -32,14 +32,23 @@ 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
|
||||
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,
|
||||
RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and cosine scheduling under DDP/FSDP.
|
||||
A focused BF16 stability analysis shows that GPT-2 residual scaling
|
||||
($\sigma = 0.02/\sqrt{2L}$) reduces per-block residual variance by a factor
|
||||
of 48, containing post-24-layer variance at 1.34 versus 17.5 under standard
|
||||
initialization. Empirically, this scaling yields a sustained loss advantage
|
||||
over Kaiming initialization, with the gap peaking at $\Delta = 0.79$ in the
|
||||
mid-training regime.
|
||||
deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA,
|
||||
SwiGLU, RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and
|
||||
cosine scheduling under DDP/FSDP. A BF16 stability analysis shows that
|
||||
GPT-2 residual scaling substantially reduces per-block residual variance
|
||||
accumulation, keeping post-training variance well below the overflow
|
||||
threshold of standard initialization; empirically this yields a sustained
|
||||
loss advantage over Kaiming initialization throughout training.
|
||||
Post-training weight distribution analysis across three
|
||||
checkpoints---varying optimizer, initialization, and training
|
||||
budget---confirms that residual-scaled projections maintain narrow
|
||||
distributions throughout training, preserving the numerical stability
|
||||
established at initialization. An SVD-based effective rank analysis
|
||||
further reveals that the model operates near its representational
|
||||
capacity, with attention Q/O projections 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}
|
||||
|
||||
% ======================================================================
|
||||
@@ -264,6 +273,46 @@ initialization over $\sim$20B tokens.}
|
||||
|
||||
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/muon_pt.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 cosine
|
||||
schedule with 2\% warmup is visible in the lower panel.}
|
||||
\label{fig:muon_pt}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:muon_pt} 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 cosine learning rate schedule
|
||||
reaches its minimum at the end of training, ensuring stable weight
|
||||
updates in the final phase.
|
||||
|
||||
\begin{figure}[H]
|
||||
\centering
|
||||
\includegraphics[width=0.85\linewidth]{data/sft_metrics.png}
|
||||
\caption{SFT training metrics over 1{,}000 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
|
||||
steps and then enters a slower decay phase. The learning rate
|
||||
follows a cosine schedule with 2\% warmup. Gradient norms stabilize
|
||||
after approximately 300 steps, indicating that the fine-tuning
|
||||
process has reached a stable optimization regime.}
|
||||
\label{fig:sft_metrics}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:sft_metrics} shows the supervised fine-tuning metrics
|
||||
for the 1K-step 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
|
||||
$\sim$1.6 over 1{,}000 steps, with the gradient norm converging to a
|
||||
stable range after the warmup phase.
|
||||
|
||||
% ======================================================================
|
||||
\section{Numerical Stability via Residual Scaling}
|
||||
\label{sec:num-stability}
|
||||
@@ -362,6 +411,60 @@ 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} (Muon, 15B tokens), \texttt{norm-15bt} (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{Training duration drives variance growth}: the
|
||||
\texttt{muon-25bt} checkpoint exhibits the largest weight std
|
||||
($\sim$0.024 for attention projections), exceeding both 15B
|
||||
checkpoints ($\sim$0.015--0.021), reflecting continued weight
|
||||
drift from $\sigma_0 = 0.02$ as training progresses.
|
||||
\item \textbf{Muon constrains early-stage drift}: at equal training
|
||||
budget (15B tokens), the Muon-trained \texttt{kami-15bt} shows
|
||||
\emph{smaller} std ($\sim$0.015) than the Normal-init
|
||||
\texttt{norm-15bt} ($\sim$0.021). Muon's orthogonalization
|
||||
step normalizes the update direction, constraining per-step
|
||||
weight movement. At 25B tokens, cumulative steps overtake
|
||||
this per-step constraint, producing the largest overall std.
|
||||
\end{itemize}
|
||||
|
||||
Critically, the residual-scaled projections remain bounded at
|
||||
$\sigma \approx 0.003$ across all checkpoints regardless of training
|
||||
duration, 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}
|
||||
|
||||
\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}
|
||||
% ======================================================================
|
||||
@@ -369,11 +472,17 @@ 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 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
|
||||
weights are available at \url{https://github.com/ViperEkura/AstrAI}.
|
||||
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. Post-training weight distribution analysis confirms that
|
||||
this scaling constraint persists throughout training, with residual-scaled
|
||||
projections maintaining narrow distributions ($\sigma \approx 0.003$)
|
||||
regardless of optimizer or training duration. 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
|
||||
@@ -553,32 +662,104 @@ This stability justifies using the base-model IFD as a data
|
||||
selection signal without re-evaluating after fine-tuning.
|
||||
|
||||
% ======================================================================
|
||||
\section{Weight Distribution by Component}
|
||||
\label{app:weight_dist}
|
||||
\section{Effective Rank Analysis}
|
||||
\label{app:eff_rank}
|
||||
% ======================================================================
|
||||
|
||||
Figure~\ref{fig:weight_dist} shows the distribution of weight
|
||||
magnitudes at initialization, grouped by component type. Embeddings
|
||||
and non-residual-scaled projections (QKV, attention output, FFN
|
||||
gate/up) follow $\mathcal{N}(0, 0.02)$, producing near-identical
|
||||
bell curves centered at zero. The residual-scaled projections
|
||||
(output projection $\mathbf{W}_o$ and FFN down-projection
|
||||
$\mathbf{W}_{\text{down}}$) use $\sigma = 0.02 / \sqrt{2L} \approx 0.0029$,
|
||||
visible as the narrow, sharply peaked distribution concentrated
|
||||
near zero. This factor-48 variance reduction is the mechanism by
|
||||
which GPT-2 residual scaling prevents BF16 underflow in deep
|
||||
Transformers (Section~\ref{sec:num-stability}).
|
||||
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}
|
||||
(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}
|
||||
|
||||
\begin{figure}[H]
|
||||
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
|
||||
\includegraphics[width=0.85\linewidth]{data/weight_dist_by_component.png}
|
||||
\caption{Weight distribution by component at initialization.
|
||||
Each panel shows the histogram of weight values for a specific
|
||||
module group (embedding, attention projections, FFN projections,
|
||||
output projections). The narrow peaks correspond to the
|
||||
residual-scaled $\mathbf{W}_o$ and $\mathbf{W}_{\text{down}}$
|
||||
projections.}
|
||||
\label{fig:weight_dist}
|
||||
\end{figure}
|
||||
\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 duration; Muon constrains drift at
|
||||
equal token count (15B) but is overtaken by longer training (25B).
|
||||
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} \\
|
||||
& (Muon, 15B) & (Normal, 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}
|
||||
|
||||
Reference in New Issue
Block a user