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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Mathematical Reasoning | CollegeMATH | -- | 337 | |
| Mathematical Reasoning | MATH 500 | Accuracy (pass@1)84.7 | 134 | |
| Mathematical Reasoning | Minerva Math | Pass@1 Accuracy29 | 51 | |
| Mathematical Reasoning | AIME 24 | Pass@1 Score16.7 | 29 | |
| Mathematical Reasoning | AIME 25 | Pass@1 Accuracy10 | 29 | |
| Text Generation | HH-RLHF and IMDB | Helpful Assistant Score6.14 | 16 | |
| Text Generation | HH-RLHF and IMDB (test) | Total Training Time (h)6.7 | 6 | |
| Mathematical Reasoning | Mathematical Reasoning Benchmarks | Forward Pass Time (s)0.82 | 4 |