Latent On-Policy Self-Distillation

In plain words

When a model learns from a teacher that is a copy of itself, the copy usually gets a hand-written hint: the right answer, some feedback, a skill note, or a past attempt. Someone has to decide what that hint looks like, and the choice limits what the model can learn.

This work lets the model learn the hint itself. Past successful attempts are pulled up for the current task and squeezed into a short run of soft, machine-only vectors. The copy reads those vectors while it judges each step of the student's own attempt, and the student is trained to match that judgment token by token. The student never sees the vectors. A rule keeps the learned vectors honest: the copy must stay better than the student at predicting successful attempts, so the vectors cannot quietly turn into noise. After training, only the student is kept, so no memory store or retriever is needed at run time.

TL;DR

• Self-distillation needs a "privileged" context for the teacher; this paper makes that context a learned latent representation of retrieved past successes §3.2.
• The student samples its own trajectories and is trained with reverse KL to match a teacher that also sees the latent context, at every visited prefix §3.3.
• A privileged-margin constraint with a dual variable keeps the teacher's log-probability advantage on successful trajectories above a threshold (m = 0.05) §3.3 tab. 4.
• Author-reported: best aggregate in all 10 backbone-benchmark comparisons across EnvScaler, BFCL-v3, ACEBench, LiveCodeBench and EvalPlus §4.2 tab. 1 tab. 2. Ablations are narrow (one benchmark, one backbone for most) §4.3.

Problem

On-policy self-distillation (OPSD) uses the same model twice: a student that sees only the task and its history, and a teacher that also sees privileged information such as answers, traces, feedback or successful rollouts §1 §2. The teacher's context is built by a fixed, designer-chosen rule, so the supervision is capped by the format the designer picked. A context format that helps in one setting can be sparse or misleading in another §4.2. The authors ask whether the privileged context can instead be learned end-to-end from experience §1.

Making the context learnable creates a new failure mode. An unconstrained composer can move the teacher toward the student, which makes the teacher uninformative while the distillation loss still goes down §3.3. The paper's answer is a constraint on the teacher's advantage over the student, plus a cold-start that begins from a useful composer §3.3 app. A.2.

Method

Setup: student, teacher and OPSD loss

The student policy \pi_\theta^S acts on the state s_t = (x, o_{\le t}, a_{<t}) and produces a multi-turn trajectory \tau §3.1 eq. 1. Standard OPSD builds a teacher by appending a fixed context c_{\mathrm{fix}} = \Phi_{\mathrm{fix}}(x, \mathcal{E}) derived from some experience source \mathcal{E}, and matches the teacher to the student on the student's visited prefixes with a token-level divergence §3.1 eq. 2 eq. 3. The paper replaces the fixed rule \Phi_{\mathrm{fix}} with a parameterised composer \Phi_\phi §3.1 eq. 4.

Composer: experience bank, encoder, compressor

• Experience bank. Only successful rollouts are stored, each as a task description plus a compact action-result trace that omits verbose observations §3.2 app. A.3.
• Retrieval. For a task x, a dense retriever returns the top J entries by cosine similarity, with J = 3 §3.2 §4.1. Document embeddings cover the task and trajectory concatenated; query embeddings cover the task only, with an instruction-style query prompt app. A.3 eq. 15. The embedding model is Qwen3-Embedding-8B, with 4,096-dimensional embeddings and an exact inner-product FAISS index app. A.3. Retrieval returns the top-scoring entries by inner product app. A.3 eq. 16; evaluation tasks are kept out of the bank app. A.3.
• Encoder. Each retrieved item m_j is encoded together with the task text by the frozen backbone plus a trainable LoRA adapter, giving hidden states \mathbf{H}_j §3.2 eq. 5. The task text is prepended in a fixed "Task to solve / Reference past trajectory" format app. A.5.
• Compressor. A QFormer-style module with K = 32 learned queries cross-attends to \mathbf{H}_j and returns \mathbf{E}_j §3.2 eq. 6. The compressor is applied to each retrieved item independently. The 8-layer block shares its weights across depth, uses a feed-forward multiplier of 4, and takes its hidden size and heads from the backbone app. A.1.
• Latent context. The J \cdot K latent tokens are concatenated, giving 96 tokens for J = 3, K = 32 §3.2 eq. 7 §4.1.

