Learning When to Think: RL for Adaptive Reasoning Computation in LLMs

Team & delivery

Joint class project for Advanced AI (Spring 2025), University of Utah, School of Computing. Contributors: Ivan Andhika, Hao Ren, and Ammon Gleason.

Tech Stack

Algorithms

• GRPO
• Adaptive Length Penalty (ALP)
• DeGRPO-Style Control Token Weighting
• LoRA Fine-Tuning

Motivation

We want the model to learn not only how to reason, but when to extend reasoning, when to self-correct via refine, and when to terminate—so marginal compute is spent where it helps most, in line with exploration–exploitation trade-offs in RL.

Problem

Fixed test-time reasoning budgets are inefficient: easy instances are over-processed while hard instances still fail for lack of exploration. Deciding how much computation to allocate per problem is a sequential decision problem—balancing further reasoning against stopping or self-correcting.

Approach

Reasoning is cast as an MDP with actions continue, refine, and terminate. Training uses Group Relative Policy Optimization (GRPO) with ALP reward shaping (token penalty scaled by per-prompt group solve rate as a difficulty proxy) and weighted loss that up-weights control-token gradients versus response tokens (DeGRPO-style). The base model is Qwen2.5-Math-7B-Instruct with LoRA on attention and MLP projections; evaluation follows a MATH-500 protocol with competition-style answer extraction (\boxed{} and ####).

Methodology

• MDP formulation: state includes prompt and reasoning trace; action space is {continue, refine, terminate}; episode ends on terminate or max-step budget.
• Reward shaping (ALP): combines exact-match correctness with adaptive token penalty conditioned on per-prompt group solve rate.
• Training objective: group-normalized advantages with weighted control-token gradients to prevent action-quality signal washout.
• Implementation setup: LoRA over attention and MLP projections with configurable rank/alpha; ablations include removing refine or restoring vanilla GRPO weighting.

Evaluation Plan

• H1: ALP plus weighted control optimization improves the accuracy-efficiency Pareto frontier against fixed-compute and vanilla baselines.
• H2: learned policy allocates more steps/tokens to harder prompts and fewer to easier prompts.
• H3: action distribution shifts with difficulty, with relatively higher continue/refine on hard items and more terminate on easy items.
• Benchmarks and metrics: MATH-500 with exact-answer extraction, measured by accuracy, tokens per problem, cost per correct, average steps, and action-count stratification.

Outcome

The paper reports accuracy and efficiency metrics (tokens, cost per correct answer), plus reasoning steps and action counts—including breakdowns intended to verify adaptive compute allocation across difficulty. The objective is a better accuracy–efficiency frontier than fixed-budget or vanilla GRPO baselines, with hypotheses on harder prompts receiving more steps/tokens and richer use of continue/refine on difficult items.

Current Status

• Documented: MDP formulation, ALP reward, DeGRPO-style training objective, MATH-500 setup, baselines, and evaluation metrics in the course paper.
• Implementation scope: LoRA adaptation, allow_refine ablation, GRPO with ALP (e.g., β=0.05, L_max=2048) and control/response weights (e.g., w_ctrl=2.0, w_resp=1.0) as described in the write-up.
• Empirical frontier and difficulty-stratified action analysis are framed as the hypotheses (H1–H3) the experiments are designed to test.

Key Contributions

• Three-action adaptive reasoning policy (continue / refine / terminate) connecting test-time compute to an explicit RL control problem, including refine as in-episode self-correction.
• ALP reward shaping that penalizes length more when group solve rate is high (easier prompts) and less when it is low (harder prompts).
• DeGRPO-style separation of control-token vs. response-token optimization to avoid washing out the action-learning signal under GRPO.

Limitations

• Current validation is focused on math reasoning benchmarks and may not transfer uniformly to other domains.
• Reward shaping introduces hyperparameter sensitivity (beta, length cap, control/response weights).
• Behavior gains depend on reliable action extraction and robust answer normalization.

Future Work

• Run larger ablation grids to isolate sensitivity to ALP and control-token weighting parameters.
• Evaluate transfer to broader reasoning tasks beyond MATH-500.
• Investigate richer control spaces and hierarchical action policies for multi-phase reasoning.

References

• Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling LLM test-time compute optimally can be more effective than scaling model parameters. ICLR, 2025.
• Zhihong Shao et al. DeepSeekMath: Pushing the limits of mathematical reasoning in open language models. arXiv:2402.03300, 2024.
• Gongfan Fang, Xinyin Ma, and Xinchao Wang. Thinkless: LLM learns when to think. NeurIPS, 2025.
• Jason Wei et al. Chain-of-thought prompting elicits reasoning in large language models. NeurIPS, 2022.
• Xuezhi Wang et al. Self-consistency improves chain of thought reasoning in language models. ICLR, 2023.
• Hunter Lightman et al. Let's verify step by step. ICLR, 2024.
• Aviral Kumar et al. Training language models to self-correct via reinforcement learning. ICLR, 2025.
• Jonathan Uesato et al. Solving math word problems with process- and outcome-based feedback. arXiv:2211.14275, 2022.
• Zhenru Zhang et al. The lessons of developing process reward models in mathematical reasoning. arXiv:2501.07301, 2025.

Research Context

Written report and experiments from our Advanced AI course project; this page summarizes the paper and hosts the PDF for readers who want the full write-up.