Waypoint-OPD: Temporal Agreement for Policy-Shift Transfer

Why People Wanted It

On-policy distillation asks the student to learn from teacher behavior at prefixes the student actually visits. Direct-OPD uses the teacher’s initial-to-final change in token log-probability as an implicit reward. That endpoint is useful, but it hides the path: two teachers can have the same start and finish while taking very different intermediate routes. So the real question is not only whether the teacher’s final preference changed, but whether its early and late changes point in compatible directions.

The Core Move

Compare the teacher’s consecutive policy shifts at each student-generated prefix, then use their positive directional agreement to scale the centered start-to-finish shift:

$d_t=C_{q_t}[\delta_t^{(1)}+\delta_t^{(2)}],\qquad A_t(a)=\operatorname{sg}[\kappa_B q_t(a)g_t d_t(a)].$

Here the waypoint checkpoint supplies consistency evidence, not a target distribution to imitate.

Quiz

What role does the Waypoint checkpoint play in Waypoint-OPD?

It replaces the initial teacher as the student’s KL reference.

It provides consistency evidence that gates the endpoint shift.

It determines the student’s top-k candidate support.

It is an additional imitation target whose full distribution the student matches.

$ check

Why the Candidate Support Follows the Student

Now the version for readers who want more of the machinery: the comparison is deliberately local to actions the student is likely to take, rather than being an unweighted comparison over the entire vocabulary. At prefix $s_t$, let $K_t=\operatorname{TopK}(\pi_\theta(\cdot\mid s_t),k)$ and normalize the student probabilities over that support to obtain $q_t(a)$. Thus $q_t(a)$ is nonnegative and sums to one over $K_t$.

The three frozen teachers score each candidate. Write their consecutive log-probability changes as $\delta_t^{(1)}(a)=\log\pi_m(a\mid s_t)-\log\pi_0(a\mid s_t)$ and $\delta_t^{(2)}(a)=\log\pi_1(a\mid s_t)-\log\pi_m(a\mid s_t)$. These telescope: their sum is the initial-to-final change. This keeps the endpoint as the transferred signal while exposing the intermediate route.

A shared shift across all candidates changes absolute log probabilities but not relative preferences. Remove that component with probability-weighted centering, $C_{q_t}[x](a)=x(a)-\sum_{b\in K_t}q_t(b)x(b)$. Then $u_t=C_{q_t}[\delta_t^{(1)}]$, $v_t=C_{q_t}[\delta_t^{(2)}]$, and $d_t=u_t+v_t$. The centering makes each shift describe changes in relative candidate preference, and the weights make likely student actions count more.

Why Agreement Is a Positive Cosine

Here is the part that gets technical: Waypoint-OPD computes cosine similarity using the same student-probability weights, $c_t=\langle u_t,v_t\rangle_{q_t}/(\|u_t\|_{q_t}\|v_t\|_{q_t})$, where $\langle x,y\rangle_{q_t}=\sum_a q_t(a)x(a)y(a)$ and $\|x\|_{q_t}=\sqrt{\langle x,x\rangle_{q_t}}$. The gate is $g_t=\max(0,c_t)$ when both norms exceed $\varepsilon=10^{-8}$; if either norm is at most that threshold, the fallback is $g_t=1$.

This is a bounded, shared-per-prefix gate: negative or zero cosine suppresses the endpoint signal, positive cosine scales it, and the low-variation fallback avoids dividing by a near-zero norm. For example, if both segments equal half the endpoint, their cosine is one; opposing segments can sum to the same endpoint yet receive a zero gate. Agreement is a consistency test, not a guarantee that the resulting supervision is beneficial.

Why Batch Scaling Does Not Change the Direction

Now the final implementation detail: per-prefix gates can have different average magnitudes across batches, so all gated rewards in a batch receive the same positive multiplier. One form is $\kappa_B=\operatorname{clip}(\tau/(|B_{\mathrm{pref}}|^{-1}\sum_{t\in B_{\mathrm{pref}}}g_t+\varepsilon_{\mathrm{num}}),\kappa_{\min},\kappa_{\max})$. This targets a chosen mean gate scale while using a numerical floor and clipping bounds.

Because $\kappa_B$ is shared and positive, it preserves relative gate weights and leaves zero gates at zero; it does not constrain the norm of the resulting reward vector. The candidate advantage is dense, $A_t(a)=\operatorname{sg}[\kappa_B q_t(a)g_t d_t(a)]$, with stop-gradient so these coefficients are treated as fixed in the student update. What remains to decide is how to use that dense signal without turning the waypoint into an imitation target.

Minimal Causal Code

# One prefix; scores are log probabilities on the student's top-k support.
with torch.no_grad():
    q = student_probs / student_probs.sum()
    early = waypoint_logp - initial_logp
    late = final_logp - waypoint_logp

    def center(x):
        return x - (q * x).sum()

    u = center(early)
    v = center(late)
    d = u + v

    dot = (q * u * v).sum()
    nu = (q * u.square()).sum().sqrt()
    nv = (q * v.square()).sum().sqrt()
    eps = 1e-8
    gate = 1.0 if (nu <= eps or nv <= eps) else (dot / (nu * nv)).clamp(min=0.0)

    # kappa_batch is shared across valid prefixes in the batch.
    candidate_advantage = kappa_batch * q * gate * d

What You Gain, and What You Lose

You gain a way to distinguish endpoint shifts with different intermediate consistency, while preserving the initial-to-final shift as the signal being transferred. The student-local top-k support concentrates the comparison on likely student actions; probability weighting and centering focus it on relative preferences; the shared batch multiplier adjusts average gate scale without changing per-batch relative weights.

You also pay for an additional frozen waypoint checkpoint and scoring pass. The gate can suppress a useful endpoint signal when segment directions disagree, while its low-norm fallback deliberately treats near-constant segment changes as ungated rather than trying to infer an unstable angle. The method does not claim that agreement alone makes supervision useful. Candidate advantages are dense on the selected support, whereas the KL penalty is estimated on sampled response tokens and anchors the student to its own initialization. The reported implementation uses dual-clipped PPO, with the same gradient as the stated surrogate at the rollout policy.

Quiz

When does the agreement gate use its low-norm fallback?

Whenever the cosine is negative.

Whenever the endpoint shift is nonzero.

Whenever either centered segment norm is at most the threshold.

Whenever the batch mean gate exceeds its target.

$ check