Representation-Based Exploration for Language Models: From Test-Time to Post-Training — part 4/5

Cleaned-text code points [47462, 63213). Bibliography omitted here; source.md includes it.

Various papers have adapted exploration techniques from deep learning to language models, including Bai et al. (2025) (count-based exploration), Gao et al. (2025) (random network distillation), and Liu et al. (2024b) ; Dwaracherla et al. (2024) (posterior sampling). 5 5 5 See also Arumugam and Griffiths (2025) , which uses a pre-trained model to simulate posterior sampling in-context for multi-turn sequential decision making tasks. These works show initial promise in terms of sample complexity benefits, but their potential to explore beyond the base model in reasoning domains has not been evaluated to our knowledge. In addition, these methods require additional learning machinery (e.g., auxiliary networks), which introduces significant complexity when scaling to language models.

#### Theoretical analysis of language model exploration

On the theoretical side, our work draws on Foster et al. (2025) , who prove that test-time exploration with representation-based diversity has provable computational benefits in a simplified linear setting. Our RepExp algorithm for test-time exploration can be viewed as a simplified, practical adaptation of their theoretical algorithm.

Other theoretical works on exploration with language models include the XPO algorithm of Xie et al. (2024) and related algorithms by Cen et al. (2024) ; Zhang et al. (2024) , 6 6 6 Cen et al. (2024) ; Zhang et al. (2024) concurrently proposed similar algorithms to XPO , but did not provide non-trivial theoretical guarantees (e.g., guarantees that indicate benefits over purely passive exploration). which augment the Online DPO objective with exploration bonuses inspired by the optimism principle. To our knowledge, these techniques have only been evaluated on RLHF tasks, and Foster et al. (2025) show that there may be computational barriers to implementing them in a way that is faithful to the theoretical guarantees.

## Appendix B Details for Inference-Time Experiments ( Section 4 )

### B.1 Details from Section 4.1

#### Hyperparameters

In Algorithm 1 , we set λ = 1.0 \lambda=1.0 . For all models and tasks, we perform a sparse projection from the respective model hidden dimension to d = 512 d=512 .

#### Preprocessing

In Algorithm 1 , after obtaining the representations h ¯ θ \bar{h}_{\theta} for every generation y y for a fixed prompt x x , we sparse project all representations down and then mean-center where the mean is taken across the response-level representations.

#### Datasets

Below we provide a brief overview of all datasets along with relevant numerical details. Note that "vanilla" sampling settings refer to τ = 1.0 \tau=1.0 , top-p = 1.0 \text{top-p}=1.0 , and min-p = 0.0 \text{min-p}=0.0 . We also do not use top-k sampling in any of the coreset experiments. Finally, we only use the test split of every dataset for all our inference-time experiments, unless specified otherwise. • MATH . This dataset contains 12.5k problems from high school math competitions, split into 7.5k training examples and 5k test examples. For each question in the test split, we generate 6400 6400 responses using vanilla settings and set the maximum response length per generation to 512 512 tokens.

• GSM8K . This dataset contains 8.79k grade school math word problems, split into 7.47k training examples and 1.32k test examples. For each question in the test split, we generate 6400 6400 responses using vanilla settings and set the maximum response length per generation to 512 512 tokens.

• MBPP+ . This dataset contains 378 basic Python programming problems that are a curated subset of the full MBPP ( Austin et al., 2021 ) dataset with more test cases. Since the dataset does not come with any train or test splits, we use the full set of questions for our experiments. For each problem, we generate 6400 6400 responses using vanilla settings, except that we set top-p = 0.95 \text{top-p}=0.95 . We set the maximum response length per generation to 768 768 tokens.

• Game of 24 . This dataset contains 1.36k questions that specify four integers that need to be combined using basic arithmetic operations ( + , − , x , / ) (+,-,x,/) to equal 24 24 . For each question, we generate 6400 6400 responses using vanilla settings and set the maximum response length per generation to 512 512 tokens. We use the version available at https://huggingface.co/datasets/nlile/24-game .

• AIME 2025 . This dataset contains the 30 30 problems taken directly from the 2025 edition of the American Invitational Mathematics Examination (AIME). For each question, we generate 8192 8192 responses using vanilla settings and set the maximum response length per generation to 8192 8192 tokens.

#### Experiment resources

We used vLLM ( Kwon et al., 2023 ) on 1 − 2 1-2 (depending on the size of the model) NVIDIA A100 40GB GPUs per model-task pair to generate the data pools for all questions in the dataset.

#### Estimating samples-to-correct

