PaperlayerView exact revision on arXiv

Source-linked AI summary

Multiple Importance Sampling for Guided Speculative Inference

Joseph Rowan, Ashish J. Khisti

arXiv:2610.33368v1cs.LGcs.ITmath.IT

TL;DR

The paper addresses GSI’s reliance on a single drafter whose support may not cover the target. It extends GSI with an MIS-weighted ensemble and establishes theoretical bounds; the approach maintains competitive pass-at-N rates with comparable inference speed.

  • Problem

    GSI relies on a single drafter whose support may not cover the target distribution, making the coverage assumption a concern.

    Source: §I · p. 1

  • Method

    The paper extends GSI to an ensemble of draft models and uses MIS weighting to aggregate proposal samples.

    Source: §I · p. 1, §III · p. 2

  • Results

    The approach maintains competitive pass-at-N rates with comparable inference speed and provides bounds on KL divergence and expected reward degradation.

    Source: A. Theoretical Analysis · p. 3, §V · p. 5

  • Takeaways & Limitations

    MIS allows the approach to bypass the requirement that one drafter’s support fully cover the target distribution.

    Source: §V · p. 5

    The fixed uniform sample allocation is considered more practical; adaptive allocation is difficult when divergence is intractable and proposal and target distributions vary across reasoning steps.

    Source: C. Optimizing the Sample Allocation · p. 4

Abstract

from arXiv · show

We study an extension of the Guided Speculative Inference (GSI) algorithm for test-time-scaling of large language models (LLMs) to a novel setting where an ensemble of smaller LLMs proposes draft reasoning steps which are later verified by a larger base model and a process reward model. Using the established framework of multiple importance sampling (MIS), we introduce weighting schemes to aggregate samples from a diverse collection of proposal distributions. We provide new bounds on the KL divergence from the optimal reward-tilted target distribution and the expected reward degradation, then demonstrate that the well-known balance heuristic is optimal. Compared to GSI with a single well-tuned drafter, combining several diverse LLMs while retaining the same total number of samples maintains competitive pass-at-$N$ downstream performance while being mathematically robust to situations where the support of a single draft model does not cover that of the target, where GSI may fail.

I. INTRODUCTION

Test-time scaling adds inference-time computation to improve already-trained models, while speculative methods use cheaper drafts with verification. This paper motivates multiple draft models and introduces the reasoning and reward setup behind that approach.

I. INTRODUCTION

  • Test-time scaling increases inference-time computation, and methods are categorized as sequential, parallel, or hybrid.Hybrid methods generate candidates at each iteration and use search to choose a promising path.Source: §I · p. 1
  • An ensemble of draft models with MIS weighting relaxes GSI’s single-drafter coverage requirement to coverage by the union of drafter supports.The paper motivates this extension as a way to improve robustness when a single drafter is poorly chosen.Source: §I · p. 1

A. Hard BoN Sampling

  • Hard BoN samples reasoning steps from the model and uses a process reward model to select the highest-reward sample at each step.Its scaling is constrained by reward hacking: an imperfect reward model can favor low-quality responses, with greater susceptibility as N grows.Source: A. Hard BoN Sampling · p. 2

B. Soft BoN Sampling

Soft BoN uses importance resampling to approximate the reward-tilted distribution, selecting among samples from the base model according to their weights. Unlike hard BoN, it has a stated KL-divergence guarantee.

B. Soft BoN Sampling

  • The target policy is the reward-tilted distribution obtained from a KL-regularized reward objective that keeps the policy close to the base model.The regularization strength is controlled by β; the paper notes that exact sampling from this distribution is generally unavailable.Source: B. Soft BoN Sampling · p. 2
  • Soft BoN draws N reasoning steps from the base model and probabilistically selects one by importance resampling.The samples receive weights, and selection probability is proportional to those weights.Source: B. Soft BoN Sampling · p. 2
  • Soft BoN’s KL-divergence guarantee distinguishes it from hard BoN, which does not have such a guarantee.The importance-resampling step can be viewed as sampling from a temperature-weighted softmax distribution over rewards.Source: B. Soft BoN Sampling · p. 2

C. Guided Speculative Inference

GSI drafts with a smaller model and verifies samples with a larger base model, but its coverage assumption can fail. The paper extends this resampling approach to multiple drafters using MIS weights.

