SofT-GRPO: Superando o Aprendizado por Reforço de LLMs com Tokens Discretos via Otimização de Políticas de Pensamento Suave com Reparametrização de Gumbel

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

November 9, 2025

Autores: Zhi Zheng, Wee Sun Lee

cs.AI

Resumo

O paradigma de pensamento suave (soft-thinking) para raciocínio em Modelos de Linguagem de Grande Porte (LLMs) pode superar o raciocínio convencional de cadeia de pensamento (Chain-of-Thought - CoT) com tokens discretos em alguns cenários, destacando seu valor para pesquisa e aplicação. No entanto, embora o padrão de raciocínio CoT com tokens discretos possa ser reforçado através de algoritmos de otimização de políticas, como a otimização de políticas relativas de grupo (group relative policy optimization - GRPO), estender o padrão de pensamento suave com Aprendizado por Reforço (RL) permanece um desafio. Essa dificuldade decorre das complexidades de injetar estocasticidade nos tokens de pensamento suave e atualizar as políticas de pensamento suave de acordo. Como resultado, tentativas anteriores de combinar pensamento suave com GRPO normalmente apresentam desempenho inferior às suas contrapartes GRPO com tokens discretos. Para liberar todo o potencial do pensamento suave, este artigo apresenta um novo algoritmo de otimização de políticas, o SofT-GRPO, para reforçar LLMs sob o padrão de raciocínio de pensamento suave. O SofT-GRPO injeta ruído Gumbel nos *logits*, emprega a técnica Gumbel-Softmax para evitar que tokens de pensamento suave fiquem fora do espaço de incorporação pré-treinado e aproveita o truque de reparametrização no gradiente de política. Realizamos experimentos com LLMs base variando de 1,5B a 7B de parâmetros, e os resultados demonstram que o SofT-GRPO permite que LLMs de pensamento suave superem ligeiramente o GRPO com tokens discretos em Pass@1 (+0,13% na precisão média), enquanto exibe um aumento substancial em Pass@32 (+2,19% na precisão média). Códigos e pesos estão disponíveis em https://github.com/zz1358m/SofT-GRPO-master.

English

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