For random sampling, we estimate the average number of samples to take (without replacement) from the data pool to find the first correct one as: samples-to-correct = N + 1 c + 1 , \text{samples-to-correct}=\frac{N+1}{c+1},

where N N indicates the size of the data pool and c c indicates the number of correct samples in the pool. Please refer to Rohatgi and Saleh (2015) for a proof. For the representation-based exploration algorithm described in Algorithm 1 , we perform 5 5 trials per question where we record the samples-to-correct for each and take their average.

#### Estimating pass@k

To compute the pass@k values plotted in Figure 6 , we follow Chen et al. (2021) and use the following unbiased estimator: pass@k = 𝔼 𝒟 ​ [ 1 − ( n − c k ) ( n k ) ] = 1 | 𝒟 | ​ ∑ i = 1 | 𝒟 | [ 1 − ( n − c k ) ( n k ) ] , \text{pass@k}=\mathbb{E}_{\mathcal{D}}\left[1-\frac{\binom{n-c}{k}}{\binom{n}{k}}\right]=\frac{1}{|\mathcal{D}|}\sum_{i=1}^{|\mathcal{D}|}\left[1-\frac{\binom{n-c}{k}}{\binom{n}{k}}\right],

where 𝒟 \mathcal{D} indicates the dataset and | 𝒟 | |\mathcal{D}| indicates its size.

### B.2 Details from Section 4.2

#### Hyperparameters

Similarly to Section 4.1 , we initialize Σ ( 0 ) − 1 = λ − 1 ​ I d \Sigma_{(0)}^{-1}=\lambda^{-1}I_{d} . We set λ = 0.1 \lambda=0.1 . For all values of β \beta , we use top-p = 0.95 \text{top-p}=0.95 and top-k = 128 \text{top-k}=128 . At every time step t t , we use a batch size of 64 to compute h θ ​ ( x , y < t , v j ) h_{\theta}(x,y_{<t},v_{j}) for all v j ∈ V v_{j}\in V where v j > − ∞ v_{j}>-\infty (note that a logit v j v_{j} is set to − ∞ -\infty if it gets filtered out by either the top-p or top-k filters mentioned earlier).

#### Computational expense

Note that computing the bonus 𝐛 ⁡ ( x , y < t , V ) \mathbf{b}(x,y_{<t},V) requires one forward pass through the model per token in the vocabulary at every time step t t . While these can be batched together since all tokens share the same prefix, this is still prohibitively time and memory intensive. To mitigate this, we combine this method with nucleus and top-k sampling such that the bonus will only need to be computed for at most k ≪ V k\ll V tokens.

#### Dataset construction

Due to the large computational cost of experiments, we only focus on MATH using Qwen-2.5-7b-Instruct . In addition, we do not evaluate on the full test split of MATH , but instead use a subset consisting of the 200 200 hardest questions as ranked by GPT-4o mini . Specifically, we sampled 1024 1024 responses for each question in the MATH test split using GPT-4o mini to estimate the per-question pass@1. We threw out all questions for which the pass@1 was 0 0 (indicating not a single response was correct among all 1024 1024 ), sorted the remaining questions, and kept the 200 200 questions with the lowest pass@1.

#### Estimating samples-to-correct

We collected up to 1024 1024 generations per question (stopping early once the correct answer was found). Since the resulting samples-to-correct value found can have high variance due to the inherent randomness in the generation process, we repeat this process 5 5 times with a different seed every time. This then results in a total of 200 × 5 = 1000 200\times 5=1000 data points, minus a few data points that didn’t finish running in time, for an effective total of 985 985 data points used in Fig. 7 .

#### Experiment resources

Experiments were run on a combination of NVIDIA A100 40GB, NVIDIA A100 80GB, and NVIDIA H100 80GB GPUs. Every run indicates one seed for a fixed question and performs up to 1024 1024 generations. We used one GPU (one of the several mentioned earlier) per run, adjusting the forward batch size to compute the elliptical bonus down from 64 64 to 32 32 when using the 40GB GPUs.

#### Computing the inverse covariance