C. Guided Speculative Inference

  • GSI uses a smaller draft model to propose samples that a larger base model verifies in parallel, with resampling weights intended to approximate the tilted distribution.The cited coverage assumption is supp πb ⊆ supp πs; without it, this approximation is not assured.Source: C. Guided Speculative Inference · p. 2

III. SOFT BEST-OF-N WITH MULTIPLE DRAFT MODELS

  • The paper extends GSI to an ensemble of draft models through multiple importance sampling (MIS).The motivation includes prompts that may favor different drafters, such as models fine-tuned on math or code.Source: §III · pp. 2–3
  • The MIS procedure samples reasoning steps from draft models, computes MIS and resampling weights, then selects a step probabilistically.The method directly extends GSI’s resampling strategy using standard MIS resampling weights.Source: §III · p. 3

A. Theoretical Analysis

The theoretical analysis establishes bounds for valid MIS schemes, addressing both closeness to the reward-tilted distribution and expected reward degradation. These results are presented as guarantees analogous to those for soft BoN and GSI.

A. Theoretical Analysis

  • The paper frames the MIS guarantees as analogous to existing bounds for soft BoN and GSI.The stated objective is to bound the output’s closeness to the optimal tilted distribution.Source: A. Theoretical Analysis · p. 3
  • For any valid MIS scheme, Theorem 1 gives a bound on KL divergence from the reward-tilted distribution.The analysis first establishes a lemma concerning the output distribution of a valid MIS scheme.Source: A. Theoretical Analysis · p. 3
  • Theorem 2 bounds expected reward degradation relative to sampling directly from the optimal tilted distribution.The bound is obtained by applying Pinsker’s inequality to Theorem 1’s KL bound.Source: A. Theoretical Analysis · p. 3

B. Choices for the MIS Weights

The section compares MIS weighting choices, including uniform weights and the balance heuristic. Uniform weights are computationally simple but require a stricter coverage assumption to ensure a valid MIS scheme.

B. Choices for the MIS Weights

2) The Balance Heuristic:

The balance heuristic, also called scheme N3, uses proposal probabilities to weight samples and yields a KL-divergence upper bound at least as tight as uniform weights. With equal sample counts from each draft model, it minimizes that bound among valid MIS weights.

2) The Balance Heuristic:

  • The balance heuristic uses each sample’s probability under its generating proposal, requiring evaluation under every proposal distribution.Source: 2) The Balance Heuristic · p. 3
  • The balance heuristic’s upper bound is at least as tight as the uniform-weight scheme’s, by convexity of the χ2 divergence.Source: 2) The Balance Heuristic · p. 4
  • With equal samples from each draft model, the balance heuristic minimizes the KL-divergence bound over all valid MIS weights.Source: 2) The Balance Heuristic · p. 4
  • Minimizing the upper bound does not establish that the balance heuristic minimizes the KL divergence itself.Source: 2) The Balance Heuristic · p. 4

C. Optimizing the Sample Allocation

The allocation strategy permits different numbers of samples from each proposal, treating the allocation as a prior on drafter effectiveness. Although the bound can be optimized over mixture weights, the authors identify fixed uniform allocation as more practically relevant in many settings.

C. Optimizing the Sample Allocation

  • Taking different numbers of samples from each proposal can encode a prior on which drafter is more effective for a prompt.A router could select a drafter suited to the question, such as one fine-tuned on mathematical problems.Source: C. Optimizing the Sample Allocation · p. 4
  • For Nj samples from proposal pj, the mixture weight is αj = Nj/N, and the KL upper bound is minimized by minimizing the mixture’s χ2 divergence from the base model.Source: C. Optimizing the Sample Allocation · p. 4
  • Fixed uniform allocation is likely more practical because the χ2 divergence is not exactly computable, distributions vary across reasoning steps, and latency often depends on the largest batch.With drafters of similar size, that largest batch typically favors uniform allocation.Source: C. Optimizing the Sample Allocation · p. 4
  • Varying sample allocation may be useful in heterogeneous environments where drafters have different computational costs.The authors leave this setting for future work.Source: C. Optimizing the Sample Allocation · p. 4

IV. EXPERIMENTS

