SofT-GRPO: Surpassing Discrete-Token LLM Reinforcement Learning via Gumbel-Reparameterized Soft-Thinking Policy Optimization (2511.06411v1)

Abstract: The soft-thinking paradigm for LLM reasoning can outperform the conventional discrete-token Chain-of-Thought (CoT) reasoning in some scenarios, underscoring its research and application value. However, while the discrete-token CoT reasoning pattern can be reinforced through policy optimization algorithms such as group relative policy optimization (GRPO), extending the soft-thinking pattern with Reinforcement Learning (RL) remains challenging. This difficulty stems from the complexities of injecting stochasticity into soft-thinking tokens and updating soft-thinking policies accordingly. As a result, previous attempts to combine soft-thinking with GRPO typically underperform their discrete-token GRPO counterparts. To fully unlock the potential of soft-thinking, this paper presents a novel policy optimization algorithm, SofT-GRPO, to reinforce LLMs under the soft-thinking reasoning pattern. SofT-GRPO injects the Gumbel noise into logits, employs the Gumbel-Softmax technique to avoid soft-thinking tokens outside the pre-trained embedding space, and leverages the reparameterization trick in policy gradient. We conduct experiments across base LLMs ranging from 1.5B to 7B parameters, and results demonstrate that SofT-GRPO enables soft-thinking LLMs to slightly outperform discrete-token GRPO on Pass@1 (+0.13% on average accuracy), while exhibiting a substantial uplift on Pass@32 (+2.19% on average accuracy). Codes and weights are available on https://github.com/zz1358m/SofT-GRPO-master

• The paper introduces SofT-GRPO, which employs Gumbel-Softmax sampling and a tailored reparameterization trick to optimize soft-thinking policy gradients in continuous token space.
• It demonstrates significant performance gains, including up to +2.19% improvement in Pass@32 metrics and enhanced token efficiency on multiple reasoning benchmarks.
• The approach addresses key challenges in soft-token embedding and training stability, paving the way for scalable reinforcement learning applications in LLMs.

SofT-GRPO: Reinforcing Soft-Thinking Policy Optimization in LLMs

Motivation and Background

Recent advances in reasoning with LLMs have emphasized the importance of fine-tuning via reinforcement learning for improved performance on complex tasks, particularly in numerical and symbolic domains. While discrete-token Chain-of-Thought (CoT) reasoning has historically been the default paradigm, emerging work on soft-thinking—where reasoning tokens are continuous weighted sums of token embeddings—has demonstrated potential to represent abstract concepts more flexibly. However, soft-thinking has historically struggled to match the performance of RLVR-optimized discrete-token reasoning, notably due to challenges in introducing stochasticity for exploration and designing losses that effectively attribute reward improvements in the continuous probability space.

SofT-GRPO Algorithm

SofT-GRPO extends Group Relative Policy Optimization (GRPO) to the soft-thinking paradigm by introducing two substantial innovations: (1) the use of Gumbel-Softmax sampling at each reasoning step, and (2) a tailored reparameterization trick for policy gradients in the soft-token space. This approach ensures that sampled soft-thinking tokens remain within the convex hull of pre-trained token embeddings, avoiding out-of-space inputs that degrade model performance.

Rollout Mechanism

• For each query, multiple parallel trajectories are sampled by injecting independent Gumbel noise into the output logits.
• Gumbel-Softmax is applied for sampling, controlling stochasticity and ensuring valid projected embeddings.
• The soft token at step $t$ is computed as a weighted sum:

$s_t = \sum_{i=1}^{|T|} y_i \, e_i$

where $$y_i = \frac{\exp((\log p_i + \epsilon_i)/T_g)}{\sum_{j}\exp((\log p_j + \epsilon_j)/T_g)}$$ for each token $i$, $p_{i}$ is the output probability, $\epsilon_i$ is Gumbel noise, and $T_g$ is the Gumbel temperature.

Policy Optimization

