PaperlayerView exact revision on arXiv

Source-linked AI summary

Privileged Information Distillation for Language Models

Emiliano Penaloza, Dheeraj Vattikonda, Nicolas Gontier, Alexandre Lacoste, Laurent Charlin, Massimo Caccia

arXiv:2602.04942v3cs.LGcs.AI

TL;DR

The paper asks how language models can transfer capabilities learned with training-time privileged information to policies that act without it, particularly when frontier-model reasoning traces are hidden. It introduces joint and on-policy self-distillation methods that use action-only privileged information, and reports strong performance across agentic and out-of-distribution evaluations. The paper also analyzes factors associated with effective privileged-information learning.

  • Problem

    The central problem is transferring capabilities learned with privileged information to test-time policies without it when distillation exposes actions but not frontier-model Chain-of-Thought reasoning.

  • Method

    The paper introduces π-Distill, which jointly trains a PI-conditioned teacher and unconditioned student with shared parameters, and OPSD, which uses on-policy RL with a reverse KL penalty.

  • Results

    π-Distill and, in some cases, OPSD outperform industry-standard SFT+RL baselines requiring full Chain-of-Thought supervision across agentic benchmarks and out-of-distribution tasks.

  • Takeaways & Limitations

    Action-only privileged information can support effective frontier-agent distillation without access to proprietary Chain-of-Thought traces, while success depends on properties of the information and teacher-student distributions.

  • Takeaways & Limitations

    The experiments use privileged information transformed from frontier-model traces, study models with at most 8B parameters, and analyze key factors observationally rather than systematically controlling all variables.

Abstract

from arXiv · show

Training-time privileged information (PI) can enable language models to succeed on tasks they would otherwise fail, making it a powerful tool for reinforcement learning in hard, long-horizon settings. However, transferring capabilities learned with PI to policies that must act without it at inference time remains a fundamental challenge. We study this problem in the context of distilling frontier models for multi-turn agentic environments, which typically hide their internal reasoning and expose only action trajectories. This breaks standard distillation pipelines, since successful behavior is observable, but the reasoning process is not. For this, we introduce π-Distill, a joint teacher-student objective that trains a PI-conditioned teacher and an unconditioned student simultaneously using the same model. Additionally, we also introduce On-Policy Self-Distillation (OPSD), an alternative approach that trains using Reinforcement Learning (RL) with a reverse KL-penalty between the student and the PI-conditioned teacher. We show that both of these algorithms effectively distill frontier agents using action-only PI. Specifically, we find that π-Distill and, in some cases, OPSD, outperform industry standard practices (Supervised finetuning followed by RL) that assume access to full Chain-of-Thought supervision across multiple agentic benchmarks, models, and forms of PI. We complement our results with extensive analysis that characterizes the factors enabling effective learning with PI, focusing primarily on π-Distill and characterizing when OPSD is competitive.

1 Introduction

The paper studies how to transfer capabilities learned with training-time privileged information to policies acting without it, especially when frontier-model reasoning traces are unavailable. It introduces two self-distillation objectives and finds strong performance across agentic benchmarks and out-of-distribution tasks.

  • Motivation: Training-time privileged information can help language models succeed on tasks they would otherwise fail, but transferring those capabilities to test-time policies without PI remains challenging.The challenge is especially relevant to reinforcement learning, where successful experiences can bootstrap learning.
  • Methods: π-Distill jointly trains a PI-conditioned teacher and an unconditioned student with shared parameters, while OPSD adds a reverse KL penalty between them during on-policy RL.The joint setup lets the teacher learn to use PI while mitigating teacher-student distribution shift; OPSD provides an on-policy alternative.
  • Motivation: Frontier-model distillation is difficult because providers often expose successful actions but hide the full Chain-of-Thought reasoning traces needed by standard pipelines.The resulting setting observes what successful agents do but not how they reason.
  • Results: The methods work on Travel Planner and τ-Bench and show strong out-of-distribution generalization across eight additional GEM and τ-Bench airline tool-use tasks.The reported evaluations include three models from two distinct model families.
  • Analysis: For π-Distill, effective learning requires useful PI, a mitigated student-teacher distributional gap, and prevention of collapse; OPSD depends more on information content.The analysis derives these factors by varying three types of PI with different information densities and induced distribution shifts.
  • Results: π-Distill consistently outperforms standard SFT+RL baselines that rely on full Chain-of-Thought data, while OPSD does so in some cases.The comparison spans multiple agentic benchmarks, models, and forms of privileged information.

2 Background

The paper formulates long-horizon agentic interaction as an MDP and describes RL optimization with GRPO. Policies generate reasoning and action tokens, receive trajectory rewards, and use group-relative clipped importance-weighted updates.

  • Agentic Interaction as an MDP: Long-horizon, multi-turn agentic environments are modeled as MDPs in which a policy interacts with the environment over extended action sequences.The evolving state aggregates the user prompt, previous model outputs, and environment responses.
  • Agentic Interaction as an MDP: A trajectory contains reasoning tokens and action tokens, allowing the formalism to distinguish internal generation from externally executed behavior.The paper denotes the trajectory as o = (z, a).
  • Reinforcement learning: RL maximizes expected return from environment interaction, with R(o, s) ∈ [−1, 1] denoting the reward assigned to a generated trajectory.The objective averages trajectory rewards over states and policy-generated outputs.
  • Policy Optimization: The paper optimizes the RL objective using GRPO, sampling groups of trajectories for each state according to the current sampling policy and transition function.Group-relative advantages compare each trajectory’s return with the group’s average return.
  • Policy Optimization: GRPO scales clipped importance-weighted policy updates by group-relative advantages computed from trajectory returns.The token-level importance ratio compares the current policy with the sampling policy, and the clipping parameter is ϵ.
  • Policy Optimization: The implementation drops the usual base-model KL penalty because recent work reports that it can hinder performance.This adjustment follows recommendations cited by the paper.

3 Methods

The paper introduces π-Distill and OPSD, two shared-parameter teacher–student objectives for transferring PI-enabled behavior to an unconditioned student. π-Distill jointly trains both policies, while OPSD uses on-policy student trajectories and reverse-KL distillation.

  • Motivation: The sequential pipeline of separately training a PI-conditioned policy and then distilling it is computationally inefficient, off-policy, checkpoint-sensitive, and empirically suboptimal.π-Distill addresses these issues through joint training in a single shared-parameter model.
  • Privileged Information Distillation (π-Distill): π-Distill jointly trains a PI-conditioned teacher and an unconditioned student within one parameter-shared model.The teacher learns from reward-optimized behavior, while shared representations support transfer to the student.
  • Privileged Information Distillation (π-Distill): The teacher objective increases reward while penalizing divergence from the student, encouraging high-reward behavior that remains learnable by the student.Shared parameters also promote transfer even without directly training the student.
  • Privileged Information Distillation (π-Distill): The student objective learns from trajectories sampled by the PI-conditioned teacher, transferring high-reward behavior without requiring PI at inference.The student is updated even though only the teacher receives privileged information.
  • Privileged Information Distillation (π-Distill): The balance parameter α controls teacher- versus student-focused learning, with α = 1 teacher-only, α = 0 student-focused, and α = 0.5 joint optimization.Parameter sharing can still improve the teacher when α = 0, while α = 0.5 optimizes both policies jointly.
  • On-Policy Self-Distillation: OPSD samples on-policy student trajectories and applies reverse KL between the student and PI-conditioned teacher as a dense per-token reward.Its expectation is taken over the student policy, distinguishing it from π-Distill’s teacher-generated trajectories.

4 Experimental Setting

The experiments evaluate PI distillation for frontier-model traces in multi-turn tool-calling environments, using several PI formats, models, and baselines. Evaluation spans τ-Bench, Travel Planner, and broader out-of-domain settings.

  • Experimental Setting: The study focuses on multi-turn tool-calling environments because weaker models struggle to sample successful trajectories independently while frontier models often hide their CoT reasoning.This makes frontier-model actions valuable as training-time PI while limiting access to full reasoning traces.
  • Environments: The evaluation uses τ-Bench airline and retail customer-service tasks, with a modified Qwen-14B user simulator and removal of the reward-hacking transfer_to_human_agents tool.The benchmark simulates booking flights and assisting shoppers through tool calls and user-information gathering.
  • Sources of Privileged Information: Frontier-model raw output trajectories are used to benchmark action-only distillation against baselines that assume access to CoT.DeepSeek-chat-v3.1 supplies the trajectories, and its accessible reasoning tokens enable the CoT comparison.
  • Sources of Privileged Information: Three PI formats vary information density: complete tool calls with arguments, tool-call names only, and self-generated hints summarizing successful frontier trajectories.Tool-call-only PI requires the model to infer arguments, while hints can filter inefficient actions and suggest improvements.
  • Sources of Privileged Information: The analysis tests how PI attributes and α affect transfer, with the optimal configuration depending on PI properties such as teacher–student KL.The paper uses these variations to identify factors associated with successful transfer.
  • Models and Baselines: Experiments cover Qwen3-4B, Qwen3-8B, and R1-Distill-Llama-8B, comparing π-Distill and OPSD with RL, SFT, and SFT-plus-RL baselines.R1-Distill-Llama-8B is warm-started with SFT with CoT because direct RL could not generate correct PI-conditioned trajectories.
  • Models and Baselines: Table 1 reports mean ± standard deviation across three random seeds for Travel Planner and τ-Bench retail and airline.The table compares shaded proposed-method rows with baselines, including those using frontier reasoning traces.

5 Main Results

π-Distill is consistently strong across the reported agentic benchmarks and often exceeds SFT with CoT followed by RL, while OPSD is competitive in selected settings. The methods also generalize across models and out-of-domain environments.

  • Main Results: π-Distill variants achieve superior performance in all but one reported setting, including cases where full CoT supervision is available.With Qwen3-8B, π-Distill consistently outperforms SFT with CoT plus RL regardless of α.
  • Main Results: 11.8% on Travel Planner, 2.08% on τ-Bench retail, and 6.00% on τ-Bench airline are the best-case π-Distill improvements.These gains support effective frontier-model distillation despite hidden CoT traces.
  • Efficiency: π-Distill requires one training phase, whereas SFT with CoT plus RL requires sweeping multiple SFT checkpoints to find peak performance.The single-phase procedure simplifies training while delivering superior performance gains in the reported comparisons.
  • OPSD: OPSD substantially outperforms base RL and SFT without CoT plus RL on τ-Bench but performs similarly to them on Travel Planner.When successful, OPSD also provides substantial out-of-domain gains on τ-Bench airline and scales with model capacity.
  • Main Results: The best α varies by setting, with no definitive universal value.This variation motivates further analysis of how the student–teacher balance interacts with task conditions.
  • Main Results: π-Distill remains useful for R1-Distill-Llama-8B even after the model has been SFT-trained on traces containing PI.This result extends the method’s reported utility to a model with prior PI-containing trace supervision.

6 Out of Domain Experiments (OOD)

π-Distill generalizes effectively to out-of-domain tool-use tasks, consistently outperforming standard RL and the base model, while OPSD is competitive for stronger models. On Qwen3-8B, both methods can also surpass SFT with CoT followed by RL.

  • Out-of-domain generalization: π-Distill consistently outperforms the base model and standard RL on the 7-dataset GEM search-tool benchmark suite for both Qwen3 models.Checkpoints are selected on τ-Bench Retail and evaluated with Pass@1 and Pass@10 across three seeds.
  • Model-size dependence: For Qwen3-8B, π-Distill variants with α = 0 or α = 0.5 and OPSD can significantly outperform SFT with CoT followed by RL.On Qwen3-4B, SFT with CoT followed by RL remains consistently strongest, while OPSD degrades substantially.
  • Generalization and forgetting: Both algorithms exhibit substantially less forgetting than standard RL, and π-Distill and OPSD generalize significantly better than SFT with CoT followed by RL for Qwen3-8B.The SFT baseline uses frontier-model CoT traces, whereas the PI methods avoid relying on them.
  • Prevention of degradation: Standard RL consistently degrades relative to the base model, whereas π-Distill avoids regression in all cases and OPSD avoids degradation in stronger reasoners.The reported results also show significant improvements for both methods as model size scales.

7 What Matters When Using Train-Time PI

The analysis identifies PI utility, teacher–student divergence, information content, and collapse prevention as key determinants of transfer. π-Distill is most robust with joint training, while OPSD benefits most from information-rich PI but can fail under high KL or negative utility.

  • Analysis framework: The analysis tracks teacher–student divergence, PI utility ∆, and maximum attainable utility ∆max to explain when PI training converts privileged access into gains.∆max is defined as the difference between the best training-task scores achieved with and without PI.
  • Divergence and utility: Higher initial KL generally predicts lower final performance, but the relationship depends on the training configuration and PI type.Figure 5 compares held-out performance with initial divergence and training-task PI utility.
  • Teacher-only training: Teacher-only training can collapse teacher and student policies, causing the teacher to ignore PI and underperform even the RL baseline.The KL falls near zero early in training, making πT approximately equal to πS.
  • Teacher learning: Training the teacher can turn initially negative PI utility into positive ∆max, showing that learning to use PI contributes to transfer.This effect occurs even when the PI initially underperforms the base model.
  • Joint training: Joint training with α = 0.5 is the most robust configuration, performing best in 7 of 16 scenarios and worst only once.Balancing teacher and student objectives mitigates failure modes of independent training, although more granular α tuning remains future work.
  • Student-only training: Student-only training succeeds with low-KL, positive-utility PI but degrades when KL rises or PI utility is negative.Negative utility can prevent the algorithm from extracting value from PI, whereas low KL can allow knowledge transfer back to the teacher.
  • OPSD: For OPSD, PI information content is the strongest predictor of transfer, with Tool Calls & Arguments generally performing best for the Qwen3 models.High-information PI is not universally successful: Qwen3-8B on τ-Bench is an exception with high KL and negative ∆max.

