PaperlayerView exact revision on arXiv

Source-linked AI summary

Trajectory-Triggered Sparse Collaborative Decoding

Amit Ranjan Trivedi, Sruthikesh Surineni, Kyong Ho Lee, Mohit Mittal, Steven Li, Rex Hill

arXiv:2610.23613v1cs.CLcs.LG

TL;DR

Sparse collaborative decoding addresses when to call a stronger target during a draft trajectory, rather than repeatedly verifying tokens or relying on fixed schedules. It uses aggregate confidence-trajectory features to predict draft failure and trigger short corrective segments. Learned triggering improves over open-loop schedules on several benchmarks, with results depending on early prediction reliability and target competence.

  • Problem

    Dense verification couples target computation to generation, while existing coarse routing does not decide whether a partial trajectory would benefit from a short correction.

    Source: §1 · p. 1

  • Method

    A lightweight predictor uses aggregate draft confidence-trajectory features to trigger bounded target-generated text segments, then returns generation to the draft.

    Source: §1 · p. 2, §5 · p. 4

  • Results

    +5.3 to 20.0 accuracy points over the best open-loop schedule across five GSM8K draft-target pairs; learned triggers also transfer to StrategyQA and help on MATH-500 when early prediction and target competence align.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    Sparse collaboration depends on placing bounded target calls when the failure signal is reliable enough and early enough to affect the trajectory.

    Source: §7 · p. 8

    The trigger is task-, draft-, and distribution-dependent, and early prediction is weaker on MATH-500, making online intervention less reliable there.

    Source: §6 · p. 6, §8 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Sparse collaborative decoding uses a draft for most tokens and calls a stronger target for short corrective segments. The paper replaces fixed schedules and local uncertainty triggers with an online predictor of failure from the evolving confidence trajectory.

1 Introduction

  • Target calls face a timing tradeoff: earlier interventions have more leverage but weaker evidence, while later calls have stronger evidence but less ability to redirect.This motivates deciding where to intervene inside the trajectory rather than using fixed schedules or waiting until generation ends.Source: §1 · p. 1, Figure 1 · p. 1
  • Local uncertainty is a weak trigger: per-token uncertainty and velocity features provide near-chance prediction of final failure.A draft may be locally confident on a wrong path or uncertain while still recoverable.Source: §1 · p. 2
  • AUC 0.86: aggregate confidence-trajectory features predict failure, enabling learned target calls when risk crosses a threshold.The predictor uses draft correctness labels; target-generated corrective segments are bounded, after which the draft resumes.Source: §1 · p. 2
  • +5.3 to 20.0 accuracy points: learned triggering beats the best open-loop schedule across five GSM8K draft-target pairs.The paper also reports transfer to StrategyQA and gains on MATH-500 when early prediction and target competence align.Source: §1 · p. 2

2 Background and Related Work

Prior work allocates inference-time compute through longer reasoning, repeated samples, search, or adaptive budgets, while small-large collaboration often uses dense verification or whole-example routing. This paper instead studies when to invoke a stronger model during a draft trajectory.

2 Background and Related Work

  • Reasoning and adaptive-compute methods spend inference resources through longer outputs, repeated samples, branching trajectories, or allocation under cost constraints.Examples include chain-of-thought prompting, self-consistency, search-based methods, and adaptive-compute methods.Source: §2 · p. 2
  • Speculative decoding verifies draft tokens throughout generation, while cascades route whole examples before generation and posthoc rewriting acts after draft completion.Sparse collaborative decoding instead decides during generation using the evolving trajectory.Source: §2 · p. 2
  • Local uncertainty is widely used for selective prediction and adaptive inference, but the paper distinguishes token-level uncertainty from trajectory-level failure prediction.It argues aggregate confidence paths are more predictive in multi-step reasoning than isolated token signals.Source: §2 · p. 2

3 Problem Setup

The setup seeks higher final-answer accuracy while limiting target-model use. A trigger maps the observable draft state to a decision about whether to call the target during generation.