The teacher is a frozen copy of the backbone. The latent tokens are inserted as ordinary positions in the first user message, replacing a block of placeholder positions §3.2 app. A.4 app. A.5. The LoRA adapter is switched off for the teacher's main forward pass, so the teacher is the plain backbone that reads the latent positions app. A.4.

Cold start

Before joint training, the composer is trained by next-token NLL on successful trajectories, with the backbone frozen and the latent tokens in the input app. A.2 eq. 14. The training trajectories are produced by the base model itself and filtered by task success app. A.2. The data size and step count are set per domain and are not stated in the paper app. A.2.

Distillation

For each supervised action token, the student and teacher next-token distributions are computed at the same prefix §3.3 eq. 8. The teacher's gradient reaches the composer through its frozen forward pass, because the teacher's activations stay in the graph §3.3.

To keep logit-level distillation affordable, the teacher's top-M entries are kept plus a single tail bucket that holds the remaining mass, with M = 20 §3.3 eq. 9 §4.1. The loss is reverse KL, \mathrm{KL}(\tilde p^S \| \tilde p^T), averaged over supervised action tokens with a 0/1 mask \omega_{t,n} §3.3 eq. 10.

Privileged-margin constraint

Per token, the paper defines a privilege \delta_{t,n} as the teacher's log-probability of the student's sampled token minus the student's log-probability of the same token, with the student side stopped from gradients §3.3 eq. 11. The trajectory verdict is A(\tau) = 2r(\tau) - 1 \in [-1, 1], with r the outcome reward §3.3. The batch margin \Delta(\phi) is the mask-weighted average of A \cdot \delta §3.3 eq. 12.

The full objective is a Lagrangian: the distillation loss, plus \beta(m - \Delta), plus an anchor pulling the latent context toward its cold-start value c_{\phi_0} §3.3 eq. 13. The dual variable is updated by projected gradient steps, \beta \leftarrow [\beta + \eta_\beta(m - \Delta)]_+ §3.3 eq. 13. The anchor adds no forward pass, according to the paper §3.3.

Algorithm

# Latent On-Policy Self-Distillation (LOPD)  [alg. 1]
# Input:  tasks; retriever; student; frozen backbone; composer (init phi_0);
#         reward verifier in [0, 1]; margin m; dual step eta_beta
# Output: trained student policy

def train_lopd(tasks, student, composer, beta=0.0):
    # Loop over batches of tasks  [alg. 1]
    for batch in sample_batches(tasks):
        # 1. Roll out the student alone, with no privileged context  [§3.3]
        trajectories = [rollout(student, task) for task in batch]
        rewards = [verify(traj) for traj in trajectories]     # in [0, 1]
        advantages = [2 * r - 1 for r in rewards]             # in [-1, 1]

        # 2. Retrieve top-J experiences and compose latent context  [§3.2]
        contexts = [composer(task, retrieve(task, k=3)) for task in batch]

        # 3. Teacher re-scores every visited prefix with its context  [eq. 8]
        teacher_dists = teacher(contexts, trajectories)

        # 4. Build top-M teacher support plus tail, M = 20  [eq. 9]
        # 5. Reverse KL(student || teacher) on supervised tokens  [eq. 10]
        distill_loss = reverse_kl_topM(student, teacher_dists)

        # 6. Per-token privilege, student side stop-gradiented  [eq. 11]
        privilege = teacher_logp - stop_grad(student_logp)

        # 7. Batch margin: mean of advantage times privilege  [eq. 12]
        margin = masked_mean(advantages * privilege)

        # 8. Joint step on student and composer, then dual update  [eq. 13]
        # The anchor is the squared distance to the cached cold-start latents.
        anchor = squared_distance(contexts, cold_start_contexts)
        loss = distill_loss + beta * (m - margin) + lam * anchor
        update(student, composer, loss)
        beta = max(0.0, beta + eta_beta * (m - margin))

    return student  # only the student is kept

