Add GQA/SwiGLU/RoPE formula blocks, Alpaca-GPT4 citation, and improve IFD-Loss Ratio interpretation
This commit is contained in:
parent
5c915514b3
commit
4a7bf7ac3f
Binary file not shown.
|
After Width: | Height: | Size: 231 KiB |
104
main.tex
104
main.tex
|
|
@ -154,7 +154,8 @@ An IFD $>1$ indicates the instruction increases the loss relative to
|
||||||
unconditional generation (the model struggles to follow it), while
|
unconditional generation (the model struggles to follow it), while
|
||||||
IFD $<1$ means the instruction provides useful guidance.
|
IFD $<1$ means the instruction provides useful guidance.
|
||||||
|
|
||||||
We compute IFD for $N=3000$ SFT samples using both the pretrained
|
We compute IFD for $N=3000$ SFT samples drawn from the
|
||||||
|
Alpaca-GPT4 dataset~\cite{alpaca} using both the pretrained
|
||||||
base model (after 15B tokens of pretraining) and a supervised
|
base model (after 15B tokens of pretraining) and a supervised
|
||||||
fine-tuned checkpoint (after 1K SFT steps).
|
fine-tuned checkpoint (after 1K SFT steps).
|
||||||
Figure~\ref{fig:ifd} shows the distribution.
|
Figure~\ref{fig:ifd} shows the distribution.
|
||||||
|
|
@ -184,13 +185,59 @@ $5.3\times$ more than unconditional loss, confirming that SFT teaches
|
||||||
instruction following rather than merely improving generic language
|
instruction following rather than merely improving generic language
|
||||||
modeling. Detailed analysis is provided in Appendix~\ref{app:ifd}.
|
modeling. Detailed analysis is provided in Appendix~\ref{app:ifd}.
|
||||||
|
|
||||||
|
\subsubsection{IFD vs.\ Loss Ratio}
|
||||||
|
|
||||||
|
We further define the \emph{loss ratio}---the fraction of
|
||||||
|
conditional loss retained after SFT---as:
|
||||||
|
\begin{equation}
|
||||||
|
\text{Loss Ratio} = \frac{L_{\text{cond}}^{\text{ckpt}}}{L_{\text{cond}}^{\text{base}}}.
|
||||||
|
\end{equation}
|
||||||
|
Table~\ref{tab:ifd_lossratio_corr} reports the pairwise correlations.
|
||||||
|
|
||||||
|
\begin{table}[H]
|
||||||
|
\centering
|
||||||
|
\caption{Pairwise correlations among IFD and Loss Ratio.}
|
||||||
|
\label{tab:ifd_lossratio_corr}
|
||||||
|
\small
|
||||||
|
\begin{tabular}{@{}lcc@{}}
|
||||||
|
\toprule
|
||||||
|
\textbf{Pair} & \textbf{Pearson $r$} & \textbf{Spearman $\rho$} \\
|
||||||
|
\midrule
|
||||||
|
IFD\textsubscript{base} vs.\ Loss Ratio & $+0.10$ & $+0.05$ \\
|
||||||
|
IFD\textsubscript{ckpt} vs.\ Loss Ratio & $+0.90$ & $+0.91$ \\
|
||||||
|
IFD\textsubscript{base} vs.\ IFD\textsubscript{ckpt} & $+0.38$ & $+0.49$ \\
|
||||||
|
\bottomrule
|
||||||
|
\end{tabular}
|
||||||
|
\end{table}
|
||||||
|
|
||||||
|
The near-perfect correlation between IFD\textsubscript{ckpt} and
|
||||||
|
Loss Ratio ($r = 0.90$) reflects a mathematical near-identity:
|
||||||
|
both are dominated by $L_{\text{cond}}^{\text{ckpt}}$ in the
|
||||||
|
numerator. Consequently, IFD\textsubscript{ckpt} is
|
||||||
|
redundant---it essentially measures how much the conditional loss
|
||||||
|
has dropped after SFT, i.e., the learning speed of each sample. In contrast, IFD\textsubscript{base} and Loss
|
||||||
|
Ratio are nearly orthogonal ($r = 0.10$), forming a complementary
|
||||||
|
two-dimensional screening space: IFD\textsubscript{base} measures
|
||||||
|
``how hard does the base model find this,'' while Loss Ratio
|
||||||
|
measures ``how much did SFT improve it.'' Samples with high
|
||||||
|
IFD\textsubscript{base} \emph{and} low Loss Ratio are the most
|
||||||
|
informative for training.
|
||||||
|
|
||||||
|
\begin{figure}[H]
|
||||||
|
\centering
|
||||||
|
\includegraphics[width=0.80\linewidth]{data/ifd_both_vs_lossratio.png}
|
||||||
|
\caption{IFD\textsubscript{base} vs.\ Loss Ratio (left),
|
||||||
|
IFD\textsubscript{ckpt} vs.\ Loss Ratio (right).}
|
||||||
|
\label{fig:ifd_lossratio}
|
||||||
|
\end{figure}
|
||||||
|
|
||||||
% ======================================================================
|
% ======================================================================
|
||||||
\section{Model Architecture}
|
\section{Model Architecture}
|
||||||
% ======================================================================
|
% ======================================================================
|
||||||
|
|
||||||
The model is a 24-layer decoder-only Transformer with Grouped Query Attention
|
The model is a 24-layer decoder-only Transformer with Grouped Query Attention
|
||||||
(GQA)~\cite{ainslie2023gqa} and SwiGLU feed-forward blocks~\cite{shazeer2020glu},
|
(GQA)~\cite{ainslie2023gqa}, SwiGLU feed-forward blocks~\cite{shazeer2020glu},
|
||||||
with Rotary Position Embedding (RoPE)~\cite{su2024roformer}.
|
and Rotary Position Embedding (RoPE)~\cite{su2024roformer}.
|
||||||
Table~\ref{tab:model_config} summarizes the configuration.
|
Table~\ref{tab:model_config} summarizes the configuration.
|
||||||
|
|
||||||
\begin{table}[H]
|
\begin{table}[H]
|
||||||
|
|
@ -210,12 +257,49 @@ Norm & RMSNorm ($\epsilon=10^{-5}$) & RoPE $\theta$ & 10,000 \\
|
||||||
\end{tabular}
|
\end{tabular}
|
||||||
\end{table}
|
\end{table}
|
||||||
|
|
||||||
Each decoder block $\ell$ computes:
|
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{equation}
|
||||||
\mathbf{h}_\ell = \mathbf{x}_\ell + \operatorname{GQA}\bigl(\operatorname{RMSNorm}(\mathbf{x}_\ell)\bigr), \qquad
|
\begin{aligned}
|
||||||
\mathbf{x}_{\ell+1} = \mathbf{h}_\ell + \operatorname{MLP}\bigl(\operatorname{RMSNorm}(\mathbf{h}_\ell)\bigr),
|
\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),
|
||||||
|
\label{eq:swiglu}
|
||||||
|
\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}
|
||||||
|
\label{eq:decoder_block}
|
||||||
\end{equation}
|
\end{equation}
|
||||||
where $\operatorname{MLP}(\mathbf{x}) = \mathbf{W}_{\text{down}}(\mathbf{W}_{\text{up}}\mathbf{x} \odot \operatorname{SiLU}(\mathbf{W}_{\text{gate}}\mathbf{x}))$.
|
|
||||||
|
|
||||||
\subsection{Initialization}
|
\subsection{Initialization}
|
||||||
|
|
||||||
|
|
@ -589,6 +673,12 @@ in the high-IFD tail.
|
||||||
% ======================================================================
|
% ======================================================================
|
||||||
\begin{thebibliography}{99}
|
\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}
|
\bibitem{ainslie2023gqa}
|
||||||
J.~Ainslie, J.~Lee-Thorp, M.~de Jong, Y.~Zemlyanskiy, F.~Lebr\'on, S.~Sanghai.
|
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
|
GQA: Training generalized multi-query transformer models from multi-head
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue