PaperlayerView exact revision on arXiv

Source-linked AI summary

RelayLLM: Efficient Reasoning via Collaborative Decoding

Chengsong Huang, Tong Zheng, Langlin Huang, Jinyuan Li, Haolin Liu, Jiaxin Huang

arXiv:2601.05167v1cs.CLcs.AIcs.LG

TL;DR

High cost and latency constrain LLM reasoning, while SLMs often lack capacity for hard tasks and query-level routing can waste computation. RelayLLM lets an SLM request targeted LLM continuations during generation. Across six benchmarks, it reports 49.52% average accuracy with only 1.07% of generated tokens from the LLM.

  • Problem

    LLMs face high computational costs and latency, SLMs typically struggle with hard reasoning, and query-level routing can waste computation by offloading tasks the SLM could mostly handle.

    Source: §1 · p. 1

  • Method

    RelayLLM makes the SLM an active controller that interleaves its reasoning with on-demand LLM assistance at critical positions.

    Source: §1 · pp. 1–2, §2 · p. 2

  • Results

    49.52% average accuracy across six benchmarks with 1.07% of generated tokens from the LLM; RelayLLM improves accuracy by 6.9% over a resource-equivalent Random Router.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    RelayLLM shows that an SLM can handle most reasoning steps with sparse expert intervention and achieve comparable reasoning accuracy to larger models with negligible computation.

    Source: §7 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

LLMs offer complex reasoning but incur high computational costs and latency, while resource-efficient SLMs often struggle with hard tasks. RelayLLM lets an SLM request targeted LLM help during generation, aiming to reduce the waste of query-level routing.

1 Introduction

  • Query-level cascading or routing can offload an entire task to the LLM, wasting computation when the SLM can handle most reasoning steps and needs help only at critical positions.The paper contrasts this all-or-nothing approach with targeted intervention.Source: §1 · p. 1
  • RelayLLM replaces coarse query routing with token-level collaboration, using the SLM as both reasoner and active controller without an additional controller.The SLM requests assistance only when necessary.Source: §1 · pp. 1–2
  • A two-stage training framework teaches call-command syntax in supervised warm-up, then uses GRPO with a context-aware reward to balance independence and help-seeking.The reward penalizes wasted costs and avoidable errors.Source: §1 · p. 2
  • 49.52% average accuracy across six benchmarks, with only 1.07% of generated tokens from the LLM; RelayLLM also improves accuracy by 6.9% over a resource-equivalent Random Router.The paper reports a 98.2% cost reduction versus a performance-matched router.Source: Figure 1 · p. 1, §1 · p. 2

2 RelayLLM Inference

RelayLLM uses an SLM as the primary reasoner and controller, interleaving its generation with bounded LLM continuations. The SLM retains the delegation trace and resumes after each intervention.

2 RelayLLM Inference

  • The inference setup pairs a resource-efficient SLM with a more powerful, computationally expensive LLM to generate a response through dynamic collaboration.Source: §2 · p. 2

2.1 Small Model Generation

  • The SLM can emit a special command pattern to request assistance, with the requested number of LLM tokens specified by n.Source: §2.1 · p. 2

2.2 Large Model Intervention

  • When the command is detected, the SLM pauses and sends the current context to the LLM, which generates up to n continuation tokens or stops at EOS.The command tokens are stripped from the context forwarded to the LLM.Source: §2.2 · p. 2

2.3 Iterative Relay of LLM and SLM

  • After the LLM responds, its tokens are added to the shared context; the SLM retains its command history and resumes reasoning from the updated context.Source: §2.3 · p. 2

3 RelayLLM Training

RelayLLM trains the SLM to decide when and how long to delegate token generation to an LLM, then resume reasoning. Its two-stage training combines supervised warm-up with GRPO and difficulty-aware rewards.

3 RelayLLM Training

  • The SLM controls token-level collaboration by invoking the LLM for a specified number of tokens, then resuming generation; training uses supervised warm-up followed by reinforcement learning.The framework also samples rollouts with and without intervention and classifies queries as solvable, teacher-dependent, or unsolvable for reward design.Source: §3 · p. 2, Figure 2 · p. 3

3.1 Cold Start via Supervised Warm-up

  • Warm-up teaches the SLM valid calling commands using self-sampled contexts with inserted command tokens, then fine-tunes those examples with cross-entropy loss.This provides a stable starting point for reinforcement learning despite a discrepancy between warm-up and inference contexts after the command.Source: §3.1 · p. 3

3.2 Policy Refinement with Reinforcement Learning

  • GRPO refines the SLM policy using rule-based verifiable rewards and group-relative comparisons of sampled responses.The group-normalized advantage encourages responses to outperform the group average.Source: §3.2.1 · pp. 3–4

