RelayGen: Intra-Generation Model Switching for Efficient Reasoning

TL;DR

RelayGen tackles the inefficiency of inference in large reasoning models by exploiting intra-output difficulty variation during long-form reasoning. It uses offline calibration to identify reliable switch cues based on post-sentence probability margins, enabling a training-free, segment-level runtime handoff from a large to a smaller unit, and full reuse with a large model later. When combined with speculative decoding, RelayGen yields substantial latency reductions with minimal accuracy loss across benchmarks. This work demonstrates that coarse-grained, empirically grounded switching can achieve deployment-friendly improvements that rival or exceed fine-grained learned routing.

Abstract

Large reasoning models (LRMs) achieve strong performance on complex reasoning tasks by generating long, multi-step reasoning trajectories, but inference-time scaling incurs substantial deployment cost. A key challenge is that generation difficulty varies within a single output, whereas existing efficiency-oriented approaches either ignore this intra-generation variation or rely on supervised token-level routing with high system complexity. We present \textbf{RelayGen}, a training-free, segment-level runtime model switching framework that exploits difficulty variation in long-form reasoning. Through offline analysis of generation uncertainty using token probability margins, we show that coarse-grained segment-level control is sufficient to capture difficulty transitions within a reasoning trajectory. RelayGen identifies model-specific switch cues that signal transitions to lower-difficulty segments and dynamically delegates their continuation to a smaller model, while preserving high-difficulty reasoning on the large model. Across multiple reasoning benchmarks, RelayGen substantially reduces inference latency while preserving most of the accuracy of large models. When combined with speculative decoding, RelayGen achieves up to 2.2$\times$ end-to-end speedup with less than 2\% accuracy degradation, without requiring additional training or learned routing components.

Figures (6)

• Figure 1: RelayGen overview. Long reasoning generation exhibits difficulty variation within a single output, enabling segment-level runtime model switching.
• Figure 2: Given a prompt, an LRM generates an output that naturally decomposes into a long reasoning stage and a subsequent answer stage, explicitly separated by special boundary tokens (e.g., <think>, </think>)
• Figure 3: Probability margin trajectories of Qwen3-32B across long reasoning examples.
• Figure 4: Post-sentence probability margin following representative discourse-level cues.
• Figure 5: Runtime model switching semantics in RelayGen.