The loss leverages the Gumbel reparameterization trick to estimate the trajectory probability in the soft domain:

• For each soft trajectory, the log probability is computed as:

$\log p(s_t) = -\epsilon_i - \exp(-\epsilon_i)$

• During policy update, reconstructed paths are compared against reward-assigned trajectories, and gradients are estimated in terms of logits rather than raw embeddings, enabling fidelity in credit assignment.
• The optimization incorporates KL regularization with respect to a reference policy, clipping, and top-p/top-k filtering for stability.

Implementation Details

• SGLang is used for the rollout agent. Policy updates are implemented via verl-0.4.x, compatible with distributed training.
• Hyperparameters: learning rate 1e-6, batch size 64, top-p = 0.95, top-k = 5, LLM temperature $T=0.6$, Gumbel temperature $T_g=0.1$. Training for a 1.5B-parameter model takes approx. 45 hours on 8×NVIDIA H200 GPUs.
• Evaluation scripts utilize Math Verify for answer assessment.

Empirical Results

Multifaceted benchmarks (AIME2024, AIME2025, AMC23, MATH-500, GSM8K, GPQA, HumanEval, MBPP) spanning in-domain and out-of-domain reasoning were used for evaluation. SofT-GRPO consistently outperforms discrete-token GRPO and prior soft-token methods (e.g., Butt et al., 2025) in Pass@32 metrics, with the following strong claims supported by data:

• Average Pass@1 Improvement: SofT-GRPO surpasses discrete-token GRPO by +0.13% on average.
• Substantial Pass@32 Gain: +2.19% on average accuracy over discrete GRPO baselines.
• Enhanced Token Efficiency: SofT-GRPO achieves shorter thinking chains, especially in LLaMA-3.2-3B-Instruct, with no penalty in token consumption compared to discrete-token approaches.
• Robustness to Hyperparameters: Ablation studies reveal training collapse when using higher top-p or Gumbel temperatures, substantiating the importance of sharp sampling and distributional alignment inherent to the proposed method.

Furthermore, majority voting amplifies gains at higher Pass@k values, indicating that ensemble soft-thinking solvers achieve superior answer coverage.

Theoretical and Practical Implications

SofT-GRPO effectively resolves the theoretical mismatch entailed by prior methods that injected Gaussian noise into embedding space. The application of Gumbel noise ensures that sampled tokens always correspond to realizable distributions over the token simplex—crucial for RLVR stability and fidelity. The Gumbel-max trick guarantees that the stochastic process respects the underlying categorical semantics, a property not afforded by naive Gaussian perturbations.

Practically, SofT-GRPO enables LLMs to reason in continuous concept space with RL fine-tuning, unlocking improved sample efficiency and accuracy for tasks where abstraction and latent search trees are pivotal. Its hardware efficiency, compatibility with standard RL frameworks, and scalable rollout strategy render it suitable for LLM deployment in production environments where both accuracy and compute utilization are essential.

Limitations and Future Directions

While gains on Pass@1 are modest, the substantial improvement at higher sample rates (Pass@16/Pass@32) highlight the value of ensemble soft-thinking—but also suggest that further work is needed on single-shot reasoning. The method is reliant on the underlying token embedding structure of pre-trained models; further work may be required to adapt SofT-GRPO to radically different architectures (e.g., vision-LLMs). Incorporating hybrid approaches that combine entropy control, curriculum learning, or compression techniques with SofT-GRPO presents promising future avenues.

Conclusion

SofT-GRPO establishes a principled framework for reinforcement learning in soft-thinking LLM reasoning, substantially increasing answer coverage and efficiency compared to discrete-token RLVR. Through robust stochastic sampling and precise gradient attribution in continuous space, it unlocks new capabilities for LLMs in complex reasoning tasks. The theoretical rigor of its approach and empirical breadth of evaluation suggests significant utility for future research and real-world large-scale deployments of LLMs employing soft reasoning chains.