SplitReason: Learning To Offload Reasoning

• Reasoning improves accuracy but creates long, memory-bound decoding sequences whose cost grows with sequence length and model size.Long CoTs require repeatedly referencing prior tokens through the KV-Cache, increasing test-time compute.Source: §1 · p. 2
• SplitReason keeps most generation on a small model while requesting difficult reasoning segments from a larger model through learned <bigmodel> boundaries.The large model processes the small model’s preceding output in parallel through prefills, enabling continuation when offloading is triggered.Source: §1 · p. 2
• The central training challenge is teaching the small model when to emit <bigmodel> and </bigmodel> at appropriate reasoning boundaries.Difficult and easy segments are annotated in OpenR1-Math-220k, then used for SFT and RLFT to regulate offloading and latency.Source: §1 · pp. 2–3
• The work releases a fine-tuning dataset and recipe for models to learn when to offload their own reasoning to larger models.Source: §1 · p. 3
• 28.3% accuracy improvement is achieved by offloading approximately 5% of the reasoning process to larger models, with 4–6× faster inference.The reported gains target small reasoning models and combine accuracy improvement with reduced inference cost.Source: §1 · p. 3
• RL4E extends fine-tuning objectives beyond human preferences by incorporating hardware-efficiency considerations into model behavior.Source: §1 · p. 3

• SFT induces reasoning behavior by imitating CoTs, while RL can refine outcomes, formatting, and intermediate reasoning behavior.GRPO estimates baselines from grouped trajectory rewards without a value network.Source: §2 · p. 3
• Sequential test-time scaling improves reasoning by generating longer or refined chains, but its token-by-token decoding cost grows quadratically with sequence length.Decode repeatedly fetches the KV-Cache and is memory-bandwidth bound, unlike highly parallel prefill.Source: §2 · p. 3
• Speculative decoding uses a small draft model and a stronger verifier, but disagreement wastes verifier prefill and draft decoding work through rollback.Its speedups rely on easy local continuations and cheaper prefill, whereas SplitReason delegates spans without token-by-token verification.Source: §2 · p. 4
• SplitReason instead trains the small model to identify spans likely to exceed its capability and delegate only those spans to the larger model.Source: §2 · p. 4

3.1 Cooperative Execution

• The cooperative loop keeps both KV-Caches current: streaming prefills prepare large-model takeover, while controlling prefills let the small model reclaim decoding with </bigmodel>.Prefill is highly parallel and cheaper than memory-bound decode, so either model can resume without a decoding delay.Source: §3.1 · p. 4
• Most (>95%) of the CoT is produced by the small model while the larger model handles selected difficult spans.The small model controls decoding, and the large model is invoked only after the small model emits <bigmodel>.Source: §3.1 · p. 4

3.2 Training Procedure

• The training pipeline annotates difficult spans in 18k OpenR1-Math-220k traces, wraps them in control tokens, and applies SFT to induce their emission.DeepSeek-chat supplies difficult-span annotations, which are aligned to the original text using fuzzy matching.Source: §3.2 · p. 5, §3.3 · p. 5
• GRPO then combines correctness, formatting, and latency-aligned offload rewards to regulate boundary quality and the desired offload ratio.The procedure does not invoke the large model during training, emphasizing latency and formatting rather than true offloaded accuracy.Source: §3.2 · p. 5, §3.3 · p. 6
• The larger model requires no fine-tuning and can be swapped among models from 7B through 32B or larger.Both SFT and GRPO are formulated without involving the large model during training.Source: §3.2 · p. 5, §3.3 · p. 6

3.3 Data Generation and Training Setup

• The data-generation setup prompts DeepSeek-chat to identify the 20% most logically complex CoT portions before fuzzy matching and tag insertion.Source: §3.3 · p. 5
• The GRPO coverage reward favors moderate offloading, penalizes excessive or minimal coverage, and penalizes unmatched control tags.Its coverage target is 40%, while rejecting requests can reduce actual offloading toward zero without noticeable AIME24 accuracy impact.Source: §3.3 · p. 6

4.1 Offloading Behavior

• 24%–28% accuracy gains result from learned 1.35%–5% offloading, exceeding random offloading, which improves accuracy by up to 20%.Learned offloading targets the hardest reasoning segments rather than handing off random decode steps.Source: §4.1 · p. 6, Figure 5 · p. 6

4.2 Offloading Across Model Sizes

• 28% higher AIME24 accuracy is achieved with SplitReason-32B while offloading only 5.54% of decoding, outperforming larger standalone-model alternatives.The gains are attributed to offloading rather than stronger fine-tuning of the 1.5B model, because disabling offloading did not improve its accuracy.Source: §4.1 · p. 6, §4.2 · p. 7
• Offloading to an 8B model raises AIME24 accuracy from 17.3% to 44%, with further improvements from 14B and 32B targets.The experiments compare a 1.5B small model with 8B, 14B, and 32B large models.Source: §4.2 · p. 7, Figure 6 · p. 7

4.3 Dataset Distribution and Inducing Offloading

• 18,500 annotated reasoning traces mostly remain below the 20% offloading target, with slightly more offloading in earlier reasoning segments.The annotation distribution is intended to identify difficult spans while limiting the fraction delegated to the large model.Source: §4.3 · p. 7, Figure 7 · p. 8
• GRPO is needed beyond SFT: SFT offloads less than 1% unreliably, whereas the final model follows formatting and offloads effectively at approximately 5%.The observed 5% rate differs from the 40% offload ratio favored by the GRPO reward, indicating that the procedure may need further tuning.Source: §4.3 · p. 7, Figure 8 · p. 8

4.4 Performance Simulation

• A 1.5B model decodes over 8× faster than a 32B model, while 32B prefill remains faster than 1.5B decoding, supporting pipelined execution.The simulator uses profiled prefill and decode throughputs to model concurrent small-model decoding and large-model prefill.Source: §4.4 · pp. 8–9, Figure 9 · p. 8

Alignment with Latency

• Latency-aware GRPO uses control tokens to optimize inference efficiency, extending reinforcement learning beyond human-preference objectives.The discussion suggests similar control-token strategies could apply to quantization, pruning, and other compression methods.Source: §5 · p. 9

Limitations

• True offloading is absent from the current GRPO reward, making its accuracy component unrepresentative of downstream accuracy; performance measurements are simulation-based.The method also requires both models to retain KV-caches and generally uses more GPUs.Source: §5 · p. 9

6 Conclusion

• 28% AIME24 accuracy improvement is achieved by the 1.5B model using SFT+GRPO to regulate <bigmodel> boundaries.Most decoding remains on the small network, while the larger model is not fine-tuned and can be swapped without retraining.Source: §6 · p. 9