Rex Hill
Author directory2026
Trajectory-Triggered Sparse Collaborative Decoding
Amit Ranjan Trivedi | Sruthikesh Surineni | Kyong Ho Lee | Mohit Mittal | Steven Li | Rex Hill
The Proceedings of the Third Workshop on Uncertainty-Aware NLP
Amit Ranjan Trivedi | Sruthikesh Surineni | Kyong Ho Lee | Mohit Mittal | Steven Li | Rex Hill
The Proceedings of the Third Workshop on Uncertainty-Aware NLP
Uncertainty estimates in language models are typically computed locally, from the predictive distribution at an individual decoding step. In multi-step reasoning, however, the outcome of interest is often trajectory-level: whether the final answer produced after many intermediate steps will be correct. We show that these two notions of uncertainty can diverge sharply, and study what an uncertainty signal should look like when the prediction target is failure of the complete reasoning trajectory. We first quantify the limitation of local uncertainty for predicting final-answer failure. On GSM8K with a Qwen-1.5B draft model, token-level entropy predicts final incorrectness at AUC 0.531, negative log-likelihood at 0.528, and logit margin at 0.474. Short-window changes in these quantities remain between 0.498 and 0.503. A model can therefore remain locally confident while following an incorrect derivation, or become locally uncertain at a step from which it subsequently recovers. Next-token uncertainty does not directly measure the risk accumulated over a reasoning trajectory. The shape of the confidence trajectory is substantially more informative. We represent each observed prefix using approximately 85 statistics describing the evolution of entropy, likelihood, logit margin, and probability concentration, including means, extrema, percentiles, slopes, variances, temporal changes, and position-dependent behavior. A lightweight logistic predictor trained only on draft trajectories labeled by final correctness reaches AUC 0.864, while the strongest individual trajectory feature reaches 0.803. Its learned weights associate failure with declining logit-margin separation, increasing temporal volatility, and weaker top-k probability concentration. A complementary wavelet analysis finds elevated mid-frequency oscillation in incorrect trajectories, suggesting that persistent instability and unresolved competition between reasoning paths carry more information than isolated uncertainty spikes.
For an uncertainty estimate used during generation, full-trajectory discrimination can be misleading because the estimate must become reliable before the trajectory is complete. We therefore evaluate prefix-conditioned uncertainty. On GSM8K, the Qwen-1.5B predictor rises from AUC 0.60 after 16 tokens to 0.82 after 128 tokens and 0.86 on the complete trajectory. Across GSM8K drafts, AUC at 128 tokens ranges from 0.69 to 0.82. On MATH-500, full-trajectory AUC reaches 0.73-0.78, yet AUC at 128 tokens remains only 0.57-0.59. Harder symbolic problems accumulate evidence more slowly. These results motivate prefix-conditioned AUC as an operational measure of uncertainty quality for systems that must act during decoding.
We test whether this distinction matters in a downstream decision problem. In sparse collaborative decoding, a small draft model generates most of a reasoning trajectory, while a stronger target model is invoked only for short textual corrections when predicted failure risk crosses a threshold. Perturbation experiments reveal a fundamental timing tradeoff: early interventions have greater influence on the eventual trajectory, precisely when uncertainty estimates are least reliable, while later estimates are more accurate but have less ability to redirect generation.
Across five GSM8K draft-target configurations, trajectory-triggered intervention improves over the best fixed insertion schedule by 5.3-20.0 accuracy points and over draft-only decoding by as much as 43.0 points. On StrategyQA it improves over the best open-loop schedule by 5.0-8.3 points, demonstrating transfer beyond arithmetic reasoning. On MATH-500 it gains 7.0-13.5 points in successful configurations but loses when early trajectory prediction is weak. The failure cases expose an important practical distinction between retrospective and actionable uncertainty: a predictor can discriminate failures well after a trajectory is complete while remaining too unreliable at the prefix where a decision must be made. Target behavior introduces a second source of uncertainty. Corrective interventions that degrade an otherwise correct draft exhibit higher entropy and negative log-likelihood and lower maximum probability and logit margin than successful interventions. Consequently, stronger standalone target accuracy does not guarantee better corrective behavior. Reliable collaborative reasoning requires both an uncertainty estimate that identifies failing trajectories early enough to act and a target continuation that is itself sufficiently reliable.
Overall, our results suggest that uncertainty for multi-step language-model reasoning should be treated as a property of an evolving trajectory rather than an isolated token. Aggregate confidence dynamics reveal failure signals that local uncertainty largely misses, while prefix-conditioned evaluation distinguishes uncertainty that predicts eventual failure from uncertainty that is available early enough.
For an uncertainty estimate used during generation, full-trajectory discrimination can be misleading because the estimate must become reliable before the trajectory is complete. We therefore evaluate prefix-conditioned uncertainty. On GSM8K, the Qwen-1.5B predictor rises from AUC 0.60 after 16 tokens to 0.82 after 128 tokens and 0.86 on the complete trajectory. Across GSM8K drafts, AUC at 128 tokens ranges from 0.69 to 0.82. On MATH-500, full-trajectory AUC reaches 0.73-0.78, yet AUC at 128 tokens remains only 0.57-0.59. Harder symbolic problems accumulate evidence more slowly. These results motivate prefix-conditioned AUC as an operational measure of uncertainty quality for systems that must act during decoding.
We test whether this distinction matters in a downstream decision problem. In sparse collaborative decoding, a small draft model generates most of a reasoning trajectory, while a stronger target model is invoked only for short textual corrections when predicted failure risk crosses a threshold. Perturbation experiments reveal a fundamental timing tradeoff: early interventions have greater influence on the eventual trajectory, precisely when uncertainty estimates are least reliable, while later estimates are more accurate but have less ability to redirect generation.
Across five GSM8K draft-target configurations, trajectory-triggered intervention improves over the best fixed insertion schedule by 5.3-20.0 accuracy points and over draft-only decoding by as much as 43.0 points. On StrategyQA it improves over the best open-loop schedule by 5.0-8.3 points, demonstrating transfer beyond arithmetic reasoning. On MATH-500 it gains 7.0-13.5 points in successful configurations but loses when early trajectory prediction is weak. The failure cases expose an important practical distinction between retrospective and actionable uncertainty: a predictor can discriminate failures well after a trajectory is complete while remaining too unreliable at the prefix where a decision must be made. Target behavior introduces a second source of uncertainty. Corrective interventions that degrade an otherwise correct draft exhibit higher entropy and negative log-likelihood and lower maximum probability and logit margin than successful interventions. Consequently, stronger standalone target accuracy does not guarantee better corrective behavior. Reliable collaborative reasoning requires both an uncertainty estimate that identifies failing trajectories early enough to act and a target continuation that is itself sufficiently reliable.
Overall, our results suggest that uncertainty for multi-step language-model reasoning should be treated as a property of an evolving trajectory rather than an isolated token. Aggregate confidence dynamics reveal failure signals that local uncertainty largely misses, while prefix-conditioned evaluation distinguishes uncertainty that predicts eventual failure from uncertainty that is available early enough.