03

Inside one fixed state

Hold the prefix and teacher still. Only then ask what forward and reverse KL do to student logits.

The same distributions, two orders

At one fixed state ss, abbreviate the teacher vector as pi=p(i∣s)p_i=p(i\mid s) and the student vector as qi=qθ(i∣s)q_i=q_\theta(i\mid s). Assume a shared vocabulary of MM tokens and positive probabilities.

MATH-EQ-FKL-001

Forward KL

DF(p∥q)=∑i=1Mpilog⁡piqiD_F(p\|q)=\sum_{i=1}^{M}p_i\log\frac{p_i}{q_i}

Teacher weights the log-ratio. Student appears in the denominator.

MATH-EQ-RKL-001

Reverse KL

DR(q∥p)=∑i=1Mqilog⁡qipiD_R(q\|p)=\sum_{i=1}^{M}q_i\log\frac{q_i}{p_i}

Student weights the log-ratio. Teacher appears in the denominator.

KL is asymmetric, so these scalars need not match. Both are dimensionless log-ratios; natural logs express them in nats. Neither formula records how ss entered the batch.

Forward KL: the familiar q − p gradient

Let the student use temperature-scaled logits, qi=softmax⁡(z/τ)iq_i=\operatorname{softmax}(z/\tau)_i. Its Jacobian is

MATH-EQ-SMX-001 · derived
∂qi∂zj=1τqi(δij−qj).\frac{\partial q_i}{\partial z_j}=\frac{1}{\tau}q_i(\delta_{ij}-q_j).

The Jacobian has shape M×MM\times M. The factor 1/τ1/\tau disappears only when τ=1\tau=1.

Teacher entropy is constant in the local derivative. Differentiate the cross-entropy term:

MATH-EQ-FKL-GRAD-001 · fixed-state derivative
DF(p∥q)=∑ipilog⁡pi−∑ipilog⁡qi,∂log⁡qi∂zj=1τ(δij−qj),∂DF∂zj=−1τ∑ipi(δij−qj)=1τ(qj−pj).\begin{aligned}D_F(p\|q)&=\sum_i p_i\log p_i-\sum_i p_i\log q_i,\\\frac{\partial\log q_i}{\partial z_j}&=\frac1\tau(\delta_{ij}-q_j),\\\frac{\partial D_F}{\partial z_j}&=-\frac1\tau\sum_i p_i(\delta_{ij}-q_j)\\&=\frac1\tau(q_j-p_j).\end{aligned}

The last line uses ∑ipi=1\sum_i p_i=1. At τ=1\tau=1, the gradient is simply student minus teacher probability.

Reverse KL: center the student log-ratio

In reverse KL, the student probability appears both as a weight and inside the logarithm. Set ri=log⁡(qi/pi)r_i=\log(q_i/p_i) and keep the state and teacher fixed:

MATH-EQ-RKL-GRAD-001 · fixed-state derivative
∂DR∂zj=∑i∂qi∂zj(ri+1)=1τ∑iqi(δij−qj)(ri+1)=1τqj[(rj+1)−∑iqi(ri+1)]=1τqj[log⁡qjpj−DR(q∥p)].\begin{aligned}\frac{\partial D_R}{\partial z_j}&=\sum_i\frac{\partial q_i}{\partial z_j}(r_i+1)\\&=\frac1\tau\sum_iq_i(\delta_{ij}-q_j)(r_i+1)\\&=\frac1\tau q_j\left[(r_j+1)-\sum_iq_i(r_i+1)\right]\\&=\frac1\tau q_j\left[\log\frac{q_j}{p_j}-D_R(q\|p)\right].\end{aligned}

The baseline being subtracted is the student-weighted mean log-ratio, which is reverse KL itself.

Both local gradient vectors sum to zero. Adding the same constant to every logit leaves softmax unchanged, so a common-shift direction must have zero derivative.

One state under the microscope

toy At x, student q=[0.75,0.25]q=[0.75,0.25] and teacher p=[0.5,0.5]p=[0.5,0.5]. At temperature 1:

Local x comparison
ChoiceLossStudent-logit gradient [A,B]
Forward KL0.143841036225890 nats[0.25, −0.25]
Reverse KL0.130812035941137 nats[0.205989804125271, −0.205989804125271]

The similar loss magnitudes hide different derivatives. More importantly for our thesis, neither row says whether x was reached through fixed data or student generation. That provenance lives outside this table.

The term the canonical toy leaves out

If the current parameters control both response sampling and the local loss, a full derivative of the expectation contains two terms:

MATH-EQ-SG-BOUNDARY-001 · general identity
∇θEy∼qθ[D‾θ(y)]=E ⁣[∇θD‾θ(y)+D‾θ(y)∇θlog⁡qθ(y)].\nabla_\theta\mathbb E_{y\sim q_\theta}[\overline D_\theta(y)]=\mathbb E\!\left[\nabla_\theta\overline D_\theta(y)+\overline D_\theta(y)\nabla_\theta\log q_\theta(y)\right].

The first term differentiates the local loss on a fixed sample. The second is the score-function contribution from changing the sampling measure itself.

The canonical GKD-style estimator in this article retains the first term and stops the second. That is why the next numerical receipt must freeze the xB branch while perturbing its local logits.

do not generalize MiniLLM begins from a sequence-level reverse-KL target and derives policy-gradient terms. Its equations are not a hidden version of this local derivative; chapter 5 keeps them separate.