| \documentclass[11pt]{article} |
| \usepackage{common} |
| \begin{document} |
| \notehead{0101}{Autoregressive likelihood, token vs.\ sequence view, SFT $=$ MLE $=$ forward KL}{Policy-optimisation lineage, batch M1. Companion notebook: \texttt{notebooks/0101-autoregressive-likelihood-sft.ipynb}.} |
|
|
| \begin{problem} |
| \question{What probability distribution over variable-length responses does a language model define, when is a response one object and when a chain of decisions, and what does SFT actually optimise?} |
| Every later note manipulates $\log\pi_\theta(o\mid q)$, ratios $\pi_\theta/\pi_{\theta'}$, KL divergences and score vectors $\nabla_\theta\log\pi_\theta$ of a \emph{variable-length} response. |
| Before any of that is meaningful we must (i) say exactly which probability distribution over responses a language model defines, and prove that it is one, (ii) say when we may treat a response as one object (sequence view) and when as a chain of decisions (token view), and (iii) identify what supervised fine-tuning (SFT) actually optimises, so that in note 0102 we can see precisely what RL changes. |
| \end{problem} |
|
|
| \section{Notation} |
| \begin{center}\small |
| \begin{tabular}{@{}ll@{}}\toprule |
| symbol & meaning\\\midrule |
| $q$ & prompt (context that is never generated, only conditioned on)\\ |
| $o=(o_1,\dots,o_{|o|})$ & response, $o_t\in\mathcal V$; $|o|\le L$ is the length\\ |
| $o_{<t}$ & the first $t-1$ tokens of $o$; $s_t=(q,o_{<t})$ is the \emph{state} at step $t$\\ |
| EOS & end-of-sequence token; a response ends at its first EOS or after $L$ tokens (truncation)\\ |
| $\pi_\theta(a\mid s_t)$ & next-token distribution (softmax of logits $z_\theta(s_t)\in\R^{|\mathcal V|}$)\\ |
| $\pi_\theta(o\mid q)$ & sequence probability $\prod_{t=1}^{|o|}\pi_\theta(o_t\mid s_t)$\\ |
| $p$ & data (teacher) distribution over $(q,o)$\\ |
| $H(p),\ \KL(p\|\pi)$ & entropy, $\KL(p\|\pi)=\E_p\log\frac{p}{\pi}$\\ |
| $e_a$ & one-hot vector of token $a$\\\bottomrule |
| \end{tabular} |
| \end{center} |
|
|
| \section{Derivation} |
| \subsection{A language model defines a normalised distribution on responses} |
| \paragraph{Step 1: chain rule.} By the product rule of probability, for any ordering of the coordinates, $\pi(o\mid q)=\prod_{t=1}^{|o|}\pi(o_t\mid q,o_{<t})$. An autoregressive LM \emph{parametrises each factor} by a softmax, so the product is a valid joint by construction; no approximation is involved. Taking logs, |
| \begin{equation}\log\pi_\theta(o\mid q)=\sum_{t=1}^{|o|}\log\pi_\theta(o_t\mid s_t).\label{eq:chain}\end{equation} |
|
|
| \paragraph{Step 2: normalisation (Lemma 1).} Let $\mathcal O_L$ be the set of responses that end at the first EOS or have length exactly $L$. Then $\sum_{o\in\mathcal O_L}\pi_\theta(o\mid q)=1$. |
|
|
| \emph{Proof.} Let $S_k$ be the total probability of all length-$k$ prefixes containing no EOS ("alive" mass), $S_0=1$. From each alive prefix $u$ of length $k$ the model emits EOS with probability $\pi(\mathrm{EOS}\mid q,u)$ and otherwise (probability $1-\pi(\mathrm{EOS}\mid q,u)$, because the conditional sums to one) stays alive. Hence the probability mass terminated at step $k+1$ is $D_{k+1}=S_k-S_{k+1}$ for $k+1\le L$. The truncated responses carry mass $S_L$. Summing, |
| \[\sum_{k=1}^{L}D_k+S_L=\sum_{k=1}^{L}(S_{k-1}-S_k)+S_L=S_0=1.\qquad\blacksquare\] |
| For $L=\infty$ one needs the extra assumption $S_k\to0$ (the model terminates almost surely); if it fails, $\sum_o\pi(o)=1-\lim S_k<1$ and "the probability of a response" silently leaks mass to infinite sequences. |
| \ver{$\sum_o\pi(o)=1$ to $1.1\times10^{-16}$ over all 121 sequences; each conditional sums to 1 to $2\times10^{-16}$.} |
|
|
| \subsection{Token view and sequence view} |
| Because the transition $s_{t+1}=s_t\oplus o_t$ is deterministic, generation is a Markov decision process whose state is the whole prefix and whose action is the next token (the formulation used in the Dr.\ GRPO analysis, arXiv:2503.20783 Sec.~3). Two facts follow directly from Eq.~\eqref{eq:chain}: |
| \begin{align} |
| \nabla_\theta\log\pi_\theta(o\mid q)&=\sum_{t=1}^{|o|}\nabla_\theta\log\pi_\theta(o_t\mid s_t)\qquad\text{(sequence score = sum of token scores)},\label{eq:scoresum}\\ |
| \nabla_{z}\log\mathrm{softmax}(z)_a&=e_a-\mathrm{softmax}(z).\label{eq:softmaxscore} |
| \end{align} |
| For Eq.~\eqref{eq:softmaxscore}: $\log\mathrm{softmax}(z)_a=z_a-\log\sum_b e^{z_b}$, whose derivative in $z_b$ is $\delta_{ab}-e^{z_b}/\sum_c e^{z_c}$. The chain rule through $z_\theta(s_t)$ then gives the token score $J_\theta(s_t)^{\!\top}(e_{o_t}-\pi_\theta(\cdot\mid s_t))$. |
| \ver{hand-written score (Eq.~\ref{eq:softmaxscore} scattered into the parameter slots) equals autograd for all 121 sequences, max error $1.1\times10^{-16}$.} |
|
|
| \paragraph{When to use which view.} A reward that is given once per response (a verifier verdict) lives naturally in the sequence view; ratios, clipping and KL penalties are applied per token (token view). The conversion between the two is Eq.~\eqref{eq:scoresum}, and \emph{every} subtlety in notes 0103--0108 is a statement about when replacing a sequence-level quantity by a product or sum of token-level ones is exact and when it is an approximation. |
|
|
| \subsection{SFT is maximum likelihood is forward KL} |
| Given pairs $(q,o)\sim p$, SFT minimises the negative log-likelihood |
| \begin{equation}\mathcal L_{\rm SFT}(\theta)=-\E_{(q,o)\sim p}\Big[\sum_{t=1}^{|o|}\log\pi_\theta(o_t\mid s_t)\Big].\label{eq:sft}\end{equation} |
| \paragraph{Step 1 (cross-entropy decomposition).} Insert $\log\pi=\log p-\log\frac{p}{\pi}$: |
| \[-\E_p\log\pi_\theta=-\E_p\log p+\E_p\log\frac{p}{\pi_\theta}=H(p)+\KL(p\,\|\,\pi_\theta).\] |
| $H(p)$ does not depend on $\theta$, so $\arg\min_\theta\mathcal L_{\rm SFT}=\arg\min_\theta\KL(p\|\pi_\theta)$: \emph{maximum likelihood is minimisation of the forward KL}. |
|
|
| \paragraph{Step 2 (Gibbs' inequality, so that the minimiser is meaningful).} By Jensen on the concave $\log$, |
| $-\KL(p\|\pi)=\E_p\log\frac{\pi}{p}\le\log\E_p\frac{\pi}{p}=\log\sum_{o:p(o)>0}\pi(o)\le0$, with equality iff $\pi=p$ on $\mathrm{supp}\,p$. If the family contains $p$, the unique minimiser reproduces it. |
| \ver{$\mathrm{CE}=H+\KL$ to $0$ (fp64 identical); L-BFGS from a seeded random start drives $\KL(p\|\pi_\theta)$ to $5\times10^{-17}$ (round-off level) and recovers the teacher sequence distribution to $2\times10^{-9}$ (max abs.\ probability error).} |
|
|
| \paragraph{Step 3 (token-level decomposition: teacher forcing).} Write $\log\frac{p(o)}{\pi(o)}=\sum_{t\le|o|}\log\frac{p(o_t\mid s_t)}{\pi(o_t\mid s_t)}$, take $\E_p$ and condition on $s_t$ (tower rule). The event $\{|o|\ge t\}$ (``the response has not ended before step $t$'') is determined by $o_{<t}$, hence by $s_t$, so the random number of terms is harmless: |
| \begin{equation}\KL(p\|\pi_\theta)=\sum_{t=1}^{L}\E_{o\sim p}\Big[\indb{|o|\ge t}\,\KL\big(p(\cdot\mid s_t)\,\big\|\,\pi_\theta(\cdot\mid s_t)\big)\Big].\label{eq:klchain}\end{equation} |
| (The weight $P_p(s_t)$ of a context is its probability of being reached, a sub-probability measure over contexts that do not terminate early; this is what ``$\E_{s_t\sim p}$'' means in the rest of the batch.) So SFT fits the next-token distribution at every context \emph{that the data visits}. Note what is \emph{not} in it: the contexts are drawn from $p$, never from $\pi_\theta$. |
| \ver{Eq.~\eqref{eq:klchain} holds to $2.2\times10^{-16}$ on the enumerated model.} |
|
|
| \paragraph{Step 4 (gradient).} Differentiating Eq.~\eqref{eq:sft} and using Eq.~\eqref{eq:scoresum}, |
| \begin{equation}\nabla_\theta\mathcal L_{\rm SFT}=-\E_{p}\Big[\sum_t\nabla_\theta\log\pi_\theta(o_t\mid s_t)\Big],\label{eq:sftgrad}\end{equation} |
| an expectation under the \emph{data}, so no samples are drawn from the model, there is no ratio and no score-function variance from sampling. At $\theta^\star$ with $\pi_{\theta^\star}=p$ it equals $-\E_{\pi_{\theta^\star}}[\nabla\log\pi]=0$ (the zero-mean score identity proved in 0102). In RL (0102) the same score vector is averaged under $\pi_\theta$ itself and weighted by a reward-derived advantage. |
| \ver{grad CE hand $=-\E_p[\text{score}]$ equals autograd to $9.7\times10^{-17}$; vanishes at the teacher to $4.9\times10^{-17}$; MC gradient from 400k samples within $1.8$ standard errors.} |
|
|
| \subsection{Loss normalisation is a modelling choice, not a detail} |
| Implementations rarely compute Eq.~\eqref{eq:sft} verbatim. Let $\ell_{t}=-\log\pi_\theta(o_t\mid s_t)$. Four conventions: |
| \begin{center}\small\begin{tabular}{@{}lll@{}}\toprule |
| name & loss per response / batch & weight on token $t$ of response $o$\\\midrule |
| token-sum & $\sum_t\ell_t$ & $1$\\ |
| sequence-mean & $\frac1{|o|}\sum_t\ell_t$ & $1/|o|$\\ |
| batch token-mean & $\frac{1}{\sum_i|o_i|}\sum_i\sum_t\ell_{i,t}$ & $1/\sum_i|o_i|$ (random)\\ |
| constant denominator & $\frac1{L}\sum_t\ell_t$ & $1/L$\\\bottomrule\end{tabular}\end{center} |
| Only token-sum and the constant denominator give Eq.~\eqref{eq:sftgrad} (the latter divided by the constant $L$). Sequence-mean reweights tokens by $1/|o|$: short responses are up-weighted, long ones down-weighted, and the resulting vector is \emph{not} a rescaling of the MLE gradient. |
| \ver{constant-denominator gradient $\times L$ equals MLE gradient exactly; sequence-mean differs from the (rescaled) MLE gradient by 44\% in relative norm on the toy model.} |
| This is irrelevant for SFT on i.i.d.\ data at moderate lengths, but it becomes a bias in RL: 0105 shows that the $1/|o|$ of GRPO is exactly such a re-weighting, and 0106 shows what DAPO changes. |
|
|
| \section{Boxed result} |
| \begin{result} |
| \[\boxed{\;\mathcal L_{\rm SFT}(\theta)=-\E_p\sum_{t=1}^{|o|}\log\pi_\theta(o_t\mid s_t)=H(p)+\KL(p\|\pi_\theta)\;}\] |
| \[\boxed{\;\nabla_\theta\mathcal L_{\rm SFT}=-\E_p\sum_t J_\theta(s_t)^{\!\top}\big(e_{o_t}-\pi_\theta(\cdot\mid s_t)\big)\;}\] |
| with $\sum_{o}\pi_\theta(o\mid q)=1$ and $\log\pi_\theta(o\mid q)=\sum_t\log\pi_\theta(o_t\mid s_t)$. |
| \end{result} |
|
|
| \section{What breaks} |
| \begin{breaks}\small |
| \begin{itemize} |
| \item \textbf{Exposure bias.} Eq.~\eqref{eq:klchain} only controls $\pi_\theta(\cdot\mid s)$ for $s\sim p$. At generation time contexts come from $\pi_\theta$, where the fit is unconstrained. \emph{Violated assumption:} train and test context distributions coincide. This is the first reason RL on own samples exists. |
| \item \textbf{Mass-covering.} $\KL(p\|\pi)=\infty$ if $\pi(o)=0<p(o)$, so a capacity-limited $\pi$ spreads mass over everything $p$ supports, including the gaps between modes. \emph{Violated assumption:} realisability ($p\in\{\pi_\theta\}$) is what makes the minimiser exact. |
| \item \textbf{Leaky normalisation.} Lemma 1 for $L=\infty$ needs almost-sure termination; truncated responses (no EOS) are legitimate members of $\mathcal O_L$ but are different events from terminated ones, which matters when a verifier needs the final answer (0110). |
| \item \textbf{Loss-normalisation bias.} Using sequence-mean or a random batch-token-mean denominator silently changes the gradient (table above); for RL this becomes the length bias of 0105. |
| \item \textbf{Sampler $\ne$ model.} Temperature, top-$p$, and low-precision inference kernels mean the distribution actually sampled is not $\pi_\theta$; see 0108. |
| \item \textbf{Empirical $\hat p$.} With finitely many samples the minimiser of $\KL(\hat p\|\pi)$ is the empirical distribution (memorisation); the population statements above hold only in expectation (the MC check shows the gradient converges at the $1/\sqrt n$ rate). |
| \end{itemize}\end{breaks} |
|
|
| \nextnote{SFT averages the score under $p$ and weights every token equally. If instead we average the \emph{same score} under the model's own samples and weight each sequence by how good it was, we get policy gradient. Note 0102 derives it from first principles, proves that the score has zero mean, and proves when a baseline is free.} |
|
|
| \section*{References} |
| {\small\begin{itemize}[leftmargin=1.2em] |
| \item Liu et al., \emph{Understanding R1-Zero-Like Training: A Critical Perspective} (Dr.\ GRPO), arXiv:2503.20783: Sec.~3 (token-level MDP, Eq.~1), Listing 1 and Table~2 (loss normalisation), Appendix A, Eqs.~4--5 (policy-gradient derivation). |
| \item Qi et al., \emph{Defeating the Training--Inference Mismatch via FP16}, arXiv:2510.26788, Sec.~2, Eqs.~1--2 (objective and REINFORCE estimator in prompt-conditional form). |
| \item Gibbs' inequality, Jensen's inequality, the KL chain rule and the softmax derivative are standard facts proved above; no external citation is needed. |
| \end{itemize}} |
| \end{document} |
|
|
