RelayGen: Intra-Generation Model Switching for Efficient Reasoning

TL;DR

RelayGen is a model switching framework that dynamically allocates smaller models to easier segments during reasoning, significantly reducing inference time while maintaining high accuracy.

Contribution

It introduces a training-free, segment-level model switching method based on difficulty estimation, enabling efficient long-form reasoning without additional training.

Findings

Up to 2.2× speedup in inference with minimal accuracy loss

Effective difficulty estimation using token probability margins

No additional training or learned routing needed

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.