Summary (Overview)

• Flux-OPD is a novel On-Policy Distillation (OPD) paradigm that uses evolving contexts as in-training supervision to capture task preferences in open-ended domains where verifiable rewards are unavailable.
• The paper provides a reverse KL decomposition showing that context-conditioned distillation targets the geometric mean of teacher distributions and includes a conflict term measuring disagreement among teachers.
• To address unstable distillation targets and conflicting distributions from evolving contexts, Flux-OPD introduces two strategies: contextual correction (anchoring to a stable context-free teacher with contextual difference signals) and contextual weighting (adjusting correction strength using the conflict term as an indicator).
• Experiments on prompt optimization for video generation and medical question answering show Flux-OPD outperforms existing OPD paradigms (OPD, OPCD, OEL) across different student-teacher settings.
• Flux-OPD demonstrates improved training stability (no loss surges), better generalization on out-of-distribution tasks (IF-Eval), and consistent gains across multiple training runs.

Introduction and Theoretical Foundation

Large language model training in open-ended domains (e.g., video generation prompt optimization, medical QA) lacks verifiable rewards, making task preferences difficult to formalize as effective supervision. Context Distillation and On-Policy Context Distillation (OPCD) use contexts as supervision signals, but contexts extracted before training provide little additional supervision once distilled into the student. This motivates evolving contexts that adapt with student performance.

However, directly using evolving contexts causes:

• Unstable distillation targets: Each update abruptly shifts the target distribution, causing loss surges.
• Conflicting distributions: Different contexts can produce conflicting target distributions.

The paper analyzes contexts through a reverse KL decomposition:

Proposition 1 (Reverse KL Decomposition): For a fixed decoding history, let $p_\theta$ and $\{q_c\}$ be distributions over vocabulary $V$, where $q_c(v) > 0$ for every context $c$ and token $v \in V$. Define the normalized geometric mean:

$$q_{\text{geo}}(v) = \frac{\exp(\mathbb{E}_c[\log q_c(v)])}{Z}, \quad Z = \sum_{v \in V} \exp(\mathbb{E}_c[\log q_c(v)])$$

Then the reverse KL objective decomposes as:

$$\mathbb{E}_c[D_{\text{KL}}(p_\theta \| q_c)] = D_{\text{KL}}(p_\theta \| q_{\text{geo}}) + (-\log Z)$$

where $-\log Z \geq 0$ is the conflict term measuring distributional disagreement among context-conditioned teachers. By Jensen's inequality:

$$Z = \sum_{v \in V} \exp(\mathbb{E}_c[\log q_c(v)]) \leq \sum_{v \in V} \mathbb{E}_c[q_c(v)] = 1, \quad -\log Z \geq 0$$

Corollary 1 (Gradient Independence): The conflict term $-\log Z$ is independent of the current student token distribution, contributing no direct gradient with respect to student parameters through $p_\theta$. Only the distillation term $D_{\text{KL}}(p_\theta \| q_{\text{geo}})$ is directly optimized.

This decomposition reveals that:

• Under reverse KL, the student is distilled toward the geometric mean of context-conditioned teachers (not arithmetic mean as under forward KL).
• The geometric mean assigns high probability only to tokens receiving consistently high probability across contexts.
• The conflict term measures distributional disagreement and can serve as an indicator for weighting.

Methodology

Flux-OPD is an iterative training process that divides a single training run into $K$ iterations, each comprising context extraction and context distillation.

Iterative Training

1. Context Extraction: At iteration $k$, collect trajectories from the previous student policy $\pi_{\theta_{k-1}}$:

   $$\mathcal{D}_{k-1} = \{(x_i, \hat{y}_i)\}_{i=1}^n$$

   The teacher extracts new experience items and accumulates them with previous items, repeated for $M$ turns with different random seeds to form a context pool:

   $$\mathcal{C}_k = \{c_{k,1}, \ldots, c_{k,M}\}$$

2. Context Distillation: Contexts sampled from $\mathcal{C}_k$ condition the teacher, allowing the target distribution to evolve with student performance.

Contextual Correction

To stabilize the distillation target, anchor it to the stable context-free teacher and inject contextual difference signals:

1. Compute normalized geometric mean of context-conditioned teachers:

   $$\log \tilde{q}_{\text{geo},k}(v) = \frac{1}{|\mathcal{C}_k|} \sum_{c \in \mathcal{C}_k} \log q_c(v)$$ $$q_{\text{geo},k}(v) = \frac{\exp(\log \tilde{q}_{\text{geo},k}(v))}{Z_k}$$

2. Compute contextual difference signal relative to context-free teacher $q_0$:

   $$\Delta_k(v) = \log q_{\text{geo},k}(v) - \log q_0(v)$$