3.2.1 GRPO Training with RLVR

  • GRPO samples multiple outputs per query, compares them within the group, and regularizes policy updates against a reference policy.The group-normalized advantage uses a small stability constant and favors responses that outperform the group average.Source: §3.2.1 · p. 4

3.2.2 Data Filtering

  • Training data excludes queries the LLM solves unreliably: only queries with at least a 50% pass rate across 10 sampled responses are retained.The authors describe this as filtering to the LLM’s competence boundary.Source: §3.2.2 · p. 4

3.2.3 Reward Design

  • Difficulty-aware rewards boost independent success on solvable queries, penalize no-call failures on teacher-dependent queries, and reward attempted calls when no rollout succeeds.For solvable queries, independent success receives 1.5; for teacher-dependent queries, failing to call receives −1.0, while correct expert-assisted answers receive the simple reward.Source: §3.2.3 · p. 4

4 Experiments

Experiments evaluate RelayLLM with Qwen3 students and an 8B teacher across six reasoning benchmarks. The reported results show accuracy gains with low token overhead and recovery of a substantial share of the student–expert performance gap.

4.1.1 Models

  • The experiments use Qwen3-0.6B and Qwen3-1.7B as students and Qwen3-8B as teacher, with all models run in non-thinking mode.The models share a family to make generation style, token distribution, vocabulary, and tokenizer more consistent for collaboration.Source: §4.1.1 · p. 5

4.2 Evaluation Setup

  • Evaluation covers six reasoning benchmarks, using avg@32 for AIME-2024 and AIME-2025 and greedy-decoding pass@1 accuracy for the other benchmarks.The benchmarks are Minerva, MATH-500, GSM8K, Olympiad-Bench, AIME-2024, and AIME-2025; GPT-4o-mini serves as a semantic judge.Source: §4.2 · p. 5

4.3 Training Details

  • Models train for one epoch on DAPO data; GRPO uses the full dataset, while RelayLLM uses the filtered subset, and the teacher is served through vLLM.The calling command halts sampling and triggers a teacher API invocation.Source: §4.3 · p. 5

4.4 Main Results

  • For Qwen3-1.7B, Difficulty-Aware-Reward achieves 49.52% average accuracy versus 49.30% for Simple-Reward, with average call ratios of 1.07% and 0.43%, respectively.The paper describes the accuracy increase as correlated with a slight increase in call ratio.Source: §4.4 · p. 5
  • Qwen3-1.7B with Difficulty-Aware-Reward recovers approximately 60% of the performance gap between the base SLM (42.50%) and expert model (54.12%).On Minerva, Qwen3-0.6B improves from 15.81% to 23.53%, invoking the large model for only 0.77% of total tokens.Source: §4.4 · p. 5

5 Analysis

Section 5 analyzes RelayLLM’s generalization, component contributions, teacher-free reasoning, delegation length, and sensitivity to the inference teacher. Together, the analyses examine when collaboration helps and how it affects performance and efficiency.

5.1 RelayLLM Generalizes to Unseen Reasoning Domains

  • 59.03% on MMLU-Pro for Qwen3-1.7B exceeds the GRPO baseline (49.76%) and CITER (53.38%) on reasoning domains unseen during training.The evaluation covers three non-mathematical benchmarks; Table 2 presents performance comparisons across reasoning and general-knowledge benchmarks.Source: §5.1 · p. 6, Table 2 · p. 6

5.2 Ablation Study

  • Removing data filtering triples the call ratio and decreases accuracy; filtering teacher-intractable queries avoids calls that yield no performance gain.Filtering too-hard data also saves training time and resources.Source: §5.2 · p. 6
  • Removing the independence incentive raises the call ratio to 4.10%, indicating that the bonus for independent success limits reliance on the teacher for solvable queries.The incentive boosts the correctness reward from 1 to 1.5 for solvable queries.Source: §5.2 · p. 6, Table 3 · p. 7
  • 47.56% accuracy without the exploration reward indicates that the reward encourages help-seeking in highly uncertain scenarios.The ablation removes the reward for seeking help in unsolvable queries.Source: §5.2 · p. 7, Table 3 · p. 7

5.3 Intrinsic Reasoning Capability

  • Without teacher access, RelayLLM (Simple-Reward) scores 61.12% on easy datasets, above GRPO, while Difficulty-Aware-Reward falls from 15.00% to 11.93% on harder datasets.Table 4 disables teacher calls; the text reports that Simple-Reward retains more independence than Difficulty-Aware-Reward, which encourages a higher call ratio.Source: §5.3 · p. 7, Table 4 · p. 7

5.4 Dynamic Token-Length Calling Minimizes Computational Cost

  • At similar accuracy, dynamic delegation reduces the call ratio from 2.87% for Fixed-100 to 1.07% for RelayLLM.Fixed-length models are retrained to request their specified token counts; RelayLLM learns to request a variable length, avoiding a set budget on simple queries.Source: §5.4 · p. 7, Table 5 · p. 7