Costs the abstract hides

• Each step runs the frozen teacher over full multi-turn trajectories, once per rollout, and back-propagates through it into the composer. The paper reports no wall-clock or GPU-hour figures §4 app. B.
• The composer must encode J retrieved trajectories per task, so every training task pays for three extra encoder passes of the backbone in addition to the teacher pass §3.2 app. A.1.

Hyperparameters

Results

Author-reported. Setup: agentic tool use on Qwen3-4B and Qwen3-8B (EnvScaler held-out 200 tasks, BFCL-v3 four multi-turn subsets, ACEBench 50 tasks); coding on Qwen3-4B and Olmo3-7B (LiveCodeBench v5 and v6, HumanEval+ and MBPP+). All methods share the training split and evaluation protocol §4.1 app. B.3. Baselines are Vanilla, GRPO, SDFT, OPSD, SDPO and Skill-SD §4.1 app. B.1. Table 1 and Table 2 reproduce the paper's numbers. The paper's green highlight is dropped.

Table 1 tab. 1: tool use, 0–100 scale.

The BFCL and ACEBench averages are consistent with a mean over the four BFCL subsets and a task-weighted mean over the 20 multi-step and 30 multi-turn ACEBench tasks tab. 1 app. B.3.

Table 2 tab. 2: code, pass@1 in percent. Averages are weighted by task count (LiveCodeBench v5 has 279 tasks and v6 has 131; HumanEval+ has 164 and MBPP+ has 378) app. B.3.

Headline claims, as written by the authors:

• LOPD is best on the aggregate in all ten backbone-benchmark comparisons §4.2.
• On Qwen3-8B it reaches 66.4 / 29.88 / 62.7 on EnvScaler / BFCL-v3 / ACEBench, against a strongest baseline of 60.2 / 29.00 / 58.0 §4.2.
• It is above the vanilla model in all ten aggregate settings, by 1.50 to 17.2 points §4.2.
• The authors also claim it surpasses GRPO and Skill-SD at a fraction of their rollout count, below the 30 percent mark §1 §5. This rests on one figure: Qwen3-4B on EnvScaler, mean reward over 1,600 generations §4.3 fig. 4 app. B.4.

Framework analysis §4.3:

• Joint optimisation (Qwen3-4B, EnvScaler, student with no context at test time) fig. 3: frozen composer 0.573; no margin (m = 0) 0.551; m ≤ 0.01 does not prevent the decline; m = 0.05 gives 0.637; m = 0.10 gives 0.626.
• Sample efficiency fig. 4: LOPD's mean reward crosses 0.61 at the 320-generation mark and hits 0.637 at generation 576, then stays between 0.63 and 0.64. GRPO and Skill-SD end at 0.611 and 0.588.
• Latent capacity fig. 5a: reward stays near 0.56 with 8 or 16 tokens per experience, rises to 0.637 at 32, and shows no consistent gain at 64 or 128.
• Retrieval count fig. 5b–e: reward goes from 0.605 with one retrieved experience to 0.637 with three. More retrievals give no monotonic gain. ACEBench at the default is 56.6 M-Step, 63.3 M-Turn, 60.6 overall.

Behaviour of the student tab. 3:

The distilled student uses more steps with fewer calls per step, and it keeps this pattern once the context is removed §4.3 tab. 3.

Paper vs code

Repository: github.com/bingreeky/LOPD, Apache-2.0, cloned at 50c3d7f (HEAD on the clone date). The paper links this repository in its author block. Scope of the release: the EnvScaler agentic pipeline only. There is no coding environment, no benchmark harness for BFCL, ACEBench, LiveCodeBench or EvalPlus, no memory-bank builder, no cold-start dataset builder, and no baseline implementations README.md. The coding and baseline numbers therefore cannot be reproduced from this release.

Discrepancies found in the code:

