R-Stitch: Dynamic Trajectory Stitching for Efficient Reasoning

Technical Summary: R-Stitch: Dynamic Trajectory Stitching for Efficient Reasoning

R-Stitch is a training-free, hybrid decoding framework designed to accelerate the inference of Large Language Models (LLMs) during complex reasoning tasks, specifically those utilizing Chain-of-Thought (CoT) prompting.

---

1. The Problem: The CoT Latency Bottleneck

While CoT prompting significantly enhances the reasoning capabilities of LLMs, it introduces a massive computational burden. Because CoT requires generating long, intermediate reasoning trajectories (often spanning thousands of tokens), the autoregressive decoding process becomes extremely slow.

The authors identify two major limitations in existing acceleration methods:

• Speculative Decoding Limitations: Traditional speculative decoding relies on "token-level consistency" between a Small Language Model (SLM) and an LLM. If the models disagree on even a single token, the system must perform a costly rollback. In reasoning tasks, consistency is often low, leading to frequent rollbacks and sometimes even slower performance than vanilla LLM decoding.
• Inefficient Use of SLMs: Existing methods fail to exploit the fact that when an SLM is correct, it often produces much shorter and more concise reasoning traces than an LLM. Speculative decoding's rigid requirement for exact token agreement prevents it from utilizing these shorter, efficient trajectories.

2. Methodology: Entropy-Guided Hybrid Decoding

The core insight of R-Stitch is that token entropy serves as an effective proxy for uncertainty. High-entropy tokens are error-prone, while low-entropy tokens are highly reliable.

A. R-Stitch (The Heuristic Approach):
R-Stitch implements a bidirectional, token-level routing strategy:

• Routing Logic: The process begins with the SLM. At each step, the system calculates the normalized entropy ($H_t$) of the predicted token.
  • If $H_t \leq \tau$ (low uncertainty): The SLM's token is accepted, and the SLM continues.
  • If $H_t > \tau$ (high uncertainty): The token is discarded, and the LLM is invoked to "correct" the trajectory and resume decoding.
• Bidirectional Flow: Conversely, if the LLM is decoding and produces a low-entropy token, control is handed back to the SLM to save computation.
• KV Cache Management: To minimize switching overhead, R-Stitch uses separate KV caches for both models and employs partial prefilling to reuse existing cache data when switching back to a previously active model.

B. R-Stitch+ (The Learned Approach):
To move beyond a fixed entropy threshold, the authors propose R-Stitch+, which uses a lightweight router (a small FFN) to decide whether to switch models.

• Reinforcement Learning (RL): The router is trained using the DAPO optimizer with a latency-aware reward function.
• Reward Design: The reward $R = r_{acc} + r_{eff}$ balances accuracy ($r_{acc}$) and efficiency ($r_{eff}$). Crucially, the efficiency reward is tied to actual predicted trajectory latency rather than just token count, allowing the model to learn the true cost of model switching.

3. Key Contributions

• New Perspective on Inconsistency: Unlike previous works that try to make SLMs and LLMs more consistent, R-Stitch treats model inconsistency as an advantage, using it to trigger high-quality LLM intervention only when necessary.
• Training-Free Efficiency: The base R-Stitch requires no retraining of the underlying models.
• Adaptive Routing: R-Stitch+ provides a data-driven way to optimize the accuracy-latency trade-off via RL.
• Conciseness Exploitation: By allowing the SLM to dictate the trajectory, the method naturally produces shorter reasoning paths, reducing the total number of tokens generated.

4. Results and Performance

The framework was validated on five mathematical reasoning benchmarks (AIME, AMC, Minerva, MATH, and OlympiadBench) using DeepSeek-R1-Distill and QWQ models.

• Substantial Speedup: R-Stitch achieves peak speedups of 3.00× (7B), 3.85× (14B), and 4.10× (32B).
• Accuracy Preservation: The method maintains accuracy levels comparable to full LLM decoding. In some cases, accuracy actually improves because the SLM's concise traces prevent the LLM from becoming overly verbose and hitting token limits.
• Superiority over Speculative Decoding: R-Stitch consistently outperforms speculative decoding, which often suffers from speedups below 1× due to frequent rollbacks in low-consistency reasoning settings.
• Low Overhead: The memory overhead of maintaining a dual KV cache is minimal (e.g., adding only ~5% memory for a 32B model), and the computational cost of entropy calculation is negligible.

5. Significance

R-Stitch represents a significant step toward making high-performance reasoning models practical for real-world, time-sensitive applications. By transforming the "problem" of model inconsistency into a "feature" for dynamic computation, it provides a principled way to balance the extreme intelligence of large models with the extreme efficiency of small models.