Latent On-Policy Self-Distillation — part 2/4
Cleaned-text code points [14206, 30189). Bibliography omitted here; source.md includes it.
Given x x and the retrieved set ℰ \mathcal{E} , the composer first encodes each item into hidden states: 𝐇 j = Enc ψ ( x , m j ) ∈ ℝ T j × d , \mathbf{H}_{j}=\operatorname{Enc}_{\psi}(x,m_{j})\in\mathbb{R}^{T_{j}\times d}, (5) where m j m_{j} is the j j -th retrieved trajectory, Enc ψ \operatorname{Enc}_{\psi} is the encoder that maps the task–experience pair into hidden states, T j T_{j} is the encoded sequence length, and d d is the hidden dimension. The variable-length states are compressed into a fixed number of latent tokens by a lightweight latent compressor. In our implementation, this compressor is QFormer-style cross-attention with learned queries, although other latent-token generators can be used under the same abstraction: 𝐄 j = Comp χ ( 𝐇 j ; 𝐐 χ ) = [ ⟨ e j , 1 ⟩ , … , ⟨ e j , K ⟩ ] ∈ ℝ K × d , \mathbf{E}_{j}=\operatorname{Comp}_{\chi}(\mathbf{H}_{j};\mathbf{Q}_{\chi})=[\langle e_{j,1}\rangle,\ldots,\langle e_{j,K}\rangle]\in\mathbb{R}^{K\times d}, (6) where Comp χ \operatorname{Comp}_{\chi} is the compressor, 𝐐 χ ∈ ℝ K × d \mathbf{Q}_{\chi}\in\mathbb{R}^{K\times d} is a learned query bank, and 𝐄 j \mathbf{E}_{j} is the latent representation of experience item m j m_{j} . We use ϕ ≔ ( ψ , χ ) \phi\coloneqq(\psi,\chi) to denote all trainable composer parameters, comprising the encoder LoRA parameters ψ \psi and compressor parameters χ \chi ; the encoder backbone itself remains frozen. Note that we do not treat the particular attention block as a contribution; its role is simply to produce a compact, differentiable context (details in Section A.1 ). The teacher receives the latent tokens as ordinary context positions: 𝒄 ϕ ( x , ℰ ) = ⨁ j = 1 J ( ⟨ e j , 1 ⟩ ⊕ ⋯ ⊕ ⟨ e j , K ⟩ ) , π θ ¯ , ϕ T ( ⋅ ∣ 𝒔 , 𝒄 ϕ ) = π θ ¯ ( ⋅ ∣ [ 𝒔 ; 𝒄 ϕ ] ) , \bm{c}_{\phi}(x,\mathcal{E})=\bigoplus_{j=1}^{J}(\langle e_{j,1}\rangle\oplus\cdots\oplus\langle e_{j,K}\rangle),\qquad\pi_{\bar{\theta},\phi}^{T}(\cdot\mid\bm{s},\bm{c}_{\phi})=\pi_{\bar{\theta}}(\cdot\mid[\bm{s};\bm{c}_{\phi}]), (7) where 𝒔 \bm{s} is an interaction state defined above and 𝒄 ϕ \bm{c}_{\phi} is the learned privileged context. Each ⟨ e j , k ⟩ \langle e_{j,k}\rangle is implemented as a continuous embedding and behaves like a special token in the teacher’s context window. The composer is cold-started on successful trajectories with the backbone frozen, yielding ϕ 0 \phi_{0} . During subsequent joint optimization, both the encoder LoRA parameters ψ \psi and compressor parameters χ \chi remain trainable as part of ϕ \phi , while the teacher backbone θ ¯ \bar{\theta} stays fixed (details in Section A.2 ).
### 3.3 Latent On-Policy Self-Distillation
Having specified how the teacher is constructed, we now describe the self-distillation process. The student first produces a multi-turn trajectory conditioned only on 𝒔 t \bm{s}_{t} . The composer then constructs 𝒄 ϕ \bm{c}_{\phi} from retrieved experience, and the teacher re-evaluates the same student prefixes with this additional context. We instantiate the teacher as a frozen reference copy θ ¯ \bar{\theta} initialized from the same backbone as the student; the two share architecture and initialization, but not weights after the student begins updating. For every turn t t and token position n n , we define: 𝒑 S t , n = π θ S ( ⋅ ∣ 𝒔 t , a t , < n ) , 𝒑 T t , n = π θ ¯ , ϕ T ( ⋅ ∣ 𝒔 t , 𝒄 ϕ ( x , ℰ ) , a t , < n ) , \displaystyle\bm{p}^{S}_{t,n}=\pi_{\theta}^{S}\big(\cdot\mid\bm{s}_{t},a_{t,<n}\big),\;\bm{p}^{T}_{t,n}=\pi_{\bar{\theta},\phi}^{T}\big(\cdot\mid\bm{s}_{t},\bm{c}_{\phi}(x,\mathcal{E}),a_{t,<n}\big), (8) where 𝒑 t , n S \bm{p}^{S}_{t,n} and 𝒑 t , n T \bm{p}^{T}_{t,n} are next-token distributions over the vocabulary 𝒱 \mathcal{V} , and θ ¯ \bar{\theta} denotes a fixed base-model checkpoint whose parameters receive no gradient updates; however, the teacher’s forward activations remain in the computation graph, so 𝒑 t , n T \bm{p}^{T}_{t,n} is differentiable with respect to 𝒄 ϕ \bm{c}_{\phi} and hence to ϕ \phi . To keep logit-level distillation efficient, we follow ( Zhao et al., 2026 ) and use the teacher top- M M vocabulary entries plus a tail bucket. This preserves the teacher’s high-probability modes while accounting for the remaining probability mass: 𝒑 ~ t , n T ( v ) = { 𝒑 t , n T ( v ) , v ∈ 𝒱 t , n M , 1 − ∑ u ∈ 𝒱 t , n M 𝒑 t , n T ( u ) , v = ⊥ , 𝒑 ~ t , n S ( v ) = { 𝒑 t , n S ( v ) , v ∈ 𝒱 t , n M , 1 − ∑ u ∈ 𝒱 t , n M 𝒑 t , n S ( u ) , v = ⊥ , , \tilde{\bm{p}}^{T}_{t,n}(v)=\begin{cases}\bm{p}^{T}_{t,n}(v),&v\in\mathcal{V}_{t,n}^{M},\\ 1-\sum_{u\in\mathcal{V}_{t,n}^{M}}\bm{p}^{T}_{t,n}(u),&v=\bot,\end{cases}\;\tilde{\bm{p}}^{S}_{t,n}(v)=\begin{cases}\bm{p}^{S}_{t,n}(v),&v\in\mathcal{V}_{t,n}^{M},\\ 1-\sum_{u\in\mathcal{V}_{t,n}^{M}}\bm{p}^{S}_{t,n}(u),&v=\bot,\end{cases}, (9) where 𝒱 t , n M \mathcal{V}_{t,n}^{M} is the teacher top- M M support at prefix ( t , n ) (t,n) and ⊥ \bot denotes the aggregated tail event. The distillation objective matches teacher and student on the student’s own multi-turn trajectory: ℒ distill ( θ , ϕ ) = 𝔼 x ∼ 𝒟 𝔼 𝝉 ∼ π θ S ( ⋅ ∣ x ) [ ∑ t = 1 | 𝝉 | ∑ n = 1 L t ω t , n D KL ( 𝒑 ~ t , n S ∥ 𝒑 ~ t , n T ) ∑ t = 1 | 𝝉 | ∑ n = 1 L t ω t , n ] , \displaystyle\mathcal{L}_{\mathrm{distill}}(\theta,\phi)=\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{\bm{\tau}\sim\pi_{\theta}^{S}(\cdot\mid x)}\left[\frac{\sum_{t=1}^{|\bm{\tau}|}\sum_{n=1}^{L_{t}}\omega_{t,n}\,D_{\mathrm{KL}}\!\left(\tilde{\bm{p}}^{S}_{t,n}\middle\|\tilde{\bm{p}}^{T}_{t,n}\right)}{\sum_{t=1}^{|\bm{\tau}|}\sum_{n=1}^{L_{t}}\omega_{t,n}}\right], (10) where 𝒟 \mathcal{D} is the task distribution and ω t , n ∈ { 0 , 1 } \omega_{t,n}\in\{0,1\} masks supervised action tokens. We use reverse KL so that the student concentrates on teacher-supported behavior. Because θ ¯ \bar{\theta} denotes a fixed checkpoint while ϕ \phi remains trainable, the distillation gradient flows through the frozen network back to the latent input (injection mechanism in Section A.4 ). This allows the composer to learn which experience features produce effective teacher supervision for the current student trajectory.
#### Privileged-Margin Constraint.
The distillation loss alone does not ensure that the teacher provides more effective supervision: an unconstrained composer can instead minimize Equation 10 by moving 𝒑 T \bm{p}^{T} toward 𝒑 S \bm{p}^{S} , yielding uninformative latent context without improving the student. We therefore prevent this collapse in two ways. ❶ First, the cold start in Section 3.2 initializes the composer to transform experience into informative latent context. ❷Second, we introduce outcome reward into the supervision signal, and more concretely, impose a privileged-margin constraint that ties the teacher’s token-level advantage to the verified outcome of the complete trajectory. For each supervised token, we define the per-token privilege: δ t , n ( ϕ ) = log π θ ¯ , ϕ T ( a t , n ∣ 𝒔 t , 𝒄 ϕ , a t , < n ) − sg [ log π θ S ( a t , n ∣ 𝒔 t , a t , < n ) ] , \delta_{t,n}(\phi)=\log\pi^{T}_{\bar{\theta},\phi}\!\left(a_{t,n}\mid\bm{s}_{t},\bm{c}_{\phi},a_{t,<n}\right)-\operatorname{sg}\!\left[\log\pi^{S}_{\theta}\!\left(a_{t,n}\mid\bm{s}_{t},a_{t,<n}\right)\right], (11) where a t , n a_{t,n} is the student’s sampled token and sg \operatorname{sg} blocks gradients to θ \theta through this path. Both log-probabilities are already computed during the teacher and student forward passes; δ t , n \delta_{t,n} requires only a gather, not an additional forward. We then use the trajectory-level verification signal A ( 𝝉 ) = 2 r ( 𝝉 ) − 1 ∈ [ − 1 , 1 ] A(\bm{\tau})=2r(\bm{\tau})-1\in[-1,1] , where r ( 𝝉 ) r(\bm{\tau}) is the outcome reward, to favor teacher advantages on successful trajectories and suppress them on unsuccessful ones: Δ ( ϕ ) = 𝔼 𝝉 [ ∑ t , n ω t , n A ( 𝝉 ) δ t , n ( ϕ ) ∑ t , n ω t , n ] . \Delta(\phi)=\mathbb{E}_{\bm{\tau}}\!\left[\frac{\sum_{t,n}\omega_{t,n}\,A(\bm{\tau})\,\delta_{t,n}(\phi)}{\sum_{t,n}\omega_{t,n}}\right]. (12) The full LOPD objective constrains ϕ \phi to maintain a minimum privilege level m > 0 m>0 : min θ , ϕ max β ≥ 0 ℒ distill ( θ , ϕ ) + β ( m − Δ ( ϕ ) ) + λ ‖ 𝒄 ϕ − sg [ 𝒄 ϕ 0 ] ‖ 2 2 , \min_{\theta,\phi}\;\max_{\beta\geq 0}\;\;\mathcal{L}_{\mathrm{distill}}(\theta,\phi)\;+\;\beta\!\left(m-\Delta(\phi)\right)+\;\lambda\left\|\bm{c}_{\phi}-\operatorname{sg}\!\left[\bm{c}_{\phi_{0}}\right]\right\|_{2}^{2}, (13) where β \beta is a dual variable updated by β ← [ β + η β ( m − Δ ( ϕ ) ) ] + \beta\leftarrow[\beta+\eta_{\beta}(m-\Delta(\phi))]_{+} , and the anchor term penalizes drift from the initialization ϕ 0 \phi_{0} in latent space, requiring no additional forward pass. If the composer degenerates to uninformative context, π T → π S \pi^{T}\to\pi^{S} , δ t , n → 0 \delta_{t,n}\to 0 , Δ → 0 < m \Delta\to 0<m , and the dual penalty activates—structurally excluding the trivial solution. Within the feasible region, outcome-weighted privilege steers the composer toward evidence that supports successful behavior, while the distillation gradient determines how that evidence is selected and encoded.
#### Algorithm Summary.
Algorithm 1 summarizes the overall workflow of LOPD . The loop collects on-policy trajectories from the student, constructs learnable latent privileged context from retrieved experience, evaluates the same visited prefixes with the teacher, and jointly updates the student and composer while adjusting the dual variable to maintain the privilege margin. At inference, only the student policy π θ S \pi_{\theta}^{S} is deployed and used.
## 4 Experiments
### 4.1 Experiment Setup
#### Training.
We evaluate LOPD under two post-training settings: ❶ agentic tool-use and ❷ coding . For tool-use training, we use the EnvScaler-derived tool-interactive corpus ( Song et al., 2026 ) , comprising 2,349 2{,}349 tasks. For coding, we use the TACO subset of DeepCoder ( TogetherAI, 2025 ) , comprising 7 7 K verified Python problems (see Section B.2 for training details). The used model backbones include Qwen3-4B , Qwen3-8B ( Yang et al., 2025 ) , Olmo3-7B ( Olmo et al., 2026 ) .
#### Evaluation.
Our evaluation mirrors the two training regimes while keeping all test sets disjoint from training data. For tool-use, we report the EnvScaler task success metric on a held-out test split of 200 200 tasks disjoint from the training pool, BFCL-v3 scores across base, missing-function, missing-parameter, and long-context subsets ( Patil et al., 2025 ) , and ACEBench multi-step and multi-turn scores ( Chen et al., 2025 ) . For coding, we report pass@1 on LiveCodeBench v5/v6 ( Jain et al., 2024 ) , HumanEval+ and MBPP+ ( Liu et al., 2023 ) (evaluation protocols in Section B.3 ).
#### Baselines.
We compare against the baselines reported in Tables 1 and 2 . Vanilla denotes the unadapted backbone. GRPO is the outcome-reward RL baseline ( Shao et al., 2024 ) . SDFT ( Shenfeld et al., 2026 ) is a demonstration-conditioned distillation baseline. The OPSD-style baselines include OPSD ( Zhao et al., 2026 ) , SDPO ( Hübotter et al., 2026 ) , and Skill-SD ( Wang et al., 2026 ) , whose privileged contexts respectively instantiate answer/trace, feedback-conditioned self-teacher, and skill-conditioned supervision. Detailed baseline configurations are provided in Section B.1 .
#### Configurations.
All trainable methods share the same backbone, training split and evaluation protocol within each setting. For LOPD , the teacher receives latent context composed from J = 3 J{=}3 retrieved experiences, each compressed into K = 32 K{=}32 latent tokens ( 96 96 total); the distillation loss uses reverse KL with SDPO-style top- M M logits plus a tail bucket, with M = 20 M{=}20 by default. The privileged-margin threshold is m = 0.05 m{=}0.05 ; the verification signal is A ( 𝝉 ) = 2 r ( 𝝉 ) − 1 A(\bm{\tau})=2r(\bm{\tau})-1 where r r is the environment reward. Tool-use rollouts are capped at 30 30 environment steps. Coding distillation uses a 16,384 16{,}384 -token response budget. We keep decoding temperature/top- p p fixed across methods within each benchmark, each baseline retains its original KL direction and loss formulation ( Section B.1 ), and within the on-policy training loop the methods differ only in privileged context construction and distillation loss. Prompt templates are provided in Section A.5 .
### 4.2 Main Results
We evaluate whether latent privileged context improves self-distillation learning across tool use and code generation domains in Tables 1 and 2 .
#### Performance.
LOPD obtains the best aggregate result in all ten backbone–benchmark comparisons. On tool use with Qwen3-4B , it improves over the strongest competing method from 61.8 61.8 to 63.7 63.7 on EnvScaler, from 25.25 25.25 to 27.38 27.38 on BFCL-v3, and from 56.0 56.0 to 60.6 60.6 on ACEBench. The advantage becomes larger with Qwen3-8B : LOPD reaches 66.4 / 29.88 / 62.7 66.4/29.88/62.7 on EnvScaler, BFCL-v3, and ACEBench, compared with the strongest baseline results of 60.2 / 29.00 / 58.0 60.2/29.00/58.0 , respectively. The same trend extends beyond interactive agents. With Qwen3-4B , LOPD improves the LiveCodeBench and EvalPlus aggregates to 48.78 48.78 and 81.36 81.36 ; with Olmo3-7B , it reaches 50.98 50.98 and 78.41 78.41 , outperforming the strongest alternatives by 2.69 2.69 and 0.55 0.55 points. These results indicate that LOPD ’s improvement transfers across task formats and model families.
#### No Hand-Crafted Context Is Universally Optimal.
A careful reader may notice that adding privileged information does not always improve upon the vanilla model in Tables 1 and 2 . For example, with Qwen3-4B , SDPO falls from 22.88 22.88 to 15.75 15.75 on BFCL-v3, from 50.6 50.6 to 38.0 38.0 on ACEBench, and from 45.61 45.61 to 38.78 38.78 on LiveCodeBench; OPSD likewise reduces the LiveCodeBench aggregate to 40.24 40.24 . More revealingly, the utility of the same context can reverse across settings: OPSD raises the BFCL-v3 long-context score from 22.00 22.00 to 29.00 29.00 , yet remains below vanilla on the overall BFCL-v3 score with Qwen3-8B ( 25.75 25.75 vs. 28.38 28.38 ) and on LiveCodeBench with both backbones ( 40.24 40.24 vs. 45.61 45.61 and 44.39 44.39 vs. 46.34 46.34 ). This behavior is consistent with how these contexts are constructed. SDPO can condition only on successful siblings produced by the current rollout group, making its privileged signal dependent on the current policy’s success coverage; OPSD instead injects a fixed oracle trace, which can be highly informative when its procedure matches the current task but need not align with the particular prefix or valid solution path visited by the student. A context format that helps in one regime can therefore become sparse, mismatched, or overly prescriptive in another. The issue is not simply whether privileged information is available, but whether its representation remains appropriate across tasks and learning states.
#### Learnable Context Matters.
LOPD avoids committing the teacher to a fixed answer, demonstration, sibling rollout, or discrete skill. Instead, it learns a compact continuous representation directly from retrieved trajectories, allowing the privileged signal to adapt to the task and the student’s visited prefixes. Unlike the reversals above, LOPD remains above the vanilla model in all ten aggregate backbone–benchmark settings, with gains ranging from 1.50 1.50 points on BFCL-v3 with Qwen3-8B to 17.2 17.2 points on EnvScaler with the same backbone. It also consistently improves over prescribed-context baselines: with Qwen3-8B , LOPD exceeds OPSD, SDFT, and Skill-SD by 14.4 / 10.2 / 6.2 14.4/10.2/6.2 points on EnvScaler and 10.0 / 8.0 / 6.7 10.0/8.0/6.7 points on ACEBench, respectively. Together, the results favor optimizing the representation of experience for supervision over searching for an increasingly elaborate hand-crafted artifact.