3 Problem Setup

  • The objective is to improve final-answer accuracy while limiting target use for a draft-generated reasoning trajectory.The setup considers a draft model Md and a stronger target model Mt.Source: §3 · pp. 2–3
  • At selected intervention points, the target generates a short text segment that is appended to the visible context, and the draft resumes from that context.No logits, hidden states, attention activations, or parameters are modified; the intervention is purely textual.Source: §3 · p. 3
  • A trigger policy maps the observable draft state, including token probabilities, entropy, margins, and their history, to a target-call decision.The policy should call when target help is likely to improve the final answer, not merely when the current token appears uncertain.Source: §3 · p. 3
  • The comparison covers seven policies and reports final-answer accuracy, target use, and—when available—wall-clock time, along with rescue and break rates.Policies include draft-only, target-only, open-loop and closed-loop SCD, cascade routing, posthoc rewriting, and self-consistency.Source: §3 · p. 3

4 Trajectory-Level Failure Prediction

Local token-level signals are weak predictors of final failure, whereas aggregate confidence-trajectory features provide a stronger signal. Prediction reliability improves with more trajectory tokens, creating a tradeoff between evidence and intervention leverage.

4 Trajectory-Level Failure Prediction

  • AUC 0.531 and 0.528: entropy and NLL are near chance for failure prediction on GSM8K with Qwen-1.5B.Velocity features are also uninformative, and a locally confident draft can still follow an incorrect path.Source: §4 · p. 3, Table 1 · p. 3
  • AUC 0.864: a lightweight logistic predictor over aggregate trajectory-shape features predicts failure from the confidence path.Its roughly 85 features summarize means, extrema, percentiles, slopes, variance, and position-dependent behavior; strong signals include margin trends and top-5 mass statistics.Source: §4 · p. 3
  • Incorrect trajectories show elevated early mid-frequency oscillations in logit-margin-derived wavelets, while correct trajectories consolidate more smoothly.This time-frequency diagnostic supports using aggregate trajectory shape rather than one-step uncertainty thresholds.Source: §4 · p. 3, Figure 2 · p. 4
  • AUC 0.60 at 16 tokens, 0.82 at 128 tokens, and 0.86 on the full trajectory: Qwen-1.5B prediction reliability rises with more observed tokens.Across drafts, GSM8K AUC@128 is near or above 0.7, while MATH-500 AUC@128 is around 0.57–0.59 despite full-trajectory AUC of 0.73–0.78.Source: §4 · p. 3, Table 2 · p. 4
  • Early prefix perturbations have greater influence on the final decoded sequence, while later positions provide stronger evidence but less ability to redirect.The measured flip rate—the fraction of examples whose final decoded sequence changes—decays with position for single-token and eight-token perturbations.Source: §4 · pp. 3–4

5 Learning Sparse Intervention Triggers

The section turns draft-trajectory features into an online target-call policy, using predicted failure risk to time bounded textual corrections. Its design balances evidence for triggering against the remaining opportunity to redirect generation.

5 Learning Sparse Intervention Triggers

  • A small draft-decoded calibration set supplies step-level statistics and final-answer correctness labels for training a logistic predictor on trajectory-shape features.At test time, the same feature extractor is applied to the observed trajectory.Source: §5 · p. 4
  • The trigger must balance evidence against intervention leverage: early calls can redirect more, while later calls have stronger evidence but less ability to change the trajectory.Segment length J trades cost for corrective reach, budget B limits repeated target use, and threshold θ controls selectivity.Source: Figure 3 · p. 4, §5 · p. 4
  • The controller calls the target when the trajectory-based risk score reaches threshold θ, subject to warmup W, cooldown C, and per-trajectory budget B.The threshold may be fixed or position-dependent.Source: §5 · p. 4
  • When triggered, the target generates a bounded text segment that is appended to the visible context, after which the draft resumes without modifying model internals.The intervention does not modify logits, hidden states, attention, or parameters.Source: §5 · p. 4
  • Table 3 reports open-loop SCD on GSM8K, with per-K marginals averaged over schedule positions P and best cells identifying the best position for each prefix length.Source: Table 3 · p. 5

6 Experiments

Experiments compare learned-trigger closed-loop sparse collaborative decoding (SCD) with fixed-schedule and other assistance policies across reasoning tasks. Learned triggering improves over open-loop on GSM8K and StrategyQA, while MATH-500 results depend on predictor reliability and target competence.

Open-loop comparison

  • Fixed target insertions already reach 86.5% for Qwen-1.5B and 75.0% for Llama-1B on GSM8K, but injecting J = 32 target tokens drops Llama-1B mean accuracy to 30.9%.Learned triggering retains the target-text insertion primitive while selecting which trajectories and positions receive help.Source: §6 · p. 5

