SofT-GRPO: Surpassing Discrete-Token LLM Reinforcement Learning via
  Gumbel-Reparameterized Soft-Thinking Policy Optimization
Authors: name: Zhi Zheng, name: Wee Sun Lee
Organization: National University of Singapore

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