3. Obtain corrected teacher distribution via log-space interpolation:

   $$q_k^{\text{flux}}(v) = \text{softmax}_v(\log q_0(v) + \lambda_k \Delta_k(v))$$

where $\lambda_k \in [0,1]$ controls correction strength. When $\lambda_k = 0$, target reduces to $q_0$; when $\lambda_k = 1$, it recovers $q_{\text{geo},k}$.

Contextual Weighting

Use the conflict term $-\log Z_k$ as an indicator to adjust correction strength:

$$\delta_k = -\log Z_k$$

Define the weight:

$$\lambda_k = \alpha \cdot \text{clip}\left(1 - \frac{\delta_k}{\tau}, \lambda_{\min}, \lambda_{\max}\right)$$

where:

• $\alpha$ is a scaling factor
• $\tau > 0$ is the conflict threshold
• $0 \leq \lambda_{\min} \leq \lambda_{\max} \leq 1$ define clipping bounds

A larger $\delta_k$ (higher conflict) reduces correction strength; smaller $\delta_k$ (consistent contexts) increases it.

Training Objective

At each iteration, train the student toward the corrected teacher using reverse KL divergence:

$$\mathcal{L}_k^{\text{flux}} = \mathbb{E}_{h_t} \left[ D_{\text{KL}} \left( p_\theta(\cdot | h_t) \| q_k^{\text{flux}}(\cdot | h_t) \right) \right]$$

Empirical Validation / Results

Experimental Setup

• Models: Qwen3-VL-Instruct (4B student, 8B teacher), Qwen2.5-VL-Instruct (7B student, 32B teacher), Qwen3 (1.7B student, 8B teacher)
• Tasks: Prompt optimization for video generation (VBench, Video-Bench) and medical question answering (HealthBench)
• Baselines: OPD, OPCD, OEL
• Context extraction: Each context from 8 student rollouts, $M=3$ random seeds, contexts updated every 300 steps

Main Results

Table 1: Prompt Optimization (Qwen3-VL-Instruct 8B Teacher, 4B Student)

Table 2: Medical QA (Qwen3 8B Teacher, 1.7B Student)

Table 3: Prompt Optimization (Qwen2.5-VL-Instruct 32B Teacher, 7B Student)

Ablation and Parameter Studies

Table 4: Ablation on HealthBench

Key findings:

• Evolving contexts + contextual correction (V2) outperforms OEL at all static $\lambda_k$ values.
• Full Flux-OPD with contextual weighting outperforms all static $\lambda_k$ variants.
• Conflict threshold $\tau = 0.1$ works best when student underperforms teacher; $\tau = 0.2$ when student outperforms teacher.

Stability and Generalization

• Training stability: Flux-OPD loss decreases steadily like OPD, while OEL exhibits abrupt loss surges at context updates. Flux-OPD also maintains stable gradient norms.
• Out-of-distribution generalization: On IF-Eval after medical QA training, Flux-OPD achieves 34.38% prompt-level strict accuracy vs. OPD's 32.90% and OEL's 31.79%.

Theoretical and Practical Implications

Theoretical Contributions:

• The reverse KL decomposition (Proposition 1) provides a principled understanding of how contexts affect distillation objectives:
  • Student is distilled toward the geometric mean of context-conditioned teachers (not arithmetic mean).
  • The conflict term $-\log Z$ measures distributional disagreement and can be used as a weighting signal.
  • The conflict term contributes no direct gradient, making it an ideal indicator for adjusting correction strength.

Practical Implications:

• Flux-OPD enables effective use of evolving contexts in a single training run, avoiding the need for multiple deployment-distillation cycles.
• The contextual correction strategy stabilizes training by anchoring to a stable context-free teacher, preventing the loss surges seen in OEL.
• The contextual weighting strategy automatically adjusts correction strength based on context consistency, preventing conflicting distributions from harming training.
• The method generalizes across different student-teacher sizes (1.7B/8B, 4B/8B, 7B/32B) and tasks (video generation, medical QA).

Conclusion

Flux-OPD is an OPD paradigm that uses evolving contexts as in-training supervision to capture task preferences in open-ended domains. The key innovations are:

1. Reverse KL decomposition showing that context-conditioned distillation targets the geometric mean of teacher distributions with a conflict term measuring disagreement.
2. Contextual correction anchoring distillation to the stable context-free teacher and injecting contextual difference signals.
3. Contextual weighting using the conflict term as an indicator to adjust correction strength.

Experimental results on prompt optimization for video generation and medical question answering demonstrate that Flux-OPD outperforms existing OPD paradigms (OPD, OPCD, OEL) across different student-teacher settings, with improved training stability and better out-of-distribution generalization.

Future Directions: The paper highlights the potential to combine teacher supervision with evolving contexts, suggesting further exploration of context evolution mechanisms and applications to other open-ended domains.