Pass the Baton: Trajectory-Relayed On-Policy Distillation

Core Methodology§

Relay On-Policy Distillation (Relay-OPD) addresses the prefix failure problem in on-policy distillation (OPD), where a student's early wrong token choices cause cascading errors. The key insight is teacher-student continuation asymmetry: given a prefix where the student has already committed an error, the teacher tends to redirect (change topic or reasoning step) while the student continues along the wrong path. This asymmetry serves as a label-free trigger for intervention.

Algorithm Overview§

During training: 1. Generate student trajectory: Sample a sequence from the student policy $\pi_s$ until a termination condition (e.g., max length, EOS). 2. Detect trigger points: Compare the student's next token logits at each step to those of the teacher $\pi_t$. A trigger is flagged when the student's top token differs from the teacher's top token and the student's confidence for its top token is high while teacher's confidence is low (empirically measured via normalized probabilities). Formally: $$\mathbb{1}_{\text{trigger}}(t) = \left[\arg\max \pi_s(\cdot|x_{<t}) \neq \arg\max \pi_t(\cdot|x_{<t})\right] \wedge \left[\pi_s(\hat{x}_t|x_{<t}) > \tau_s \right] \wedge \left[\pi_t(\hat{x}_t|x_{<t}) < \tau_t\right]$$ where $\hat{x}_t$ is student's top token, $\tau_s$ and $\tau_t$ are thresholds. 3. Relay intervention: At the first trigger point (or earliest in a limited budget), switch generation to the teacher for a fixed small number of steps $K$ (e.g., 2–4 tokens), producing a teacher leg. Then the student resumes from where the teacher left off, continuing the trajectory. 4. Train on relay trajectory: Use the entire trajectory (student prefix + teacher leg + student continuation) as the on-policy data. The student is trained with token-level KL divergence against the teacher on its own tokens, but only on 'student' portions (not on teacher leg). The loss is: $$\mathcal{L}_{\text{Relay-OPD}} = \sum_{t \in \text{student tokens}} D_{KL}(\pi_t(\cdot|x_{<t}) \,||\, \pi_s(\cdot|x_{<t}))$$

Relay Budget§

To avoid excessive deviation from the student's policy, a relay budget $B$ limits the total number of teacher tokens inserted per trajectory. $B$ is typically small (e.g., 10 tokens total) and spread across early positions to maximize effect. This concentrates intervention where prefix failures are most harmful.

Implementation Details§

• Models: Teacher is Qwen3-4B-Instruct-2507; students are Qwen3-0.6B and 1.7B (non-thinking variant, which lacks explicit chain-of-thought)
• Datasets: 8 mathematical reasoning benchmarks (GSM8K, MATH, etc.)
• Hyperparameters: Thresholds $\tau_s=0.7$, $\tau_t=0.3$; teacher leg length $K=3$; relay budget $B=9$.
• Training: Standard supervised fine-tuning setup with batch size 128, learning rate 5e-5, cosine schedule.

Pseudocode / PyTorch Snippet§

import torch
import torch.nn.functional as F

def relay_opd_step(student, teacher, input_ids, max_len, budget, K, tau_s, tau_t):
    # input_ids: initial prompt (batch, seq_len)
    student.eval()  # for generation, student is trained afterwards
    teacher.eval()
    
    batch_size = input_ids.shape[0]
    current_ids = input_ids
    trigger_count = torch.zeros(batch_size, dtype=torch.long)
    teacher_leg_active = torch.zeros(batch_size, dtype=torch.bool)
    
    while current_ids.shape[1] < max_len:
        # Forward student and teacher
        with torch.no_grad():
            s_logits = student(current_ids).logits[:, -1, :]  # batch x vocab
            t_logits = teacher(current_ids).logits[:, -1, :]
        
        s_probs = F.softmax(s_logits, dim=-1)
        t_probs = F.softmax(t_logits, dim=-1)
        s_top_token = s_probs.argmax(dim=-1)
        t_top_token = t_probs.argmax(dim=-1)
        
        # Determine trigger: student top != teacher top AND student confidence high AND teacher confidence low
        mask_match = (s_top_token != t_top_token)
        mask_s_high = (s_probs.max(dim=-1).values > tau_s)
        mask_t_low = (t_probs.gather(1, s_top_token.unsqueeze(1)).squeeze() < tau_t)
        trigger = mask_match & mask_s_high & mask_t_low & ~teacher_leg_active
        
        # For triggered samples, start teacher leg if budget remains
        can_trigger = trigger_count < budget
        start_leg = trigger & can_trigger
        
        # Update trigger count for those starting leg
        trigger_count[start_leg] = trigger_count[start_leg] + K
        teacher_leg_active[start_leg] = True
        
        # Sample next token: if teacher leg active, use teacher; else use student
        if teacher_leg_active.any():
            # For teacher leg, sample from teacher logits
            t_probs_for_sample = F.softmax(t_logits / 0.7, dim=-1)  # temperature 0.7
            next_token = torch.multinomial(t_probs_for_sample, 1).squeeze(1)
            # Decrement leg counter (a separate counter for each sample, omitted for brevity)
        else:
            s_probs_for_sample = F.softmax(s_logits / 1.0, dim=-1)
            next_token = torch.multinomial(s_probs_for_sample, 1).squeeze(1)
        
        current_ids = torch.cat([current_ids, next_token.unsqueeze(1)], dim=1)
        
        # Update teacher_leg_active: decrement counter per sample (not shown)
    
    # Now train on all student tokens (not teacher leg) with KL div
    # Separate trajectory into student and teacher segments; compute loss only on student tokens.
    # ... (omitted for brevity)
    return loss

Results§

Relay-OPD achieves state-of-the-art results on all 8 benchmarks. For the 1.7B student, average accuracy improves by +5.73% over standard OPD and +1.49% over FastOPD. Training trajectory length reduces by over 50%, demonstrating both efficiency and effectiveness.

The method is particularly effective on harder reasoning tasks where prefix failures are more frequent, and the relay budget concentrates intervention on early tokens, minimizing deviation from the student's policy.