Representation-Based Exploration for Language Models: From Test-Time to Post-Training — part 2/5
Cleaned-text code points [15794, 31707). Bibliography omitted here; source.md includes it.
We focus our experiments on an exploration strategy that avoids these shortcomings: An adaptation of elliptic bonuses and sampling—a de facto standard for linear bandits and active learning ( Abbasi-Yadkori et al., 2011 ; Chu et al., 2011 ; Ash et al., 2021 ; Henaff et al., 2022 ; Saran et al., 2023 ; Foster et al., 2025 ) 2 2 2 Indeed, elliptical bonuses are ubiquitous in linear bandits and reinforcement learning, the simplest non-tabular RL setting, where they have strong provable guarantees. Beyond this, elliptic bonuses and iterative schemes such as Algorithm 1 have a long history in the theory of optimal experimental design ( Kiefer and Wolfowitz, 1960 ; Pukelsheim, 2006 ; Allen-Zhu et al., 2021 ) and active learning ( Cesa-Bianchi et al., 2009 ; Agarwal, 2013 ; Gu et al., 2014 ; Chaudhuri et al., 2015 ) —with a representation derived from the language model’s hidden states. This approach is arguably the simplest principled strategy that is appropriate for language models, and already yields significant performance improvements in our experiments. Our use of elliptic bonuses is particularly inspired by Foster et al. (2025) , who prove that test-time exploration with such bonuses has provable computational benefits in a simplified language model setting with frozen features.
At a high level, elliptical bonus methods operate over a d d -dimensional feature space and adopt a linear-algebraic measure of novelty: given previously seen feature vectors h 1 , … , h i − 1 h_{1},\ldots,h_{i-1} the novelty (or bonus) of a new feature vector h h is defined as div ( h ∣ h 1 : i − 1 ) = h ⊤ Σ i − 1 h Σ i = λ I d + ∑ j < i h j h j ⊤ \displaystyle\texttt{div}(h\mid{}h_{1:i-1})=h^{\top}\Sigma_{i}^{-1}h\qquad\Sigma_{i}=\lambda I_{d}+\sum_{j<i}h_{j}h_{j}^{\top} (2)
These bonuses are grounded in the theory of linear regression: If we fit a linear model f θ ( h ) = ⟨ θ , h ⟩ f_{\theta}(h)=\langle\theta,h\rangle on features h 1 , … , h i − 1 h_{1},\ldots,h_{i-1} (with associated regression targets), the prediction error on h h will be bounded by div ( h ∣ h 1 : i − 1 ) \texttt{div}(h\mid{}h_{1:i-1}) ( Lattimore and Szepesvári, 2020 ) . Thus, div ( h ∣ h 1 : i − 1 ) \texttt{div}(h\mid{}h_{1:i-1}) reflects novelty, as it will be large for features h h that are poorly represented by the training dataset.
To adapt elliptical bonuses to language models, we use representations extracted from the model itself as the feature vectors. Formally, given a prompt x x and a response y i = y i 1 , … , y i T y_{i}=y_{i}^{1},\ldots,y_{i}^{T} of T T tokens, we form the feature vector as h ¯ θ ( x , y i ) := 1 T ∑ t = 1 T h θ ( x , y i 1 : t ) \bar{h}_{\theta}(x,y_{i}):=\frac{1}{T}\sum_{t=1}^{T}h_{\theta}(x,y_{i}^{1:t}) where h θ ( x , y i 1 : t ) ∈ ℝ d h_{\theta}(x,y_{i}^{1:t})\in\mathbb{R}^{d} is the last-layer hidden state of the model on input ( x , y i 1 : t ) (x,y_{i}^{1:t}) (the activation prior to the unembedding matrix). In Fig. 4 , we ablate this choice by comparing it to the effectiveness of using representations at the last token h θ ( x , y i 1 : T ) h_{\theta}(x,y_{i}^{1:T}) or penultimate token h θ ( x , y i 1 : T − 1 ) h_{\theta}(x,y_{i}^{1:{T-1}}) instead. We reduce dimensionality to 512 using a sparse random projection ( Li et al., 2006 ) ; see Appendix B for details.
#### Representation-based exploration for inference-time selection
Fig. 3 presents RepExp , our main algorithm for inference-time selection using representation-based elliptical bonuses. Here, given a single prompt x x and a set of candidate generations 𝒴 = { y 1 , … , y N } \mathcal{Y}=\left\{y_{1},\ldots,y_{N}\right\} , we iteratively select the generation that maximizes the elliptical bonus via y t + 1 = arg max y ∈ 𝒴 h ¯ θ ( x , y ) Σ t − 1 h ¯ θ ( x , y ) y_{t+1}=\argmax_{y\in\mathcal{Y}}\bar{h}_{\theta}(x,y)\Sigma_{t}^{-1}\bar{h}_{\theta}(x,y) , leveraging the representations h ¯ θ ( x , y ) \bar{h}_{\theta}(x,y) described above. We efficiently update the inverse covariance matrix Σ − 1 \Sigma^{-1} using the Woodbury identity for O ( d 2 ) O(d^{2}) time per step ( Vetterling and Press, 1992 ) . We formally present our procedure in Algorithm 1 .
Algorithm 1 RepExp 1: input: Embeddings h ¯ θ \bar{h}_{\theta} (abbrv. h ¯ \bar{h} ) , generations 𝒴 \mathcal{Y} for prompt x x , budget k k , regularization param. λ \lambda . 2: Initialize L ← { y 1 } , y 1 ∼ Unif ( 𝒴 ) L\leftarrow\{y_{1}\},y_{1}\sim\text{Unif}(\mathcal{Y}) . 3: Initialize inverse covariance Λ 0 = λ − 1 I d \Lambda_{0}=\lambda^{-1}I_{d} . 4: for t = 1 t=1 to k − 1 k-1 do 5: Λ t ← Λ t − 1 − Λ t − 1 h ¯ ( x , y t ) h ¯ ( x , y t ) ⊤ Λ t − 1 1 + h ¯ ( x , y t ) ⊤ Λ t − 1 h ¯ ( x , y t ) \Lambda_{t}\leftarrow\Lambda_{t-1}-\frac{\Lambda_{t-1}\bar{h}(x,y_{t})\bar{h}(x,y_{t})^{\top}\Lambda_{t-1}}{1+\bar{h}(x,y_{t})^{\top}\Lambda_{t-1}\bar{h}(x,y_{t})} . 6: y t + 1 = argmax y ∈ 𝒴 h ¯ ( x , y ) ⊤ Λ t h ¯ ( x , y ) y_{t+1}=\operatorname*{argmax}\limits_{y\in\mathcal{Y}}\bar{h}(x,y)^{\top}\,\Lambda_{t}\,\bar{h}(x,y) . 7: L ← L ∪ { y t + 1 } L\leftarrow L\cup\{y_{t+1}\} . 8: return: L L . Figure 4: Representation ablation. We compare averaging all token representations to using those from the penultimate or final token. Averaging is over 2x more sample efficient.
#### Representation-based exploration for RL post-training
For our post-training experiments, we use the same representations h ¯ θ ( x , y ) \bar{h}_{\theta}(x,y) as above, but directly augment the rewards with elliptic bonuses instead of performing coreset selection. Concretely, given the current iterate π θ \pi_{\theta} in GRPO, we first sample a group of k k responses y 1 , … , y k ∼ i.i.d. π θ ( ⋅ ∣ x ) y_{1},\ldots,y_{k}\overset{\textrm{i.i.d.}}{\sim}\pi_{\theta}(\cdot\mid{}x) for each prompt x x . Letting Σ := λ I d + ∑ i = 1 k h ¯ θ ( x , y i ) h ¯ θ ( x , y i ) ⊤ \Sigma\vcentcolon=\lambda{}I_{d}+\sum_{i=1}^{k}\bar{h}_{\theta}(x,y_{i})\bar{h}_{\theta}(x,y_{i})^{\top} , we define the reward for response y i y_{i} as 3 3 3 The bonus here can be interpreted as a leverage score for y i y_{i} ( Drineas et al., 2006 ; Cohen et al., 2015 ) . r ⋆ ( x , y i ) + β ⋅ h ¯ θ ( x , y i ) ⊤ Σ − 1 h ¯ θ ( x , y i ) r^{\star}(x,y_{i})+\beta\cdot\bar{h}_{\theta}(x,y_{i})^{\top}\Sigma^{-1}\bar{h}_{\theta}(x,y_{i}) , where β > 0 \beta>0 is a bonus parameter. While one could also imagine performing inference-time coreset selection in the loop with GRPO, this approach is more practical and efficient, and it achieves significant improvements in performance. We refer the reader to Section 5 and Appendix C for further details.
#### Why representation-based elliptical bonuses?
We summarize several desirable properties of these bonuses. First, by leveraging the hidden state of the model in featurization, the bonuses capture rich information about the generations, thereby incorporating the language model’s prior knowledge. Second, the method is history-aware 4 4 4 In the RL setting, we do not let the covariance matrix persist across multiple iterations of the same question, and hence there it is only group-aware . : the covariance matrix summarizes all previously selected generations, and redundancy with previous selection (in representation space) is penalized. Finally, the method is simple and scalable, involving no additional learning machinery and using rank-one updates to avoid costly matrix inversions.
## 4 Inference-Time Exploration: Experimental Results
In this section, we investigate the performance of representation-based exploration for the inference-time selection problem. We detail the experimental setup in Section 4 , present main findings in Section 4.1 , and present additional experiments with a “token-level” variant in Section 4.2 .
#### Datasets
We use the test splits of the following five datasets: MATH ( Hendrycks et al., 2021 ) , GSM8K ( Cobbe et al., 2021 ) , MBPP+ ( Liu et al., 2023 ) , Game-of-24 ( Yao et al., 2023 ) , and AIME 2025 . We chose these tasks as they cover easy ( GSM8K ), medium ( MATH ), and harder ( Game-of-24 , AIME ) difficulty levels in math. In addition, we include MBPP+ to verify that our findings transfer to the coding domain. For a more detailed overview of these datasets, please refer to Section B.1 .
#### Models
We consider a range of model families and sizes: Phi-3-Medium ( Abdin et al., 2024a ) and Phi-4 , Llama-3.2-3B-Instruct and Llama-3.1-8B-Instruct ( Dubey et al., 2024 ) , ( Abdin et al., 2024b ) , Qwen-2.5-X-Instruct ( Qwen et al., 2024 ) for X ∈ { 0.5B , 3B , 7B , 14B , 32B } \texttt{X}\in\left\{\texttt{0.5B},\texttt{3B},\texttt{7B},\texttt{14B},\texttt{32B}\right\} , and Mistral-7B ( Jiang et al., 2023 ) .
#### Algorithms
In our experiment protocol, we initially draw a pool of N N candidate generations from the base model, where unless otherwise specified we use temperature τ = 1.0 \tau=1.0 and top-p = 1.0 \text{top-p}=1.0 , which we refer to as vanilla settings (for MBPP+ , we set top-p = 0.95 \text{top-p}=0.95 ). Then we compare RepExp with budget k k with the baseline of random sampling (without replacement) of k k responses from this pool. We consider generating the pool using different samplers such as nucleus and min-p sampling in Fig. 6 , but always use random sampling without replacement as the baseline. See Section B.1 for further details.
### 4.1 Benefits of representation-based exploration
We present our results as a series of Research Findings (RF), expanding on the findings in Fig. 1 .
#### RF1 : RepExp improves verifier efficiency across models and tasks
In Fig. 1 , we plot the samples-to-correct , defined as the expected number of samples k k with which we query the verifier before finding a correct answer, for all model-task pairs. We compare RepExp , which picks responses to a fixed question according to Algorithm 1 , with the random sampling baseline. For both algorithms, we average the samples-to-correct across all questions in the dataset. Our results show the bulk of the data fall below the line y = x y=x , indicating exploration improves over random sampling in most cases. For example, we find RepExp obtains a 50% improvement in samples-to-correct for Qwen-2.5-14b-Instruct in MATH , GSM8K , MBPP+ , and Game-of-24 .
#### RF2 : The benefits of RepExp grow with model strength
Since RepExp relies on the model’s internal representations, it is natural to hypothesize that weaker models might have worse representations and thus benefit less from exploration. To validate this hypothesis, we expand the collection of models in Fig. 1 to include additional weaker models (e.g., Qwen-2.5-0.5B-Instruct and Mistral-7b ). For each task, we rank models according to their pass@1 performance, and plot the relative improvement of representation-based exploration over random sampling in Fig. 5 . We indeed observe a strong correlation between model strength and the benefit from representation-based exploration: weaker models (e.g., Qwen-2.5-0.5B ) experience no benefit or even degradation, while the strongest models (e.g., Qwen-2.5-32B ) almost uniformly benefit.
#### RF3 : RepExp provides more improvement for harder questions
Beyond RF2 —which provides insight into the benefits of RepExp across models —we also evaluate the benefits across question difficulty , for a fixed model and task. To this end, we sort all questions for a given task by their samples-to-correct under random sampling with a reference model ( GPT-4o-mini ). We then group the questions in bins, each containing 10% of the dataset, and plot the average samples-to-correct for each bin for both RepExp and random sampling. As displayed in Fig. 5 (right), RepExp matches or improves verifier efficiency across all bins, with the largest improvements on the hardest bins (e.g., the hardest 20% of questions on MATH ). Concretely, on the hardest Game-of-24 questions, we find that RepExp with Phi-4 provides a 3x improvement in verifier efficiency.
#### RF4 : RepExp improves verifier efficiency over standard generation modifications
We now investigate the effect of alternative base sampling strategies that might already induce diversity. Using Qwen-2.5-7B-Instruct on MATH , we change the underlying generation strategy to use one of five different generation settings: vanilla (no changes), low temperature ( τ = 0.6 \tau=0.6 ), high temperature ( τ = 1.5 \tau=1.5 ), min-p ( Minh et al., 2025 ) ( τ = 1.5 , p = 0.05 \tau=1.5,p=0.05 ), and nucleus sampling ( Holtzman et al., 2020 ) ( top-p = 0.9 \text{top-p}=0.9 ). In Fig. 6 , we find that RepExp improves verifier efficiency in all settings, except for when paired with high-temperature sampling. We suspect this is because high-temperature sampling tends to produce less coherent responses, which may look novel in representation space, yet do not necessarily contain correct answers.
### 4.2 Extension: Representation-based exploration at the token level
While useful in its own right as a testbed for benchmarking the viability of exploration methods, one drawback of the inference-time selection setting is that compute—as measured by N N , the size of the per-question data pool—may need to be rather large relative to k k for selection to yield improvements. As an extension, we conduct a preliminary investigation into algorithms that use elliptic bonuses to guide the autoregressive generation process itself, removing the need to generate such a pool at all.
#### Representation-based exploration for autoregressive generation
To guide sampling for improved diversity, given a budget k k , we use features from responses 1 , … , i − 1 1,\ldots,i-1 to guide the generation of the i th i\text{th} response by modifying the logits at every generation step. Specifically, consider the i th i\text{th} generation for a given prompt x x . At each position t t within the generation, we perturb the | V | |V| -dimensional token-level logit vector as: 𝐳 ~ ( i ) ( x , y < t ) = 𝐳 ( i ) ( x , y < t ) + β ⋅ 𝐛 ( i ) ( x , y < t , V ) , 𝐳 ~ ( i ) ( x ) ∈ ℝ | V | , \tilde{\mathbf{z}}^{(i)}(x,y_{<t})\;=\;\mathbf{z}^{(i)}(x,y_{<t})+\beta\cdot\mathbf{b}^{(i)}(x,y_{<t},V),\quad\tilde{\mathbf{z}}^{(i)}(x)\in\mathbb{R}^{|V|}, where the bonus 𝐛 ( i ) ( x , y < t , V ) \mathbf{b}^{(i)}(x,y_{<t},V) is a token-level elliptic bonus, defined as: 𝐛 j ( i ) ( x , y < t , V ) = h ~ θ ( x , y < t , v j ) ⊤ Σ ( i ) − 1 h ~ θ ( x , y < t , v j ) , \mathbf{b}_{j}^{(i)}(x,y_{<t},V)=\sqrt{\tilde{h}_{\theta}(x,y_{<t},v_{j})^{\top}\Sigma_{(i)}^{-1}\tilde{h}_{\theta}(x,y_{<t},v_{j})}, for v j ∈ V v_{j}\in V . Here h ~ θ ( x , y < t , v j ) \tilde{h}_{\theta}(x,y_{<t},v_{j}) is a mean-centered Transformer representation for sequence ( y < t , v j ) (y_{<t},v_{j}) ; see Section B.2 for further details.
RF5 : RepExp for autoregressive generation improves solve rate . In Fig. 7 , we visualize the effect of token-level representation-based exploration for Qwen-2.5-7B-Instruct on the MATH task. We use two values for the bonus parameter β \beta (0.5, 1.0) and compare the pass@k to vanilla autoregressive generation for the 200 hardest (but solvable) questions in MATH , as judged by GPT-4o-mini . While token-level exploration tends to solve fewer problems compared to vanilla generation when given a small budget, this trend reverses when the budget exceeds 512 − 640 512-640 (depending on choice for β \beta ). Fig. 10 further shows that the improvement in solve rate is largest on the hardest questions. These results are encouraging, though further research is required (e.g., on more tasks and models) before one can draw a definitive conclusion. Further, our implementation is not optimized for efficiency and requires at least one additional forward pass per generation step compared to naive sampling.
### 4.3 Application: Protein sequence generation