Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Token-Efficient RL for LLM Reasoning

About

We propose reinforcement learning (RL) strategies tailored for reasoning in large language models (LLMs) under strict memory and compute limits, with a particular focus on compatibility with LoRA fine-tuning. Building on early policy gradient methods with baseline subtraction, we design critic-free methods that operate on a small, informative subset of output tokens to reduce memory usage and stabilize training. We introduce S-GRPO, a stochastic variant of Group Relative Policy Optimization, and T-SPMO, a token-level prefix matching approach for fine-grained credit assignment. Applied to Qwen2-1.5B, our methods raise accuracy on the SVAMP benchmark from 46% to over 70% and show strong performance on multi-digit multiplication. Surprisingly, full-token GRPO under LoRA fails to improve over the base model, suggesting that selective token-level optimization may act as an implicit regularizer in low-parameter training regimes.

Alan Lee, Harry Tong• 2025

Related benchmarks

TaskDatasetResultRank
Mathematical ReasoningCollegeMATH--
337
Mathematical ReasoningMATH 500
Accuracy (pass@1)84.7
134
Mathematical ReasoningMinerva Math
Pass@1 Accuracy29
51
Mathematical ReasoningAIME 24
Pass@1 Score16.7
29
Mathematical ReasoningAIME 25
Pass@1 Accuracy10
29
Text GenerationHH-RLHF and IMDB
Helpful Assistant Score6.14
16
Text GenerationHH-RLHF and IMDB (test)
Total Training Time (h)6.7
6
Mathematical ReasoningMathematical Reasoning Benchmarks
Forward Pass Time (s)0.82
4
Showing 8 of 8 rows

Other info

Follow for update