Temperature-grouped reinforcement learning speeds exploration in llms
TGRL: Temperature-Grouped Reinforcement Learning for Efficient Exploration in LLMs
Machine LearningComputation and Language
Summary
Getting large language models (LLMs) to explore different possibilities efficiently is a big challenge in training them with rewards. The authors propose a method called Temperature-Grouped Reinforcement Learning (TGRL) that groups model outputs based on temperature settings to measure how much trying different options helps improve results. Instead of just increasing the number of tries, TGRL uses differences in reward outcomes to guide learning more effectively. This approach lets models reach high accuracy faster and improves performance across various tasks like math problems, coding contests, and interactive environments.
What this means in practice
- •For ai engineers: Improve large language model training speed and accuracy by using temperature-based reward contrasts to better guide exploration without extra rollout cost.
- •For software developers: Enhance performance of AI-powered coding assistants by integrating TGRL to generate higher quality code solutions faster.
Authors
Zihan Lin, Xiaohan Wang, Jie Cao, Jiajun Chai, Wei Lin, Guojun Yin, Ran He
Abstract
Efficient exploration often remains a central bottleneck in reinforcement learning with verifiable rewards (RLVR). Although temperature control and test-time scaling strategies can increase rollout diversity of large language models (LLMs), they either expand the sample budget at rollout time or leave the benefit of exploration unquantified. To this end, we propose Temperature-Grouped Reinforcement Learning (TGRL), which turns temperature-induced diversity into an explicit training signal. For each prompt, TGRL partitions its rollout group into low- and high-temperature subsets, estimates exploration gain through their reward contrast, and allocates this group-level signal as token-level credit using Jensen--Shannon (JS) divergence between the corresponding temperature-scaled next-token distributions induced by the same logits. Notably, TGRL reaches equivalent accuracy up to 36% faster than strong RLVR baselines without expanding the rollout budget. Across 11 benchmarks from diverse domains, TGRL broadly improves over strong RLVR baselines: it improves the six-benchmark math average by 1.6% at 32B, raises CodeForces rating by 196.7 points and LiveCodeBench Pass@16 by 4.4%, and improves ALFWorld/WebShop success rates by 6.3%/4.9%. Comprehensive ablations and wall-clock analysis confirm the efficacy of all proposed components. Code is available at https://github.com/1229095296/TGRL/tree/main.