R-Stitch Dynamic Trajectory Stitching for Efficient Reasoning

Extensive Reading

01Author Info#

Existing acceleration methods like Speculative Decoding have limitations:

Rigid Consistency: They require the Small Language Model (SLM) to match the LLM's tokens exactly. If the SLM phrases a correct reasoning step differently, speculative decoding rejects it, wasting computation.
Low Agreement: In complex reasoning tasks, token-level agreement between SLMs and LLMs is often low, leading to frequent rollbacks and minimal speed gains.

Use entropy as an uncertainty proxy

The paper analyzes entropy patterns on AMC using DeepSeek-R1-Distill-Qwen-7B (LLM) and L1-1.5B-Short (SLM), revealing three key observations:

Incorrect answers show higher entropy. Reasoning traces leading to wrong answers have significantly higher mean entropy (Figure 3a).
Most tokens are near-deterministic. Over 89% of SLM tokens have entropy ≤ 0.1, indicating high prediction confidence (Figure 3b).
High-entropy tokens trigger errors. Harmful tokens are often preceded by locally elevated entropy, making entropy a useful routing signal (Figure 3c).

Entropy-Guided Routing

The system monitors the normalized entropy () of the probability distribution at each decoding step.
- SLM LLM: If the SLM's uncertainty exceeds a fixed threshold (), control is passed to the LLM. The uncertain token is discarded, and the LLM regenerates it.
- LLM SLM: Once the LLM stabilizes and produces a token with low uncertainty (), control is handed back to the SLM to save computational cost.

To overcome the limitations of a fixed entropy threshold, the authors introduce R-Stitch+, which replaces the heuristic rule with a learned policy.

Learned Router: A lightweight router (a small MLP head) takes the model's hidden states as input and outputs a binary decision: continue with the current model or switch.
Latency-Aware Reward Function:The router is trained using Reinforcement Learning (specifically the DAPO algorithm). The reward function () explicitly balances accuracy and real-world latency:
Accuracy (): Binary reward for correct final answers.
Efficiency (): Modeled as , where is the estimated trajectory latency. Crucially, latency is only penalized if the answer is correct, preventing the model from learning to output short, incorrect nonsense just to be fast.
Latency Estimation:Since measuring wall-clock time during RL training is impractical, the authors use a regression model to estimate latency based on the number of tokens processed (prefill length) and the current KV cache size.