Learning to Reason with Mixture of Tokens

## Abstract

Abstract Reinforcement learning with verifiable rewards (RLVR) has become a leading approach for improving large language model (LLM) reasoning capabilities. Most current methods follow variants of Group Relative Policy Optimization, which samples multiple reasoning completions, scores them relative to each other, and adjusts the policy accordingly. However, these approaches invariably sample discrete tokens at each reasoning step, discarding the rich distributional information in the model’s probability distribution over candidate tokens. While preserving and utilizing this distributional information has proven beneficial in non-RL settings, current RLVR methods seem to be unnecessarily constraining the reasoning search space by not using this information. To address this limitation, we investigate mixture-of-token generation (MoT-G) in RLVR. We present a unified framework that generalizes existing MoT-G approaches, including existing training-free methods that construct mixture embeddings as weighted sums over token embeddings, and extend RLVR to operate directly in this continuous mixture space for generating chain-of-thought. Evaluating two MoT-G variants on Reasoning-Gym, a suite of reasoning-intensive language tasks, we find that MoT-G methods achieve substantial improvements (5–35% gains on 7/10 tasks) compared to standard decoding with the Qwen2.5-1.5B model, while reaching comparable accuracy with half the number of trajectories, suggesting improved training efficiency. Through comprehensive hidden-state and token-level analyses, we provide evidence that MoT-G’s benefits may stem from its ability to maintain higher hidden-state entropy throughout the reasoning process and promote exploration in token space.

## 1 Introduction

Reinforcement learning with verifiable rewards (RLVR) has emerged as a dominant paradigm for enhancing the reasoning capabilities of large language models. In this approach, models are provided with a “scratchpad” to develop their reasoning step-by-step before producing final answers, with training focused solely on outcome-based rewards while allowing the model to autonomously discover effective reasoning strategies (Shao et al., 2024; Kambhampati et al., 2025). Recent advances in RLVR have demonstrated remarkable success across challenging domains, from achieving gold medal-level performance on International Mathematical Olympiad problems Huang and Yang (2025) to saturating previously difficult benchmarks in logical reasoning Guo et al. (2025). These achievements highlight the potential of allowing models to develop their own reasoning strategies rather than constraining them to human-prescribed approaches, suggesting that RLVR represents a fundamental advancement in how we train LLMs for complex reasoning tasks.

Most current RLVR implementations rely on group relative policy optimization (GRPO) and similar techniques. Which all basically operate by, for a given question, using discrete token sampling to generate some $G$ number of completions or trajectories (reasoning scratchpad + final answer), scoring those final answers, and using the relative scores of each rollout to train the LLM to give better answers (Zheng et al., 2025; Liu et al., 2025b). These methods have proven to work on a variety of reasoning intensive tasks (Stojanovski et al., 2025; Liu et al., 2025a; Guo et al., 2025), and there has been substantial research effort to improve this flavor of RL algorithm, from better normalization methods, to expanding it to multi-turn and tool use trajectories.

However, much less attention has been paid to a fundamental assumption underlying these approaches: the commitment to discrete tokens at each reasoning step. This discrete sampling paradigm forces models to make early, irreversible decisions in token space, limiting their ability to maintain uncertainty and explore alternative reasoning paths during critical intermediate steps of complex problem-solving.

Figure: (1) Mixture of Tokens based Generation (MoT-G)
Refer to caption: figs/performance_differences_relative_to_vanilla.png

Although some prior work has explored continuous or mixture token representations in language modeling (Zhang et al., 2025b; Deng et al., 2024), there has been limited large-scale studies systematically investigating the integration of mixture-of-token generation within the RLVR training framework (Butt et al., 2025).
Furthermore, existing work has not focused on understanding the mechanistic benefits that mixture token representations might provide for reasoning tasks, particularly in terms of exploration capabilities and training efficiency. This represents a significant gap in our understanding of how token-level design choices impact effectiveness of reinforcement learning for reasoning and this paper addresses the following question:

To study this research question, we introduce mixture-of-token generation (MoT-G) for RLVR, a framework that enables models to maintain distributions over multiple tokens at each reasoning step rather than committing to discrete selections. Our main contributions can be summarized as:

1. Generalized framework for mixture of tokens based generation (MoT-G).
We introduce MoT-G for RLVR, a unified framework that enables mixture-of-token generation within reinforcement learning paradigms for reasoning, generalizing training-free approaches and more importantly extending RLVR to operate in continuous mixture spaces rather than discrete token selections.

