RelayLLM: Efficient Reasoning via Collaborative Decoding

TL;DR

RelayLLM tackles the efficiency gap in complex LLM reasoning by enabling a small language model to selectively delegate critical tokens to a larger model through an interleaved, token-level collaboration scheme. A two-stage training pipeline—supervised warm-up to learn the <call> command and GRPO-based reinforcement learning with verifiable rewards—teaches the SLM to balance independence and strategic help-seeking. Empirical results on six benchmarks show RelayLLM achieving $49.52\%$ average accuracy while invoking the LLM for only $1.07\%$ of tokens, delivering up to $98.2\%$ cost reductions relative to performance-matched routers. The approach generalizes to unseen reasoning domains, and ablations confirm the importance of data filtering, independence incentives, and exploration rewards in driving both efficiency and accuracy.

Abstract

Large Language Models (LLMs) for complex reasoning is often hindered by high computational costs and latency, while resource-efficient Small Language Models (SLMs) typically lack the necessary reasoning capacity. Existing collaborative approaches, such as cascading or routing, operate at a coarse granularity by offloading entire queries to LLMs, resulting in significant computational waste when the SLM is capable of handling the majority of reasoning steps. To address this, we propose RelayLLM, a novel framework for efficient reasoning via token-level collaborative decoding. Unlike routers, RelayLLM empowers the SLM to act as an active controller that dynamically invokes the LLM only for critical tokens via a special command, effectively "relaying" the generation process. We introduce a two-stage training framework, including warm-up and Group Relative Policy Optimization (GRPO) to teach the model to balance independence with strategic help-seeking. Empirical results across six benchmarks demonstrate that RelayLLM achieves an average accuracy of 49.52%, effectively bridging the performance gap between the two models. Notably, this is achieved by invoking the LLM for only 1.07% of the total generated tokens, offering a 98.2% cost reduction compared to performance-matched random routers.

Figures (4)

• Figure 1: Results are averaged across six mathematical benchmarks. The "Random Router" baseline randomly directs questions to either the small or large model. The "Perfect Router" baseline directs only questions SLM cannot solve to large model. The x-axis represents the Call Ratio (percentage of tokens generated by the teacher model), and the y-axis denotes the average accuracy.
• Figure 2: Overview of the RelayLLM framework.(Left) Collaborative Inference: The Small Language Model acts as a central controller. During generation, it can actively trigger an "intervention" by generating a <call> command. The Large Language Model then generates the specified number of tokens, after which control returns to the SLM to complete the reasoning. (Right) Difficulty-Aware Reward Design: During GRPO training, we sample a group of rollouts (both with and without LLM intervention) and classify the query difficulty into three distinct scenarios: Solvable, Teacher-Dependent, and Unsolvable, to guide the policy optimization.Distinct reward designs are applied to each scenario to align the model's behavior with the optimal strategy.
• Figure 3: Impact of teacher model size on student performance. We evaluate two student models (Difficulty-Aware) across six benchmarks. The x-axis represents the size of the teacher model used during the collaborative inference process, ranging from "None" (the same as Sec. \ref{['sec:intri']}, prevent model to call teacher model) to 14B. The reported scores are averaged across all six datasets.
• Figure 4: A qualitative case study showcases that our model invoke LLM to show how to validate the final answer, then solve the problem by itself.