8 Ablation on β

The β ablation shows that KL regularization is generally important for π-Distill, especially when the teacher is updated, although its optimal value varies by setting.

  • Ablation findings: β is important for achieving the best performance in 17/21 ablated configurations across π-Distill and OPSD settings.The sweep evaluates β = {0, 0.1, 0.25, 0.5} over three random seeds.
  • π-Distill regularization: For π-Distill, β > 0 generally improves or matches performance, particularly when teacher updates are enabled with α > 0.The best β remains sensitive to the task, PI type, and training configuration.

9 Related Work

The paper relates PI distillation to latent-reasoning, self-bootstrapping, privileged-signal, and context-distillation approaches. Its teacher–student setup shares parameters and avoids requiring oracle answers or frontier-model CoT.

  • Latent reasoning: STaR uses hints after incorrect answers and supervised fine-tuning on generated reasoning traces, resembling π-Distill student-only training without teacher updates.The paper distinguishes its KL-regularized off-policy RL objective from STaR’s SFT procedure.
  • Variational reasoning: Variational reasoners iteratively train separate teacher and student parameters with oracle answers, whereas this work shares parameters and does not assume ground-truth solutions.The paper compares against a modified similar setup because its environments lack ground-truth solutions.
  • Self-bootstrapping and privileged signals: Related self-bootstrapping work uses self-generated hints or privileged oracle solutions as exploration signals for hard reinforcement-learning and reasoning tasks.These approaches are framed as training the teacher to use privileged information.
  • Context distillation: Both proposed methods are forms of context distillation, which aims to internalize benefits from rich contexts into model weights.Related work also studies adapters that retain context-dependent knowledge for use at test time.

10 Limitations & Future Work

The study demonstrates effective PI-based distillation without CoT, but its evidence remains bounded by the source of PI, model scale, and observational analysis.

  • All privileged information is transformed from frontier model traces, leaving PI acquisition without frontier actions or ground-truth answers for future work.
  • Experiments are limited to models with ≤8B parameters, so larger-scale behavior remains untested.
  • The factor analysis is observational rather than systematically controlled, limiting conclusions about how configurations cause failure or success.

11 Conclusion

The paper introduces PI-based teacher–student algorithms that improve policies even when privileged information is unavailable at test time. Across training domains and out-of-domain datasets, the methods are reported as highly effective, while the study also identifies base-model statistics that predict success.

  • π-Distill and OPSD use training-time PI to produce improved policies without requiring that information at test time.
  • Evaluations across two training domains and eight OOD datasets find both methods highly effective in all explored settings.
  • Distillation success can often be predicted using only base-model statistics.

13 Author Contributions

The paper credits its contributors across project leadership, algorithm and environment implementation, experiments, advising, research direction, technical guidance, writing, and figures.

  • Emiliano Penaloza proposed leveraging PI, led the project, implemented the algorithms and RL environments, conducted most experiments, and primarily wrote the paper.
  • Dheeraj Vattikonda supported experiments, implemented and conducted OOD experiments, and contributed substantially to writing and figures.
  • Nicolas Gontier advised the project and provided feedback that improved experimental rigor.
  • Alexandre Lacoste advised the project, contributed insights on PI, improved rigor, and assisted with writing.
  • Laurent Charlin influenced the framing and perspective, and contributed to writing, including much of the introduction.
  • Massimo Caccia shaped the research direction, proposed OPSD, guided ablations, and contributed technical guidance, writing, and figures.

D.2 Reference Model Ablation

The appendix studies reference-policy choice and related training behavior, while documenting failure modes and evaluation details for the agentic experiments.

  • Reference Model Ablation: Using the student as the KL reference yields stable learning across α ∈{0, 0.5, 1}, whereas a fixed base reference can degrade or collapse performance.
  • Reference Model Ablation: For π-Distill, β > 0 generally improves or matches performance and is especially useful when training the teacher.
  • The reported 31.11% averages each seed’s peak, so it can exceed the peak of the curve that averages seeds at each step.
  • The original Travel Planner reward encourages reward hacking, with policies exploiting the easy-constraint score and converging to the same response structure.
  • The revised reward checks each easy constraint individually before its corresponding hard constraint, rather than requiring all easy constraints first.
Loading 2602.04942v3…