2. Empirical Results. We numerically analyze the efficacy of our work on a variety of reasoning tasks from Reasoning-Gym (Stojanovski et al., 2025) in Section [4](https://arxiv.org/html/2509.21482v1#S4). We provide comprehensive empirical evaluation demonstrating significant performance improvements across reasoning tasks, with MoT-G achieving 5–35% gains on 7 out of 10 tasks while requiring half the trajectories to reach comparable performance, suggesting improved training efficiency. We present experiments ablation studies on different model sizes, hyperparameters etc. in the Appendix [B](https://arxiv.org/html/2509.21482v1#A2).

3. Effect of MoT-G on information in the hidden states and exploration. We offer mechanistic insights into why MoT-G benefits reasoning through hidden-state and token-level analyses, showing that the approach maintains higher entropy and promotes exploration in token space, providing a deeper understanding of how continuous representations enhance reasoning capabilities.

Notation. We denote vectors by bold-face small case letters, $\mathbf{x},\mathbf{y}..$, matrices and tensors by bold-face capital case letters $\mathbf{X},\mathbf{Y}..$ and sets by curly letters $\mathcal{X},\mathcal{Y}..$. Denote $\mathcal{P}(\cdot)$ to be the probability simplex of the set. $\mathbb{P}$ and $\mathbb{E}$ denote the probability and expectation over the appropriate measure.

## 2 Background: Related Work and Preliminaries

Chain-of-Thought and Scratchpad Reasoning. It has long been established that large language models demonstrate substantially improved performance when provided with intermediate reasoning steps, commonly referred to as “scratchpads” or chain-of-thought prompting (Wei et al., 2023).

Formally, for an input sequence of embeddings (prompt), $\mathbf{X}_{0}=(\mathbf{x}_{-m+1},\dots\mathbf{x}_{0})^{\mathsf{T}}\in\mathbb{R}^{d\times m}$, where $m$ is the prompt length, rather than directly predicting a final answer $y$, the models generate an intermediate reasoning trajectory $\tau=(\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{n})$ where each $\mathbf{x}_{t}$ represents the embedding corresponding to the token of the $t$-th reasoning step. This is followed by the final answer $y$. This can be expressed as: $\mathbb{P}(y|\mathbf{X}_{0})=\sum_{\tau}\mathbb{P}(y|\tau,\mathbf{X}_{0})\mathbb{P}(\tau|\mathbf{X}_{0})$ where $\mathbb{P}(\tau|\mathbf{X}_{0})=\prod_{t=1}^{T}\mathbb{P}(\mathbf{x}_{t}|\mathbf{x}_{<t},\mathbf{X}_{0})$ represents the probability of generating the reasoning chain $\tau$. This approach has proven remarkably effective across diverse reasoning tasks, from mathematical problem-solving to logical inference, fundamentally changing how we approach complex reasoning with language models.

Reinforcement Learning with Verifiable Rewards (RLVR).
GRPO was first successfully demonstrated in DeepSeekMath (Shao et al., 2024) and has since improved reasoning across models including DeepSeek R1, Nemotron, and Phi-4.
The key insight is that models can learn to develop their own effective reasoning strategies through outcome-based training alone. Rather than requiring human supervision of intermediate reasoning steps, RLVR enables models to discover optimal reasoning pathways by sampling multiple completion trajectories and learning from final outcomes.

The objective in RLVR is to optimize a policy $\pi_{\theta}$ (parameterized by $\theta$) to maximize expected return,

$$ $\displaystyle\pi^{*}=\underset{\mathbf{\theta}}{\arg\min}\ \mathcal{J}(\mathbf{\theta})=\underset{\mathbf{\theta}}{\arg\min}\ \mathbb{E}_{\tau\sim\pi_{\mathbf{\theta}}}\left\{r(\tau)\right\}$ $$

where $\tau$ represents a complete reasoning trajectory and $r(\tau):\mathcal{T}\rightarrow\mathbb{R}$ is the reward function defined on the space of trajectories $\tau$. Most current implementations employ Group Relative Policy Optimization (GRPO), which computes advantages relatively within groups of sampled trajectories:

$$ $\displaystyle\mathcal{J}_{\text{GRPO}}=\mathbb{E}_{(\tau_{1},\tau_{2},\dots,\tau_{G})\sim\pi_{\theta_{\text{old}}}}\left\{\frac{1}{G}\sum_{g=1}^{G}\frac{1}{|\tau_{g}|}\sum_{t=1}^{|\tau_{g}|}A_{g,t}-\beta\mathbb{D}_{\text{KL}}[\pi_{\theta}||\pi_{\text{ref}}]\right\}$ $$

where the approximate advantage for trajectory $g$ at step $t$ is computed as $A_{g,t}=\frac{r_{g}-\mu}{\sigma}$, with $r_{g}$ being the reward of trajectory $g$, $\mu$ the mean group reward, and $\sigma$ the standard deviation.

This approach has catalyzed the current wave of multi-turn reasoning systems, tool-use capabilities, and agentic AI applications by demonstrating that sophisticated reasoning behaviors can emerge from simple outcome optimization.
Although there has been research critically analyzing RLVR’s complete utilization of base model potential (Yue et al., 2025) and reasoning consistency (Shojaee et al., 2025), recent empirical (Liu et al., 2025a) and theoretical (Wen et al., 2025) results demonstrate that GRPO incentivizes correct reasoning paths and logical integrity.

Continuous and Mixture of Token based Generation. Parallel to advances in reinforcement learning for reasoning, a separate line of research has questioned the fundamental assumption of discrete token sampling in autoregressive generation. COCONUT (Hao et al., 2024) uses distillation to replace chain-of-thought with continuous embeddings, with theoretical work showing polynomial-time solutions through token superposition (Zhu et al., 2025a).
Training-free methods include soft-thinking (Zhang et al., 2025b), which averages token embeddings weighted by probabilities, and approaches using VQ-VAE compression (Su et al., 2025), speculative soft tokens (Xu et al., 2025), embedding optimization (Zhu et al., 2025b), and Bayesian posterior mixing (Zhuang et al., 2025). The Coconut paper (Hao et al., 2024) and related mixture-of-tokens work (Zhang et al., 2025b) have demonstrated that such approaches can improve generation quality and enable more nuanced reasoning in standard (non-RL) settings. However, none have systematically studied mixture generation within reinforcement learning frameworks for reasoning tasks.

At each generation step, language models produce a rich probability distribution $\mathbb{P}(\cdot|z_{<t})$ over the vocabulary $\mathcal{D}$, yet standard generation discards this distributional information by sampling a single token $z_{t}\sim\mathbb{P}(\cdot|z_{<t})$ and feeding only its one-hot representation forward. Recent work has explored preserving this distributional information through mixture-of-token approaches. In soft-thinking (Zhang et al., 2025b), mixture embeddings are constructed as weighted combinations over all vocabulary tokens: $\mathbf{x}_{\text{mix}}=\sum_{z\in\mathcal{D}}\mathbb{P}(z|\mathbf{x}_{<t})\mathbf{x}^{z}$ where $\mathbf{x}^{z}$ represents the embedding of token $z$. More generally, mixture-of-token generation (MoT-G) methods maintain representations that preserve uncertainty and alternative paths rather than committing to discrete selections. Formally, instead of standard discrete generation, MoT-G uses continuous representations: $\mathbf{x}_{\text{mix}}=\texttt{aggregate}(\texttt{sample}(\mathbf{x}_{\text{mix},<t}))$, where sample and aggregate are defined later.

Combining RLVR with Mixture Generation.
Despite the demonstrated benefits of both RLVR for reasoning and mixture-of-token approaches for generation quality, no prior work has systematically investigated their combination. This represents a significant missed opportunity: while RLVR methods excel at discovering effective reasoning strategies through outcome optimization, they may be unnecessarily constraining the exploration space by forcing discrete commitments at each reasoning step. Conversely, mixture approaches have shown promise for preserving reasoning flexibility but have not been studied in the context of reinforcement learning where models must learn to balance exploration and exploitation across extended reasoning trajectories. The central question we address is whether extending RLVR to operate in continuous mixture spaces can enhance both reasoning performance and training efficiency, potentially allowing models to defer commitments and maintain consideration of alternative reasoning paths throughout the learning process.

## 3 Mixture Of Tokens Generation (MoT-G)

We now describe one of our key contribution, a generalized framework for mixture of token generation. Put simply, we replace the embedding of the sampled token with a mixture of $k$ tokens. Consider a large language model $\mathcal{L}$ which has vocabulary $\mathcal{D}$ and context length $\mathsf{N}$. Let (input) token embeddings be mapped by the function $l:\mathcal{D}\to\mathcal{X}$ and let
$\mathbf{X}_{t}=(\mathbf{x}_{1},\dots,\mathbf{x}_{t})\in\mathcal{X}^{t}$ denote a tensor of token embeddings corresponding to a sequence of length $t$. Denote the transformer function of $\mathcal{L}$ by $h:\mathcal{X}^{\mathsf{N}}\to\mathcal{P}(\mathcal{D})$ which takes in input as token embedding tensor and outputs a probability distribution over the vocabulary and denote $\mathbf{p}_{t}\in\mathcal{P}(\mathcal{D})$ to be the output at generation step $t$.

Recently a technique titled “soft-thinking” was introduced which uses a weighted average of token embeddings during the generation phase (Zhang et al., 2025b). We provide recipe for generating using a mixture of token embeddings which generalizes soft-thinking and allows for a wide-variety of choices to design a generation approach. Note that any standard generation technique can be derived as a special case of this framework. There are four key steps in the generation process - initialize - for initialization, sample - sampling, aggregate - for aggregation and endcriteria - for end criteria. We now describe the mixture of token generation:

Mixture of Token Generation (MoT-G): initialize $\mathbf{X}_{0}=(\mathbf{x}_{-m+1},\dots\mathbf{x}_{0})^{\mathsf{T}}\in\mathbb{R}^{d\times m}$. Do the following for $t=1,2,\dots$ until $\texttt{endcriteria}(\mathbf{X}_{t})$.

Step 1. Sample $k$ tokens $\mathcal{S}_{t}=\texttt{sample}(\mathbf{p}_{t}=h(\mathbf{X}_{t-1}))$.

Step 2. Get embedding $\mathbf{x}_{t}=\texttt{aggregate}(\mathcal{S}_{t},h(\mathbf{X}_{t-1}))$. Append $\mathbf{x}_{t}$ to end of $\mathbf{X}_{t}$.

After the loop terminates, do standard generation (described above) till the end of sequence token.

We now discuss $\texttt{initialize},\texttt{endcriteria},\texttt{sample}$ and aggregate in detail.

Possible Initialization and End Criteria:
One approach is to start answering directly using the chat template (if available) that the model was instruction finetuned on. Similar to soft-thinking (Zhang et al., 2025b), one can add a think token (e.g. <think>) at the beginning of the generation. Further akin to the think step by step prompt (Wei et al., 2023), one can add prompts which encourage the model to consider contrasting options (e.g. think about different approaches to solve).

The end-criteria decides when the mixture of token generation ends and standard generation resumes. One way is to generate till a end think token (e.g. </think> token) is the most-likely token. Another technique which the soft-thinking takes is to stop when the entropy of $\mathbf{p}_{t}=h(\mathbf{X}_{t})$ is below a certain threshold (chosen hyperparameter) for consecutive rounds (another hyperparameter) and perform parameter sweeps for the optimal value of the entropy threshold (Zhang et al., 2025b).
Note that setting the hyperparameters is a challenging task dependent on the problem instance, which can be subsided if using RL which allows the LLM to learn when to stop using MoT-G.

Sample Function: The sample function $\texttt{sample}:\mathcal{P}(\mathcal{D})\to 2^{\mathcal{D}}$ takes in as input the probability distribution and returns $k$ tokens. Previous research in the training-free regime (Zhang et al., 2025b; Zhuang et al., 2025) consider all tokens and change the aggregation mechanism (described next), however for ensuring diversity across trajectories we consider a subset of tokens as the input.
There are different sampling techniques one can come up with and we discuss a few.

- 1.
Top-$k$ sampling This is extension of greedy sampling to the mixture domain, where the tokens with the $k$ highest probabilities are selected.
- 2.
min-$p$ sampling: Sample $k$ tokens without replacement from the subset of tokens which have probabilities greater than $p$ (hyperparameter), with normalized probabilities.
- 3.
Nucleus sampling: Sample $k$ tokens from the subset of tokens with cumulative probability greater than some threshold (hyperparameter).
- 4.
SWR $k$ sampling: Sample $k$ tokens without replacement (SWR) with some temperature (similar to normal sampling in standard autoregressive generation).

Aggregate Function: The aggregate function, $\texttt{aggregate}:2^{\mathcal{D}}\times\mathcal{P}(\mathcal{D})\to\mathbb{R}^{d}$ aggregates embeddings of the sampled vectors into a single mixture embedding.
There are many different approaches that one can take to aggregate $k$ tokens into a single token embedding. A more general framework would map $k$ tokens into $m\ (<k)$ tokens, and future work can look at such an approach bridging the gap between this paper and COCONUT where embeddings are learned to compress think tokens across generation steps. We now highlight a few obvious choices for aggregation.

The most obvious choice is to consider a weighted sum of the sampled tokens with different weighing choices, $\sum_{i\in|\mathcal{S}_{t}|}w_{i}\mathcal{S}_{t}[i]$, where $w_{i}$ are the weights for $i$-th token in the sampled set $\mathcal{S}_{t}$. One can weigh the embedding uniformly by $\nicefrac{{1}}{{k}}$, effectively taking the bary-centre of the sampled tokens. Further one can weigh the embeddings by normalized probabilities (Zhang et al., 2025b), which allows more representation of tokens with higher probabilities.  Zhuang et al. (2025) comes up with a more sophisticated aggregation mechanism where the Bayesian posterior on the weighted sum to be used is updated based on the sampled token and the probability distribution. One could also weigh the vectors by randomly sampled weights from a suitable distribution, e.g. Dirichlet distribution with the normalized probabilities as the support parameter. Such an approach allows for better representation of more likely tokens on average but with added randomness.

Although we consider weighted aggregation, there are other options for aggregation, for e.g., element-wise maximum of the sampled embeddings, which lies on the convex hull. We benchmark aggregation and sampling schemes on GSM-8K in the training-free regime in Appendix [B](https://arxiv.org/html/2509.21482v1#A2).

### 3.1 Adapting GRPO for MoT-G

We now discuss how GRPO can be modified for generating reasoning tokens with MoT-G.

Formatting and Rewards. Similar to standard generation, we consider using specific tokens for letting the model produce its chain of thought and once the end think token is the most-likely token to stop the mixture generation. The answer of the model is extracted from a specified format. Since the answer is generated after the mixture based generation ends, one can straightforwardly give exact-match or LLM-as-a-judge correctness rewards. Similarly the length of reasoning can be rewarded. However, any other reward which is based on rewarding the content of the reasoning chain (for e.g. using a LLM judge) and is sensitive to coherence might need some additional engineering difficult to reward in practice, which is a limitation of this method. One possible workaround is sampling a token from the set of tokens being mixed and reward this auxillary chain of thought.

Loss Computation. However the computing loss for such a generation technique is more subtle since at each step there are multiple tokens over which the loss can be propagated. We disambiguate between two different ways one can compute loss for the mixture based generation, each of them is significant in our opinion and we study a variant of each of them experimentally.

Single token losses: One can propagate the loss using the log probabilities of a single sampled token from the set of tokens used for generating mixture embeddings, $z_{t}\sim\mathcal{S}_{t}$. The log probabilities of this token can either be left unweighted or weighed by their respective probabilities.

Multiple token losses: One can propagate the loss is to each of the tokens in $\mathcal{S}_{t}$. For on-policy RL without stabilization ($\pi_{\text{old}}=\pi_{\mathbf{\theta}}$), the multiplicative factor does not matter for the advantage term, as long as the appropriate tokens are gathered. For the KL regularization loss it does create a difference theoretically, however empirically we don’t observe one can use a single token approximation.

For on-policy RL with stabilization ($\pi_{\text{old}}$ in Shao et al. (2024)) one needs a plug-in estimate for $\log\pi_{\mathbf{\theta}}(\tau)$. In standard autogeneration when the trajectory is considered the sequence of tokens ($\tau=(z_{1},\dots,z_{t})$), one can substitute this as $\log\pi_{\mathbf{\theta}}(\tau)=\log(\Pi_{n=1}^{t}\mathbb{P}(z_{n}|(z_{1},\dots,z_{n-1})))=\log(\Pi_{n=1}^{t}\mathbf{p}_{n}(z_{n}))=\sum\log\mathbf{p}_{n}(z_{n})$, where we denote $\mathbf{p}_{n}=\mathbb{P}(z_{n}|(z_{1},\dots,z_{n-1})))$.
Since now the sequence of tokens is replaced by a sequence of mixture of tokens given by the matrix with entries $(z_{m,n})_{k\times t}$. For each generation step $t$, one ideally needs to compute the probability $\mathbb{P}(\{z_{1,t},\dots,z_{1,t}\}|\mathbf{x}_{1},\dots,\mathbf{x}_{t-1})$ which can be combinatorially expensive to compute (especially if sampling without replacement). One way to approximate is to sum the log probability of each token weighted by the probability, $\log(\pi_{\mathbf{\theta}}(\tau))\approx\sum_{m,n}\mathbf{p}_{n}(z_{m,n})\log(\mathbf{p}_{n}(z_{m,n}))$.

### 3.2 Exploration in GRPO with Mixture of Tokens

Directly applying soft thinking (Zhang et al., 2025b) to GRPO based RL in LLMs is challenging because soft thinking does not allow for any randomness, since the probability distribution $\mathbf{p}_{t}$ over the dictionary and the embedding map $l$ are deterministic (fixed) given the same sequence of embedding vectors, $\mathbf{X}_{t}$. This leads to zero exploration when doing multiple trajectories in a group. We now describe and motivate two generation techniques which we analyze in our experiments.

A. Dirichlet stochastic weighing with top-$k$ tokens (Dirichlet). We sample the top$-k$ tokens greedily and the weights are sampled from a Dirichlet distribution (described in Appendix [C.2](https://arxiv.org/html/2509.21482v1#A3.SS2)) with the normalized probability distribution as the parameter. There are two primary reasons for considering this, (a) such a prior allows for exploration and increases the entropy of the mixture input embedding (b) it does so in a controlled fashion, i.e., because the Dirichlet distribution has higher entropy around the centroid of the probability simplex and low entropy near the boundary, the resulting mixture weights are aligned with the model’s confidence and do not alter them considerably (see Figure [6](https://arxiv.org/html/2509.21482v1#A3.F6) in Appendix). Note that as $k$ increases the weights are more likely to be attributed to unhelpful tokens.

B. Sampling $k$ tokens without replacement and weighing with normalized probability (Different Tokens). We now state the following result, which illustrates the fundamental tradeoff involved in the choice of $k$ for this MoT-G generation method,

###### Proposition 1 .

Consider the LLM $\mathcal{L}$ described above with dictionary $\mathcal{D}$, token embedding space $\mathcal{X}$. We consider mixture of token generation using sampling $k$ tokens without replacement and aggregating them using a weighted sum with their normalized probability. Let $G$ be the number of trajectories. Let $L_{t}$ be the number of unique tokens across the trajectories at generation step $t$ and $\mathbf{x}_{g}\in\mathcal{X}$ to be the generated mixture embedding for trajectories $g\in[G]$. With the expectation taken over the generation process, $\mathbb{E}\{L_{t}\}$ is a monotonically increasing function in $k$. Further, $\mathbb{E}_{g,g^{\prime}\sim[G]}\left\{\|\mathbf{x}_{g}-\mathbf{x}_{g^{\prime}}\|^{2}\right\}$ is a monotonically decreasing function in $k$.

The proof is in Appendix [C](https://arxiv.org/html/2509.21482v1#A3) and follows straightforwardly using arguments from probability theory.
This shows the tradeoff inherent to sampling $k$ tokens - we improve the coverage of tokens across mixtures as $k$ increases but increases the overlap in the resulting embedding mixture. Therefore it is suitable to choose a reasonably small value of $k$, and we do our experiments primarily with $k=2$.

###### Proposition 1 .

## 4 Numerical Experiments with Reasoning-Gym

We demonstrate how our method using the two generation methods described in the last section perform better than Single-Token GRPO in a variety of different environments. Out of the $10$ environments, our methods improve performance in $5$ environments and perform at par with just $5$ chains when compared to standard generation with $10$ chains, hence significantly improving diversity and allowing better scaling. We fix the model to be Qwen2.5 with 1.5 billion parameters to understand how our method compares to Single-Token GRPO in an isolated setting. We conclude the section with analysis of the hidden states using recent methods in interpretability to infer why our methods allow for better exploration. We provide supplementary experiments and ablations in Appendix [B](https://arxiv.org/html/2509.21482v1#A2).

To analyze the performance of model’s reasoning performance when trained with MoT-G, we need a evaluation harness which contains examples which are not in-distribution of the training data so that there is enough room for improvement and one can illustrate improvement in sampling correct reasoning paths. We benchmark Dirichlet standard math datasets including Math-500 and GSM-8K in Appendix [B](https://arxiv.org/html/2509.21482v1#A2) where our methods perform competitively but gains are minimal ($1-2\%$). However we primarily demonstrate the efficacy of our method on a recent released collection of procedurally generated reasoning tasks, titled, Reasoning-Gym (Stojanovski et al., 2025). Our main results are benchmarked on $10$ tasks across $5$ different categories, prime factorization (math), simple integration (math), family relationships (graphs), acre (induction), color cube rotation (cognition), graph color (algorithmic), syllogism (logic), self reference (logic) and number sequence (logic) and mixed (color cube rotation, prime factorization, shortest path, acre, graph color, family relationships).

Experiment Methodology We first fixed the primary experimental setup: We used the instruction finetuned version of Qwen2.5 with 1.5 billion parameters. Appendix [B](https://arxiv.org/html/2509.21482v1#A2) has experiments on other model sizes of the same family. The number of learning steps was fixed to $1000$ steps, the learning rate to $10^{-6}$, the KL regularization parameter to $0.2$, the batch size was $1$, the maximum generation steps and prompt length were both fixed to $1024$ for all experiments. We evaluate on a holdout set every $50$ iterations with $100$ samples. We selected these $10$ tasks because the other tasks we tried on (specifically maze, sudoku, shortest path, dice, course schedule) do not show learning with the fixed parameters (with our method or standard generation). We benchmark 3 methods: (a) standard generation and standard GRPO loss (single-token) (b) Sampling top-$k$ tokens, aggregating them with weights sampled from Dirichlet distribution with the normalized probability as parameter and loss on single token (Dirichlet) (c) Sampling $k$ tokens and aggregating tokens with normalized probability weights and propagating loss to all the sampled tokens (Different Tokens).

### 4.1 Aggregate Analysis on Reasoning Gym Tasks with Qwen2.5-1.5B

**Table 1: Accuracy (%) $\pm$ standard deviation by task, number of chains, and method. Dir=Dirichlet, DT=Different Tokens, Single=single-token. All values are the reported pass@1, for further significance testing we ran each experiment three times, with three different random seeds, and report the mean and standard deviation.**
| Task | 5 Chains | 10 Chains |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | DT | Single | Dir | DT | Single |  |
| Acre | 42.7 $\pm$ 3.4 | 42.0 $\pm$ 2.4 | 38.5 $\pm$ 0.2 | 42.5 $\pm$ 0.5 | 41.5 $\pm$ 9.1 | 38.6 $\pm$ 1.2 |
| Color Cube Rotation | 24.5 $\pm$ 1.0 | 21.3 $\pm$ 1.9 | 21.2 $\pm$ 0.2 | 25.5 $\pm$ 0.8 | 25.2 $\pm$ 1.5 | 21.2 $\pm$ 0.8 |
| Graph Color | 33.1 $\pm$ 1.3 | 21.4 $\pm$ 12.1 | 23.8 $\pm$ 1.1 | 36.1 $\pm$ 1.5 | 23.2 $\pm$ 14.2 | 25.2 $\pm$ 1.3 |
| Mixed | 29.3 $\pm$ 1.6 | 32.8 $\pm$ 1.6 | 24.9 $\pm$ 2.5 | 31.5 $\pm$ 2.3 | 33.9 $\pm$ 1.9 | 25.1 $\pm$ 0.9 |
| Family Relationships | 40.3 $\pm$ 4.3 | 46.8 $\pm$ 3.6 | 39.3 $\pm$ 1.5 | 47.0 $\pm$ 0.6 | 55.7 $\pm$ 8.3 | 43.2 $\pm$ 2.9 |
| Simple Integration | 44.0 $\pm$ 2.6 | 59.3 $\pm$ 1.9 | 25.6 $\pm$ 2.1 | 57.5 $\pm$ 3.4 | 73.8 $\pm$ 3.2 | 43.6 $\pm$ 4.9 |
| Prime Factorization | 78.2 $\pm$ 0.6 | 79.7 $\pm$ 4.4 | 72.3 $\pm$ 1.5 | 83.6 $\pm$ 3.9 | 82.9 $\pm$ 2.9 | 76.3 $\pm$ 0.5 |
| Syllogism | 82.5 $\pm$ 5.0 | 73.3 $\pm$ 12.3 | 82.5 $\pm$ 4.5 | 84.6 $\pm$ 3.1 | 78.2 $\pm$ 13.1 | 84.4 $\pm$ 4.1 |
| Number Sequence | 46.3 $\pm$ 13.2 | 44.9 $\pm$ 4.3 | 54.9 $\pm$ 3.4 | 38.9 $\pm$ 8.6 | 50.7 $\pm$ 7.8 | 61.0 $\pm$ 5.9 |
| Self Reference | 13.6 $\pm$ 3.3 | 7.9 $\pm$ 4.1 | 30.7 $\pm$ 0.8 | 24.5 $\pm$ 7.6 | 16.2 $\pm$ 13.2 | 30.4 $\pm$ 0.8 |

We summarize the results in Table [1](https://arxiv.org/html/2509.21482v1#S4.T1). We present the pass@1 accuracies of the best checkpoint for the different tasks for the different methods with different number of trajectories (groups or chains).

Insight 1. We observe that there are three different types of performance behaviors exhibited by our methods.
Type A (First 6 tasks): Tasks where both the methods perform significantly better with gains of around $5-40\%$ in the pass@1 accuracy. Further on these tasks the performance using mixture methods perform competitively or better with half the number of chains with modest ($2-10\%$) gains. This is a desired consequence since it allows for more exploration with the same number of chains which allows for improvement in policy. Type B (Syllogism) : Tasks where one of the methods perform competitively. Type C (Number Sequence and Self Reference): Tasks where both the mixture underperform the single-token model by roughly $5-20\%$. We speculate that Type C tasks may require more precise, deterministic reasoning steps where maintaining uncertainty through mixture representations introduces harmful noise rather than beneficial exploration.

### 4.2 Hidden state and token-level analysis of the MoT-G methods

Figure: Figure 2: Difference in entropy evolution with respect to single-token over generation steps for different environment. The entropy is computed using the eigenvalue distribution of the Gram matrix of the hidden state vectors of a particular layer. The entropy is representative of the informativeness and diversity of features (Skean et al., 2025). RLVR in LLMs with MoT-G generation generally leads to a higher entropy across layers and generation steps than with single-token generation.
Refer to caption: figs/eigenvalue_plots.png

Analysis of Hidden States.
We use techniques from Skean et al. (2025) to analyze the entropy of eigenvalue distribution of the Gram matrix of the sequence of hidden layer vectors for different generation steps of a particular layer. The entropy has shown to be indicative of the information representation quality, higher entropy indicates more information and more diverse features Skean et al. (2025). Consider the hidden layer $l$ of the LLM. For generation step $t$ let $\mathbf{z}^{(l)}_{t}\in\mathbb{R}^{d_{h}}$ be the hidden state vector. Then for the generation steps $t=1,\dots,n$, denote the sequence of hidden state vectors $\mathbf{Z}_{n}=[\mathbf{z}^{(l)}_{1},\dots,\mathbf{z}^{(l)}_{n}]\in\mathbb{R}^{n\times d_{h}}$. We consider the eigenvalue distribution of the Gram matrix, $\mathbf{K}_{n}=\mathbf{Z}_{n}\mathbf{Z}_{n}^{\mathsf{T}}\in\mathbb{R}^{n\times n}$. Specifically we consider the von Neumann entropy of the eigenvalue distribution, $S(\mathbf{Z}_{n})=-\sum_{i}\frac{\lambda_{i}}{\sum_{j}\lambda_{j}}\log\left(\frac{\lambda_{i}}{\sum_{j}\lambda_{j}}\right)$, where $\lambda_{i}$ is the $i$-th largest eigenvalue of $\mathbf{Z}_{n}$.

We plot the difference in the entropy $S(\mathbf{Z}_{n})$ between our mixture methods (Dirichlet and Different-Tokens) and single-token generation in Figure [2](https://arxiv.org/html/2509.21482v1#S4.F2) for different total generation steps $n=1,\dots,200$, for $4$ different environments for the $l=1,\dots,12$ hiddent layers for Qwen2.5-1.5B model with the the checkpoint of training step $800$ with $10$ trajectories.

Insight 2: The MoT-G Dirichlet and different tokens methods generally have a higher entropy than the single-token generation. Further the gap increases with increasing generation step (tokens) and is present across different layers. The entropy usually plateaus or decreases with generation steps when the latter tokens become less informative. This indicates that MoT-G is able to retain more information in the hidden-layer embedding space than single-token generation, and gives one possible explanation of the improved downstream performance observed when using MoT-G with RLVR.

Figure: Figure 3: Average unique token counts used, by each method, to generate the next token embedding across trajectories. The unique tokens are indicative of the exploration in the token space.
Refer to caption: figs/token_counts.png

Token Diversity Analysis. We further analyze the number of unique tokens sampled for each training step across the trajectories averaged across the generation steps for different methods in Figure [3](https://arxiv.org/html/2509.21482v1#S4.F3).

Insight 4. MoT-G methods explore (consider) more unique tokens on average than single-token generation. The Different tokens method has the highest exploration but also has higher variance. Note that both the token count and hidden state entropy is unusually high for the self-reference task (where MoT-G methods perform worst than standard generation; example in Appendix) showing that the model using MoT-G can potentially get confused in more esoteric tallying intensive tasks.

Improvement in Reasoning with Mixture of Tokens. Figure [4](https://arxiv.org/html/2509.21482v1#S4.F4) shows the sampled tokens and corresponding coefficients of mixing for a specific problem of the dataset family-relationships.

Insight 5. RLVR with MoT-G enables the LLM to learn how to accordingly weigh different options (e.g. son and father in Step 10 of middle plot in Figure [4](https://arxiv.org/html/2509.21482v1#S4.F4)) and also keep a consistent belief of the different possibilities (e.g. Victoria v/s Robert, Margaret v/s Matthew, child and son).

Figure: Figure 4: Mixture token weights for a family relationship reasoning problem across training checkpoints. Each row shows a generation step, with the left column displaying the highest-weighted token and its probability, and the right column showing the second-highest weighted token. The model learns to simultaneously consider multiple plausible reasoning paths, such as weighing both“father” and “son” as potential relationships, rather than committing prematurely to a single discrete choice.
Refer to caption: figs/token_choice_heatmap.png

## 5 Conclusion

This paper investigates mixture-of-token generation as an alternative to discrete sampling in reinforcement learning for reasoning. We present a unified framework for incorporating mixture representations into RLVR and demonstrate the necessary modifications to standard training pipelines. Our evaluation on Reasoning-Gym reveals that MoT-G methods achieve substantial improvements on most tasks (5-35% gains on 7/10 tasks) while requiring half the trajectories to reach comparable performance, though modest regressions occur on some tasks. Analysis suggests these improvements stem from enhanced hidden state informativeness and increased token level exploration during reasoning. Two promising directions for future work emerge: (a) Curriculum Learning: Investigating whether curricula that progressively increase both task difficulty and mixture size K can yield sustained performance gains across increasingly challenging reasoning problems, and (b) Transfer Learning: Understanding how mixture token representations learned on specific reasoning tasks generalize to new domains and if cross-task knowledge transfer can be improved through MoT-G.
The primary limitations are that mixture-of-thought decoding makes it hard to reward the quality of intermediate reasoning and requires approximate, expensive likelihood/KL calculations during training. The added uncertainty can sometimes hurt performance on precise, deterministic tasks.

## References

- Butt et al. [2025]
Natasha Butt, Ariel Kwiatkowski, Ismail Labiad, Julia Kempe, and Yann Ollivier.
Soft tokens, hard truths, 2025.
URL [https://arxiv.org/abs/2509.19170](https://arxiv.org/abs/2509.19170).
- Cobbe et al. [2021]
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman.
Training verifiers to solve math word problems.
*arXiv preprint arXiv:2110.14168*, 2021.
- Deng et al. [2024]
Yuntian Deng, Yejin Choi, and Stuart Shieber.
From explicit cot to implicit cot: Learning to internalize cot step by step, 2024.
URL [https://arxiv.org/abs/2405.14838](https://arxiv.org/abs/2405.14838).
- Geiping et al. [2025]
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein.
Scaling up test-time compute with latent reasoning: A recurrent depth approach, 2025.
URL [https://arxiv.org/abs/2502.05171](https://arxiv.org/abs/2502.05171).
- Guo et al. [2025]
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al.
Deepseek-r1 incentivizes reasoning in llms through reinforcement learning.
*Nature*, 645(8081):633–638, 2025.
- Hao et al. [2024]
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian.
Training large language models to reason in a continuous latent space, 2024.
URL [https://arxiv.org/abs/2412.06769](https://arxiv.org/abs/2412.06769).
- Hou et al. [2025]
Zhenyu Hou, Xin Lv, Rui Lu, Jiajie Zhang, Yujiang Li, Zijun Yao, Juanzi Li, Jie Tang, and Yuxiao Dong.
T1: Advancing language model reasoning through reinforcement learning and inference scaling.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=tnxONP8zTE](https://openreview.net/forum?id=tnxONP8zTE).
- Huang and Yang [2025]
Yichen Huang and Lin F. Yang.
Gemini 2.5 pro capable of winning gold at imo 2025, 2025.
URL [https://arxiv.org/abs/2507.15855](https://arxiv.org/abs/2507.15855).
- Ji et al. [2025]
Yifan Ji, Zhipeng Xu, Zhenghao Liu, Yukun Yan, Shi Yu, Yishan Li, Zhiyuan Liu, Yu Gu, Ge Yu, and Maosong Sun.
Learning more effective representations for dense retrieval through deliberate thinking before search, 2025.
URL [https://arxiv.org/abs/2502.12974](https://arxiv.org/abs/2502.12974).
- Kambhampati et al. [2025]
Subbarao Kambhampati, Kaya Stechly, and Karthik Valmeekam.
(how) do reasoning models reason?
*Annals of the New York Academy of Sciences*, 1547(1):33–40, 2025.
doi: https://doi.org/10.1111/nyas.15339.
URL [https://nyaspubs.onlinelibrary.wiley.com/doi/abs/10.1111/nyas.15339](https://nyaspubs.onlinelibrary.wiley.com/doi/abs/10.1111/nyas.15339).
- Liu et al. [2025a]
Mingjie Liu, Shizhe Diao, Ximing Lu, Jian Hu, Xin Dong, Yejin Choi, Jan Kautz, and Yi Dong.
Prorl: Prolonged reinforcement learning expands reasoning boundaries in large language models, 2025a.
URL [https://arxiv.org/abs/2505.24864](https://arxiv.org/abs/2505.24864).
- Liu et al. [2025b]
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin.
Understanding r1-zero-like training: A critical perspective, 2025b.
URL [https://arxiv.org/abs/2503.20783](https://arxiv.org/abs/2503.20783).
- Shao et al. [2024]
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo.
Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024.
URL [https://arxiv.org/abs/2402.03300](https://arxiv.org/abs/2402.03300).
- Shojaee et al. [2025]
Parshin Shojaee, Iman Mirzadeh, Keivan Alizadeh, Maxwell Horton, Samy Bengio, and Mehrdad Farajtabar.
The illusion of thinking: Understanding the strengths and limitations of reasoning models via the lens of problem complexity, 2025.
URL [https://arxiv.org/abs/2506.06941](https://arxiv.org/abs/2506.06941).
- Skean et al. [2025]
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Nikul Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv.
Layer by layer: Uncovering hidden representations in language models.
In *Forty-second International Conference on Machine Learning*, 2025.
URL [https://openreview.net/forum?id=WGXb7UdvTX](https://openreview.net/forum?id=WGXb7UdvTX).
- Stojanovski et al. [2025]
Zafir Stojanovski, Oliver Stanley, Joe Sharratt, Richard Jones, Abdulhakeem Adefioye, Jean Kaddour, and Andreas Köpf.
Reasoning gym: Reasoning environments for reinforcement learning with verifiable rewards, 2025.
URL [https://arxiv.org/abs/2505.24760](https://arxiv.org/abs/2505.24760).
- Su et al. [2025]
DiJia Su, Hanlin Zhu, Yingchen Xu, Jiantao Jiao, Yuandong Tian, and Qinqing Zheng.
Token assorted: Mixing latent and text tokens for improved language model reasoning, 2025.
URL [https://arxiv.org/abs/2502.03275](https://arxiv.org/abs/2502.03275).
- Wang et al. [2025]
Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, Yuqiong Liu, An Yang, Andrew Zhao, Yang Yue, Shiji Song, Bowen Yu, Gao Huang, and Junyang Lin.
Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for llm reasoning, 2025.
URL [https://arxiv.org/abs/2506.01939](https://arxiv.org/abs/2506.01939).
- Wei et al. [2023]
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou.
Chain-of-thought prompting elicits reasoning in large language models, 2023.
URL [https://arxiv.org/abs/2201.11903](https://arxiv.org/abs/2201.11903).
- Wen et al. [2025]
Xumeng Wen, Zihan Liu, Shun Zheng, Zhijian Xu, Shengyu Ye, Zhirong Wu, Xiao Liang, Yang Wang, Junjie Li, Ziming Miao, Jiang Bian, and Mao Yang.
Reinforcement learning with verifiable rewards implicitly incentivizes correct reasoning in base llms, 2025.
URL [https://arxiv.org/abs/2506.14245](https://arxiv.org/abs/2506.14245).
- Xu et al. [2025]
Yige Xu, Xu Guo, Zhiwei Zeng, and Chunyan Miao.
Softcot: Soft chain-of-thought for efficient reasoning with llms, 2025.
URL [https://arxiv.org/abs/2502.12134](https://arxiv.org/abs/2502.12134).
- Xu et al. [2024]
Zifan Xu, Haozhu Wang, Dmitriy Bespalov, Xian Wu, Peter Stone, and Yanjun Qi.
LaRS: Latent reasoning skills for chain-of-thought reasoning.
In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, *Findings of the Association for Computational Linguistics: EMNLP 2024*, pages 3624–3643, Miami, Florida, USA, November 2024. Association for Computational Linguistics.
doi: 10.18653/v1/2024.findings-emnlp.206.
URL [https://aclanthology.org/2024.findings-emnlp.206/](https://aclanthology.org/2024.findings-emnlp.206/).
- Yao et al. [2023]
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik R Narasimhan.
Tree of thoughts: Deliberate problem solving with large language models.
In *Thirty-seventh Conference on Neural Information Processing Systems*, 2023.
URL [https://openreview.net/forum?id=5Xc1ecxO1h](https://openreview.net/forum?id=5Xc1ecxO1h).
- Yao et al. [2024]
Wenlin Yao, Haitao Mi, and Dong Yu.
Hdflow: Enhancing llm complex problem-solving with hybrid thinking and dynamic workflows, 2024.
URL [https://arxiv.org/abs/2409.17433](https://arxiv.org/abs/2409.17433).
- Yue et al. [2025]
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Yang Yue, Shiji Song, and Gao Huang.
Does reinforcement learning really incentivize reasoning capacity in llms beyond the base model?, 2025.
URL [https://arxiv.org/abs/2504.13837](https://arxiv.org/abs/2504.13837).
- Zhang et al. [2025a]
Junyu Zhang, Runpei Dong, Han Wang, Xuying Ning, Haoran Geng, Peihao Li, Xialin He, Yutong Bai, Jitendra Malik, Saurabh Gupta, and Huan Zhang.
Alphaone: Reasoning models thinking slow and fast at test time, 2025a.
URL [https://arxiv.org/abs/2505.24863](https://arxiv.org/abs/2505.24863).
- Zhang et al. [2025b]
Zhen Zhang, Xuehai He, Weixiang Yan, Ao Shen, Chenyang Zhao, Shuohang Wang, Yelong Shen, and Xin Eric Wang.
Soft thinking: Unlocking the reasoning potential of llms in continuous concept space, 2025b.
URL [https://arxiv.org/abs/2505.15778](https://arxiv.org/abs/2505.15778).
- Zhao et al. [2025]
Andrew Zhao, Yiran Wu, Yang Yue, Tong Wu, Quentin Xu, Yang Yue, Matthieu Lin, Shenzhi Wang, Qingyun Wu, Zilong Zheng, and Gao Huang.
Absolute zero: Reinforced self-play reasoning with zero data, 2025.
URL [https://arxiv.org/abs/2505.03335](https://arxiv.org/abs/2505.03335).
- Zheng et al. [2025]
Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, Jingren Zhou, and Junyang Lin.
Group sequence policy optimization, 2025.
URL [https://arxiv.org/abs/2507.18071](https://arxiv.org/abs/2507.18071).
- Zhu et al. [2025a]
Hanlin Zhu, Shibo Hao, Zhiting Hu, Jiantao Jiao, Stuart Russell, and Yuandong Tian.
Reasoning by superposition: A theoretical perspective on chain of continuous thought, 2025a.
URL [https://arxiv.org/abs/2505.12514](https://arxiv.org/abs/2505.12514).
- Zhu et al. [2025b]
Qinglin Zhu, Runcong Zhao, Hanqi Yan, Yulan He, Yudong Chen, and Lin Gui.
Soft reasoning: Navigating solution spaces in large language models through controlled embedding exploration.
In *Forty-second International Conference on Machine Learning*, 2025b.
URL [https://openreview.net/forum?id=4gWE7CMOlH](https://openreview.net/forum?id=4gWE7CMOlH).
- Zhuang et al. [2025]
Yufan Zhuang, Liyuan Liu, Chandan Singh, Jingbo Shang, and Jianfeng Gao.
Text generation beyond discrete token sampling, 2025.
URL [https://arxiv.org/abs/2505.14827](https://arxiv.org/abs/2505.14827).

## Appendix A Concurrent Submission

Concurrent to preparing our arXiv submission,  [Butt et al., 2025] (posted September 23, 2025 on arXiv) study ‘soft’ token chains of thought optimized with RL, introducing noise-injected soft/fuzzy tokens and deriving a REINFORCE training objective, they report gains on math and reasoning benchmarks with Llama/Qwen up to 8B. Our work instead formulates MoT-G inside RLVR, operating on explicit k-token mixtures (e.g., top-k + Dirichlet weights) and training with GRPO on Reasoning-Gym. We also provide mechanistic analyses of hidden state entropy/exploration and show efficiency gains in trajectories. Given the timing, we provide a brief comparison below in the spirit of research integrity and transparency. We will try our best to update the arXiv submission to numerically compare a version of their approach within our setup.

- 1.
Mixture object.
*Soft Tokens, Hard Truths:* uses continuous “soft” tokens formed as probability-weighted mixtures in embedding space (plus “fuzzy” variants).
*This paper:* proposes a generalized *mixture-of-token* generation mechanism over a discrete subset of $k$ tokens (e.g., top-$k$/nucleus). Their approach can be seen as a special case of our approach with weighted aggregation over the entire dictionary ($k=|\mathcal{D}|$) and added Gaussian noise. We study two completely different MoT-G instances.
- 2.
RL objective and training loop.
*Soft Tokens, Hard Truths:* derives and trains with a REINFORCE/RLOO-style objective on trajectories containing soft/fuzzy tokens. Their loss function is different and based on the embeddings.
*This paper:* instantiates mixture-of-token generation *inside RLVR* using a GRPO/RLVR-style objective tailored to that framework.
- 3.
Decoding at inference.
*Soft Tokens, Hard Truths:* trains with soft/fuzzy tokens, but reports best results with *hard* token decoding at test time.
*This paper:* uses standard hard-token decoding; mixtures are primarily a *training-time* device to enrich exploration within RLVR.
- 4.
Evaluation focus and metrics.
*Soft Tokens, Hard Truths:* evaluates on math-reasoning benchmarks (e.g., GSM8K, MATH, OlympiadBench) with pass@$k$ metrics on open models.
*This paper:* evaluates within the RLVR setting (e.g., Reasoning-Gym) and emphasizes *trajectory efficiency* (accuracy vs. number of trajectories) and analysis of entropy/hidden-state dynamics.
- 5.
Goal and contribution style.
*Soft Tokens, Hard Truths:* emphasizes feasibility and gains of RL with soft/fuzzy tokens plus guidance on when hard decoding helps.
*This paper:* is more expository: it formalizes mixture-of-token generation inside RLVR and provides mechanistic (hidden-state) analyses and efficiency trade-offs.
- 6.
Hyperparameters:
*Soft Tokens, Hard Truths:* full-probability mixtures with temperature/noise schedules in embedding space.
*This paper:* subset size $k$, candidate selection rule (top-$k$ or nucleus), Dirichlet concentration for mixture weights, and GRPO/RLVR hyperparameters.

## Appendix B Supplementary Experiments

Additional Experiment Details: Most of the experiments were done on a single H100 GPU with 80GB of VRAM, however some experiments were done on multiple such GPUs. We attach slurm scripts and code to run our experiments efficiently.

In general we run all our experiments with default parameters of reasoning gym.
Note that for the 3B and 7B models, task we increased the difficulty of color cube rotation (min rotations = 8, max rotations 20), graph coloring (num colors=4,min num vertices=10,max num vertices=20,edge probability = 0.4), family relationship (min family size=10,max family size=20) because the tasks were too easy otherwise.

Table [2](https://arxiv.org/html/2509.21482v1#A2.T2): Qwen2.5-3B. The MoT-G variants consistently outperform single-token decoding on most tasks, with the Different-Tokens method showing the largest gains on Family Relationships ($57.60\pm 3.11$ vs. $41.27\pm 2.10$) and solid lifts on Acre ($58.40\pm 2.09$ vs. $51.80\pm 0.20$). Improvements are modest but positive on Mixed and Graph Color, while Color Cube Rotation is roughly on par across methods. Overall, these results indicate that preserving token uncertainty during the “think” phase benefits relational and inductive reasoning more than low-signal, visually grounded puzzles, aligning with the paper’s argument that MoT-G enhances exploration without sacrificing coherence.

Table [3](https://arxiv.org/html/2509.21482v1#A2.T3): Qwen2.5-7B. Scaling the base model preserves (and often amplifies) MoT-G’s advantages: Different-Tokens again leads on Family Relationships ($79.30\pm 2.12$ vs. $73.27\pm 3.32$) and improves Acre ($67.40\pm 1.83$ vs. $59.87\pm 3.25$), with both MoT-G variants boosting the Mixed suite as well. Absolute accuracies remain low on Graph Color for all methods, suggesting algorithmic bottlenecks rather than decoding alone. These results show that mixture-based thinking complements model scale, translating uncertainty-aware hidden states into higher pass@1 without additional trajectories.

Table [4](https://arxiv.org/html/2509.21482v1#A2.T4): Trajectory count (15 vs. 20 chains). With fewer chains, MoT-G matches or exceeds Single-Token, highlighting sample-efficiency gains. For Acre, Different-Tokens improves the 15-chain maximum ($47.1\pm 3.2$ vs. $42.1\pm 3.7$) and remains ahead at 20 chains ($46.0\pm 3.7$ vs. $38.2\pm 1.6$). Prime Factorization shows strong Dirichlet performance at 15 chains ($86.5\pm 0.5$ vs. $81.7\pm 1.4$) and sustained advantage at 20 chains ($85.3\pm 0.2$ vs. $81.5\pm 3.0$). In Number Sequence, MoT-G is competitive but not uniformly superior, underscoring task dependence. Collectively, the table supports the claim that MoT-G can reach comparable or better accuracy with fewer trajectory by encouraging broader exploration per step.

Table [5](https://arxiv.org/html/2509.21482v1#A2.T5), Table [6](https://arxiv.org/html/2509.21482v1#A2.T6), Table [7](https://arxiv.org/html/2509.21482v1#A2.T7): Temperature sensitivity. When decoding temperature rises, Dirichlet-weighted MoT-G remains notably robust while Different-Tokens and Single-Token degrade sharply. At $T=1.6$ (10 chains) on Family Relationships, Dirichlet attains $36.9$ vs. $7.1$ (Different-Tokens) and $6.2$ (Single-Token); at $T=2.0$ (10 chains) on Acre, Dirichlet still reaches $32.7$ vs. $0.2$ and $8.9$, respectively. Similar patterns hold on Color Cube Rotation. Because Dirichlet perturbs weights over a stable top-$k$ set (Figure [6](https://arxiv.org/html/2509.21482v1#A3.F6)), it preserves structure under higher stochasticity, whereas resampling tokens or committing to single tokens collapses.

**Table 2: Accuracy (%) with Qwen2.5-3B Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | Dirichlet | Different Tokens | Single Token |
| --- | --- | --- | --- |
| Acre | 55.00 ± 2.27 | 58.40 ± 2.09 | 51.80 ± 0.20 |
| Color Cube Rotation | 26.30 ± 5.80 | 21.80 ± 5.37 | 22.40 ± 5.39 |
| Family Relationships | 44.47 ± 1.63 | 57.60 ± 3.11 | 41.27 ± 2.10 |
| Graph Color | 2.53 ± 0.12 | 3.20 ± 0.72 | 1.70 ± 0.42 |
| Mixed | 41.89 ± 3.91 | 40.61 ± 2.67 | 39.44 ± 2.30 |

**Table 3: Accuracy (%) with Qwen2.5-7B Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | Dirichlet | Different Tokens | Single Token |
| --- | --- | --- | --- |
| Acre | 63.33 ± 1.40 | 67.40 ± 1.83 | 59.87 ± 3.25 |
| Color Cube Rotation | 22.60 ± 7.81 | 20.65 ± 8.17 | 23.27 ± 7.09 |
| Family Relationships | 72.40 ± 2.11 | 79.30 ± 2.12 | 73.27 ± 3.32 |
| Graph Color | 2.93 ± 0.81 | 4.00 ± 1.60 | 2.80 ± 0.85 |
| Mixed | 51.67 ± 1.61 | 50.89 ± 2.87 | 45.61 ± 1.51 |

**Table 4: Accuracy (%) $\pm$ standard deviation by task, number of chains, and method. Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | 15 Chains | 20 Chains |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | DT | Single | Dir | DT | Single |  |
| Acre | 44.3 $\pm$ 3.0 | 47.1 $\pm$ 3.2 | 42.1 $\pm$ 3.7 | 42.7 $\pm$ 1.6 | 46.0 $\pm$ 3.7 | 38.2 $\pm$ 1.6 |
| Color Cube Rotation | 27.7 $\pm$ 1.4 | 28.0 $\pm$ 2.6 | 22.3 $\pm$ 2.1 | 13.2 $\pm$ 0.7 | 30.5 $\pm$ 3.7 | 22.3 $\pm$ 2.8 |
| Mixed | 32.5 $\pm$ 0.6 | 33.8 $\pm$ 1.7 | 23.0 $\pm$ 1.4 | 33.6 $\pm$ 0.6 | 30.6 $\pm$ 0.6 | 25.3 $\pm$ 1.2 |
| Number Sequence | 60.7 $\pm$ 3.8 | 50.7 $\pm$ 10.3 | 59.8 $\pm$ 3.5 | 58.3 $\pm$ 6.8 | 60.9 $\pm$ 7.7 | 62.2 $\pm$ 4.9 |
| Prime Factorization | 86.5 $\pm$ 0.5 | 84.1 $\pm$ 5.1 | 81.7 $\pm$ 1.4 | 85.3 $\pm$ 0.2 | 84.4 $\pm$ 4.0 | 81.5 $\pm$ 3.0 |

**Table 5: Accuracy (%) with temperature T=1.2. Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | 5 Chains | 10 Chains |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | DT | Single | Dir | DT | Single |  |
| Acre | 42.4 | 26.8 | 34.2 | 40.9 | 29.4 | 33.4 |
| Color Cube Rotation | 24.6 | 12.4 | 14.6 | 23.0 | 14.4 | 14.9 |
| Family Relationships | 38.6 | 33.6 | 19.4 | 42.4 | 48.3 | 28.6 |
| Graph Color | 23.8 | 5.8 | 9.6 | 28.4 | 5.0 | 15.7 |
| Number Sequence | 46.0 | 8.2 | 9.4 | 50.2 | 17.8 | 8.6 |

**Table 6: Accuracy (%) with temperature T=1.6. Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | 5 Chains | 10 Chains |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | DT | Single | Dir | DT | Single |  |
| Acre | 36.8 | 3.8 | 24.8 | 36.8 | 4.3 | 22.9 |
| Color Cube Rotation | 11.2 | 1.0 | 1.8 | 20.8 | 1.0 | 9.1 |
| Family Relationships | 17.8 | 2.0 | 0.6 | 36.9 | 7.1 | 6.2 |
| Graph Color | 13.2 | 0.2 | 0.2 | 15.1 | 0.5 | 0.8 |
| Number Sequence | 5.2 | 2.0 | 0.4 | 10.0 | 1.0 | 0.1 |

**Table 7: Accuracy (%) with temperature T=2.0. Dir=Dirichlet, DT=Different Tokens, Single=Single-Token.**
| Task | 5 Chains | 10 Chains |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | DT | Single | Dir | DT | Single |  |
| Acre | 17.0 | 0.0 | 7.8 | 32.7 | 0.2 | 8.9 |
| Color Cube Rotation | 6.2 | 0.0 | 3.2 | 17.6 | 0.1 | 3.7 |
| Family Relationships | 2.4 | 0.0 | 2.4 | 28.1 | 0.2 | 2.0 |
| Graph Color | 0.8 | 0.0 | 0.0 | 5.2 | 0.0 | 0.0 |
| Number Sequence | 1.0 | 0.0 | 1.0 | 2.1 | 0.0 | 0.5 |

**Table 8: Accuracy (%) by task, $k$, number of chains for the different tokens method.**
| Task | k=4 | k=8 |
| --- | --- | --- |
| Acre | 47.4 | 48.6 |
| Color Cube Rotation | 22.2 | 26.0 |
| Family Relationships | 41.8 | 43.6 |
| Graph Color | 22.2 | 14.2 |
| Prime Factorization | 86.0 | 68.8 |

**Table 9: Accuracy (%) by task and method for Math benchmarks. k=k, C=chains, Dir=Dirichlet, Single=Single-Token.**
| Task | k=2 | k=4 | k=8 |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dir | Single | Dir | - | Dir | - |  |
| GSM8K | 75.6 | 71.7 | 71.5 | - | 71.7 | - |
| MATH500 | 50.5 | 46.2 | 48.2 | - | 44.4 | - |

Figure: Figure 5: Progress of training for different algorithms under different temperatures (0.9 versus 1.2). Our methods can sometimes slow in the beginning (learning the mixture of tokens generation) but catch up quite quickly. The gains from the different-tokens technique is lesser with higher temperature possibly because of more irrelevant. Dirichlet based techniques are more robust to the temperature since top-$k$ tokens are the same, hence only the weights get more random with increasing temperature. We note that the different tokens based technique show improvement on some tasks with increasing $k$ on some tasks and Dirichlet based technique (not plotted) detoriates with increasing $k$, illustrating Proposition [1](https://arxiv.org/html/2509.21482v1#Thmproposition1) that the choice of $k$ is an important task-dependent hyperparameter.
Refer to caption: figs/temporal_accuracy_grid_combined.png

In Figure [5](https://arxiv.org/html/2509.21482v1#A2.F5) we plot the evaluation accuracy across evaluation steps for different $k$ for the different tokens approach and for different temperatures (on different columns). Increasing $k$ to $4$ helps improve the overall accuracy especially in acre and family relationships while doesn’t improve in graph color and color-cube rotation. This can be intuitively understood by the result of Proposition [1](https://arxiv.org/html/2509.21482v1#Thmproposition1) which claims that the choice of $k$ is task-sensitive for the different token method. For higher temperature it is evident that Dirichlet based sampling is the most robust to increasing temperatures (since the tokens used for sampling is not different but only the weights). The single-token and different tokens methods are more sensitive.

Table [8](https://arxiv.org/html/2509.21482v1#A2.T8) k-sensitivity (Different-Tokens).
Increasing the number of mixed tokens from $k=4$ to $k=8$ helps on relational/inductive tasks,Acre (47.4 to 48.6), Color Cube Rotation (22.2 to 26.0), and Family Relationships (41.8 to 43.6),but hurts more algorithmic ones,Graph Color (22.2 to 14.2) and Prime Factorization (86.0 to 68.8). This pattern underscores the predicted trade-off: larger $k$ broadens exploration but can dilute signal when precise token commitments matter, so a smaller $k$ is often safer for brittle, algorithmic settings.

Table [9](https://arxiv.org/html/2509.21482v1#A2.T9) Math benchmarks (Dirichlet vs. Single-Token).
On GSM8K and MATH500, Dirichlet mixing with a small set ($k=2$) yields consistent gains over Single-Token decoding (75.6 vs. 71.7 on GSM8K; 50.5 vs. 46.2 on MATH500). As $k$ increases, accuracy declines (e.g., GSM8K 71.5 at $k=4$; MATH500 48.2 at $k=4$ and 44.4 at $k=8$), suggesting that larger mixtures can over-soften the signal required for deterministic math. Dashes indicate settings not evaluated for the baseline. Together with Table 8, this supports the guidance that the optimal $k$ is task-dependent, with smaller $k$ favored for math-heavy tasks.

### B.1 Benchmarking Extensions of Mixture of Tokens based Generation on GSM-8K

This section presents the numerical experiments which illustrate efficacy of the different variations of the generalized mixture of tokens method. We do this for different small models of the Qwen family on GSM-8K [Cobbe et al., 2021] to a) see if different sampling and aggregation techniques affect the downstream performance b) which techniques form a good basis for selecting as generation mechanisms for our GRPO experiments.

**Table 10: Performance of extensions of Training-Free Soft Thinking on GSM-8K. The sampling and aggregation methods have an impact on the downstream performance. For example, when doing element wise maximum (a method which is more sensitive to embeddings) gives 0 accuracy when used with top-k sampling but gives $44\%$ with nucleus sampling.**
|  |  | Qwen3 1.7B | Qwen3 4B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Sampling Method | Aggregation Method | k=2 | k=4 | k=8 | k=2 | k=4 | k=8 |
| Baselines |  |  |  |  |  |  |  |
| Standard (Greedy) | — | $75\pm 3$ | $78\pm 3$ |  |  |  |  |
| Soft Thinking (Original) | — | $77\pm 3$ | $79\pm 3$ |  |  |  |  |
| Other mixture of token methods |  |  |  |  |  |  |  |
| Top-k | Prob. Weighted Avg. | $76\pm 3$ | $75\pm 3$ | $76\pm 3$ | $80\pm 3$ | $79\pm 3$ | $79\pm 3$ |
| Top-k | Dirichlet Weighted Avg. | $80\pm 3$ | $80\pm 3$ | $80\pm 3$ | $79\pm 3$ | $75\pm 3$ | $80\pm 3$ |
| Top-k | Element Wise Maximum | $3\pm 1$ | 0 | 0 | 0 | 0 | 0 |
| Min-p | Prob. Weighted Avg. | $78\pm 2$ | $81\pm 2$ |  |  |  |  |
| Min-p | Dirichlet Weighted Avg. | $81\pm 2$ | $83\pm 2$ |  |  |  |  |
| Min-p | Element Wise Maximum | 0 | $3$ |  |  |  |  |
| Nucleus Sampling | Prob. Weighted Avg. | $78\pm 2$ | $78\pm 2$ |  |  |  |  |
| Nucleus Sampling | Dirichlet Weighted Avg. | $78\pm 2$ | $77\pm 2$ |  |  |  |  |
| Nucleus Sampling | Element Wise Maximum | $18\pm 2$ | $44\pm 2$ |  |  |  |  |
| k-sampling | Prob. Weighted Avg. | $74\pm 3$ | $78\pm 2$ | $78\pm 3$ | $78\pm 3$ | $79\pm 3$ | $79\pm 3$ |
| k-sampling | Dirichlet Weighted Avg. | $79\pm 3$ | $79\pm 3$ | $85\pm 3$ | $79\pm 3$ | $76\pm 3$ | $?$ |
| k-sampling | Element Wise Maximum | $3\pm 1$ | 0 | 0 | $5$ | $5$ | 0 |

## Appendix C Proofs

### C.1 Proof of Proposition 1

###### Proof.

Let $p_{t}=h(\mathbf{X}_{t-1})$ denote the probability distribution over dictionary $\mathcal{D}$ at step $t$. For each rollout $g\in[G]$, we sample $k$ tokens $S_{g}=\{z_{g,1},\ldots,z_{g,k}\}$ without replacement.

Part 1: For the expected number of unique tokens $L_{t}$ across $G$ trajectories, let $q_{j}(k)$ denote the probability that token $j$ is included when sampling $k$ tokens. For sampling without replacement:

$$ $q_{j}(k)=1-\prod_{i=0}^{k-1}\frac{1-p_{j}-\sum_{m\in S_{i}}p_{m}}{1-\sum_{m\in S_{i}}p_{m}}\geq q_{j}(k-1)$ $$

where $S_{i}$ denotes the first $i$ sampled tokens. Since $q_{j}(k)$ increases with $k$, the expected number of unique tokens can be computed using indicator variables. Let $I_{j}=1$ if token $j$ appears in at least one rollout, 0 otherwise. Then $L_{t}=\sum_{j=1}^{|\mathcal{D}|}I_{j}$ and by linearity of expectation:

$$ $\mathbb{E}[L_{t}]=\sum_{j=1}^{|\mathcal{D}|}\mathbb{P}(I_{j}=1)=\sum_{j=1}^{|\mathcal{D}|}\left[1-(1-q_{j}(k))^{G}\right]$ $$

is monotonically increasing in $k$.

Part 2:
Let $Z\sim p$ and write $Y:=e_{Z}$. For any set $S\subseteq D$ with $p(S):=\sum_{j\in S}p_{j}>0$,
define the conditional mean

$$ $m(S)\;:=\;\mathbb{E}[Y\mid Z\in S]\;=\;\sum_{j\in S}\frac{p_{j}}{p(S)}\,e_{j}.$ $$

Under the standard sequential PPS-without-replacement construction, we can couple the random
sampled sets so that $S_{k}\subset S_{k+1}$ almost surely. In this scheme, one rollout’s mixture
embedding with sample size $k$ is exactly

$$ $X_{k}\;=\;m(S_{k}).$ $$

Let $\mu:=\mathbb{E}Y=\sum_{j\in D}p_{j}e_{j}$. For any random vector $U$ with an independent copy $U^{\prime}$,

$$ $\mathbb{E}\|U-U^{\prime}\|^{2}\;=\;2\,\mathrm{tr}\big(\mathrm{Cov}(U)\big)\;=\;2\big(\mathbb{E}\|U\|^{2}-\|\mathbb{E}U\|^{2}\big).$ $$

Hence it suffices to prove that

$$ $\mathrm{tr}\big(\mathrm{Cov}(X_{k+1})\big)\;\leq\;\mathrm{tr}\big(\mathrm{Cov}(X_{k})\big).$ $$

Fix a realization of the nested sets $S_{k}\subset S_{k+1}$. Consider the two-atom
$\sigma$-algebras $\mathcal{G}_{k}:=\sigma(\{Z\in S_{k}\})$ and
$\mathcal{G}_{k+1}:=\sigma(\{Z\in S_{k+1}\})$.
By construction, conditioning on $\{Z\in S_{k+1}\}$ is strictly *coarser* information than
conditioning on $\{Z\in S_{k}\}$ (because $S_{k}\subset S_{k+1}$).
Since $m(S)=\mathbb{E}[Y\mid\sigma(\{Z\in S\})]$, the variance of these $L^{2}$-projections is
monotone in the information:

$$ $\mathrm{Var}\!\left(\,\mathbb{E}[Y\mid\mathcal{G}_{k+1}]\,\right)\;\leq\;\mathrm{Var}\!\left(\,\mathbb{E}[Y\mid\mathcal{G}_{k}]\,\right).$ $$

(Equation $(3)$ is the standard monotonicity of $\mathrm{Var}(\mathbb{E}[Y\mid\mathcal{H}])$ with respect
to the partial order of $\sigma$-algebras: coarser $\mathcal{H}$ yields smaller variance.)

Taking expectation over the randomness of the sampled sets yields

$$ $\mathbb{E}\,\mathrm{tr}\big(\mathrm{Cov}(X_{k+1})\big)=\mathbb{E}\,\mathrm{tr}\Big(\mathrm{Cov}\big(\mathbb{E}[Y\mid\mathcal{G}_{k+1}]\big)\Big)\;\leq\;\mathbb{E}\,\mathrm{tr}\Big(\mathrm{Cov}\big(\mathbb{E}[Y\mid\mathcal{G}_{k}]\big)\Big)=\mathbb{E}\,\mathrm{tr}\big(\mathrm{Cov}(X_{k})\big),$ $$

which is $(2)$. Combining $(2)$ with the identity $(1)$ for $U=X_{k}$ and $U^{\prime}=X_{k}^{\prime}$ gives the claim.
∎

###### Proof.

### C.2 Background on Dirichlet distribution

Figure: (1) Center ( $\mathbf{p}=[\tfrac{1}{3},\tfrac{1}{3},\tfrac{1}{3}]$, $c=1$). Symmetric spread around the simplex barycenter; high diversity across tokens.
Refer to caption: figs/dir_center_c1.png

The Dirichlet distribution $\mathrm{Dir}(\boldsymbol{\alpha})$ is defined on the $(K\!-\!1)$-simplex
$\Delta^{K-1}=\{\mathbf{x}\in\mathbb{R}^{K}:\,x_{i}\geq 0,\ \sum_{i}x_{i}=1\}$ with density,

$$ $f(\mathbf{x}\mid\boldsymbol{\alpha})=\frac{1}{B(\boldsymbol{\alpha})}\prod_{i=1}^{K}x_{i}^{\alpha_{i}-1},\qquad B(\boldsymbol{\alpha})=\frac{\prod_{i=1}^{K}\Gamma(\alpha_{i})}{\Gamma\!\big(\sum_{i=1}^{K}\alpha_{i}\big)}.$ $$

where $\Gamma$ is the gamma function. Parameterizing as $\boldsymbol{\alpha}=c\,\mathbf{p}$ with concentration $c>0$ and base probabilities
$\mathbf{p}\in\Delta^{K-1}$ gives
$\mathbb{E}[X_{i}]=p_{i}$, $\mathrm{Var}(X_{i})=\tfrac{p_{i}(1-p_{i})}{c+1}$, and
$\mathrm{Cov}(X_{i},X_{j})=-\tfrac{p_{i}p_{j}}{c+1}$ for $i\neq j$.
Thus $c$ controls spread (smaller $c$ $\Rightarrow$ broader draws) while $\mathbf{p}$ sets the center.
Figure [6](https://arxiv.org/html/2509.21482v1#A3.F6) shows $K=3$ cases at $c=1$: a symmetric center
($\mathbf{p}=[\tfrac{1}{3},\tfrac{1}{3},\tfrac{1}{3}]$), an asymmetric middle
($[0.6,0.3,0.1]$), and a corner ($[0.95,0.04,0.01]$). As the center approaches a vertex,
samples concentrate near that corner, illustrating the intuition that “points in the corner stay in the corner.”

## Appendix D Related Work on Other Test-Time Scaling Methods in LLMs

Other Test-Time Scaling Methods in LLMs:
Yao et al. [2024] combines fast and slow thinking modes and switches between them based on instance difficulty, so the model spends extra steps only when useful.
Zhang et al. [2025a] provides a single view that covers many prior “more steps help” methods, and lets the system control how much slow vs fast reasoning to use in a fine-grained way.
Ji et al. [2025] introduce DEBATER, a retriever that learns better document embeddings by running a simple step-by-step thinking process, giving stronger retrieval when more test-time thinking is allowed.
[Yao et al., 2023] Tree of Thoughts scales by exploring branches of partial thoughts, tuning breadth and depth to trade off cost and accuracy.
Geiping et al. [2025] scales compute by repeating a recurrent block to any depth at inference, growing inner iterations rather than only producing more tokens.
Xu et al. [2024] LaRS learns a latent “reasoning skill” space and a policy that picks the skill level for a question, then chooses in-context examples that match it, which adapts test-time effort to the task.
Wang et al. [2025] improves RLVR by updating only on forking tokens, the points where the path can split, which focuses extra samples where they matter most.
Hou et al. [2025] pretrains with trial-and-error plus self-check CoT, then scales RL by over-sampling to raise sample diversity, which supports stronger test-time sampling without training collapse.
Zhao et al. [2025] uses self-play to generate diverse trajectories at inference, which can be ramped up for harder inputs.
One can similarly extend our techniques to other recent algorithms including GSPO and DR-GRPO.

## Appendix E Example Problems

For MoT-G methods we only show the token with the highest probability.

### E.1 ACRE

### E.2 Color Cube Rotation

### E.3 Family Relationships

### E.4 Graph Color

### E.5 Prime Factorization

### E.6 Self Reference

### E.7 Simple Integration

### E.8 Syllogism

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: 1. Does the method harm the coherence of the final outputs? I noticed that there are some garbled example outputs. I also think it should be a drawback of "soft-thinking" methods.

2. How does k affect memory and speed? Any instability at larger k?

3. Clarify end criteria and whether the model learns to stop mixing vs. using a fixed heuristic. Any ablations on stopping criteria?
soundness: 2
strengths: 1. Strong motivation to preserve distributional information in RLVR

2. Two practical instantiations with intuitive exploration properties. Dirichlet weighting is a straightforward way to align exploration with model confidence while controlling variance.

3. Improvements on a diverse suite of reasoning tasks, including relational/inductive/cognitive tasks in Reasoning-Gym.
summary: The paper proposes Mixture-of-Token Generation (MoT-G) for RL with verifiable rewards, replacing hard token commitments during chain-of-thought with continuous “mixture embeddings” formed by aggregating k sampled tokens per step. It unifies sampling (top-k, nucleus, min-p, k-sampling) and aggregation (probability-weighted, Dirichlet-weighted, etc.) and adapts GRPO to train in this space with single- vs. multi-token loss variants and approximate KL/log-prob handling. Two concrete instantiations, including Dirichlet (top-k + Dirichlet weights) and Different-Tokens (sample without replacement + normalized weights), are evaluated. On Reasoning-Gym with Qwen2.5-1.5B, MoT-G improves pass@1 by 5–35% on 7/10 tasks and often matches 10-chain accuracy using 5 chains; gains on GSM8K and MATH500 are modest. Analyses show higher hidden-state entropy and greater token-level diversity, suggesting better exploration; Dirichlet is more robust to higher decoding temperatures.
weaknesses: 1. Baseline coverage is too narrow. Comparisons are primarily against single-token GRPO. Missing comparisons to strong RLVR variants (e.g., DAPO/GSPO/DR-GRPO updates), and to established test-time scaling methods or exploration enhancements. No end-to-end comparison against training-free soft-thinking at inference within the same RLVR training regime

2. Evaluation scope is limited, and the improvements are not consistent and significant. Main gains are on Reasoning-Gym; math benchmarks show small gains and two tasks (Number Sequence, Self-Reference) regress notably. Lack of harder, or widely adopted reasoning benchmarks (e.g., AIME24/25, OlympiadBench, BIG-bench subsets) under the RLVR setting limits external validity.

3. The paper highlights “half the trajectories,” but mixture steps increase per-step compute (k soft-embeddings, aggregation, KL approximations). I think the method should be more efficient. However, there is no wall-clock training time, FLOPs, or GPU-hours analysis to demonstrate these points.

4. The overall writing quality of the work can be improved. For example, in lines 128 and 129, the GRPO objective and notation contain inconsistencies, i.e., arg min vs. maximize expected return. The overall presentation is not good. It is a bit difficult for readers to directly understand and follow the work.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 2
questions: See Weakness
soundness: 2
strengths: - This paper is easy to follow and has a reasonable motivation. Authors identifies the core limitations of current RLVR methods that rely on discrete token sampling which discardes distributional information and limiting the search space, and then proposes a logically clear improvement framework (named MoT-G).
- Clear experimental results. The experimental section of the paper is clearly designed, and through comprehensive testing on Reasoning-Gym, it strongly demonstrates that the MoT-G method achieved significant performance improvements of 5-35% on 7 out of 10 tasks.
summary: This paper points out a limitation of current Reinforcement Learning with Verifiable Rewards (RLVR) methods, which rely on discrete token sampling and thus discard valuable distributional information, constraining the reasoning search space. Thus, authors propose a "Mixture-of-Token Generation" (MoT-G) framework where a weighted mixture of multiple token embeddings is generated at each reasoning step, allowing the model to maintain uncertainty and explore alternative paths. Experiments show that MoT-G achieves significant performance gains of 5-35% on 7 out of 10 tasks in the Reasoning-Gym benchmark. The paper's analysis suggests these benefits may stem from MoT-G's ability to maintain higher hidden-state entropy and promote token-space exploration during reasoning.
weaknesses: - The specific computational overhead is unclear, even though the authors used fewer trajectories. The computational time and cost of MoT-G seem to be much higher than the standard method because it must aggregate $k$ tokens and maintain gradients for all of them during training, rather than just for a single token.
- The performance of MoT-G is not universally applicable. The method performs poorly on certain tasks. Moreover, why does the MoT-G method's performance degrade when increasing trajectories from 15 to 20, as shown in Table 4, contradicting the standard RL principle that more samples should improve the policy estimate and final accuracy?

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: - How sensitive are the results to the number of sampled tokens (*k*) and the choice of aggregation (Dirichlet vs. uniform)?  
- Can the method generalize to larger models or other reasoning domains?
soundness: 3
strengths: - Clear formulation of mixture-based generation under the RLVR framework.  
- Includes interpretability analyses (entropy, diversity, visualization).  
- Shows potential for more efficient exploration with fewer rollouts.  
- Presentation and figures are clean and understandable.
summary: The paper introduces **Mixture-of-Token Generation (MoT-G)** for reinforcement learning with verifiable rewards (RLVR, e.g., GRPO). Instead of sampling a single token, the model aggregates multiple token embeddings into a continuous mixture to enhance exploration. Two variants—**Dirichlet weighting** and **Different Tokens sampling**—are evaluated on the Reasoning-Gym benchmark using Qwen2.5-1.5B. The approach achieves modest gains (5–35%) on half of the tasks and minimal improvements (≈1–2%) on math datasets. Analyses of hidden-state entropy and token diversity suggest that MoT-G improves exploration, though causal links remain unclear.
weaknesses: - **Low novelty:** Essentially applies prior soft-token or mixture-of-token ideas to RLVR without new objectives or theory.  
- **Limited experiments:** Only one base model (Qwen2.5-1.5B) and fixed GRPO setup; no scaling or robustness tests.  
- **Marginal gains:** Minimal improvement on core reasoning benchmarks (e.g., GSM8K, Math-500).  
- **Weak ablations:** Lacks analysis isolating the contribution of mixture components or sampling parameters.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: 1.	How sensitive is MoT-G’s performance to the choice of mixture size k and Dirichlet concentration parameters across different reasoning domains?
2.	Can the proposed method be combined with other exploration-enhancing techniques such as temperature scaling or entropy regularization?
3.	Does the higher hidden-state entropy correlate causally with reasoning correctness, or could it reflect noise?
4.	How would MoT-G perform when scaled to larger models (e.g., ≥14B) or used with curriculum or transfer learning setups?
soundness: 2
strengths: The proposed method itself is coherent and easy to follow.  The authors conduct experiments on several benchmarks to validate the effectiveness of the proposed method.
summary: The paper proposes Mixture-of-Token Generation (MoT-G), a framework that extends Reinforcement Learning with Verifiable Rewards (RLVR) to operate in a continuous token space rather than discrete sampling. By aggregating multiple candidate tokens into mixture embeddings, MoT-G preserves uncertainty and promotes exploration during reasoning. Experiments on Reasoning-Gym show 5–35% accuracy gains on most tasks using the Qwen2.5-1.5B model, achieving comparable results with half the trajectories.
weaknesses: While MoT-G presents an effective way to extends RLVR to operate in a continuous token space, several limitations temper its overall contribution.

**(1) Model scale limitation:** all experiments are conducted primarily on the Qwen2.5-1.5B model, with brief extensions to 3B and 7B variants in the appendix. Such relatively small models may not fully capture the complex reasoning dynamics that large-scale RLVR systems (e.g., ≥30B parameters) exhibit. The scalability and stability of MoT-G in larger models remain unclear. 

**(2) Limited dataset coverage:** the evaluation focuses almost entirely on Reasoning-Gym, a procedurally generated benchmark. While diverse, it may not represent the full variety of real-world reasoning tasks such as mathematical proofs, code reasoning, or commonsense multi-hop inference. The absence of widely used public datasets (e.g., GSM8K, MATH, or ARC-Challenge) weakens the generalizability of results.

**(3) Computational complexity:** mixture-based generation introduces higher memory and compute costs per step, yet the paper provides no quantitative efficiency analysis beyond “half the trajectories.”

**(4) Performance regressions:** on deterministic tasks (e.g., self-reference, number sequence), MoT-G underperforms, indicating that the added uncertainty may harm precise reasoning rather than help exploration.


Decision: Reject