• Anchor scale. The paper's anchor is the squared L2 norm of c_\phi - \mathrm{sg}[c_{\phi_0}] §3.3 eq. 13. The code takes the mean over all elements, which divides by J \cdot K \cdot d worker.py. The effective anchor weight is therefore much smaller than \lambda = 0.2 implies, and the size of the difference depends on the backbone's hidden size d.
• Two updates, not one. Algorithm 1 updates (\theta, \phi) jointly on one loss alg. 1. The code takes an AdamW step on the student worker.py, then computes the composer loss from the student's log-probabilities collected before that step and takes a separate AdamW step on the composer with its own learning rate worker.py. The composer's gradient therefore comes from the student's pre-update log-probabilities, so the two updates are sequential, not one joint step.
• Loss normalisation differs between the two sides. Student loss sums over all supervised tokens and divides by the global token count worker.py. Composer loss averages each record's KL and margin over that record's tokens, then averages over records worker.py worker.py. Long trajectories therefore weigh less per token in the composer gradient. The reported batch margin \Delta is token-weighted, matching eq. 12 worker.py.
• Teacher forwards per step. The paper says the per-token privilege needs only a gather from the teacher and student forward passes, with no extra forward §3.3. The code runs the teacher twice per step: once under no-grad for the top-M support worker.py, and once with gradients in the composer phase worker.py. Each phase also loads the backbone from disk again worker.py worker.py worker.py.
• Anchor cost. The paper says the anchor needs no extra forward pass §3.3. The code runs the cold-start adapter once per unseen retrieved text and caches the result to disk teacher.py. Steady-state steps need no extra pass, but the first encounter does.
• Retrieval embeddings. The paper embeds documents as task plus trajectory and queries as task with an instruction prompt app. A.3 eq. 15. The shipped bank indexes only the task instruction string and encodes queries without any prompt bank.py bank.py. The paper's concatenated documents may have been built elsewhere; the build script is not in the repository.
• Self-retrieval. The training, cold-start and latent-evaluation configs all set allow_self_retrieval: true lopd_qwen3_4b_envscaler.yaml compressor_sft_qwen3_4b_envscaler.yaml envscaler_latent_eval200_qwen3_4b.yaml. The class default is false config.py, and the retrieval filter skips the query's own entry only when the flag is off bank.py. The paper says the bank holds only training-split successes and that evaluation tasks are excluded app. A.3, but says nothing about a training task retrieving its own trajectory. Whether the released bank contains the training tasks' own successes cannot be checked from the repository.
• Query-side encoder. The config leaves the embedding model to the user bank.py. The paper's Qwen3-Embedding-8B is not wired in.
• Weight decay and schedule. The paper states neither app. A.2 app. B.1. The code constructs AdamW without a weight-decay argument, so PyTorch's default applies worker.py worker.py. No learning-rate scheduler appears in the repository.
• Cold-start clipping. The paper gives one clipping value, 3.0 app. A.2. The SFT config clips the QFormer at 30.0 and the LoRA at 3.0 compressor_sft_qwen3_4b_envscaler.yaml.
• Latent initialisation. The paper writes \mathcal{N}(0, 1/\sqrt{d}) app. A.1. The code draws from a normal with standard deviation 1/\sqrt{d}, that is \mathrm{randn}/\sqrt{d} qformer.py. The paper's notation can be read either way.
• Compressor details not in the paper. The attention keys and values include the latent queries themselves qformer.py. The block adds linear input and output projections and a final LayerNorm qformer.py qformer.py.
• Silent truncation. Sequences longer than 40,000 tokens are cut and their supervised tokens are dropped serialization.py lopd_qwen3_4b_envscaler.yaml.
• Failed rollouts. Rollouts that error are excluded from both losses and from the margin, and the student takes a zero-gradient step to keep ranks in sync rollout.py worker.py worker.py. The paper does not mention this.
• Tail clamp. The tail log-probability is clamped to at most −1e-7 before exponentiation loss.py. This keeps the tail finite when the top-M mass rounds to one.
• Model name. The paper's coding backbone is "Olmo3-7B" §4.1. The README and the released checkpoint name OLMo-3-7B-Think README.md.
• Figure 4 budget. The shipped config runs 100 steps of 32 tasks lopd_qwen3_4b_envscaler.yaml. The paper's sample-efficiency figure stops at 1,600 generations §4.3 fig. 4. The number of steps in the paper's run is not stated.

