🤖 machine learning

Representation-Based Exploration for Language Models: From Test-Time to Post-Training

This paper demonstrates that incorporating a simple, representation-based diversity bonus derived from pre-trained hidden states into both inference-time sampling and post-training reinforcement learning significantly enhances language model performance and sample efficiency by promoting the discovery of novel behaviors rather than merely refining existing ones.

Original authors: Jens Tuyls, Dylan J. Foster, Akshay Krishnamurthy, Jordan T. Ash

Published 2026-07-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jens Tuyls, Dylan J. Foster, Akshay Krishnamurthy, Jordan T. Ash

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 a world where computers don't just follow instructions but actually learn to think for themselves, solving puzzles and writing code by trying, failing, and trying again. This is the realm of Reinforcement Learning (RL), a branch of artificial intelligence where an agent learns by interacting with its environment to maximize rewards. Think of it like training a dog: if it sits, it gets a treat; if it jumps, it gets nothing. Over time, the dog learns the best behaviors to get the most treats.

In the world of large language models (the super-smart AI chatbots we use today), this process is called "post-training." Scientists feed the model a math problem or a coding task, let it generate an answer, and check if it's correct. If it is, the model gets a "reward" and learns to do that specific thing better next time. However, there's a catch. Current methods often act like a student who just memorizes the textbook answers. They get really good at the specific tricks they already know, but they rarely discover new ways to solve problems. They just get sharper at what they already do. The big question scientists are asking is: Can we teach these AI models to be curious explorers instead of just memorizers? Can we encourage them to try weird, novel, and diverse approaches to find solutions that their base training never showed them?

This paper, titled "Representation-Based Exploration for Language Models," dives into exactly that question. The researchers, Jens Tuyls, Dylan J. Foster, Akshay Krishnamurthy, and Jordan T. Ash, propose a clever new way to make AI models explore. Instead of just letting the model guess randomly or simply repeating what it knows, they give the model a "diversity bonus." Imagine you are looking for a hidden treasure in a giant forest. A random walker might just stumble around, while a smart walker might try to visit every distinct type of tree to make sure they don't miss a clue. This paper suggests using the AI's own internal "thoughts" (its hidden states) to measure how different a new idea is from the ones it has already tried. If an idea feels "fresh" or "novel" compared to its previous attempts, it gets a bonus reward, encouraging the model to keep exploring that new path.

The team tested this idea in two main ways. First, they looked at "inference-time" exploration, which is like a one-off test. They generated thousands of answers to a single math problem and used their "diversity bonus" to pick the most interesting ones to check. They found that this method was incredibly efficient. For example, on a difficult math dataset called MATH, using their method with a model called Qwen-2.5-14b-Instruct improved the efficiency of finding a correct answer by over 50% compared to just picking answers at random. In fact, for some tasks, they needed 3 to 13 times fewer attempts to find a correct solution when using this exploration strategy.

Second, they integrated this exploration strategy into the actual training process (post-training). They wanted to see if it could help the model learn new behaviors rather than just sharpening old ones. The results were promising. When they trained the model with this exploration bonus, it didn't just get better at the same old tricks; it started generating answers that looked much more novel and creative. On a tough math competition dataset called AIME 2024, their exploration-enhanced model performed as well as a standard training method that used four times as many samples. This suggests that by deliberately encouraging diversity, they can help AI models discover new capabilities without needing massive amounts of extra data or computing power.

However, the authors are careful to note that this isn't a magic bullet that solves everything. They found that the method works best with stronger, more capable models; weaker models sometimes didn't benefit or even got worse. They also showed that while this approach helps the model find new solutions, it doesn't necessarily mean the model is "smarter" in a general sense, but rather that it is better at navigating the vast space of possible answers to find the right one. The paper suggests that this "deliberate exploration" is a practical path forward, offering a way to move beyond simply sharpening existing skills and toward discovering genuinely new behaviors in language models.

Technical Summary: Representation-Based Exploration for Language Models

Problem Statement

Reinforcement learning (RL) has shown promise in enhancing language models (LMs) for tasks with verifiable rewards, such as mathematical reasoning and code generation. However, a critical limitation of current RL post-training techniques is the uncertainty regarding whether they genuinely discover novel behaviors or merely "sharpen" (amplify) capabilities already present in the base model with modest probability. Existing evidence suggests that standard RL recipes often fail to unlock capabilities absent in the pre-trained model, instead converging on known solutions. Furthermore, data curation and algorithmic interventions to mitigate this are becoming bottlenecks, particularly in complex, open-ended domains.

The paper posits that deliberate exploration—explicitly incentivizing the model to discover truly novel and diverse behaviors—is essential to realizing the full potential of RL for language model reasoning. The central challenge is quantifying novelty and diversity within the combinatorially large space of language in a scalable manner, while leveraging the prior knowledge embedded in pre-trained models to guide this search efficiently.

Methodology

The authors propose a representation-based exploration strategy, utilizing elliptical bonuses derived from the hidden states of the pre-trained language model. The methodology is evaluated through a two-pronged approach:

1. Inference-Time Selection

To isolate the role of diversity from optimization and generalization, the authors first define an inference-time selection problem. Given a prompt $x$ and a large pool of $N$ candidate responses generated by a fixed model $\pi$ , the goal is to select a subset of $k$ responses that is maximally diverse and has a high probability of containing a correct answer.

Mechanism (RepExp): The algorithm iteratively selects responses to maximize an elliptical bonus. For a response $y$ , the feature vector $\bar{h}_\theta(x, y)$ is constructed by averaging the last-layer hidden states of the model over the tokens of the response.
Diversity Metric: The novelty of a candidate $y$ given a set of previously selected responses is defined as $div(y) = \bar{h}_\theta(x, y)^\top \Sigma^{-1} \bar{h}_\theta(x, y)$ , where $\Sigma$ is the covariance matrix of the selected features. The inverse covariance is updated efficiently using the Woodbury identity.
Goal: This approach selects a "coreset" of diverse responses to query a verifier, aiming to improve verifier efficiency (reducing the number of samples needed to find a correct answer).

2. RL Post-Training

The authors integrate this exploration strategy into the GRPO (Group Relative Policy Optimization) post-training pipeline.

Reward Augmentation: Instead of selecting a subset, the elliptical bonus is added directly to the extrinsic reward. For a rollout $y_i$ , the reward becomes $r(x, y_i) + \beta \cdot \bar{h}_\theta(x, y_i)^\top \Sigma^{-1} \bar{h}_\theta(x, y_i)$ .
Implementation: The covariance matrix $\Sigma$ is re-initialized for each batch of training rollouts to measure novelty relative to the current batch. The feature vectors are projected to a lower dimension (e.g., 512) via sparse random projection to ensure scalability.

3. Token-Level Extension (Preliminary)

As a proof of concept, the authors also explore modifying the autoregressive generation process itself. By perturbing token-level logits with an elliptical bonus at each generation step, they aim to guide diversity during the generation process rather than post-hoc selection.

Key Contributions and Findings

1. Improved Verifier Efficiency (Inference-Time)

The paper demonstrates that representation-based exploration significantly improves verifier efficiency across a wide range of models (Phi, Llama, Qwen) and tasks (MATH, GSM8K, MBPP+, Game of 24).

Result: For models like Qwen-2.5-14b-Instruct, the method achieves over 50% improvement in the number of samples required to find a correct answer compared to random sampling.
Correlation with Model Strength: The benefits of exploration scale with model strength; stronger models (e.g., Qwen-32B) benefit significantly, while weaker models may see no improvement or degradation.
Difficulty Sensitivity: The method provides the greatest efficiency gains on the hardest questions (e.g., the top 20% hardest in MATH).

2. Enhanced Post-Training Performance and Diversity

Integrating representation-based exploration into RL post-training yields superior performance compared to standard GRPO and "Unlikeliness" baselines.

Pass@k Improvement: The method maintains or improves pass@k rates for large values of $k$ , whereas standard GRPO often degrades performance relative to the base model for large $k$ (a phenomenon known as "diversity collapse").
Sample Efficiency: On the AIME 2024 benchmark, the post-trained Qwen-2.5-7b-Instruct using RepExp achieves a pass@80 that matches the pass@256 of standard GRPO, representing a 3x improvement in test-time sample efficiency.
Anti-Sharpening: Analysis of log-likelihoods under the base model reveals that RepExp generates responses that are significantly more novel (lower likelihood) than those from standard GRPO, which tends to sharpen the base model's distribution. This suggests the method successfully escapes the "sharpening regime."

3. Application to High-Cost Verification

In a preliminary experiment on protein sequence generation (where verification via wet-lab work is expensive), the method improved verifier efficiency by up to 56.9x compared to random sampling, demonstrating potential utility in domains where verification is costly.

Significance and Claims

The paper claims that deliberate exploration, guided by a principled, representation-based notion of diversity, is a practical path toward discovering new behaviors in language models beyond mere sharpening.

Novelty of Approach: Unlike prior work that relies on auxiliary networks (e.g., Random Network Distillation) or complex data curation, this method leverages the pre-trained model's own hidden states, requiring no additional learning machinery.
Validation Strategy: The authors emphasize the value of their two-stage evaluation (inference-time selection followed by post-training) as a way to disentangle the benefits of diversity from other RL mechanisms.
Modest Scope: The authors acknowledge that their current implementation is a "simple, principled" scheme. They do not claim to have solved all exploration challenges but suggest that this approach provides a strong baseline and a framework for understanding the tradeoffs of more sophisticated strategies. They note that while token-level exploration shows promise, it currently incurs higher computational costs that require optimization.

In summary, the paper argues that by explicitly incentivizing the discovery of diverse behaviors through representation-based bonuses, language models can be guided to find novel solutions that standard RL post-training misses, thereby expanding their reasoning capabilities more effectively.