PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhi Zheng, Yu Gu, Wei Liu, Yee Whye Teh, Wee Sun Lee

arXiv:2511.06411v2cs.AIcs.LG

TL;DR

Soft-thinking may represent abstract ideas that discrete tokens cannot, but applying RLVR to it is challenging because its tokens lack stochasticity and policy-gradient signals. SofT-GRPO uses Gumbel-based sampling and reparameterized policy updates; its soft-thinking results surpass discrete-token GRPO on averaged Pass@1 and show substantial Pass@32 gains.

  • Problem

    Soft-thinking can represent abstract concepts beyond discrete tokens, but applying RLVR is challenging because deterministic soft tokens lack the stochasticity and gradients needed for exploration and policy updates.

    Source: §1 · pp. 1–2

  • Method

    SofT-GRPO samples soft-thinking paths with Gumbel noise and Gumbel-Softmax, then uses Gumbel reparameterization for policy updates.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    SofT-GRPO soft-thinking surpasses discrete-token GRPO on averaged Pass@1 and yields substantial gains on Pass@32.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper concludes that its results demonstrate advantages of soft-thinking for LLM reasoning, with reported gains in numerical, scientific, and code reasoning.

    Source: §6 · p. 8

    Under the soft-thinking pattern, inputs may fall outside the pretrained discrete-token embedding space.

    Source: §5.2 · p. 8

Abstract

from arXiv · show

