TLDR: SofT-GRPO is a novel policy optimization algorithm that enhances Large Language Models (LLMs) by reinforcing their ‘soft-thinking’ reasoning pattern. Unlike traditional discrete-token reasoning, soft-thinking uses continuous representations, but has been challenging to optimize with Reinforcement Learning (RL). SofT-GRPO addresses this by injecting Gumbel noise and employing Gumbel-Softmax and reparameterization tricks. Experiments show it surpasses discrete-token GRPO in accuracy, especially at higher sample rates (Pass@16, Pass@32), improves token efficiency, and generalizes well to out-of-domain tasks, demonstrating a significant step forward for soft-thinking LLMs.
Large Language Models (LLMs) have shown remarkable abilities in various tasks, especially in reasoning. Traditionally, LLMs reason using a method called ‘discrete-token Chain-of-Thought’ (CoT), where they generate a sequence of distinct words or tokens to arrive at an answer. However, a newer approach, known as ‘soft-thinking,’ has emerged, offering a more nuanced way for LLMs to process information.
Soft-thinking allows LLMs to reason using continuous representations rather than discrete tokens. Imagine it as thinking in shades of gray instead of just black and white. Instead of picking a single word, the model considers a weighted blend of many possible words’ meanings, represented as a continuous vector. This can help LLMs express more abstract concepts and potentially explore a wider range of reasoning paths.
While soft-thinking shows great promise, combining it with Reinforcement Learning (RL) – a powerful technique used to train models by rewarding desired behaviors – has been a significant challenge. Existing RL methods, like Group Relative Policy Optimization (GRPO), work well for discrete-token reasoning but have struggled to effectively enhance soft-thinking. The main difficulties lie in introducing controlled randomness into these continuous ‘soft-thinking tokens’ and updating the model’s decision-making process accordingly.
This is where a new algorithm, SofT-GRPO, comes into play. Developed by Zhi Zheng and Wee Sun Lee, SofT-GRPO is designed specifically to overcome these challenges and unlock the full potential of soft-thinking in LLMs. The core idea behind SofT-GRPO is to inject a special kind of randomness, called Gumbel noise, into the model’s output probabilities. It then uses a technique called Gumbel-Softmax to ensure that these continuous ‘soft-thinking tokens’ remain within a valid range that the LLM understands. Finally, it employs a ‘reparameterization trick’ to efficiently update the model’s soft-thinking policies based on the rewards it receives.
The process involves generating groups of soft-thinking reasoning paths, each with a bit of Gumbel noise, and then optimizing the LLM to favor paths that lead to better answers. This approach allows for effective exploration of diverse reasoning strategies while maintaining the stability needed for training.
Experiments were conducted across various LLMs, ranging from 1.5 billion to 7 billion parameters, on a variety of reasoning tasks, including numerical, scientific, and code-related problems. The results were compelling: SofT-GRPO enabled soft-thinking LLMs to slightly outperform discrete-token GRPO in immediate accuracy (Pass@1) and showed substantial improvements in scenarios where multiple attempts are allowed (Pass@16 and Pass@32). This means that with SofT-GRPO, soft-thinking LLMs are more likely to find the correct answer when given a few chances.
Beyond accuracy, SofT-GRPO also demonstrated benefits in token efficiency, meaning the models could arrive at solutions using fewer ‘thinking’ steps, especially noticeable in smaller LLMs. It also showed good generalization to tasks outside its primary training domain, such as scientific and code reasoning. Furthermore, combining SofT-GRPO with a ‘majority voting’ technique, where the most common answer from multiple runs is chosen, further boosted its performance, making it an even more robust problem-solver.
Also Read:
- Enhancing AI’s Math Skills: A Self-Evolving Approach to Multimodal Reasoning
- Boosting Reasoning in Smaller AI Models: A New Approach to Label-Free Learning
In essence, SofT-GRPO provides a robust framework for enhancing the soft-thinking capabilities of LLMs, pushing them beyond the limitations of traditional discrete-token reasoning. This research highlights a promising direction for developing more capable and efficient AI models. For more technical details, you can refer to the original research paper here.
