🤖 machine learning

Exploration-Driven Optimization for Test-Time Large Language Model Reasoning

The paper proposes Exploration-Driven Optimization (EDO), a novel framework that integrates diversity-promoting exploration objectives into iterative post-training methods like iDPO and GRPO to resolve the tension between inference-time sampling diversity and RL-induced distribution sharpening, thereby enhancing LLM reasoning performance and stability across both in-distribution and out-of-distribution tasks.

Original authors: Changhao Li, Yuchen Zhuang, Chenxiao Gao, Haotian Sun, Rushi Qiang, Chao Zhang, Bo Dai

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Changhao Li, Yuchen Zhuang, Chenxiao Gao, Haotian Sun, Rushi Qiang, Chao Zhang, Bo Dai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant student (a Large Language Model) who is learning to solve complex math puzzles. You want them to become a master problem-solver.

Usually, when we train these students, we use a method called Reinforcement Learning. Think of this like a strict coach who only rewards the student for finding the single best answer. Over time, the student gets very good at finding that one specific path. But here's the problem: they stop exploring. They become "stuck" in a rut, only knowing one way to solve a problem. If that one way fails, they have no backup plan.

This creates a conflict. To solve really hard puzzles, we often use a technique at test-time called "Self-Consistency." This is like asking the student to solve the same problem 10 different times and then picking the answer that appears most often. This works great if the student is generating 10 different approaches. But if the student has been trained to only know one approach, asking them to try 10 times just gives you 10 copies of the same (potentially wrong) answer.

The paper introduces a new training method called EDO (Exploration-Driven Optimization) to fix this.

The Core Idea: The "Curious Explorer" vs. The "Safe Specialist"

The authors realized that to make the "Self-Consistency" trick work, the model needs to be a Curious Explorer during training, not just a Safe Specialist.

The Old Way (The Specialist): The coach says, "If you get the answer right, great! If you get it wrong, try to be more like the last time you got it right." The student learns to repeat the same successful pattern over and over. The result is a very narrow, "sharp" way of thinking.
The New Way (EDO - The Explorer): The coach says, "Get the answer right, but also, don't just repeat what you did last time. Try a slightly different path. If you keep doing the exact same thing, I'm going to gently nudge you to try something new."

The Creative Analogy: The Maze and the Flashlight

Imagine the student is in a giant, dark maze (the problem space) trying to find the exit (the correct answer).

Standard Training: The student finds a path to the exit. They get a reward. Next time, they are trained to follow that exact same path with laser focus. They stop looking at the walls or other corridors. If that one path gets blocked later, they are stuck.
EDO Training: The student finds a path to the exit and gets a reward. But the coach adds a rule: "You must also wander into a few side corridors you haven't tried before." The student learns to keep their "flashlight" (probability distribution) wide, illuminating many different paths, not just the one they know works.

How It Works in Practice

The paper takes two existing training methods (called iDPO and GRPO) and adds this "curiosity" rule to them. They call the new versions ED-iDPO and ED-GRPO.

The Mechanism: Mathematically, they add a "penalty" if the model gets too comfortable with its previous answers. It forces the model to stay slightly "uncomfortable" and diverse, keeping its options open.
The Result: When they test these new models, they don't just get one good answer; they generate a wide variety of different solutions.

Why This Matters (The "Test-Time" Magic)

Because the model is now generating diverse solutions, the "Self-Consistency" trick (asking for 10 answers and voting) suddenly works much better.

Before: Ask for 10 answers $\rightarrow$ Get 10 identical wrong answers $\rightarrow$ Vote $\rightarrow$ Still wrong.
With EDO: Ask for 10 answers $\rightarrow$ Get 10 different approaches (some right, some wrong) $\rightarrow$ Vote $\rightarrow$ The correct answer wins because it appeared multiple times in different forms.

The Results

The authors tested this on hard math competitions (like the AIME and MATH datasets).

Accuracy: The new methods (ED-iDPO and ED-GRPO) solved more problems correctly than the previous best methods.
Diversity: The models produced much more varied answers (measured by how different the words and steps were).
Stability: Unlike other methods that sometimes cause the model to "collapse" (forget everything and become repetitive), EDO kept the model stable and creative throughout the training process.

Summary

In simple terms, EDO teaches AI models to be less obsessed with being "perfectly consistent" and more willing to be "creatively diverse." By forcing the model to explore different paths during training, it becomes much better at solving hard problems when asked to generate multiple solutions at test time. It's the difference between a student who memorizes one solution and a student who understands the landscape of the problem well enough to find the answer from many different angles.

Technical Summary: Exploration-Driven Optimization for Test-Time Large Language Model Reasoning

1. Problem Statement

Recent advancements in Large Language Models (LLMs) have demonstrated that combining post-training techniques (specifically Reinforcement Learning, RL) with inference-time scaling significantly enhances reasoning capabilities. However, a fundamental tension exists between these two paradigms:

Inference-Time Scaling: Techniques such as Self-Consistency, Best-of-N (BoN) sampling, and Monte Carlo Tree Search rely on diverse sampling from a relatively flattened probability distribution to discover optimal solutions.
RL Post-Training: Conventional RL-based post-training (e.g., PPO, DPO, GRPO) inherently sharpens output distributions. By fine-tuning the model toward high-confidence, high-reward solutions, these methods often lead to "mode collapse," reducing the diversity of generated trajectories.

This distributional collapse limits the effectiveness of test-time scaling, as the model fails to generate the diverse candidate solutions necessary for aggregation methods to identify the correct answer. Existing approaches often treat inference-time scaling as a post-hoc procedure, failing to align the training objective with the requirements of diverse exploration.

2. Methodology: Exploration-Driven Optimization (EDO)