Across MMLU-Pro and MATH-500, MIS strategies using one-third as many samples as single-model GSI achieve comparable pass-at-N rates, and MIS-N3 improves tilted reward on MMLU-Pro. MIS-N3 with the balance heuristic adds 3–10% runtime overhead compared with GSI, while the conclusion reports theoretical guarantees and robustness to incomplete draft support.

IV. EXPERIMENTS

  • Experiments use Qwen2.5-Math 7 B as target, three 1.5 B draft models for MIS, and a fine-tuned 7 B PRM on 100 prompts from each reasoning dataset.Responses are capped at 40 steps and 4096 tokens; GSI uses Qwen2.5-Math 1.5 B as its drafter.Source: §IV · p. 4

A. Hyperparameter Selection

  • The experiments fix KL regularization at β = 20 and use a rejection threshold α = 0.5 for GSI, MIS-N1, and MIS-N3.Draft steps below the threshold are rejected in favor of vanilla soft BoN sampling from the target model.Source: A. Hyperparameter Selection · p. 5

B. Results on Reasoning Tasks

C. Comparative Runtime

  • 3–10% is the reported range of per-step runtime overhead for MIS compared with GSI in the most computationally expensive settings.The comparison covers N = 96 and N = 192; the authors describe MIS-N3 with the balance heuristic as viable despite additional proposal evaluations.Source: C. Comparative Runtime · p. 5

V. CONCLUSION

  • The conclusion reports new KL-divergence guarantees and says the ensemble approach avoids requiring a single drafter to cover the target distribution’s support.It identifies routers, other sources of drafter diversity, and domains beyond multi-step reasoning as directions for future work.Source: §V · p. 5

APPENDIX

The proof of Lemma 1 computes the probability that a sampled candidate is selected, focusing first on candidates with positive proposal probability. It then generalizes across candidates under the conditions for a valid MIS scheme.

A. Proof of Lemma 1

  • The proof begins by computing the probability that X1 = x is selected as output, then focuses on the event S = 1.S(x) is introduced as the set of proposal indices with positive probability at x.Source: A. Proof of Lemma 1 · p. 7
  • The argument generalizes from S = 1 to any proposal index and substitutes the result into equation (A4).It invokes Jensen’s inequality and the valid-MIS conditions to handle indices outside S(x).Source: A. Proof of Lemma 1 · p. 7
  • For indices outside S(x), the proof sets the corresponding MIS terms to zero, allowing the summation to extend over all proposal indices without changing the result.The text states this extension yields the next expression in the proof.Source: A. Proof of Lemma 1 · p. 7

B. Proof of Theorem 1

The proof of Theorem 1 bounds the KL divergence by applying the Lemma 1 representation and Jensen’s inequality. It then uses the valid-MIS conditions and a weighted arithmetic-harmonic mean inequality to obtain the bound.

B. Proof of Theorem 1

  • The proof starts from the KL divergence and restricts the sum to proposal indices with positive probability at each x.The passage identifies this set as S(x) and then applies Lemma 1 and Jensen’s inequality.Source: B. Proof of Theorem 1 · p. 8
  • The proof expands a nonnegative inner sum over all x and invokes the first canonical requirement for a valid MIS scheme.These steps provide the bound used in the KL-divergence argument.Source: B. Proof of Theorem 1 · p. 8
  • A weighted arithmetic-harmonic mean inequality yields the bound even when some MIS weights are zero, provided the positive weights form a partition of unity on the target support.The proof uses e^βr(x) ≥ 1 and adopts a zero-over-zero convention for the final expression.Source: B. Proof of Theorem 1 · p. 8

C. Proof of Theorem 2

The proof of Theorem 2 bounds total variation using its variational representation and a function built from the bounded reward. Pinsker’s inequality then converts the Theorem 1 KL bound into the desired result.

C. Proof of Theorem 2

  • The proof begins with the variational representation of total variation between measures P and Q.It defines the function class as functions on X* with supremum norm at most 1.Source: C. Proof of Theorem 2 · p. 9
  • Pinsker’s inequality converts the total-variation bound into the desired bound.The proof then bounds D_KL(π*∥π̃) using Theorem 1.Source: C. Proof of Theorem 2 · p. 9

Related papers

Loading 2610.33368v1…