The inverse covariance matrix Σ ( i ) − 1 \Sigma^{-1}_{(i)} includes all mean-centered hidden representations for all generated tokens in all i − 1 i-1 complete sequences generated so far for a fixed prompt x x . Note that the mean μ ( i ) \mu^{(i)} of the raw hidden representations h θ h_{\theta} is computed as: μ ( i ) = 1 H ​ ∑ j = 0 i − 1 ∑ t = 1 T j h θ ​ ( x , y < t ( j ) ) , H = ∑ j = 0 i − 1 T j , \ \mu^{(i)}=\frac{1}{H}\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}h_{\theta}(x,y_{<t}^{(j)}),\quad H=\sum_{j=0}^{i-1}T_{j}, where T j T_{j} indicates the length (in number of tokens) of response j j . Because of this, the mean changes after every generation i i . This means some care is required when computing the inverse covariance matrix Σ ( i ) − 1 \Sigma^{-1}_{(i)} with mean-centered hidden representations. To account for this, we will separately keep track of the inverse covariance matrix Σ ~ ( i ) − 1 \tilde{\Sigma}^{-1}_{(i)} with non -mean-centered hidden representations as well as the mean μ ( i ) \mu^{(i)} . Then, we compute the mean-centered inverse covariance matrix Σ ( i ) − 1 \Sigma^{-1}_{(i)} as

Σ ( i ) − 1 = Σ ~ ( i ) − 1 − ( Σ ~ ( i ) − 1 ​ μ ( i ) ​ μ ( i ) T ​ Σ ~ ( i ) − 1 − 1 / H + μ ( i ) T Σ ~ ( i ) − 1 μ ( i ) ) . \Sigma^{-1}_{(i)}=\tilde{\Sigma}_{(i)}^{-1}-\left(\frac{\tilde{\Sigma}_{(i)}^{-1}\mu^{(i)}\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}}{-1/H+\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}\mu^{(i)}}\right).

This result is immediate from Proposition B.2 .

.

{proposition} Given the current generation step i i , the non-mean-centered inverse data covariance matrix Σ ~ ( i ) − 1 \tilde{\Sigma}^{-1}_{(i)} and the current mean μ ( i ) \mu^{(i)} , the mean-centered inverse data covariance matrix Σ ( i ) − 1 \Sigma^{-1}_{(i)} is given as Σ ( i ) − 1 = Σ ~ ( i ) − 1 − ( Σ ~ ( i ) − 1 ​ μ ( i ) ​ μ ( i ) T ​ Σ ~ ( i ) − 1 − 1 / H + μ ( i ) T Σ ~ ( i ) − 1 μ ( i ) ) , H = ∑ j = 0 i − 1 T j . \Sigma^{-1}_{(i)}=\tilde{\Sigma}_{(i)}^{-1}-\left(\frac{\tilde{\Sigma}_{(i)}^{-1}\mu^{(i)}\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}}{-1/H+\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}\mu^{(i)}}\right),\quad H=\sum_{j=0}^{i-1}T_{j}. Here, T j T_{j} indicates the length (in number of tokens) of response j j .

