Representation-Based Exploration for Language Models: From Test-Time to Post-Training — part 3/5
Cleaned-text code points [31707, 47462). Bibliography omitted here; source.md includes it.
As pointed out in Section 2.1 , we believe inference-time exploration can be useful in its own right for domains where verification is expensive. While an extensive study of this merits further study, we provide some preliminary experiments in the domain of protein sequence generation. In this domain, verification requires lab work, which is much more time consuming than sampling protein sequences from a generative protein model. Following prior work, we use two proxy metrics described below in lieu of expensive wet-lab verification.
For our experimental setup, we follow the unconditional generation setup in Hayes et al. (2025) and sample 2048 protein sequences for fully masked sequence prompts ranging in length from 64 to 916 with increments of size 4 (i.e. 64, 68, …, 916) using ESM3-open . Then, we perform structure prediction for all sampled sequences using ESMFold ( Lin et al., 2023 ) , which returns pLDDT and pTM scores per sequence. We count a sampled sequence as plausible when pLDDT > 0.8 and pTM > 0.8. To get representations for every sequence, we use esmc-600m , the 600M parameter version of ESM C ( Hayes et al., 2025 ; ESM Team, 2024 ) , and average the last-layer hidden representations. Finally, we perform inference-time selection for every prompt and its 2048 corresponding sampled sequences, using both random selection and RepExp . We find the average samples-to-correct for random to be 240.3 and for RepExp to be 6.7. This corresponds to a 35.9x verifier efficiency improvement. We also plot the corresponding pass@k curves in Fig. 8 , for which we find up to 56.9x improvements in verifier efficiency.
We find these results to be very promising and suggestive that our method can be practical in a domain where verification is genuinely expensive.
## 5 Exploration for RL Post-Training
Following the methodology in Section 2 , we now investigate the use of representation-based exploration to guide the RL post-training process.
#### Tasks and models
We use Qwen-2.5-7b-Instruct evaluated on MATH , GSM8K , and AIME 2024 . Because there are only 30 30 questions in AIME 2024 , we follow Yu et al. (2025) and use the DAPO-Math-17K dataset for training, leaving AIME 2024 for evaluation only. Please refer to Appendix C for exact details on train, validation, and test splits for all tasks.
#### Baselines
We compare our method with three baselines: (1) Unlikeliness ( He et al., 2025 ) modifies GRPO by scaling the extrinsic rewards by a value inversely related to the likelihood of a generation under the current policy. (2) GRPO is simply an unmodified version of the original GRPO algorithm. (3) Base Model is the original untrained model included as a reference point to see if methods can improve upon it, especially at high values of k k .
Representation-based Exploration ( RepExp ). We augment the rewards in GRPO with representation-based bonuses as described in Section 3 . Concretely, we add sequence-level elliptic bonuses to the binary extrinsic rewards provided by the verifier: r i = R ( x , y i ) + b ( x , y i ) r_{i}=R(x,y_{i})+b(x,y_{i}) for the i th i^{\text{th}} rollout y i y_{i} of a given prompt x x . As mentioned in Section 3 , we specifically use leverage score-like elliptic bonuses, which allow easier control over their scale as they are bounded in [ 0 , 1 ] [0,1] . The covariance matrix Σ \Sigma used to compute bonuses is re-initialized for each batch of RL training. This way, the bonus b ( x , y i ) b(x,y_{i}) measures the novelty of y i y_{i} only with respect to the other rollouts in the batch—previously generated sequences for x x are not considered. To better maximize the bonus along all relevant directions in representation space, we draw a new random projection of h ¯ θ ( x , y i ) \bar{h}_{\theta}(x,y_{i}) at each optimization step. For further details, please refer to Appendix C .
RF6 : RepExp improves pass@k. Fig. 2 compares pass@ k k curves after training for all methods. In line with earlier work, we find that all instantiations of GRPO improve the pass@k for small values of k k , and that standard GRPO degrades performance relative to the base model for large values of k k ( Yue et al., 2025 ) . Exploration appears to be an essential part of mitigating the latter effect: policies fit using RepExp preserve or improve pass@ k k for large values of k k with limited reductions for small k k . This phenomenon is more pronounced for RepExp than for Unlikeliness.
#### RF7 : RepExp responses look novel
To provide further insight into whether RepExp is able to move beyond merely sharpening the base model, we run an additional experiment inspired by Figure 4 in Karan and Du (2025) . Specifically, we sample a single response from the base model, the GRPO post-trained model, and the RepExp post-trained model for each question in the full test set of MATH . We then score all responses under the base model in terms of log likelihood. We plot the resulting histogram in Fig. 9 . We find that responses from RepExp tend to be less likely under the base model, indicating that it generates more novel responses and is therefore not merely performing sharpening. In contrast, notice that standard GRPO exhibits sharpening behavior, as demonstrated by the movement of probability mass towards the right with respect to the base model. Taken together with RF6, these results suggest that with the right exploration strategy, we may be able to escape the sharpening regime and discover novel model behaviors.
## 6 Discussion
#### Related work
Several recent works aim to encourage exploration in language models, either by adapting exploration techniques from deep reinforcement learning, or by augmenting PPO or GRPO in ways that are more specialized to language models ( He et al., 2025 ; Cheng et al., 2025 ; Chen et al., 2025b ; Zhou et al., 2025 ; Setlur et al., 2025 ; Liu et al., 2025 ) . Examples of the former include count-based exploration via pseudo-counts ( Bai et al., 2025 ) and at the outcome level ( Song et al., 2025 ) , random network distillation ( Liu et al., 2024b ; Gao et al., 2025 ) , and posterior sampling ( Dwaracherla et al., 2024 ) . Examples of the latter include rewarding unlikely-but-correct responses ( He et al., 2025 ) , entropy bonuses ( Cheng et al., 2025 ) , adapting the number of rollouts based on question difficulty ( Yang et al., 2025 ) , training the model to explore in-context ( Setlur et al., 2024 ) , using a learned classifier to jointly optimize diversity and quality ( Li et al., 2025 ) , adding a diversity term based on determinantal point processes to the RL objective ( Chen et al., 2025a ) , and reformulating post-training to more-directly maximize the pass@ N N objective ( Balashankar et al., 2024 ; Chow et al., 2025 ; Chen et al., 2025b ; Walder and Karkhanis, 2025 ; Tang et al., 2025 ) . Among these, we experiment with the unlikeliness reward approach of He et al. (2025) as a baseline, due to its robust performance and clean implementation. More generally, our work is unique in (1) the specific representation-based objective, and (2) our focus on inference-time as a means to validate methods with minimal confounding factors. See Appendix A for a detailed overview.
#### Final remarks
Our work shows that deliberate exploration is a viable path toward expanding the reasoning capabilities of language models, offering the possibility of discovering novel behaviors that would be unlikely under naive sampling. While our results show that representation-based diversity is effective at incentivizing exploration, the algorithm design space for exploration techniques is vast, and there is still much to understand regarding how to best use the knowledge encoded in foundation models to guide exploration. Along these lines, natural directions for future work include: 1. Scaling up RL compute, and combining exploration with other techniques known to improve reasoning behavior in RL post-training, such as prolonged reinforcement learning ( Liu et al., 2025 ) .
2. Exploration for autoregressive generation. Our results in Section 4.2 show that incentivizing diversity during autoregressive generation is a promising approach to reducing the computational burden of exploration, but much remains to be done in terms of (1) understanding which diversity metrics are most helpful, and (2) optimizing the implementation to close the compute gap.
3. Beyond verifiable rewards. How can we deliberately incentivize exploration in domains without verifiable rewards, while simultaneously mitigating reward hacking?
## Acknowledgements
We thank Adam Block, Qinghua Liu, and Max Simchowitz for valuable feedback on this work. We thank Manan Tomar, Audrey Huang, Spencer Whitehead, Prabhat Nagarajan, and Karthik Narasimhan for support and encouragement throughout the project.
## Reproducibility Statement
To ensure reproducibility of our results, we have listed all relevant details (hyperparameters, experiment resources, etc.) for the inference-time experiments in Appendix B , and those for RL post-training in Appendix C . In addition, the code for all our experiments and plots can be found at https://rep-exp.github.io . We used Weights & Biases for experiment tracking and visualizations to develop insights for this paper.
## Appendix A Additional Related Work
#### Exploration at test time
Test-time alignment techniques for language models are an active area of research with many complementary threads ( Khanov et al., 2024 ; Chen et al., 2024 ; Shi et al., 2024a ; Liu et al., 2024a ; Jinnai et al., 2024 ; Shi et al., 2024b ) , but exploration has not typically been the focus of this line of work.
Most closely related to our work, Setlur et al. (2025) , propose a test-time exploration approach based on the idea of learning to explore in-context . They propose to encourage exploration within a long chain of thought by training the LLM to chain operations such as generation, verification, and refinement together in search of a solution. This is somewhat complementary to our inference-time exploration framework, which aims to improve diversity across parallel generations once the model is fixed; these techniques could potentially be combined.
Also related, Xu et al. (2025) consider the problem of learning from language (non-verifiable) feedback, and propose an iterative prompting approach to enable exploration at test time; their work focuses on simpler exploration domains, but with more difficult, implicit feedback.
#### Exploration in RL post-training
Exploration in RL post-training for reasoning is a growing area of research, motivated by the observation that standard techniques tend to simply sharpen responses already covered by the base model ( Yue et al., 2025 ; Gandhi et al., 2025 ; Wu et al., 2025 ) . A number of recent works, discussed below, aim to improve diversity and expand the reasoning frontier by incorporating bonuses into the GRPO objective or by otherwise augmenting it. Briefly, our work is unique in terms of (1) the specific representation-based diversity objective we focus on, and (2) our focus on inference-time exploration as a means to validate diversity metrics before applying them to post-training.
He et al. (2025) introduce an unlikeliness reward to GRPO, which reweights the reward by ranking generations according to their unlikeliness under the sampling policy. Unlikeliness reward is a form of diversity metric, similar to our representation-based diversity metrics. Cheng et al. (2025) observe that high-entropy (high uncertainty) tokens in the model’s output often correspond to critical reasoning steps, and augment the GRPO objective with entropy bonuses to encourage exploration at these high-entropy steps. Entropy can be seen as another form of diversity metric in our setup. Another option is to learn the diversity metric as in Li et al. (2025) , who use a learned classifier to determine whether a pair of responses is semantically equivalent. They then use the classifier score to scale the reward in GRPO to joinly optimize quality and diversity. Our work, in contrast, does not require training any auxiliary models for computing diversity.
Various works ( Balashankar et al., 2024 ; Chow et al., 2025 ; Chen et al., 2025b ; Walder and Karkhanis, 2025 ; Tang et al., 2025 ) formulate the problem of directly post-training to maximize the pass@ N N objective, deriving approximate gradient estimators and using them for policy optimization. As discussed in Chow et al. (2025) ; Chen et al. (2025b) , these gradient estimators implicitly encourage exploration, since they allow the model to distribute probability mass across a more diverse range of responses when it is uncertain about the correct answer. Our work instead focuses on using the language model representations to deliberately incentivize novel behaviors, including in a novel inference-time setting.
Zhou et al. (2025) consider a setting where ground truth answers are available (as opposed to just rewards), and propose to encourage exploration by prompting the model to generate self-explanations for the ground truth answers and incorporating this as an SFT term in the GRPO loss. Unlike our method, their approach does not directly optimize for diversity, and cannot be used in settings where ground truth answers are unavailable (e.g. coding).
Yang et al. (2025) investigate the role of “breadth" (batch size) and “depth” (number of rollouts) in RLVR. They show that increasing breadth through full batch updates and increasing depth through more rollouts for harder questions has complementary benefits and overall improves pass@1 and pass@k performance. We view this as orthogonal and potentially complementary to our work.
Lanchantin et al. (2025) introduce diverse preference optimization (DivPO), an alignment method to optimize for both quality and diversity. In contrast to our work, their method is designed for the RLHF setting and applied to non-reasoning tasks (e.g. creative writing).
Concurrent work of Song et al. (2025) adapts tabular UCB-style bonuses to language model post-training with GRPO, but their approach—unlike representation-based exploration—is only suitable for domains with a small, discrete set of possible outcomes. Other concurrent work of Chen et al. (2025a) optimizes a diversity term along with the rewards, where the diversity term captures the volume of the responses in representation space by computing the determinant of the gram matrix. While related, our method instead adds a leverage-based score to the reward of each individual response.
Lastly, Liu et al. (2025) take a complementary approach and aim to incentivize reasoning beyond the base model through (1) prolonged RL training (increasing the overall amount of training steps), and (2) periodically resetting the reference model; they show that this can increase pass@ N N performance beyond the base model in a variety of reasoning tasks. This approach is complementary, and could likely be combined with our techniques.
#### Representation-based diversity
Our findings regarding benefits of inference-time exploration with representation-based diversity parallel the findings of Ivison et al. (2025) , who evaluated the effectiveness of various data selection schemes for instruction tuning, and found a similar representation-based scheme to be the most effective when normalized for compute. In addition, there is a long line of work relying on representation-based exploration for RL through elliptic bonuses in non-LLM settings ( Agarwal et al., 2020a ; Agarwal et al., 2020b ; Henaff et al., 2022 ; Ash et al., 2021 ) .
#### Adapting exploration techniques from deep reinforcement learning