DistIL: Learning from Rich Feedback via Distributional Imitation

Paper: 2606.05152 Authors: Rishabh Agrawal, Jacob Fein-Ashley, Paria Rashidinejad Categories: cs.LG, cs.AI, cs.CL

The Gap

Current reasoning model training relies on Reinforcement Learning from Verifiable Rewards (RLVR): generate many candidate solutions, check if the final answer is correct, assign a binary reward. This recipe wastes rich intermediate feedback like execution traces, tool outputs, and step-by-step corrections. Existing approaches using self-distillation (reverse KL, Jensen-Shannon divergence) fail to guarantee that incorporating “better” trajectories actually improves the policy—they can increase probability on worse actions even when the expert has higher reward.

Problem: Binary rewards waste rich feedback
   |
   v
Assumption: Expert distribution available on policy-visited states
   |
   v
Method: Forward cross-entropy with expert distribution (DistIL)
   |
   v
Evidence: Monotonic improvement guarantees + empirical gains
   |
   v
Conclusion: Rich credit assignment > binary final-answer signals

The Increment

One sentence: Before this paper, reasoning models trained on final-answer correctness; after, they can propagate disagreement with expert trajectories backward through entire reasoning chains while guaranteeing policy improvement.

Core Mechanism

DistIL builds on DAgger (Dataset Aggregation) but operates in the distributional regime. At each training iteration, the current policy π generates trajectories. For states visited by π, an expert provides a distribution over what actions it would take. The learner minimizes forward cross-entropy between its action distribution and the expert’s distribution at each decision point.

The key operation is credit assignment: when the expert and student disagree at a future state, that disagreement signal propagates backward through the sequence-level gradient. If the expert places high probability on a different action at timestep t+5, the gradient adjusts not just that future decision but also earlier choices at t, t+1, t+2 that led to that divergence.

Unlike standard RL objectives that only see a final scalar reward, DistIL receives distributional feedback at every intermediate state. The forward cross-entropy direction matters: it pulls the student’s distribution toward the expert’s, rather than pushing the expert toward the student (reverse KL) or meeting halfway (Jensen-Shannon). This directional choice ensures that when the expert is better, updates provably increase probability on higher-reward actions.

Trajectory Generation:
   Current Policy π
        |
        v
   [s0] --a0--> [s1] --a1--> [s2] --a2--> [final]
        |           |           |
        v           v           v
   Expert dist  Expert dist Expert dist
   P*(·|s0)     P*(·|s1)    P*(·|s2)
        |           |           |
        +-----+-----+-----+
              |
              v
    Forward Cross-Entropy: H(P*, π)
              |
              v
    Gradient propagates disagreement backward
    through sequence: ∇ affects a0, a1, a2

Think of this as a driving instructor who sits in the passenger seat while you practice. Standard RL is like only telling you “crash” or “no crash” at the end. DistIL’s instructor continuously shows you what they would do at each intersection—turn signal, speed adjustment, lane positioning. When you deviate from their distribution of actions at mile 5, the learning signal reaches back to correct your steering at mile 2 that set up the deviation. The “forward” direction means you’re adjusting your habits toward the instructor’s, not averaging between your mistakes and their expertise. The instructor’s distribution (not single action) means they can say “70% brake gently, 30% change lanes”—capturing uncertainty and multiple valid choices.

Key Concepts

• Forward vs Reverse KL Divergence: KL divergence measures how different two probability distributions are, but direction matters. Forward KL D_KL(Expert || Student) heavily penalizes the student for placing probability mass where the expert doesn’t—forcing coverage of all expert-approved options. Reverse KL D_KL(Student || Expert) does the opposite: it’s fine if the student ignores some expert options, as long as where it does place mass, the expert also approves. For learning from a better teacher, forward KL (used here) ensures you don’t miss good options; reverse KL lets you cherry-pick and can lock onto suboptimal modes even when better data is available.