Checked and consistent: reverse-KL direction loss.py eq. 10; top-M with tail loss.py eq. 9; A = 2r - 1 with r in [0, 1] worker.py; m = 0.05, \lambda = 0.2, \eta_\beta = 0.5, \mathrm{lr}_\phi = 10^{-5}, \beta_0 = 0 lopd_qwen3_4b_envscaler.yaml; J = 3, K = 32, M = 20 lopd_qwen3_4b_envscaler.yaml compressor_sft_qwen3_4b_envscaler.yaml; the teacher's encoder adapter is disabled during the main pass lora.py app. A.4; latent injection by a contiguous span replacement via torch.cat serialization.py app. A.4; framing text prompt_protocol.py app. A.5.

Pitfalls

• Self-retrieval may leak the answer into training. If a training task can retrieve its own successful trajectory, the teacher is shown the solution it is being distilled toward. The paper's bank description does not rule this out, and the configs enable self-retrieval lopd_qwen3_4b_envscaler.yaml. Build the bank so that each training task's own entry is excluded, or confirm that the paper's bank does not include training-split successes for the same tasks.
• The margin is load-bearing. Without it the student falls below the frozen-composer baseline (0.551 against 0.573) §4.3 fig. 3. Weak margins below 0.02 do not prevent the decline. Do not drop the constraint to save complexity.
• Gradients through a 40k-token teacher. The composer's gradient flows through the full frozen backbone over every visited prefix. The paper gives no memory or time budget, and the code relies on gradient checkpointing and FSDP across eight ranks lopd_qwen3_4b_envscaler.yaml lopd_qwen3_4b_envscaler.yaml worker.py.
• Placeholder contract. The latent span must be contiguous and must match J \cdot K exactly, and the same framing text must be used at train and test time. The code raises an error if the span runs past the length limit serialization.py. The paper reports that token-ID and embedding-tensor paths give identical greedy output on a sanity check, not on the full pipeline app. A.4.
• Interpretability is not established. Decoding the latent tokens through the LM head gives fragmented, multilingual or code-like tokens, not a readable procedure §4.4 fig. 6. Do not read the latent tokens as a readable memory.
• Baseline context is not a like-for-like comparison. OPSD and SDFT receive ground-truth trajectories that the authors sometimes synthesised with stronger external models (DeepSeek-V4-Pro and Qwen3.7-Max), when the base model's own pool was insufficient app. B.1. Those baselines' privileged signal is therefore not purely self-generated.
• Statistical support is thin. The paper does not report seeds, variance or repeated runs. The efficiency claim, the margin sweep and the sensitivity sweeps are each on one benchmark, and most use Qwen3-4B §4.3 app. B.4.
• Reward is continuous on EnvScaler. The reward is the fraction of subtasks completed, in [0, 1], so A = 2r - 1 takes values across the range, not only ±1. The margin is then not a pure success signal app. B.2 worker.py.
• Anchor weight is not transferable as written. Because the code averages over elements, a \lambda tuned against the paper's squared norm will not reproduce the same constraint strength [see Paper vs code].
• Experience is minimal by design. The bank stores action-result traces with observations removed. The paper argues richer sources could be added without changing the framework §3.2, but that has not been tested here.

Lineage

• Builds on on-policy self-distillation with a fixed privileged context: OPSD (arXiv:2601.18734), SDPO (arXiv:2601.20802), Skill-SD (arXiv:2604.10674) and SDFT (arXiv:2601.19897) §2 app. B.1. The change is that the privileged context is a learned latent composer rather than a designer-chosen artifact §3.2.
• Uses a QFormer-style compressor and a LoRA encoder, with latent tokens injected at placeholder positions app. A.1 app. A.4.
• Compared against GRPO (DeepSeekMath, arXiv:2402.03300) as the outcome-reward RL baseline §4.1 app. B.1.
• The agentic training data is EnvScaler (arXiv:2601.05808) §4.1.
• Superseded by: not known.