The authors propose Exploration-Driven Optimization (EDO), a framework designed to promote diversity during the post-training phase while maintaining alignment with task objectives. EDO extends reward-biasing exploration objectives from prior work (e.g., VPO, XPO) to iterative post-training settings.

Core Mechanism

EDO introduces a regularization term that explicitly encourages the current policy ( $\pi_\theta$ ) to diverge from the previous policy iterate ( $\pi_{t-1}$ ), rather than solely anchoring to a static reference policy ( $\pi_{ref}$ ).

Reward Biasing Formulation:
The method reinterprets test-time computation as an implicit Q-function. To encourage exploration, EDO maximizes an objective that includes a term promoting the divergence from the previous policy.
- The optimization target is reformulated to maximize $J^*(r^{(t)})$ , which, under specific constraints (zero expectation under the previous policy), translates to minimizing the reverse Kullback-Leibler (KL) divergence: $-\beta D_{KL}(\pi_{t-1} \parallel \pi_\theta)$ .
- This term acts as a "repulsive" force, preventing the policy from collapsing into the dominant modes of the previous iteration, while the standard forward KL term ( $D_{KL}(\pi_\theta \parallel \pi_{ref})$ ) maintains stability and prevents catastrophic forgetting.
Integration with Existing Algorithms:
EDO is seamlessly integrated into two popular iterative RL paradigms:
- ED-iDPO (Exploration-Driven Iterative Direct Preference Optimization): Modifies the iDPO loss by adding the exploration term. The objective becomes:
  $L_{ED-iDPO} = L_{iDPO} - \alpha J^*(r^{(t)})$
  where $\alpha$ controls the strength of the exploration bias.
- ED-GRPO (Exploration-Driven Group Relative Policy Optimization): Adapts the GRPO objective to include the token-level reverse KL divergence between the current and previous policies, smoothing the output probability distribution without requiring a separate critic model.
Iterative Training Pipeline:
The framework operates in an iterative loop:
- Generate $N$ candidate solutions from the current policy.
- Construct preference pairs (for iDPO) or group rewards (for GRPO) based on correctness.
- Update the policy using the EDO-augmented objective.
- Repeat, using the updated policy as the reference for the next iteration's exploration term.

3. Key Contributions

The paper outlines three primary contributions:

Principled Framework: Introduction of EDO, a framework that promotes a flattened output distribution to balance exploration and exploitation, specifically facilitating more effective inference-time scaling.
Seamless Integration: Demonstration that EDO can be integrated into standard post-training pipelines (iDPO and GRPO) to produce enhanced variants (ED-iDPO and ED-GRPO) that significantly improve solution diversity and reasoning performance.
Empirical Validation: Comprehensive evaluation showing consistent improvements across various backbones (Qwen, LLaMA) and datasets, highlighting EDO's robustness in preventing over-optimization collapse.

4. Experimental Results

In-Distribution Performance

Experiments were conducted on three reasoning benchmarks: GSM8K (grade school math), Math (high school competition), and s1K (challenging competition-level questions).

Accuracy Gains: Both ED-iDPO and ED-GRPO outperformed strong baselines (ETO, GRPO, DAPO).
- On s1K, ED-GRPO achieved a 38.6% accuracy (greedy) and 48.4% (Self-Consistency), surpassing the best baseline (ETO+SC) by 1.1% and 1.3% respectively.
- On Math and GSM8K, EDO variants consistently improved accuracy over standard CoT and Self-Consistency baselines.
Synergy with Test-Time Scaling: The performance gains were most pronounced when combined with Self-Consistency, indicating that the diverse solutions generated by EDO are effectively aggregated by inference-time methods.

Out-of-Distribution (OOD) Generalization

Models trained on s1K were evaluated on five OOD datasets (AIME24, AIME25, MATH500, Minerva Math, Olympiad Bench).

Plug-and-Play Capability: EDO-trained models showed strong cross-dataset generalization without retraining.
Improvements: ED-GRPO delivered an average gain of 2.4% over the GRPO baseline across all OOD tasks, demonstrating that enhanced exploration improves transferability to heterogeneous data distributions.

Ablation and Analysis

Training Stability: Unlike standard GRPO or ETO, which exhibited performance degradation or flatlining after a few epochs (signs of over-optimization), EDO variants maintained steady performance gains across multiple training iterations.
Entropy Preservation: Figure 4 in the paper shows that while standard methods rapidly drive policy entropy to low levels, ED-GRPO maintains significantly higher entropy throughout training, preventing mode collapse.
Diversity Metrics: EDO variants produced substantially more diverse outputs (measured by Distinct-4). For instance, on s1K, ED-GRPO improved Distinct-4 by 47.0% compared to the baseline.
Case Study: A qualitative analysis on a geometry problem revealed that standard models produced homogeneous, incorrect reasoning paths, whereas ED-iDPO generated diverse strategies (coordinate geometry, region decomposition, symmetry arguments), allowing test-time voting to select the correct answer.

5. Significance and Claims

The paper positions EDO as a practical framework for resolving the tension between RL-based alignment and the diversity requirements of test-time scaling.

Preventing Collapse: The authors claim EDO effectively prevents "over-optimization collapse," a common failure mode in RLHF where models lose the ability to explore alternative solutions.
Compatibility: The method is highlighted as compatible with various policy optimization algorithms and test-time inference techniques (Self-Consistency, BoN, Tree Search).
Scalability: By enabling models to leverage diverse inference pathways, EDO allows for more effective utilization of test-time compute, a critical factor for scaling reasoning capabilities without increasing model parameters.

The authors conclude that EDO offers a principled approach to balancing exploration and exploitation, establishing a new baseline for LLM reasoning that is particularly effective in settings relying on test-time scaling.