Sparse collaborative decoding improves accuracy by triggering target calls smartly
Trajectory-Triggered Sparse Collaborative Decoding
Computation and LanguageMachine Learning
Summary
Generating text or answers often uses a main model that guesses tokens and a bigger, more accurate model to check them. The problem is deciding when to ask the bigger model for help. The authors show that just looking at uncertainty at each step is a poor way to decide when to call the bigger model. Instead, they use patterns over time in the guess quality, which predicts errors much better. This smart timing lets them call the bigger model only when needed, improving accuracy on math and reasoning tasks while saving effort.
What this means in practice
- •For ai model engineers: Improve large language models' output accuracy by calling expensive target models only at important points identified from trajectory patterns.
- •For automated math solver developers: Enhance math problem solvers by using trajectory-triggered corrections for better accuracy without excessive computation.
Authors
Amit Ranjan Trivedi, Sruthikesh Surineni, Kyong Ho Lee, Mohit Mittal, Steven Li, Rex Hill
Abstract
Dense assistance methods such as speculative decoding repeatedly verify draft tokens with a target model, keeping target computation coupled to generation. We study sparse collaborative decoding, where a draft model generates most tokens and a target model is invoked only for short corrective segments. The key problem is deciding when to call the target. We show that local uncertainty is a weak trigger: entropy, negative log-likelihood, logit margin, top-$k$ mass, and velocity features give near-chance failure prediction, with AUC 0.50-0.53. In contrast, aggregate trajectory-shape features over the draft confidence path reach AUC 0.86 using a lightweight logistic predictor trained only from draft correctness labels. This induces a trigger-timing tradeoff: early calls have more leverage but weaker evidence, while later calls have stronger evidence but less ability to redirect the trajectory. Used online, learned sparse collaborative decoding improves over the best open-loop schedule by +5.3 to +20.0 accuracy points across five GSM8K draft-target pairs. On MATH-500, it gains 7.0-13.5 points when early prediction and target competence align, but loses when either is weak. On StrategyQA, it improves over open-loop by 5.0-8.3 points, showing transfer beyond arithmetics.