Accuracy-cost frontier

  • Closed-loop beats the interpolated open-loop accuracy-cost frontier in 16 of 21 Qwen-1.5B cells and 13 of 21 Llama-1B cells; losses occur mainly when segments are too short to repair trajectories.Figure 4's cumulative BLEU traces show draft-only drift without recovery and rapid alignment recovery after triggered insertions.Source: §6 · p. 5, Table 4 · p. 5, Figure 4 · p. 6

Threshold sweeps

  • Position-dependent thresholds help only slightly within standard error for Qwen-1.5B, while aggressive early firing yields larger gains for weaker Llama-1B drafts.The thresholds are deployment knobs rather than universal constants; the sweep tests early and late thresholds at J = 96.Source: §6 · p. 6, Table 5 · p. 6

Generalization across tasks

  • +5.3 to +20.0 points over the best open-loop schedule across five GSM8K triples; on MATH-500, learned triggering wins three of five triples by +7.0 to +13.5 points but loses when prediction or target competence is weak.The two MATH losses occur for Qwen drafts paired with gpt-4o-mini; operational AUC@128 is 0.57–0.59, versus 0.76–0.82 on GSM8K.Source: §6 · p. 6, Table 6 · p. 6

Generalization to commonsense reasoning

  • +5.0 to +8.3 points over the best open-loop cell on all three StrategyQA triples, with closed-loop SCD improving Llama-1B from 48.3% draft-only accuracy to 68.7%.StrategyQA uses shorter trajectories than GSM8K, so the warmup window and baseline threshold sequence are scaled to its maximum sequence length.Source: §6 · p. 6, Table 7 · p. 7

Policy comparisons

  • On Llama-1B GSM8K with gpt-4o-mini, closed-loop SCD reaches 88.8% accuracy versus 87.8% for cascade and 93.8% for target-only, at 7.8s versus target-only's 6.8s.Cascade uses one 184.8-token call on average; closed-loop uses 2.27 calls averaging 86 tokens each, retaining draft computation while intervening in-trajectory.Source: Table 8 · p. 7, §6 · p. 7

Failure modes and deployment rule

  • Closed-loop SCD requires both a reliable early trigger and a bounded, reliable target continuation; swapping gpt-4o-mini for gpt-4o flips the same Qwen-0.5B MATH setup from a −5.5-point loss to a +7.0-point gain.Higher target-only accuracy alone is insufficient: longer or more uncertain interventions can overwrite draft context and degrade correct drafts.Source: §6 · pp. 7–8

Relation to speculative decoding

  • Speculative decoding is comparable only for open-weight Qwen pairs with shared tokenization and full target-logit access; SCD instead offers sparse text corrections for API targets that return text without full logits.The seven-policy evaluation also includes draft-only, target-only, self-consistency, posthoc rewrite, cascade, and open-loop SCD.Source: §6 · pp. 5, 8, Table 12 · p. 8

7 Conclusion

Sparse collaborative decoding uses a draft for most generation and short target corrections, with learned trajectory-shape signals guiding target-call placement. The results emphasize that success depends on reliable early prediction and bounded target continuations.

7 Conclusion

  • Aggregate trajectory-shape features provide a stronger failure signal than local uncertainty and support online placement of sparse target calls.The paper concludes that this is a trajectory-aware alternative to dense verification.Source: §7 · p. 8
  • Learned triggering transfers to StrategyQA when early prediction and target competence align, alongside gains over fixed schedules on GSM8K.The conclusion does not claim transfer is unconditional.Source: §7 · p. 8
  • Sparse collaboration requires both a reliable early trigger and bounded target continuations.Target computation is allocated only when it can change the reasoning trajectory.Source: §7 · p. 8

8 Limitations

The limitations span transfer across tasks and distributions, target-intervention behavior, and the scope of evaluation. The appendix also documents implementation settings and decoding procedures, including how open-loop and learned-trigger variants differ.

8 Limitations

  • The trigger is task-, draft-, and distribution-dependent, and exact-match reasoning benchmarks may not transfer directly to other domains.MATH-500 early prediction is weaker; answer extraction and normalization can affect symbolic datasets.Source: §8 · p. 9
  • In high-stakes settings, the authors call for external validation, calibrated abstention, and task-specific safety checks rather than treating SCD as a correctness guarantee.Interventions may increase trust in outputs that remain incorrect, and black-box target behavior can be difficult to audit.Source: §8 · p. 9

