TLDR: A new research paper introduces Representation-Based Exploration (RepExp), a method that uses an LLM’s internal hidden states to incentivize diverse and novel responses. This approach significantly improves language models’ ability to find correct answers more efficiently during inference and prevents “diversity collapse” during reinforcement learning post-training, helping models discover new behaviors beyond just refining existing ones.
Reinforcement Learning (RL) holds immense promise for advancing the capabilities of large language models (LLMs), enabling them to discover new and valuable behaviors autonomously. However, a key challenge has been whether current RL methods truly help LLMs uncover novel solutions, or if they merely refine and improve upon behaviors the model already possesses, a phenomenon often referred to as “sharpening.” A new research paper delves into this question, proposing a novel approach to encourage deliberate exploration in LLMs.
The paper, titled “Representation-Based Exploration for Language Models: From Test-Time to Post-Training,” introduces a method that significantly boosts the diversity of behaviors and the success rates of LLMs, both during inference (test-time) and post-training with RL. The core idea revolves around using a simple, yet powerful, representation-based bonus derived from the LLM’s internal hidden states to incentivize the discovery of novel and varied responses.
Understanding the Challenge: Beyond Sharpening
Traditional RL for LLMs, particularly in tasks like mathematical reasoning or code generation, has led to impressive breakthroughs. Yet, there’s growing evidence that these methods often enhance existing capabilities rather than unlocking entirely new ones. This limitation becomes particularly apparent in complex, open-ended domains where current techniques struggle to elicit truly desired, novel behaviors. The researchers argue that deliberate exploration, which explicitly encourages the model to find diverse solutions, is crucial for realizing the full potential of RL in language models.
The Solution: Representation-Based Exploration (RepExp)
The proposed method, RepExp, tackles this by quantifying novelty and behavior diversity in a scalable way. It adapts a technique called “elliptic bonuses,” commonly used in other machine learning fields, to language models. Essentially, RepExp uses the hidden states (internal representations) of the pre-trained language model as “feature vectors.” When the model generates responses, RepExp assigns a bonus to responses that are “novel” or “diverse” in this representation space, meaning they are not well-represented by previously seen or selected responses. This encourages the model to explore different solution paths.
The researchers evaluated RepExp in two main settings:
- Inference-Time Selection: In this scenario, the goal is to select a small, diverse set of high-quality responses from a larger pool of candidates generated by a fixed model. RepExp iteratively selects responses that maximize this diversity bonus. This setting helps isolate the impact of diversity from other complex RL mechanisms.
- RL Post-Training: Here, the exploration strategy is integrated directly into the RL training pipeline. Instead of just using the standard reward for correctness, the model also receives an additional “exploration bonus” based on the novelty of its generated responses within the training batch.
Key Findings and Impact
The results are compelling:
For inference-time exploration, RepExp significantly improved “verifier efficiency” – the average number of samples needed to find a correct answer. For instance, on the Qwen-2.5-14b-Instruct model, it achieved over 50% improvement in efficiency across various reasoning tasks like GSM8K, MATH, MBPP+, and Game-of-24. The benefits were particularly pronounced for stronger models and on harder questions, suggesting that better internal representations lead to more effective exploration.
For RL post-training, integrating RepExp into the RL pipeline (specifically GRPO) led to improved reasoning performance. Crucially, it addressed the “diversity collapse” phenomenon, where standard RL methods often degrade the model’s ability to produce diverse correct answers for larger budgets (pass@k for large k). RepExp maintained or even improved pass@k rates across all values of k, demonstrating that it helps discover new behaviors rather than just sharpening existing ones. For example, on AIME 2024, their post-trained Qwen-2.5-7b-Instruct’s pass@80 matched the pass@256 of GRPO, showing a 3x improvement in test-time sample efficiency.
The paper also explored an extension where representation-based exploration guides the autoregressive generation process at the token level, showing promising initial results for improving solve rates on harder questions for larger budgets.
Also Read:
- AI’s Hidden Language: Decoding Reasoning Success Through Latent Signals
- Dynamic Temperature Control Enhances LLM Reasoning in Reinforcement Learning
A Path Towards True Discovery
Overall, the findings suggest that deliberate exploration, guided by the right understanding of diversity, offers a practical way for LLMs to move beyond merely refining existing knowledge and instead discover genuinely novel behaviors. This could be a significant step towards unlocking the full potential of reinforcement learning for language models, especially in complex reasoning domains where data curation is a bottleneck. For more in-depth information, you can read the full research paper here.