5.5 Distributional Alignment

  • 49.52% is the 1.7B student’s peak accuracy when the inference teacher matches the training teacher; using a larger teacher slightly lowers performance.Figure 3 reports scores averaged across six benchmarks as inference-teacher size ranges from None to 14B.Source: Figure 3 · p. 8, §5.5 · p. 8
  • Even teachers weaker than the student outperform the “None” baseline; excluding the training-matched 8B teacher, performance positively correlates with teacher size.The authors attribute the larger-teacher exception to distribution shift outweighing its stronger reasoning capability.Source: §5.5 · p. 8, Figure 3 · p. 8

6 Related Work

Related work places RelayLLM among model-collaboration approaches and reinforcement-learning methods for LLM reasoning.

6.1 Model Collaboration

  • Model collaboration includes weight-level merging and logits-level ensembling, while efficiency-focused work uses speculative reasoning or interleaved collaborative decoding.Speculative approaches use judges, verifiers, or large-model guidance; collaborative decoding interleaves generation through a controller or joint policy.Source: §6.1 · p. 8

6.2 RL for LLM Reasoning

  • Reinforcement learning has been applied to mathematical, code-generation, multimodal, and structured-data tasks, alongside methods such as self-play, DAPO, VAPO, and high-entropy optimization.The cited work presents these as training paradigms and algorithmic techniques for supporting diverse applications and complex behaviors.Source: §6.2 · p. 8

7 Conclusion

RelayLLM addresses all-or-nothing routing by using the large model as an on-demand tool at critical positions, and its GRPO-based training learns help-seeking behavior. The reported results include a 6.9% advantage over resource-equivalent random routers and comparable reasoning accuracy to larger models with negligible computation.

7 Conclusion

  • RelayLLM treats the large model as an on-demand tool at critical positions, addressing the inefficiency of all-or-nothing routing.The conclusion says small models can handle most reasoning steps when supported at those positions.Source: §7 · p. 8
  • GRPO-based training effectively learns and optimizes help-seeking behavior.Source: §7 · p. 8
  • RelayLLM outperforms resource-equivalent random routers by 6.9% and achieves comparable reasoning accuracy to larger models with negligible computation.Source: §7 · p. 8

A Case Study

The case studies use Qwen3-1.7B with Difficulty-Aware training to illustrate how invoking the LLM can help solve questions.

A Case Study

  • The case studies use Qwen3-1.7B with Difficulty-Aware training to illustrate how LLM invocation can help solve questions.Source: A Case Study · p. 12

B Detailed Analysis of Delegation Length Strategies

RelayLLM matches Fixed-100 on standard benchmarks with a lower call ratio, while delegation-length strategies show different accuracy and cost trade-offs on Minerva. The appendices also specify the inference prompt and answer-judging procedure.

B Detailed Analysis of Delegation Length Strategies

C.1 Inference Prompt

  • The inference prompt requests step-by-step reasoning and a final answer within \boxed{}, with {problem_statement} standing for the actual math problem.Source: §C.1 · p. 12

C.2 GPT-4o-mini Judge Prompt

  • GPT-4o-mini judges mathematical answers whose final forms can be complex, using a correctness check against benchmark ground truth.Source: §C.2 · p. 12
  • The judge prompt uses {answer} for the model-generated solution and {response} for the benchmark ground truth.Source: §C.2 · p. 12

D Other Related Work

The supplied material connects RelayLLM to efficient-reasoning approaches and includes a qualitative case study and a benchmark-table description.

D Other Related Work

  • RelayLLM is positioned alongside efficient chain-of-thought, speculative best-of-N decoding, in-context learning, and non-myopic generation methods.Source: D Other Related Work · p. 12
  • Figure 4’s case study shows LLM invocation for validating the final answer, followed by solving the problem independently.Source: Figure 4 · p. 13
  • Table 6 covers six benchmarks, reports mean@32 for AIME datasets and pass@1 for others, and compares retrained Fixed-k strategies.The caption says best accuracy values are bolded and the most efficient call ratio is highlighted.Source: Table 6 · p. 14

E Hyperparameter

GRPO training uses AdamW with a constant learning rate and weight decay, alongside settings for long reasoning outputs and rollout stability.

  • GRPO post-training uses AdamW with a constant learning rate of 1 × 10−6, weight decay of 1 × 10−2, and global batch size 32.Source: E Hyperparameter · p. 14
  • Training allows prompts up to 4096 tokens and responses up to 8192 tokens for extended chain-of-thought generation.Source: E Hyperparameter · p. 14
  • GRPO uses group size G = 8, KL coefficient β = 0.01, and rollout temperature 1.0.The group size samples eight outputs per prompt to estimate the baseline; the KL coefficient is used to limit policy deviation from the reference model.Source: E Hyperparameter · p. 14
Loading 2601.05167v1…