Technical Summary: Knowledge-Enhanced Preference Optimization (KEPO) for Multimodal Reasoning

1. Problem Statement

Reinforcement Learning (RL) has emerged as a primary paradigm for inducing explicit reasoning behaviors in Large Multimodal Vision-Language Models (VLMs). However, applying RL to reasoning-intensive tasks, particularly in high-stakes domains like healthcare, faces two fundamental failure modes:

1. Sparse Reward and Ambiguous Credit Assignment: Most existing methods rely on sparse, trajectory-level rewards (e.g., binary correct/incorrect). In long-context reasoning, a single local error can invalidate an otherwise coherent chain of thought. This sparsity leads to ambiguous credit assignment, where the model cannot distinguish which intermediate steps contributed to the final outcome, resulting in high-variance gradient estimates and unstable optimization.
2. Exploration Collapse (The "Learning Cliff"): Under standard on-policy RL, naive exploration often fails to generate any reward-bearing trajectories when rewards are rare or binary. This traps the policy in a "learning cliff" regime where the model receives no informative signals to improve, particularly in complex multimodal settings where errors can stem from both visual perception and language generation.

Recent attempts to mitigate sparse rewards using on-policy knowledge distillation (leveraging a strong teacher model for dense, token-level supervision) have shown promise. However, these methods typically apply distillation uniformly across all generated trajectories. The authors argue this is ill-suited for reasoning tasks: low-quality trajectories often originate from early logical errors or hallucinations. Distilling these flawed trajectories forces the student to mimic the teacher's distribution conditioned on an incorrect context, injecting noisy and misaligned gradients that hinder learning.

2. Methodology: KEPO

To address these challenges, the authors propose Knowledge-Enhanced Preference Optimization (KEPO), a unified post-training framework that integrates two tightly coupled components: Quality-Gated On-Policy Distillation and Knowledge-Enhanced Exploration.

2.1 Quality-Gated On-Policy Distillation

KEPO modifies the standard RL objective by introducing a dense teacher supervision term that is selectively activated based on trajectory quality.

• Mechanism: Instead of distilling all student-generated trajectories, the method applies a reward-based threshold ($\tau$). The distillation loss is only computed for trajectories where the reward $r_i \geq \tau$.
• Rationale: This ensures that dense teacher guidance is aligned with the reinforcement learning objective. It prevents the student from learning from "flawed contexts" where the reasoning chain has already broken down.
• Objective Function: The KEPO objective ($J_{KEPO}$) combines:
  1. An advantage-weighted preference optimization term (driving improvement based on relative trajectory quality).
  2. A quality-gated distillation term ($I_{r_i \geq \tau} D(\pi_T \parallel \pi_\theta)$) providing dense supervision only on reward-aligned paths.
  3. A KL regularization term to constrain policy updates.

2.2 Knowledge-Enhanced Exploration via Hint-Aware Rejection Sampling

To address the "learning cliff" where naive exploration fails to find reward-bearing trajectories, KEPO introduces an adaptive exploration mechanism.

• Trigger: During training, if a standard on-policy rollout of $G$ trajectories yields zero rewards (i.e., $\max r(y_i) = 0$), the exploration mechanism is triggered.
• Teacher-Guided Hints: A teacher model ($\pi_T$) generates a reasoning hint ($h$) conditioned on the input ($x$) and the ground-truth answer ($y$). Note: The ground-truth answer is used only as privileged information during this training-time failure recovery and is never used at inference.
• Hint-Aware Sampling: A self-enhanced policy ($\pi_E$), sharing parameters with the student but conditioned on the hint $h$ and the privileged answer, generates a proposal trajectory.
• Rejection Sampling: The proposal is accepted only if it yields a positive reward. If not, sampling is repeated up to a budget $B$.
• Injection: Accepted hint-aware trajectories are injected into the rollout buffer and jointly optimized with standard on-policy trajectories. This allows the policy to escape the learning cliff without relying on static offline supervision.

3. Key Contributions

The paper outlines four primary contributions:

1. KEPO Framework: A unified post-training framework integrating quality-gated teacher supervision with adaptive exploration to enhance RL for multimodal VLMs.
2. Identification of Failure Modes: The authors identify and address two specific failure modes in reasoning-oriented RL: (a) contextual noise caused by uniform distillation over flawed trajectories, and (b) exploration stagnation under sparse rewards. They demonstrate that selective gating and hint-aware injection are key to resolving these.
3. Experimental Benchmarking: The work instantiates recent on-policy distillation frameworks (specifically MM-GKD) in the multimodal vision-language setting to provide a principled basis for comparing distillation-based and RL-based post-training approaches.
4. Empirical Validation: Extensive experiments on a low-resource, single-source medical VQA benchmark demonstrate that KEPO yields robust reasoning and superior out-of-distribution (OOD) generalization compared to strong baselines.

4. Experimental Results

The authors evaluate KEPO on OmniMedVQA, a medical VQA benchmark spanning eight imaging modalities.

• Setup: Models are trained solely on MRI data (In-Domain) and tested on MRI plus seven other modalities (CT, X-Ray, Ultrasound, Dermoscopy, Fundus, OCT, Microscopy) as Out-of-Distribution (OOD) tests. This isolates cross-modality generalization capabilities.
• Baselines: Comparisons include general-purpose VLMs (Qwen3-VL), medical-specific VLMs (HuatuoGPT), standard post-training methods (SFT, GRPO, MM-DAPO), and multimodal on-policy distillation baselines (MM-GKD).
• Performance:
  • KEPO achieves the highest average accuracy (80.25%) and average OOD accuracy (78.00%) among all compared methods.
  • It significantly outperforms GRPO and uniform on-policy distillation baselines, particularly in the OOD setting.
  • Ablation Studies: Removing the exploration component (KEPO-KE) reduces OOD accuracy, confirming the value of guided exploration. Removing the quality gate ($\tau=0$) also degrades performance compared to the gated version ($\tau=1$), validating that restricting distillation to high-quality trajectories prevents noisy gradient injection.
  • Training Dynamics: KEPO demonstrates faster convergence and greater stability than GRPO, particularly in the early stages of training where sparse rewards typically cause stagnation.

5. Significance and Claims

The paper claims that KEPO provides a principled framework for integrating dense teacher guidance into reinforcement-based post-training. Its significance lies in:

• Solving the "Learning Cliff": By actively steering exploration toward informative trajectories when naive sampling fails, KEPO enables learning in regimes where standard RL would stagnate.
• Stabilizing Reasoning: By filtering out low-quality trajectories for distillation, KEPO avoids the injection of noisy gradients that typically plague reasoning-intensive tasks, allowing for the stable emergence of coherent chain-of-thought behaviors.
• Generalization: The framework enables models to learn transferable reasoning patterns that generalize across modalities, even when trained on a single source modality with limited data.

The authors position KEPO not as a replacement for RL, but as a structured augmentation that leverages teacher models dynamically to guide exploration and provide dense credit assignment specifically when standard on-policy mechanisms fail.