Proof of Section B.2 . We can write Σ ( i ) \displaystyle\Sigma^{(i)} = ∑ j = 0 i − 1 ∑ t = 1 T j ( h t ( j ) − μ t ( i ) ) ​ ( h t ( j ) − μ t ( i ) ) T \displaystyle=\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}(h^{(j)}_{t}-\mu^{(i)}_{t})(h_{t}^{(j)}-\mu^{(i)}_{t})^{T} = ∑ j = 0 i − 1 ∑ t = 1 T j h t ( j ) ​ ( h t ( j ) ) T − ∑ j = 0 i − 1 ∑ t = 1 T j h t ( j ) ​ ( μ ( i ) ) T − ∑ j = 0 i − 1 ∑ t = 1 T j μ ( i ) ​ ( h t ( j ) ) T + ∑ j = 0 i − 1 ∑ t = 1 T j μ ( i ) ​ ( μ ( i ) ) T \displaystyle=\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}h^{(j)}_{t}(h^{(j)}_{t})^{T}-\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}h^{(j)}_{t}(\mu^{(i)})^{T}-\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}\mu^{(i)}(h_{t}^{(j)})^{T}+\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}\mu^{(i)}(\mu^{(i)})^{T} = Σ ~ ( i ) − ( ∑ j = 0 i − 1 ∑ t = 1 T j h t ( j ) ) ​ ( μ ( i ) ) T − μ ( i ) ​ ( ∑ j = 0 i − 1 ∑ t = 1 T j ( h t ( j ) ) T ) + ∑ j = 0 i − 1 ∑ t = 1 T j μ ( i ) ​ ( μ ( i ) ) T \displaystyle=\tilde{\Sigma}^{(i)}-\left(\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}h_{t}^{(j)}\right)(\mu^{(i)})^{T}-\mu^{(i)}\left(\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}(h_{t}^{(j)})^{T}\right)+\sum_{j=0}^{i-1}\sum_{t=1}^{T_{j}}\mu^{(i)}(\mu^{(i)})^{T} = Σ ~ ( i ) − H ​ μ ( i ) ​ ( μ ( i ) ) T − μ ( i ) ​ H ​ ( μ ( i ) ) T + H ​ μ ( i ) ​ ( μ ( i ) ) T \displaystyle=\tilde{\Sigma}^{(i)}-H\mu^{(i)}(\mu^{(i)})^{T}-\mu^{(i)}H(\mu^{(i)})^{T}+H\mu^{(i)}(\mu^{(i)})^{T} = Σ ~ ( i ) − 2 ​ H ​ μ ( i ) ​ ( μ ( i ) ) T + H ​ μ ( i ) ​ ( μ ( i ) ) T \displaystyle=\tilde{\Sigma}^{(i)}-2H\mu^{(i)}(\mu^{(i)})^{T}+H\mu^{(i)}(\mu^{(i)})^{T} = Σ ~ ( i ) − H ​ μ ( i ) ​ ( μ ( i ) ) T . \displaystyle=\tilde{\Sigma}^{(i)}-H\mu^{(i)}(\mu^{(i)})^{T}. Inverting Σ ( i ) \Sigma^{(i)} , we conclude that Σ ( i ) − 1 \displaystyle\Sigma_{(i)}^{-1} = ( Σ ~ ( i ) − H ​ μ ( i ) ​ μ ( i ) T ) − 1 \displaystyle=\left(\tilde{\Sigma}^{(i)}-H\mu^{(i)}\mu_{(i)}^{T}\right)^{-1} = Σ ~ ( i ) − 1 − ( Σ ~ ( i ) − 1 ​ μ ( i ) ​ μ ( i ) T ​ Σ ~ ( i ) − 1 − 1 / H + μ ( i ) T Σ ~ ( i ) − 1 μ ( i ) ) , \displaystyle=\tilde{\Sigma}_{(i)}^{-1}-\left(\frac{\tilde{\Sigma}_{(i)}^{-1}\mu_{(i)}\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}}{-1/H+\mu_{(i)}^{T}\tilde{\Sigma}_{(i)}^{-1}\mu^{(i)}}\right), where we used the Woodbury matrix identity lemma 7 7 7 https://en.wikipedia.org/wiki/Woodbury_matrix_identity in the last step with U = μ ( i ) U=\mu^{(i)} , C = − H C=-H , and V = μ ( i ) T V=\mu_{(i)}^{T} . ∎

We exclude representations from the current sequence i i in Σ ( i ) \Sigma^{(i)} to keep the generation from veering off topic, and update Σ ( i ) − 1 \Sigma^{-1}_{(i)} after each generation i i using T i T_{i} consecutive applications of the Sherman-Morrison update in Algorithm 1 , one for each hidden representation of the output tokens. Finally, we found it necessary for numerical stability to perform all covariance related computations in double precision.

#### Additional plots

In Fig. 10 , we provide a revised version of Fig. 7 where we add shaded areas indicating one standard error to the left plot, and we additionally add β = 0.5 \beta=0.5 to the right plot.

## Appendix C Details for RL Post-Training Experiments ( Section 5 )

#### Hyperparameters

We use verl for training ( Sheng et al., 2024 ) , and provide a full overview of all common hyperparameters in Table 3 , all hyperparameters specific to unlikeliness in Table 3 , and all hyperpameters specific to RepExp in Table 3 . Note that for AIME 2024 , we adjusted the maximum prompt length to 2048, the maximum response length to 8192, the train batch size to 512, the ppo mini batch size to 128, and the ppo micro batch size per gpu to 8. Also, for RepExp on AIME 2024 , we increased the sparse projection dimension from 32 to 128.

#### Algorithm details

We note that we mean-center the representations h ¯ θ \bar{h}_{\theta} that are used to compute the elliptic bonuses as described in Section 3 , where the mean is taken over all the response-level representations of the current group of rollouts for a fixed prompt x x . In addition, we do not add a bonus for questions where all rollouts in the batch are incorrect, as we found this to empirically hurt performance.

#### Dataset splits

For MATH , we use the original 7.5k train split for training, MATH-500 ( Lightman et al., 2023 ) for validation, and the 4.5k (originally 5k minus the problems in MATH-500 )) test split for testing. For GSM8K, we use the original 8.79k train split for training, the first 512 512 examples from the test split for validation, and the remaining 807 807 examples from the test split for testing. Finally for AIME 2024 , we use 4096 4096 examples randomly chosen from the full DAPO-Math-17K dataset for training and use the full AIME 2024 dataset both for validation and testing.

#### Modifications to unlikeliness baseline