Add optimizer/initialization comparison figure (ckpt_comparison)

This commit is contained in:
ViperEkura 2026-07-08 21:59:56 +08:00
parent ff13a45d95
commit 5726bf6c84
2 changed files with 67 additions and 0 deletions

BIN
data/ckpt_comparison.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 243 KiB

View File

@ -330,6 +330,15 @@ initialization over 15B tokens.}
Figure~\ref{fig:loss} shows both loss curves; GPT-2 residual scaling (lower 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. 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] \begin{table}[H]
\centering \centering
\caption{Loss at 0.125B-interval milestones, 0--1B tokens.} \caption{Loss at 0.125B-interval milestones, 0--1B tokens.}
@ -491,6 +500,64 @@ Response length & $-0.36$ & $-0.48$ & $-0.56$ & $-0.79$ & $+0.58$ & $+0.72$ \
\end{tabular} \end{tabular}
\end{table} \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.
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
$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.
\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{Weight Distribution by Component} \section{Weight Distribution by Component}
\label{app:weight_dist} \label{app:weight_dist}