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 , abbreviate the teacher vector as and the student vector as . Assume a shared vocabulary of tokens and positive probabilities.
Forward KL
Teacher weights the log-ratio. Student appears in the denominator.
Reverse KL
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 entered the batch.
Forward KL: the familiar q − p gradient
Let the student use temperature-scaled logits, . Its Jacobian is
The Jacobian has shape . The factor disappears only when .
Teacher entropy is constant in the local derivative. Differentiate the cross-entropy term:
The last line uses . At , 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 and keep the state and teacher fixed:
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 and teacher . At temperature 1:
| Choice | Loss | Student-logit gradient [A,B] |
|---|---|---|
| Forward KL | 0.143841036225890 nats | [0.25, −0.25] |
| Reverse KL | 0.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:
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.