SimpleOPD Simple Tokenizer-Agnostic On-Policy Distillation
of Long-Context Reasoning

1.1 Setup

Let \(\mathbf{x}\) denote the input conversation as a list of messages. The student and teacher with different tokenizers may use different chat templates, denoted by \(\mathcal{C}_{\theta}\) and \(\mathcal{C}_{\phi}\). We first construct the student input context \(\mathbf{c}_{\theta}\) and sample a response-token sequence with the student policy executed by the rollout engine \(\pi_{\theta_{\mathrm{roll}}}\):

\[ y_{1:n} = (y_1,\ldots,y_n) \sim \pi_{\theta_{\mathrm{roll}}}(\cdot \mid \mathbf{c}_\theta). \]

Decoding \(y_{1:n}\) via the student tokenizer gives the response surface string \(s\). Instead of passing the student context to the teacher, we reconstruct the teacher context text \(u_\phi\) using its own chat template and append the student response:

\[ c_\phi = \mathcal{C}_\phi(\mathbf{x}), \qquad u_\phi = c_\phi \oplus s, \]

where \(\oplus\) denotes string concatenation. The complete text \(u_\phi\) is then tokenized with the teacher encoder, giving \(z_{1:m}\). This lets the teacher evaluate the student response under its native tokenizer and chat template — although \(z\) and \(y\) may differ, the response text being evaluated is identical on both sides. The incremental text spans contributed by response tokens satisfy:

\[ \bigoplus_{t=1}^{n} \tau_\theta(y_t) \;=\; \bigoplus_{i=1}^{m} \tau_\phi(z_i) \;=\; s. \]

1.2 Cross-Tokenizer Alignment

Define the cumulative response text preceding each student and teacher token as

\[ P_\theta(t) = \bigoplus_{k=1}^{t-1} \tau_\theta(y_k), \qquad P_\phi(i) = \bigoplus_{k=1}^{i-1} \tau_\phi(z_k), \]

A teacher position \(i\) is aligned with a student position \(t\) if both tokenizations have consumed the same response prefix and the current tokens contribute the same text span:

\[ \mathcal{M} = \big\{\,(i,t) : P_\phi(i) = P_\theta(t) \;\wedge\; \tau_\phi(z_i) = \tau_\theta(y_t)\,\big\}. \]

Equivalently, an aligned teacher–student pair covers the same start and end offsets in the shared response string \(s\). Tokens that overlap only partially are not aligned — the log-probability of one teacher token cannot be uniquely assigned to multiple student tokens, nor can the log-probabilities of multiple teacher tokens be merged into one student token.

1.3 Aligned Teacher Log-Probabilities

For each student position \(t\), define the alignment indicator

\[ a_t = \mathbf{1}\big[\exists\, i \text{ such that } (i,t) \in \mathcal{M}\big]. \]

Let \(\log\pi_\theta\) and \(\log\pi_\phi\) be the log-probabilities of the student and teacher policies. We construct a student-length teacher target:

\[ \widetilde{\ell}_t^{\phi} = \begin{cases} \log \pi_{\phi}\big(z_{i} \mid c_{\phi},\, z_{<i}\big), & a_t = 1,\\[2pt] \log \pi_\theta\big(y_t \mid c_{\theta},\, y_{<t}\big), & a_t = 0. \end{cases} \]

Thus aligned positions inherit the teacher log-probability, while unmatched positions fall back to the student's own log-probability. We report the lexical overlap ratio

\[ \rho = \frac{|\mathcal{M}|}{n}, \]

which measures the fraction of student response tokens that receive teacher supervision. The complete alignment procedure is shown in Algorithm 1.

1.4 On-Policy Distillation Objective

The cross-tokenizer distillation objective is defined over aligned response positions:

\[ \mathcal{L}_{\mathrm{Distill}}(\theta) = \mathbb{E}_{y \sim \pi_\theta}\left[ \sum_{t=1}^{n} \log \pi_\theta(y_t \mid c_\theta, y_{<t}) - \widetilde{\ell}_t^{\phi} \right]. \]

This objective is a token-aligned surrogate for the reverse KL divergence. It compares teacher and student probabilities only at positions where the two tokenizers induce the same local segmentation of the response string. When the tokenizers are identical, every student token is aligned and the objective reduces to

\[ \mathcal{L}_{\mathrm{Distill}}(\theta) = \mathbb{E}_{y \sim \pi_\theta}\left[\log \frac{\pi_\theta(y \mid c_\theta)}{\pi_\phi(y \mid c_\phi)}\right] = D_{\mathrm{KL}}\big(\pi_\theta(\cdot \mid c_\theta) \parallel \pi_\phi(\cdot \mid c_\phi)\big). \]

When the tokenizers differ, the same response string is factorized into different token sequences, so the exact token-level KL is not directly computable. To enable multiple policy updates on the same rollout batch, we substitute the online log-probability with its pre-update counterpart for unmatched positions and use fixed policy advantages with the PPO clipped loss:

\[ \widehat{A}_t = \widetilde{\ell}_t^{\phi} - \log \pi_{\theta_{\mathrm{old}}}(y_t \mid c_\theta, y_{<t}), \qquad r_t = \frac{\pi_\theta(y_t \mid c_\theta, y_{<t})}{\pi_{\theta_{\mathrm{old}}}(y_t \mid c_\theta, y_{<t})}. \]

The resulting objective is

\[ \mathcal{L}_\theta = -\mathbb{E}\left[ \sum_{t=1}^{n} \min\Big( r_t \widehat{A}_t,\; \mathrm{clip}(r_t, 1-\epsilon, 1+\epsilon)\, \widehat{A}_t \Big) \right]. \]

1.5 Training Stabilization

Naive OPD makes students verbose and truncation-prone. Two simple fixes keep training steady: