PaperlayerView exact revision on arXiv

Source-linked AI summary

Route-and-Generate: KL-Regularized Router Improvement for Reasoning Routing in LLMs

Tien-Phat Nguyen, Truong Nguyen, Thanh-Thien Le, Hoang Phan, Thien Huu Nguyen, Linh Ngo Van, Trung Le

arXiv:2610.30351v1cs.AIcs.LG

TL;DR

Large language models can overthink simple inputs and underthink harder ones, while accuracy-focused training can favor verbose reasoning. Route-and-Generate separates strategy selection from answer generation and allocates a fixed generation budget across formats; across reasoning benchmarks, it improves the accuracy–efficiency trade-off.

  • Problem

    Reasoning models can spend costly extra tokens on routine inputs yet still under-allocate depth on harder questions, while accuracy-focused training couples answer quality with reasoning strategy.

    Source: Introduction · p. 1

  • Method

    Rouge uses a format router to allocate a fixed generation budget among reasoning formats and a format-conditioned generator, with concise-correct rewards and KL-regularized routing improvement.

    Source: Introduction · p. 1, Methodology · p. 2, Closed-form KL-Regularized Router · p. 4

  • Results

    Across diverse reasoning benchmarks, Rouge achieves a better accuracy–efficiency trade-off, producing shorter generations while maintaining strong performance.

    Source: Introduction · p. 2, Conclusion · p. 7

  • Takeaways & Limitations

    Rouge explicitly controls reasoning cost, favoring cheaper strategies on easy inputs and heavier reasoning on harder ones.

    Source: Conclusion · p. 7

    The approach assumes a small set of reasoning formats can cover most problem instances in a supervised setting with reference signals such as answers, unit tests, or verifier inputs.

    Source: Motivation · p. 2

Abstract

from arXiv · show

Large language models often apply complex reasoning indiscriminately, leading to overthinking on simple tasks and underthinking on harder ones. We present Route-and-Generate (Rouge), a hierarchical training framework that enables LLMs to adaptively decide when and how much reasoning is required for each input. Our approach decouples reasoning effort allocation from answer generation by introducing a routing policy that allocates a fixed sampling budget across multiple reasoning formats, while the base model is optimized with GRPO using per-sample rewards. To promote efficient reasoning, we employ a rank-based reward that favors shorter correct solutions, and update the routing policy via a closed-form KL-regularized policy improvement step distilled into a neural router. This design encourages concise answers on easy problems while reserving longer reasoning for harder ones, achieving improved accuracy-efficiency trade-offs across multiple reasoning benchmarks.

Introduction

Reasoning models can overthink routine inputs while under-allocating depth on hard ones, and accuracy-focused training can favor verbosity. Route-and-Generate separates strategy selection from answer generation; across multiple reasoning benchmarks, it improves efficiency without sacrificing accuracy.

  • Reasoning models can overthink routine inputs while still allocating too little depth to the hardest questions.Additional reasoning tokens add latency and cost, and the gap between enough and too much reasoning widens as models scale.Source: Introduction · p. 1
  • Accuracy-focused training asks one policy to learn both correct answers and reasoning strategy, encouraging verbosity as extra tokens yield diminishing returns.Source: Introduction · p. 1
  • The paper frames the routing problem as allocating a fixed compute budget across alternative strategies to maximize correctness.This separates strategy selection from answer execution and makes reasoning effort controllable.Source: Introduction · p. 1
  • Rouge decouples reasoning-strategy selection from answer generation and derives closed-form soft routing targets through an optimization problem.This makes reasoning effort an explicit routing decision rather than part of a single generation policy.Source: Introduction · p. 1
  • Across multiple reasoning benchmarks, Rouge improves efficiency without sacrificing accuracy.Source: Introduction · p. 2

Related Work

Prior work studies explicit reasoning, reinforcement learning for reasoning, and adaptive efficiency methods. The paper situates its approach against methods that favor verbose reasoning or regulate cost without explicitly allocating computation among strategies.

  • Explicit intermediate reasoning improves mathematical, logical, and symbolic reasoning but increases generation length and computational cost, often uniformly across inputs.The cited approaches include chain-of-thought prompting, programmatic reasoning, and iterative self-reflection (Wei et al. 2022).Source: Related Work · p. 2
  • Reinforcement learning for reasoning can encourage complex behaviors, but jointly optimizing reasoning structure, length, and generation tends to favor verbose trajectories.The cited work reports higher success probability for longer trajectories under accuracy-driven objectives (Cobbe et al. 2021; Shao et al. 2024).Source: Related Work · p. 2
  • Adaptive methods reduce average computational cost through concise reasoning, format selection, or efficiency-aware objectives, but generally do not allocate resources across strategies in one decision process.Related work also uses global length or computation penalties that couple reasoning effort to execution (Wu et al. 2025; Hou et al. 2026; Aggarwal and Welleck 2025; Luo et al. 2025).Source: Related Work · p. 2

Methodology

Rouge separates input-specific reasoning-format selection from answer generation, allocating a fixed sampling budget across formats and training the generator and router with distinct objectives.

Motivation

  • Rouge separates choosing a semantic reasoning mode from executing it, treating output length as an outcome rather than a target token count.The stated setting assumes a small set of formats covers most instances, with questions paired to reference signals such as answers, unit tests, or verifier inputs.Source: Motivation · pp. 2–3

Format Routing and Candidate Generation

  • A lightweight encoder-based router predicts a format distribution and allocates a fixed generation budget across formats, while format-specific templates condition candidate generation.The resulting mixture can assign concise formats to simpler inputs and heavier reasoning to harder ones.Source: Format Routing and Candidate Generation · p. 3

Reward Design

  • The rank-based reward favors the shortest correct candidate, decays exponentially with rank, and gives incorrect responses a lower reward.Its batch-relative efficiency signal is invariant to absolute length scales.Source: Reward Design · p. 3
  • Format-level utility aggregates candidate rewards within each reasoning format to guide the router’s allocation policy.The aggregation is the mean reward across candidates generated under a format.Source: Reward Design · p. 3

Generator Optimization

  • The generator is optimized with GRPO-style policy gradients using rewards normalized across the rollout set and KL regularization to a frozen reference policy.The normalization uses the rollout rewards’ mean and standard deviation.Source: Generator Optimization · pp. 3–4

Closed-form KL-Regularized Router

  • The router’s KL-regularized improvement reweights a reference routing policy by format utility, with β controlling exploitation versus exploration.The improved policy is distilled into the parametric router; β is annealed from large to small to encourage broad early exploration and sharper later routing.Source: Closed-form KL-Regularized Router · p. 4

Training Algorithm

  • The training pipeline shares a question’s sampling budget across formats, scores generated candidates, then updates the generator with GRPO and the router with closed-form policy improvement.The framework uses a format router and a format-conditioned generator.Source: Methodology · p. 2, Training Algorithm · p. 4

Experiments

Across mathematical and general reasoning benchmarks, Rouge is more token-efficient than GRPO and ARM while maintaining strong accuracy. The experiments compare models, datasets, and evaluation measures under a setup designed to isolate adaptive format allocation from length control.

Experimental Setup

  • The evaluation uses Qwen2.5 3B and 7B backbones, three training datasets, six mathematical benchmarks, two general-reasoning benchmarks, and Pass@1, Maj@8, and average completion tokens.The training mixture is CommonsenseQA, GSM8K, and MATH; evaluation spans arithmetic through Olympiad-style problems and includes BBH and CSQA.Source: Experimental Setup · p. 4, Mathematical Reasoning · p. 4, General Reasoning · p. 5
  • Comparisons hold the backbone, training data, format inventory, and rollout budget fixed; SFT, GRPO, O1-Pruner, and ARM test distinct competing explanations for efficiency gains.The controls distinguish adaptive format allocation from ordinary accuracy-driven RL and token-length reduction.Source: Experimental Setup · p. 4

Mathematical Reasoning

  • Across Qwen2.5-3B and Qwen2.5-7B, Rouge produces the shortest solutions under Pass@1 and Maj@8 while maintaining, and sometimes exceeding, baseline accuracy.SFT and O1-Pruner can use fewer tokens by under-thinking, with large accuracy drops on difficult tasks such as AIME.Source: Mathematical Reasoning · p. 4, Table 1 · p. 5
  • Rouge’s efficiency gains are largest on GSM8K and SVAMP, while it remains highly competitive on harder AIME and AMC benchmarks.The overall results show improved token efficiency without materially sacrificing correctness, despite a few isolated regressions.Source: Mathematical Reasoning · p. 5

General Reasoning

  • Against GRPO and ARM, Rouge is the most token-efficient across both model sizes while preserving strong general-reasoning performance.On CSQA it matches the best baselines with far fewer tokens; on BBH it remains substantially more efficient and competitive on Pass@1 and Maj@8.Source: General Reasoning · p. 5

Analysis

Rouge uses less underthinking waste than ARM and GRPO, and allocates more reasoning to harder tasks. Its training accuracy rises as response length falls, while reward-margin sensitivity remains stable across tested settings.

Underthinking

  • Underthinking measures wasted computation when an early correct intermediate thought is followed by an incorrect final answer; lower scores are better.Source: Underthinking · p. 5

Training Stability

Underthinking

  • Rouge’s underthinking waste is markedly below ARM/GRPO, while incorrect outputs are not drastically shorter, arguing against early stopping as the explanation.GRPO shows the strongest underthinking behavior, consistent with ineffective strategy switching rather than deepening a viable line of thought.Source: Underthinking · pp. 5–6, Figure 2 · p. 6

Routing Behavior

  • Rouge favors cheaper formats on easier problems and shifts toward Code and Long-CoT on harder benchmarks, while GRPO largely collapses to Long-CoT.The allocation is more graded by task difficulty than GRPO’s and avoids unnecessary long reasoning on easy inputs compared with ARM.Source: Routing Behavior · p. 6

Ablations

  • Larger ρ more strongly favors concise correct solutions, and performance changes smoothly across settings without collapse, though the preferred balance depends on task difficulty.Source: Ablations · p. 7

Conclusion

Rouge separates reasoning strategy selection from answer generation to control reasoning cost. Across diverse reasoning benchmarks, it produces shorter generations while maintaining strong performance and a better accuracy–efficiency trade-off.

  • Rouge decouples reasoning strategy selection from answer generation, allowing explicit control of reasoning cost in large language models.Source: Conclusion · p. 7
Loading 2610.30351v1…