# Flux-OPD: On-Policy Distillation with Evolving Contexts

> Flux-OPD stabilizes on-policy distillation with evolving contexts by anchoring to a context-free teacher and weighting correction by teacher disagreement.

- **Source:** [arXiv](https://arxiv.org/abs/2607.28022)
- **Published:** 2026-08-01
- **Permalink:** https://picx.dev/p/DRyKXl
- **Whiteboard:** https://picx.dev/p/DRyKXl/image

## Summary

## 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)**

| Method | VBench (Wan2.1-VACE-1.3B) | VBench (CogVideoX-2B) | Video-Bench (Wan2.1) | Video-Bench (CogVideoX) |
|--------|---------------------------|----------------------|---------------------|----------------------|
| Teacher | 81.48 | 77.68 | 3.74 | 3.44 |
| Student | 81.40 | 77.97 | 3.73 | 3.42 |
| + OPD | 81.13 | 77.42 | 3.73 | 3.34 |
| + OPCD | 81.18 | 76.85 | 3.74 | 3.40 |
| + OEL | 80.26 | 73.46 | 3.60 | 3.01 |
| **+ Flux-OPD** | **82.25** | **78.11** | **3.77** | **3.45** |

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

| Method | Completeness | Context Awareness | Accuracy | Communication | Instruction Following | Total |
|--------|-------------|------------------|----------|---------------|---------------------|-------|
| Teacher | 30.70 | 32.40 | 46.61 | 63.65 | 52.96 | 37.38 |
| Student | 11.34 | 20.14 | 29.40 | 40.84 | 40.23 | 19.06 |
| + OPD | 13.01 | 21.13 | 29.10 | 39.79 | 38.80 | 19.63 |
| + OPCD | 12.16 | 20.84 | 29.78 | 41.09 | 37.06 | 19.53 |
| + OEL | 11.88 | 21.27 | 29.67 | 41.07 | 38.68 | 19.66 |
| **+ Flux-OPD** | 12.86 | **21.30** | **31.50** | **41.99** | 39.67 | **20.61** |

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

| Method | VBench (Wan2.1) | VBench (CogVideoX) | Average |
|--------|----------------|-------------------|---------|
| Teacher | 81.70 | 78.87 | 80.29 |
| Student | 81.38 | 77.83 | 79.61 |
| + OPD | 82.00 | 77.78 | 79.89 |
| + OPCD | 81.51 | 77.78 | 79.65 |
| + OEL | 81.71 | 77.78 | 79.75 |
| **+ Flux-OPD** | **82.26** | **78.21** | **80.24** |

### Ablation and Parameter Studies

**Table 4: Ablation on HealthBench**

| Variant | Evolving Contexts | Contextual Correction | Contextual Weighting | $\lambda_k$ | Score |
|---------|------------------|---------------------|---------------------|-------------|-------|
| V1 | ✗ | ✓ | ✓ | Eq. (20) | 19.63 |
| OEL | ✓ | ✗ | ✗ | ≡ 1.0 | 19.66 |
| V2 | ✓ | ✓ | ✗ | ≡ 0.9 | 19.73 |
| V2 | ✓ | ✓ | ✗ | ≡ 0.7 | 20.03 |
| V2 | ✓ | ✓ | ✗ | ≡ 0.5 | 19.97 |
| **Full** | ✓ | ✓ | ✓ | Eq. (20) | **20.61** |

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.

---

_Markdown view of https://picx.dev/p/DRyKXl, served by PicX — AI-generated visual whiteboard summaries of research papers._