A Appendix Roadmap

  • The appendix covers implementation and decoding, predictor reliability, schedule sensitivity, cross-task generalization, target-side intervention quality, and evaluation, compute, and limitations.Source: A Appendix Roadmap · p. 10

B Implementation Details

  • Default settings use greedy draft decoding, bounded target segments, and a lightweight trigger trained on draft-side trajectory statistics and final-answer correctness labels.Table 13 lists the default decoding, triggering, and training settings.Source: B Implementation Details · pp. 10–11, Table 13 · p. 11

C Decoding Procedures and Baselines

  • OPEN-LOOP fixes a global insertion schedule, while CLOSED-LOOP evaluates a learned failure predictor online; both append target text to the visible context before draft generation resumes.Cascade and posthoc instead use the predictor after draft completion.Source: C Decoding Procedures and Baselines · p. 11
  • CLOSED-LOOP records draft-side signals, computes trajectory features and risk after warmup and cooldown checks, and calls the target when risk crosses a threshold and budget remains.The target generates a bounded segment, after which draft decoding resumes.Source: C Decoding Procedures and Baselines · p. 11

D Trajectory Features and Predictor Training

The appendix details how the predictor is built and how its reliability and intervention outcomes vary with trajectory length, schedule, and intervention frequency. Aggregate trajectory features are trained from draft traces and correctness labels, while calls that are too short, too late, or too frequent can limit gains.

D Trajectory Features and Predictor Training

  • The predictor summarizes 85-dimensional aggregate features from draft prefixes and trains an L2-regularized logistic model using draft traces and final-answer correctness labels.Failure signatures include declining separation, increased volatility, and weaker probability concentration rather than isolated uncertainty spikes.Source: D Trajectory Features and Predictor… · pp. 11–12

E Predictor Reliability Over Trajectory Length

F Open-Loop and Closed-Loop Sweeps

G Non-Monotonic Accuracy Under Increasing Target Interventions

  • Accuracy is non-monotonic in intervention count: a small number of well-timed calls can help, while additional calls may arrive too late or overwrite useful draft state.TAP prefix length shapes trajectory initialization; closed-loop SCD controls when target computation is reallocated during decoding.Source: G Non-Monotonic Accuracy Under… · pp. 12–13, Figure 7 · p. 13

H Attention Dynamics under Target Prefixing

The attention analysis examines how target-generated prefixes affect draft attention dynamics on GSM8K. Longer prefixes show a later shift toward concentrated attention, and the observed effect is not explained by persistent attention to prefix tokens.

H Attention Dynamics under Target Prefixing

I Full Cross-Task Generalization Grid

The cross-task grid finds the largest gains for weaker drafts and selective MATH-500 failures. Intervention quality also depends on target behavior and how well the draft continues from injected target text.

J Target Behavior and Intervention Reliability

I Full Cross-Task Generalization Grid

  • The full grid finds the largest gains for weaker drafts, while MATH-500 failures are selective and depend on target repair competence.The same draft can lose with a weaker or less bounded target yet recover with a target able to repair the fired context.Source: I Full Cross-Task Generalization Grid · p. 14, Table 21 · p. 15

J Target Behavior and Intervention Reliability

  • Degrading interventions can overwrite a partly correct plan, arrive after commitment to a wrong setup, or mismatch symbolic answer formats.The analysis motivates shorter target continuations, position-dependent thresholds, and task-specific answer normalization.Source: J Target Behavior and Intervention… · p. 15

K Draft Tuning for Target Followability

M Datasets, Prompting, and Evaluation

K Draft Tuning for Target Followability

L Speculative Decoding Comparison

  • Where speculative decoding is feasible, it approaches target-level accuracy through dense verification; SCD instead offers sparse text-correction points, including for black-box targets without logits.The comparison does not imply that SCD should replace speculative decoding when its assumptions hold.Source: L Speculative Decoding Comparison · p. 16, Figure 11 · p. 17

N Statistical Reporting and Compute

The section reports evaluation sample sizes, uncertainty guidance, and target-token cost reporting for the extended speculative-decoding comparison.

N Statistical Reporting and Compute

Loading 2610.23613v1…