🤖 machine learning

Learning from Language Feedback via Variational Policy Distillation

This paper introduces Variational Policy Distillation (VPD), a framework that overcomes the limitations of passive teacher-student self-distillation by co-evolving both policies through a Variational Expectation-Maximization process, thereby enabling continuous improvement in extracting actionable signals from language feedback to enhance performance on complex reasoning and code generation tasks.

Original authors: Yang Li, Erik Nijkamp, Semih Yavuz, Shafiq Rayhan Joty

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yang Li, Erik Nijkamp, Semih Yavuz, Shafiq Rayhan Joty

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Silent Teacher"

Imagine you are learning to solve complex puzzles (like math problems or coding challenges) using a very smart but slightly clumsy robot (the AI model).

In the old way of training these robots (RLVR), the robot tries to solve a puzzle, and a strict judge gives it a simple "Yes" or "No" at the very end.

The Problem: If the robot makes a tiny mistake in step 3 of a 10-step solution, it gets a "No." If it hallucinates nonsense and fails completely, it also gets a "No."
The Result: The robot doesn't know why it failed. It's like trying to learn to drive by only getting a red light at the end of the trip if you crash, with no explanation of whether you hit a pedestrian or ran a stop sign. This makes learning incredibly slow and inefficient.

The Current "Fix": The Passive Mentor

To fix this, researchers tried a new method where the robot gets language feedback. Instead of just "No," it gets a note saying, "You forgot to divide by zero in step 3."

The Idea: The robot reads this note and tries to fix its next attempt.
The Flaw: In existing methods, the robot acts as its own teacher. It reads the note and tries to guess the right answer.
The Limit: The robot is only as good as its current ability to read. If the note is confusing, or if the robot is too dumb to understand the note, the "teacher" version of the robot gives bad advice. Once the robot hits a ceiling in its ability to understand the notes, it stops improving. The teacher is passive; it never gets smarter at explaining things.

The New Solution: VPD (The "Co-Evolving Duo")

The authors propose Variational Policy Distillation (VPD). Instead of one robot trying to teach itself, they create a two-person team that grows together. They use a mathematical framework called Expectation-Maximization (EM), which is like a two-step dance:

Step 1: The "E-Step" (The Teacher Gets Smarter)

The Analogy: Imagine a Coach (the Teacher) and a Player (the Student).
What happens: The Player tries the puzzle and fails. The Coach looks at the failure and the feedback note.
The Magic: In this new method, the Coach doesn't just read the note once. The Coach actively trains to understand the note better. The Coach learns, "Ah, when the note says 'syntax error,' it usually means the Player missed a semicolon." The Coach gets better at diagnosing the problem.
The Safety Net: Crucially, the Coach is trained to stay close to the Player's current skill level. The Coach doesn't try to teach quantum physics if the Player is still learning arithmetic. This keeps the advice reachable and prevents the Player from getting overwhelmed.

Step 2: The "M-Step" (The Student Learns)

What happens: Now that the Coach has a sharper, clearer diagnosis, the Player tries to copy the Coach's thinking.
The Result: The Player internalizes the Coach's dense, detailed advice. The Player learns why it failed, not just that it failed.

Why This is Better (The "Shared Brain" Trick)

Usually, having a Coach and a Player means you need two separate computers (or a lot of memory), which is expensive.

VPD's Trick: They put the Coach and the Player in the same brain (a single neural network).
How it works: When the network needs to be the "Coach," it looks at the puzzle plus the feedback note. When it needs to be the "Player," it looks at the puzzle without the note.
Benefit: They co-evolve. As the Player gets better, the Coach gets better at explaining things. As the Coach gets better, the Player gets better at understanding. They pull each other up, avoiding the "ceiling" where other methods get stuck.

What the Paper Actually Found

The authors tested this on three main things:

Coding: Where the computer gives error messages (like a compiler).
Science: Where the computer checks if the answer is right or wrong.
Math: Where the answer must be perfectly correct.

The Results:

Coding & Science: VPD was the clear winner. It learned faster and more stably than the old methods. It could handle the feedback notes much better because the "Coach" kept getting sharper at reading them.
Math & "Cold Starts" (Starting from scratch): Here, the paper found a limit. If the robot starts with zero knowledge (a "cold start") or if the math is extremely rigid, VPD still struggled a bit compared to the old "Yes/No" method.
- Why? Sometimes, the feedback notes are too noisy or confusing for a beginner to understand. In these specific cases, the old method of just trying millions of times until it gets a "Yes" (Pure RL) is still the most powerful, though it is very slow.

Summary

Think of VPD as upgrading from a silent referee (who just says "Out!") to a dynamic coaching duo.

Old Way: The student tries, fails, and guesses why.
New Way (VPD): The student tries, fails, and a Coach (who is actively learning how to interpret the failure notes) gives a tailored lesson. Then the student practices that lesson. They repeat this, with the Coach getting smarter at teaching and the Student getting smarter at learning, all while sharing the same brain to save space.

The paper proves this works great for complex tasks where feedback is available, but it admits that for the hardest, most rigid math problems or total beginners, the old "trial and error" method still has the edge.

Technical Summary: Learning from Language Feedback via Variational Policy Distillation

Problem Statement

Reinforcement Learning from Verifiable Rewards (RLVR) has driven recent advances in the reasoning capabilities of Large Language Models (LLMs) by optimizing against objective, outcome-based correctness. However, standard RLVR methods (e.g., GRPO) rely on sparse, binary outcome signals (correct/incorrect). This sparsity creates a severe credit assignment bottleneck: a minor arithmetic error in a complex derivation receives the same zero-reward as a completely nonsensical hallucination. Consequently, on-policy algorithms face extreme exploration bottlenecks on hard problems where initial success rates are near zero, wasting latent information in "near-miss" trajectories.

To address this, recent on-policy self-distillation methods (e.g., SDPO, OPSD) utilize rich textual feedback (e.g., compiler errors, automated critiques) to generate dense, token-level supervision. However, these approaches suffer from a critical flaw: they treat the feedback-conditioned "self-teacher" as a fixed, passive function. The teacher relies entirely on its zero-shot ability to interpret feedback. As the student policy improves, the teacher's ability to distinguish between increasingly subtle reasoning errors plateaus because it is never explicitly trained to refine its diagnostic interpretation. This leads to a ceiling effect where the student is starved of meaningful gradients.

Methodology: Variational Policy Distillation (VPD)

The authors propose Variational Policy Distillation (VPD), a framework that formalizes learning from language feedback as a Variational Expectation-Maximization (EM) problem. Unlike passive distillation, VPD co-evolves both the student and the teacher policies within a single, shared-weight neural network.

The Variational Formulation

The framework treats the optimal policy $\pi^*$ (the reward-tilted distribution) as an intractable target. VPD introduces a parameterized teacher network $q_\phi(y|x, C)$ , conditioned on dense diagnostic feedback $C$ , to serve as a tractable approximate posterior for $\pi^*$ . This formulation decomposes training into two alternating steps:

E-Step (Teacher Refinement):
- Goal: Actively train the teacher to interpret language feedback and approximate the optimal reward-tilted distribution.
- Mechanism: The teacher is optimized to distinguish between successful and failed trajectories given the textual critique. Since the feedback $C$ is unique to each trajectory, standard paired preference optimization (like DPO) is inapplicable. Instead, VPD employs Binary Classifier Optimization (BCO), an unpaired preference optimization framework.
- Dynamic Reference Prior: Crucially, the teacher is not anchored to a static base model. To prevent distribution shift and ensure the teacher remains a useful critic for the student's current capabilities, the reference prior is dynamically anchored to the current student policy ( $\pi_{ref} \leftarrow \pi_\theta$ ). This enforces an adaptive trust region, ensuring the teacher's target distribution remains reachable by the student.
M-Step (Student Optimization):
- Goal: Internalize the refined, dense guidance from the teacher into the unconditioned student policy.
- Mechanism: The student minimizes the token-level KL divergence against the updated teacher $q_\phi$ on its own on-policy rollouts. Since the student operates without the privileged feedback $C$ at inference time, it must implicitly learn the reasoning corrections discovered by the teacher.

Architecture and Efficiency

Shared-Weight Network: Both the student ( $\pi_\theta$ ) and teacher ( $q_\phi$ ) are instantiated within a single network ( $\phi = \theta$ ). The distinction is purely contextual: the teacher is invoked by appending diagnostic feedback $C$ to the prompt, while the student relies solely on the input $x$ . This eliminates the massive memory overhead of multi-model paradigms.
Asymmetric Updates: The framework supports asymmetric update frequencies (e.g., multiple M-steps per E-step), acting similarly to a target network in RL to stabilize the distillation process.

Key Contributions

Formalization of Variational EM: The paper formalizes on-policy learning from language feedback as a Variational EM procedure. This introduces an explicit, feedback-aware teacher update (absent in prior self-distillation methods) implemented via unpaired preference optimization (BCO).
Dynamic Trust Region: By dynamically anchoring the teacher to the current student policy, VPD enforces an adaptive trust region that ensures stable on-policy KL distillation, avoiding the catastrophic distribution shifts seen in methods using static priors.
Comprehensive Empirical Study: The authors evaluate VPD across three distinct sources of diagnostic feedback:
- Deterministic environment verifiers (LiveCodeBench).
- Contrastive sibling rollouts (SciKnowEval).
- Autonomous self-critique via an LLM judge.
Characterization of Limits: The paper rigorously stress-tests the framework on base-model cold-start scenarios and rigid mathematical reasoning tasks to establish the fundamental bounds of feedback-driven self-distillation compared to pure sparse RL.

Experimental Results

Evaluated on benchmarks spanning competitive programming (LiveCodeBench), scientific reasoning (SciKnowEval), and mathematical reasoning (Math500, AIME), VPD consistently outperforms standard RLVR (GRPO) and existing self-distillation baselines (SDPO).

Code Generation (LiveCodeBench): VPD achieved a state-of-the-art pass rate of 49.62% on Qwen3-8B, outperforming GRPO (45.61%) and SDPO (47.33%). The method successfully maintained a growing reward margin between correct and incorrect responses, whereas SDPO's guidance signal degraded over time.
Scientific Reasoning (SciKnowEval): VPD established new state-of-the-art results across multiple models (Qwen3-1.7B, 8B, and OLMo3-7B). For instance, on Qwen3-8B, VPD achieved 77.15% average accuracy compared to SDPO's 74.44%. Notably, VPD eliminated the late-stage training degradation (validation accuracy drops) frequently observed in SDPO and hybrid baselines.
Self-Critique: Even when using autonomous self-critique (where the model judges its own failures without positive pairs), VPD significantly outperformed SDPO (78.14% vs. 74.87% on Qwen3-8B).

Fundamental Limits and Cold Start

The paper identifies specific regimes where self-distillation struggles relative to pure sparse RL:

Base Model Cold Starts: When applied to base foundation models (e.g., Qwen3-4B-Base) lacking instruction-following alignment, self-distillation methods (including VPD) struggle because the model cannot effectively parse the diagnostic feedback. While VPD significantly delays the training collapse observed in SDPO, pure RL (GRPO) remains the most effective paradigm for eliciting emergent reasoning in these cold-start scenarios.
Mathematical Reasoning: On strict mathematical benchmarks (DAPO-Math, Math500), VPD again delays training collapse but ultimately underperforms pure GRPO. The authors hypothesize that the strict, non-forgiving nature of mathematical derivations makes self-distillation vulnerable to "poisoning" by noisy or imprecise intermediate teacher guidance, whereas sparse RL encourages broader exploration to find rigorous paths.

Significance and Claims

The paper claims that VPD represents a principled shift from treating language feedback as a static heuristic to treating it as a learnable, co-evolving component of the optimization process. By embedding self-distillation within a rigorous EM framework with a dynamic trust region, VPD extracts significantly more value from language feedback than passive methods, offering stable optimization dynamics and consistent performance gains.

However, the authors maintain a modest stance regarding the ultimate limits of the approach. They explicitly state that while VPD mitigates the limitations of self-distillation, pure sparse RL remains the most effective paradigm for the most challenging logical domains and base-model initialization. The work illuminates the fundamental bounds of language-driven self-distillation, suggesting that while it is a powerful tool for refining reasoning, it may not yet fully replace the exploratory power of outcome-based RL in all contexts.