The soft-thinking paradigm for Large Language Model (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

1. Introduction

Discrete-token reasoning can struggle to represent abstract concepts, while extending RLVR to soft-thinking is challenging because soft tokens lack stochasticity and gradients. SofT-GRPO addresses these challenges with Gumbel-based rollouts and policy updates.

1. Introduction

  • Across three 1.5B–7B LLMs and five numerical benchmarks, SofT-GRPO soft-thinking surpasses discrete-token GRPO on averaged Pass@1 and yields substantial Pass@32 gains.The contribution summary also reports outperformance across eight in- and out-of-domain benchmarks, especially at Pass@16 or Pass@32.Source: Figure 1 · p. 1, §1 · p. 2
  • Soft-thinking replaces each discrete CoT token with a weighted sum of token embeddings, a representation proposed to express abstract ideas that discrete tokens may not capture.Earlier soft-thinking approaches can outperform discrete-token CoT without fine-tuning on a range of tasks.Source: §1 · pp. 1–2
  • Extending RLVR to soft-thinking is difficult because deterministic soft tokens lack the randomness and gradients needed for exploration and policy updates.Prior Gaussian-noise approaches struggled to match discrete-token RLVR, largely because they had difficulty attributing RLVR advantages to logits.Source: §1 · p. 2
  • SofT-GRPO injects Gumbel noise into output token probabilities and uses Gumbel-Softmax to obtain diverse soft-thinking paths within the pretrained embedding space.This addresses the rollout challenge of sampling diverse paths while avoiding invalid soft-token inputs.Source: §1 · p. 2
  • SofT-GRPO uses Gumbel reparameterization for gradient estimation to attribute policy improvements to the LLM’s output probability distributions.The paper presents this as precise attribution in policy optimization.Source: §1 · p. 2

2. Preliminaries

The preliminaries contrast discrete-token CoT with soft-thinking and review how prior work applied sampling and RLVR to each. They describe soft-thinking’s representational promise alongside unresolved challenges in reinforcing it.

2.1. Discrete-Token CoT Reasoning

  • Discrete-token CoT generates reasoning tokens R from question Q before predicting answer tokens A.The question, reasoning, and answer are token sequences from the vocabulary.Source: §2.1 · p. 2
  • SFT fine-tunes LLMs on high-quality CoT labels, which are difficult to obtain for complex datasets, and raises concerns about out-of-domain generalization.The cited discussion contrasts this reliance on labeled reasoning with RLVR methods that do not require CoT labels.Source: §2.1 · pp. 2–3
  • Discrete-token GRPO samples multiple CoTs per question and optimizes toward higher-reward examples, but updates only one token per reasoning step.Its rewards depend on answer quality, and RLVR can improve reasoning without CoT labels.Source: §2.1 · p. 3

2.2. Soft-Thinking Reasoning Paradigm

  • Soft-thinking replaces discrete reasoning tokens with real-valued vectors formed by weighting token embeddings with output probabilities.These vectors lie in the convex hull of the embeddings, preserving validity while transmitting more information.Source: §2.2 · p. 3
  • Prior soft-thinking sampling methods use Gumbel noise with Gumbel-Softmax or Dirichlet resampling to add randomness to the reasoning process.The Gumbel-Softmax temperature and Dirichlet scaling parameter control their respective sampling procedures.Source: §2.2 · p. 3
  • Without fine-tuning, sampled soft-thinking can outperform conventional discrete-token CoT across numerical, code, and scientific reasoning tasks.The cited result uses Gumbel noise and Gumbel-Softmax sampling.Source: §2.2 · p. 3

2.3. Related Work on RLVR Soft-Thinking

  • A prior soft-thinking GRPO approach adds Gaussian noise to soft tokens, restores rollout values, and calculates path log probabilities using Gaussian reparameterization.It leaves the answer-generation update unchanged.Source: §2.3 · pp. 3–4
  • Prior soft-thinking GRPO improves Pass@32 over discrete-token GRPO but severely degrades average accuracy; existing attempts also fail to retain soft-thinking’s no-finetune advantage over discrete-token CoT.The cited work attributes the weaknesses to adding noise to soft tokens rather than logits and not using advanced sampling methods.Source: §2.3 · p. 4

3. SofT-GRPO: Reinforcing Soft-Thinking Policy with Gumbel Reparameterization

SofT-GRPO combines stochastic soft-thinking rollouts with a Gumbel-reparameterized policy update to assign reward credits to token probabilities.

3. SofT-GRPO: Reinforcing Soft-Thinking Policy with Gumbel Reparameterization

  • SofT-GRPO uses Gumbel-Softmax rollouts and a Gumbel-reparameterized loss to assign reward credits to token probabilities.The method samples soft-thinking paths with controllable randomness, then updates the policy using the reparameterization trick.Source: §3 · p. 4

3.1. Group Rollout with Gumbel Noise

  • For each query, SofT-GRPO stores a group of soft-thinking CoTs in a replay buffer and uses Gumbel-Softmax resampling to diversify deterministic rollouts.The method uses Gumbel-Softmax because stochasticity should preserve soft tokens near the pre-trained discrete-token embedding space, mitigating out-of-vocabulary risk.Source: §3.1 · pp. 4–5

3.2. Gumbel Reparameterization for Loss Function

  • SofT-GRPO treats rollout Gumbel noise as RLVR actions and uses its density to estimate soft-token probabilities for policy updates.For off-policy clipping, it reconstructs soft-thinking paths and noise from the rollout; the loss encourages noise associated with higher advantages and updates probabilities of several tokens per step.Source: §3.2 · p. 5

4. Experiments

Across three LLMs and numerical, scientific, and code benchmarks, SofT-GRPO improves soft-thinking performance over discrete-token GRPO, particularly at higher sample rates, while majority voting can further strengthen results.

4.1. Implementation Details

  • The experiments use three LLMs from 1.5B to 7B parameters, with five in-domain numerical benchmarks and three out-of-domain benchmarks.Training uses 40,315 DeepScaler queries; inference settings include top-p 0.95 and top-k 5 for SofT-GRPO.Source: §4 · p. 5, §4.1 · pp. 5–6
  • Baselines include discrete-token GRPO, unfine-tuned models, and soft-thinking with Gumbel-Softmax at τg = 0.5.The reported metrics are Mean@32, Pass@16, and Pass@32.Source: §4.1 · p. 6

4.2. Main Result

  • SofT-GRPO shows clear average improvements in Pass@16 (+1.80%) and Pass@32 (+2.19%) compared with GRPO, and its superiority also holds across different temperatures.Applying soft-thinking on top of discrete-token-trained GRPO gives inconsistent improvement.Source: §4.2 · p. 6

4.3. Comparison in Out-of-Domain Datasets

  • On GPQA Diamond, HumanEval, and MBPP, SofT-GRPO models trained on numerical questions retain advantages from Pass@1 through Pass@32.Table 2 reports average accuracies on these out-of-domain datasets.Source: §4.3 · p. 7, Table 2 · p. 7

4.4. Boosting SofT-GRPO with Majority Voting

  • With majority voting, SofT-GRPO outperforms No-Finetune and discrete-token GRPO on Major@16 and Major@32 across the reported benchmarks.The comparison covers AIME2024, AIME2025, AMC23, and GSM8K.Source: §4.4 · p. 7, Table 3 · p. 7

4.5. Comparison to Butt et al. (2025)

  • SofT-GRPO clearly improves on the reported Butt et al. (2025) soft-thinking RLVR results under the same training dataset and LLM.The compared method also requires transferring d-dimensional vectors between rollout and RLVR workers, making implementation more difficult.Source: §4.5 · p. 7, Table 4 · p. 7

4.6. Comparison in the Token Efficiency

  • SofT-GRPO reduces thinking length versus No-Finetune and does not greatly increase tokens versus discrete-token GRPO.The LLaMA-3.2-3B-Instruct model shows a clear reduction in thinking length.Source: §4.6 · p. 7
  • Figure 3 tracks training or validation curves for Gumbel-noise and top-p/τg settings, while Table 5 compares Dirichlet- and Gaussian-noise variants.The Figure 3 caption identifies top-p = 0.95 and τg = 0.1 as the settings discussed in panels (b,c).Source: Figure 3 · p. 8, Table 5 · p. 8

5. Discussion

The ablations compare noise and hyperparameter choices, and the post-training visualization shows high-probability paths retaining interpretability alongside distributions that preserve multiple reasoning branches.

5.1. Ablation on the Added Noise

  • Replacing Gumbel noise with Dirichlet noise brings no improvement, while Gaussian noise causes poorer initial performance.The study compares training reward curves, validation curves, and final variant performance.Source: §5.1 · p. 8

5.2. Ablation on Hyper-Parameters

  • Changing top-p to 1.0 or τg to 0.25 collapses training, which the authors attribute to increased KL divergence and soft-token inputs outside the pretrained embedding space.The stated settings for SofT-GRPO are top-p = 0.95 and τg = 0.1.Source: §5.2 · p. 8

5.3. Visualization: Soft-Thinking after SofT-GRPO

  • After SofT-GRPO fine-tuning, high-probability reasoning paths remain interpretable, and distributions with mass on a token and its antonym indicate multiple branches.The example antonyms are “Yeah” and “No.”Source: §5.3 · p. 8

6. Conclusion

SofT-GRPO reinforces soft-thinking and is reported to outperform discrete-token GRPO across Pass@1, Pass@16, and Pass@32. The discussion distinguishes this accuracy goal from latent reasoning’s efficiency focus and limits its RLVR comparison to GRPO.

6. Conclusion

  • SofT-GRPO reinforces soft-thinking with Gumbel-based stochasticity and reparameterized policy updates, outperforming discrete-token GRPO on Pass@1, especially Pass@16 and Pass@32.The conclusion also reports that majority voting can boost it into a better solver.Source: §6 · p. 8

A.1. RLVR on Discrete-Token CoT Reasoning

  • Discrete-token RLVR includes GRPO, Dr. GRPO, DAPO, and Lite PPO, with other work targeting response efficiency or entropy control.These approaches optimize discrete-token CoT reasoning.Source: §A.1 · p. 13
  • The paper considers GRPO as its sole RLVR pipeline, leaving a comprehensive investigation of other modifications such as entropy control and DAPO for future work.The authors describe those modifications as potentially orthogonal to SofT-GRPO’s sampling and credit-assignment contributions.Source: §A.1 · p. 13

A.2. Latent Reasoning

  • Token-wise latent reasoning can improve token efficiency but shows a performance drop, whereas SofT-GRPO targets soft-thinking accuracy against discrete-token GRPO on general reasoning tasks.Latent reasoning also includes auxiliary strategies that can boost the original LLM while sacrificing running efficiency.Source: §A.2 · p. 13

A.3. The Working Mechanism of Soft-thinking

Soft-thinking represents each reasoning step with a probability-weighted embedding mixture, recursively approximating a sum over discrete reasoning paths. The section also describes limits of that representation, the prompts used, and a mismatch in Gaussian-noise treatments of sparse soft tokens.

A.3. The Working Mechanism of Soft-thinking

  • Soft-thinking recursively replaces sums over discrete paths with expected-embedding soft tokens, a first-order linear approximation that retains multiple possible reasoning trajectories.The paper describes these mixtures as concept tokens that can represent abstract ideas; it also says LLMs can process them without additional fine-tuning because they lie in the embedding vectors’ convex hull.Source: §A.3 · p. 14
  • Soft-token expressiveness is limited because some points in the embedding convex hull are not equally interpretable or accessible, and out-of-distribution tokens may be linked to reasoning degradation or repetition.The authors connect this limitation to the approximation failing for particularly smooth or spread-out distributions.Source: §A.3 · p. 14

B. Prompt of Language Reasoning and Latent Generation

C.1. Motivation: Mismatch in Butt et al. (2025)

  • Gaussian noise in embedding space does not define a true likelihood on token-probability mixtures because the embedding map is non-injective.Different mixtures can map to the same embedding, losing information about the original token distribution.Source: §C.1 · p. 16
  • For sparse top-k mixtures, Gaussian perturbations almost surely leave the convex hull of any k token embeddings, so the perturbed embedding may not represent a valid top-k soft token.The text gives k=10 to 30 with k << d as a general setting and says most perturbed embeddings are not realizable by such distributions.Source: §C.1 · p. 16

C.2. Proof for Theorem 3.1

The appendix invokes the Gumbel-max trick and explains how SofT-GRPO tracks Gumbel perturbations for rollout and policy correction. It treats the perturbations as the stochastic actions and the resulting soft tokens as deterministic functions of them.

C.2. Proof for Theorem 3.1

  • The proof section develops the Gumbel-max trick from nonnegative weights and independent samples from a standard Gumbel distribution.The supplied theorem statement identifies these inputs but does not state its conclusion.Source: §C.2 · p. 17

C.3. Derivation of Eq. (8), Eq. (11) and Eq. (12)

  • Given logits or probabilities and Gumbel noise, the rollout produces a deterministic soft token, so SofT-GRPO attributes stochasticity and probability mass to the noise.The derivation treats the sampled perturbation as the RL action and the soft-token generation as an environment transition.Source: §C.3 · p. 18
  • For off-policy correction, SofT-GRPO reuses realized rollout perturbations and evaluates their density under the current policy rather than defining a sampling density directly over soft tokens.The paper notes that a perturbation maps to one deterministic soft token, while a soft token can map to multiple perturbations.Source: §C.3 · p. 18

D. Detailed Parameters for Training & Testing

The supplementary experiments report training cost, token-efficiency comparisons, and temperature robustness. SofT-GRPO takes longer to run than discrete-token GRPO, while the reported efficiency and accuracy comparisons are qualitative in these passages.

D. Detailed Parameters for Training & Testing

E. Supplementary of Experiments

  • Table 9 reports token counts across all queries and, separately, across correct queries for the baselines and SofT-GRPO.The table caption specifies the meaning of the two token-count measures but provides no values in the supplied passage.Source: Table 9 · p. 20

E.1. Supplementary of Token Efficiency

  • SofT-GRPO improves token efficiency over No-Finetune on both measures and, for LLaMA-3.2-3B-Instruct, maintains or improves efficiency during training unlike GRPO.The text says it does not cause a severe token improvement relative to discrete-token GRPO.Source: §E.1 · p. 20, Figure 4 · p. 20

E.2. Experiments under Different Temperatures

  • Across five numerical benchmarks on DeepSeek-R1-Distill-Qwen-1.5B, SofT-GRPO reports stronger average Pass@K results than GRPO and No-Finetune across temperatures 0.6 to 1.4, from Pass@1 to Pass@32.The tested temperatures are 0.6, 0.8, 1.0, 1.2, and 1.4.Source: §E.2 · p. 21, Figure 5 · p. 21

E.3. Supplementary Ablation on Hyper-Parameters

The supplementary analyses connect hyperparameter instability to possible soft-token inputs outside the pretrained embedding space, then test Pass@K up to 1024 and assess significance for selected Pass@16 and Pass@32 comparisons.

E.3. Supplementary Ablation on Hyper-Parameters

  • Higher top-p or τg can cause training collapse, attributed to incomprehensible soft-tokens; variants at top-p 1.0 or τg 0.25 also show higher KL divergence.The divergence may indicate inputs outside the pretrained embedding space, while PPO-policy KL below 10^-5 makes precision issues less likely in these variants.Source: §E.3 · p. 22

E.4. Pass@K up to K=1024

  • SofT-GRPO consistently outperforms GRPO and No-Finetune from Pass@16 through Pass@1024; GRPO beats No-Finetune only at small K.The AIME2024 curve fluctuates with only 30 instances, while all algorithms eventually reach complete correctness on AMC23.Source: Figure 7 · p. 23, §E.4 · p. 23

E.5. P-values for Significance on Pass@K

  • SofT-GRPO significantly outperforms GRPO on AIME2024 Pass@16 (0.0067) and AMC23 Pass@32 (0.022), with p-values of 0.13 on the other two pairings.The tests use 64 groups for Pass@16 and 32 groups for Pass@32, with 0.05 as the significance threshold.Source: §E.5 · p. 24, Figure 8 · p. 24, Figure 9 · p. 24
  • AMC23 yields p-values of 0.081 for Pass@16 and 0.054 for Pass@32.Both are described as quite small, but they exceed the stated 0.05 threshold.Source: §E.5 · p. 24

F. Discussion & Analysis on Results

The discussion contrasts entropy retention under SofT-GRPO with entropy collapse under discrete-token GRPO and relates soft-token updates to broader Pass@K strategies. The following sections specify the baselines, datasets, inference framework, and licenses used in the study.

F. Discussion & Analysis on Results

  • SofT-GRPO updates the main and low-probability components of soft-tokens on both positive and negative instances, unlike sampled-token-only updates.The discussion places this alongside Pass@K methods that redistribute group rewards or avoid over-concentration on top-1 selection.Source: F. Discussion & Analysis on Results · p. 25
  • Token-probability entropy converges near zero with discrete-token GRPO but remains stable and can improve with SofT-GRPO.Figure 10 presents the entropy curve for DeepSeek-R1-Distill-Qwen-7B during training.Source: F. Discussion & Analysis on Results · p. 25, Figure 10 · p. 25

G. Baselines & Datasets & Licenses

G.1. Baselines

  • Baselines include No-Finetune base LLMs, discrete-token GRPO, and Butt et al. (2025) Soft Token; the authors report that Soft Token requires substantial communication between rollout and policy-update workers.They therefore report that method’s results from Butt et al. (2025) for comparison in Table 4.Source: §G.1 · p. 26

G.2. Datasets

  • The benchmarks comprise five in-domain numerical-reasoning datasets, one out-of-domain scientific-reasoning dataset, and two out-of-domain code-reasoning datasets.The listed datasets are AIME2024, AIME2025, AMC23, MATH-500, GSM8K, GPQA-Diamond, HumanEval, and MBPP.Source: §G.2 · p. 26

G.3. Inference Framework

  • The inference implementation uses SGLang with continuous batching, RadixAttention, and compressed finite state machines; its optional overlap scheduler stores soft probabilities and indices in VRAM future maps.A custom kernel resolves slot addresses to fetch those states just in time.Source: §G.3 · p. 26

G.4. Licenses

  • The paper lists licenses for the base LLMs, datasets, and frameworks in Table 10.The table is captioned as a summary of licenses.Source: §G.4 · p. 26, Table 10 · p. 26

H. Visualization

The visualization examples describe preserved multi-path soft-thinking and interpretability, and report a blended token distribution that the authors consider a likely product of SofT-GRPO.

H. Visualization

Related papers

Loading 2511.06411v2…