• Distributional Expert Access: Instead of querying “what’s the single best action here?”, the learner accesses “what’s your probability distribution over actions?”. This is weaker than needing a perfect oracle (expert might be uncertain) but stronger than scalar rewards (captures relative preferences). Practically, this could be an ensemble of models, a trained value function estimating action quality, or self-evaluation scores. The key property: local access on states the current policy visits, not requiring the expert to demonstrate optimal trajectories from scratch.

• Sequence-Level Credit Assignment: In language model training, we make many sequential decisions (choosing tokens). A mistake at token 5 might stem from a bad choice at token 2. Standard supervised learning treats each token independently. Sequence-level gradients flow through the entire chain: when the objective measures disagreement at the sequence end, backpropagation attributes responsibility to earlier decisions. DistIL’s forward cross-entropy naturally creates this: expert-student distributional mismatch at future states produces gradients that adjust earlier policy decisions that led there.

Framework Shift

Before (RLVR):                       After (DistIL):
                                     
Policy π generates N samples         Policy π generates samples
        |                                    |
        v                                    v
Final answer checked                 Expert distribution at each state
    [Correct=1, Wrong=0]                     |
        |                                    v
        v                            Forward CE: H(P*, π)
RL objective (e.g. PPO)                      |
Maximize reward signal                       v
                                    Gradient propagates backward
Single bit at end                   Rich signal at every step

[One sentence: From end-of-sequence binary signals to every-step distributional alignment, the core shift is continuous multi-step credit assignment.]

Expert Assessment

Problem choice: Real gap. RLVR’s wastefulness is well-known—practitioners manually engineer reward shaping to inject intermediate signals. Formalizing how to use rich feedback with theoretical guarantees is a natural next step, not manufactured. This sits at the intersection of imitation learning (classical DAgger) and modern LLM training (RLHF/RLVR), bridging two mature areas.

Method maturity: Clever adaptation rather than invention. DAgger is 2011 tech; applying it distributionally and analyzing the objective choice (forward KL) is the insight. The reverse-KL failure modes (Theorem 3.1) and monotonic improvement guarantee (Theorem 4.1) provide real value—prior work used reverse KL or JS without realizing they lack these properties. However, the practical expert construction (Section 6: Best-of-N sampling, self-evaluation) feels more like “what we tried” than principled design. The gap between “assume expert distribution” and “here’s how to build one” could be tighter.

Experimental integrity: Baselines are fair (RLVR, self-distillation variants) and domains diverse (math, code, reasoning). The improvements are consistent but not massive (Table 1: +2-8% over RLVR). A concern: expert quality isn’t isolated. When DistIL outperforms baselines, is it the forward-KL objective or simply having access to distributional feedback? An ablation where baselines also get expert distributions would clarify. The Pass@N results (Table 2) are interesting but feel secondary—optimization for different metrics should be explicit.

Writing quality: Strong motivation and theory sections. The related work undersells the connection to reward shaping and curriculum learning, which solve similar problems differently. Section 5 (regret bounds) is dense—moving proof sketches to appendix and leading with intuition would help. The biggest miss: Figure 1 should show a worked example (actual reasoning trace, expert distribution, gradient flow), not just algorithmic pseudocode. Readers will skim the algorithm but remember a concrete instance.

Verdict: Weak accept — Solid theoretical contribution (monotonic improvement guarantee) and consistent empirical gains, but expert construction remains underspecified and experiments don’t fully isolate the objective’s contribution from the feedback type.

Takeaways

For practitioners: If you have access to intermediate feedback (tool outputs, execution traces, model ensembles), use forward cross-entropy to align with it rather than reverse KL or JS divergence. The directional choice matters—forward CE ensures you cover all expert-preferred options, not just mode-collapse onto one.

For researchers: The failure mode of reverse KL (Theorem 3.1) transfers beyond this setting. Any time you’re distilling from a better distribution into a worse one, check whether your objective guarantees improvement or could increase probability on worse actions. The proof technique (bounding policy improvement via distributional shift) is reusable.

Concrete technique: Best-of-N sampling as a poor-man’s expert distribution (Section 6.1). Generate N trajectories, score them, construct a categorical distribution over actions weighted by trajectory success. Simple, requires no new models, and empirically works. This is immediately deployable in any setting with verifiable outcomes.