Technical Summary: When Does Online Imitation Learning Help in LLM Post-Training?

Problem Statement

The paper addresses a fundamental question in Large Language Model (LLM) post-training: Why and when does online imitation learning (IL), specifically on-policy distillation, outperform offline supervised fine-tuning (SFT)?

While empirical studies suggest online IL is superior, the prevailing theoretical explanation—that online interaction mitigates "error accumulation" over long horizons (as in DAgger-style algorithms)—is challenged by recent findings. Foster et al. [15] demonstrated that in realizable settings (where the expert policy lies within the student's hypothesis class), offline IL with log-likelihood objectives achieves the same sample complexity guarantees as online methods. This raises the question: if error accumulation is not the sole driver, what is the source of online IL's advantage in practical LLM training, where student models are often smaller and less capable than the experts they imitate?

The authors propose that non-realizability (misspecification)—the scenario where the student policy class $\Pi$ cannot represent the expert policy $\pi_e$—is the critical factor. They investigate whether online interaction provides benefits in non-realizable settings even in the absence of long-horizon compounding errors (i.e., in contextual bandits with horizon $H=1$).

Methodology

Theoretical Framework

The authors abstract LLM post-training as a contextual bandit problem with horizon $H=1$.

• Contexts ($X$): Prompts.
• Actions ($Y$): Generated response sequences.
• Policy ($\pi$): A conditional distribution over responses.
• Realizability: Defined as $\pi_e \in \Pi$. If $\pi_e \notin \Pi$, the setting is non-realizable.

They analyze two learning paradigms:

1. Offline IL (SFT): Learns from a static dataset of i.i.d. expert demonstrations $(x_i, y_i)$ by minimizing behavior cloning loss (negative log-likelihood).
2. Online IL (On-policy Distillation): The learner samples from its current policy $\pi$, queries the expert for the probability $\pi_e(y|x)$, and optimizes a reverse-KL divergence objective (or a general shaping function $f(\pi_e(y|x))$).

Key Analytical Tools

1. Discrepancy Analysis: The paper critiques traditional bounds that rely on distributional discrepancies (e.g., Total Variation, Hellinger distance) between the expert and student. They argue these bounds are overly pessimistic in non-realizable settings and misaligned with the true reward objective.
2. Information-Theoretic Lower Bounds (Offline): The authors derive a lower bound for offline IL sample complexity in non-realizable settings, showing it scales with a coverage coefficient $C_\infty^e$. This coefficient measures the mismatch between the expert's distribution and the best achievable policy within the restricted student class.
3. Structural Condition for Online IL (Misspecification–Reward Alignment): To explain online IL's success, they propose a condition where the expert's signal (used for distillation) aligns with the reward function in a way that the student class can exploit, even if the student cannot fully represent the expert.
   • Assumption 1 (Misspecification–Reward Alignment): The expert signal $f(\pi_e(y|x))$ is decomposed into a reward-aligned component $\alpha r(x,y)$ and a residual $h(x,y)$. The condition requires that within the set of policies near-optimal for the online objective, the residual term $h$ does not significantly distinguish between policies that differ in reward performance. Essentially, the "misspecification" must be orthogonal to the reward or unrepresentable by the student class.

Key Contributions and Results

1. Realizable Settings: Offline IL Suffices

The authors empirically verify that in realizable settings, online IL offers no advantage over offline IL.

• Experiments: They conducted experiments on synthetic (Countdown) and real-world math reasoning (GSM8K, DeepScaleR) tasks. In these setups, the "expert" was generated by an RL-trained model from the same family as the student, guaranteeing realizability.
• Findings: SFT (offline) fully recovered the expert's performance. On-policy distillation (online) provided no gains in final accuracy or training speed compared to SFT. This confirms that error accumulation is not the primary driver of online IL's success in these regimes.

2. Non-Realizable Settings: The Information Bottleneck

In non-realizable settings (where the expert is a much larger or different model), the paper identifies a fundamental limitation for offline IL:

• Theorem 1: Offline IL algorithms learning from i.i.d. expert data face an information-theoretic barrier. If the expert places low probability mass on the specific responses that the restricted student class can learn (high $C_\infty^e$), offline IL cannot achieve vanishing suboptimality regardless of sample size. The expert's demonstrations may simply not contain the "right" examples for the student to learn the optimal policy within its class.

3. Online IL's Advantage via Structural Alignment

The paper provides a theoretical justification for why online IL succeeds in non-realizable settings:

• Theorem 2: Under the Misspecification–Reward Alignment Condition, online IL can achieve low reward suboptimality even with large distributional discrepancies.
• Mechanism: Online interaction allows the student to query the expert on states (prompts) it actually visits. If the expert's signal is structured such that the reward-relevant information is accessible to the student (even if the expert's full distribution is not), online IL can converge to a high-reward policy.
• Empirical Validation: Using DeepSeek-R1 models, the authors showed that distributional discrepancies (e.g., in response length or keyword frequency) can be massive (e.g., 30% TV distance) while the performance gap remains small (5%). This supports the claim that discrepancy-based bounds are loose. Furthermore, they showed that distilling from an expert with better "alignment" (Skywork-7B) yielded significantly better student performance than a less aligned expert (R1-Distill-7B), consistent with their theoretical bound estimates.

Significance and Claims

The paper claims to shift the theoretical understanding of LLM post-training from an error-accumulation perspective to a misspecification perspective.

1. Reframing the Advantage: The primary advantage of online IL in LLMs is not merely correcting compounding errors over long horizons, but rather overcoming the information-theoretic bottleneck caused by non-realizability. In realizable settings, offline SFT is theoretically and empirically sufficient.
2. Necessity of Reward-Aware Analysis: The authors argue that analyzing IL solely through distributional discrepancies (e.g., KL divergence) is insufficient and often misleading. A reward-dependent analysis is required to explain why online methods succeed when the student cannot perfectly mimic the expert.
3. Practical Implications: The results suggest that the performance gap observed between SFT and online methods (like RL or on-policy distillation) in practice often stems from the non-realizability of the expert relative to the student. If the expert is too strong or structurally different, offline data may be insufficient, and online interaction becomes necessary to guide the student toward high-reward regions of its own policy space.

The paper concludes that while their framework explains the offline-online gap, it does not fully explain why on-policy distillation is sometimes more sample-efficient than RL (which also uses online interaction), suggesting that the multi-step structure of language generation and dense reward signals remain open areas for future investigation.