Where the measure lives

A sampled trajectory induces the prefix states. The local loss arrives only after those states exist.

Turn trajectories into a distribution over visited prefixes

The autoregressive product tells us how a response is generated. To expose the measure beneath the loss, we need one more object: the distribution over prefixes induced by a response source.

Let $\rho$ be any distribution over prompt–response pairs. Its normalized prefix occupancy is

MATH-EQ-OCC-001 · derived

d_\rho(s)=\mathbb E_{(x,y)\sim\rho}\!\left[\frac{1}{T(y)}\sum_{t=1}^{T(y)}\mathbf 1\{s=(x,y_{<t})\}\right].

Sample a sequence from $\rho$ , give each of its $T(y)$ visited positions weight $1/T(y)$ , then ask how much mass lands on state $s$ .

length boundary This preserves GKD's per-sequence mean. Pooling all tokens across variable-length sequences would weight long responses more heavily. The toy fixes two evaluated positions and does not define a universal EOS rule.

Do not let five operations collapse into one verb

“Distill on student outputs” compresses an entire computation. Expanding it reveals where the source measure stops and the local derivative starts.

MATH-FIG-001 · sourced + teaching decompositionFive operations, two kinds of arrow

01SampleStudent emits a response token. The draw is discrete.
02VisitThe emitted token becomes part of the next prefix.
03EvaluateTeacher and student are queried at that same prefix.
04CompareA chosen local divergence returns a scalar in nats.
05UpdateDifferentiate local student logits; the sampled path stays fixed.

solid data and state flow dashed gradient/update

Read the edge types. The sampled token travels along the solid path. The gradient begins only after a prefix has been fixed and the local loss has been formed. This five-label decomposition is ours; its operations follow GKD Eq. 2, Eq. 4, and Algorithm 1.

The separation is useful even if the real implementation fuses teacher and student forward passes. Conceptually, collection decides which states exist in the batch; the local loss decides how probability mass is compared at each state.

The toy makes one branch reproducible

OPD-TOY-001 At prefix x, the student distribution is [0.75, 0.25] in vocabulary order [A,B]. The contract supplies a uniform value of 0.8, so inverse-CDF sampling selects B and the next visited prefix is xB.

MATH-FIG-003 · toyA prefix string does not carry its own provenance

prefixxstudent [0.75, 0.25]

A · probability 0.75xAsupplied fixed occurrenceCould also be student-generated under another draw.

B · probability 0.25xBrealized student occurrenceu_token = 0.8 lands in [0.75, 1).

Exact branch intervals
Token	Student interval at x	This fixture
A	[0, 0.75)	not selected by u = 0.8
B	[0.75, 1)	selected → xB

The contract labels the supplied xA occurrence as the off-policy comparison. That label belongs to how this occurrence was collected—not to the characters “xA.”

The draw is fixed so we can audit every later number. It is still a sample realization, not the exact expectation over all student responses. At the same time, the displayed xA comparison should not be treated as intrinsically off-policy: another student draw could reach the same string.

Now place a divergence on the selected path

For each state $s_t=(x,y_{<t})$ on the selected response, query teacher $p$ and student $q_\theta$ at that identical context. GKD's token-distribution reduction is

MATH-EQ-REDUCE-001 · sourced form

\overline D(\theta;x,y)=\frac{1}{T(y)}\sum_{t=1}^{T(y)}D\!\left(p(\cdot\mid s_t)\,\middle\|\,q_\theta(\cdot\mid s_t)\right).

The response selects the states. The symbol $D$ then compares two full next-token distributions at each one and returns a scalar.

The author manuscript's Eq. 2 uses exactly this position-average structure. It does not force $D$ to be reverse KL.

Cross the source measure with the local choice

Write $\mathcal D_{\mathrm{fix}}$ for a fixed prompt–response distribution and $q_{\bar\theta}$ for a behavior-policy snapshot. At collection, $\bar\theta$ may equal the current student. During the canonical GKD-style update, the sampled response is held fixed:

MATH-EQ-MIX-001 · GKD factorization

\begin{aligned}L_\lambda(\theta;\bar\theta)={}&(1-\lambda)\,\mathbb E_{(x,y)\sim\mathcal D_{\mathrm{fix}}}[\overline D(\theta;x,y)]\\&+\lambda\,\mathbb E_{x\sim X,\,y\sim q_{\bar\theta}(\cdot\mid x)}[\overline D(\theta;x,y)].\end{aligned}

$\lambda$ is unitless and controls the expected contribution of student-generated output data. It does not occur inside the definition of forward or reverse KL.

Source special cases in GKD
λ	Output source	Careful reading
0	fixed / supervised	Not automatically SeqKD: fixed `Y` may be ground truth or teacher-generated, and the local objective is declared separately.
0.5	mixed	Equal expected source weight. Algorithm 1 chooses a source for a batch, so a finite batch history need not contain an exact half.
1	student-generated	On-policy for the fresh behavior policy at collection; prompts still come from `X`.

The paper evaluates these three settings and crosses them with forward KL, reverse KL, and several JSD coefficients. See §4, “GKD Variants.”

Detach the path, not the local student

Agarwal et al. state after Eq. 4 that they do not backpropagate through the student sampling distribution. In the toy, that means the 0.8 → B → xB branch is frozen during the update. It does not mean the student distribution at xB is constant.

HELD FIXED

prompt $x$ ; realized prefixes; teacher probabilities; vocabulary alignment; temperature; categorical outcome

DIFFERENTIATED

student logits $z_\theta(s_t)$ ; then student parameters through each state's network Jacobian

The next chapter performs precisely that fixed-state derivative.