transformers
753de4f6 - [Generation] Add static ensemble verification for lossy speculative decoding (#45979)

Commit
17 hours ago
[Generation] Add static ensemble verification for lossy speculative decoding (#45979) * feat(generation): Add static ensemble verification for lossy speculative decoding Add assistant_ensemble_weight parameter to GenerationConfig that enables static ensemble verification, a training-free method that increases draft token acceptance rates by relaxing the verification distribution. The verification distribution becomes v(x) = w * p(x) + (1-w) * q(x), where w is the ensemble weight. This provably achieves the Pareto-optimal tradeoff between acceptance rate and distributional bias. Changes: - Add assistant_ensemble_weight to GenerationConfig - Modify _speculative_sampling to use ensemble acceptance ratio - Support greedy decoding with argmax(v) in Case 2 - Add error for incompatible candidate generators (no logits) - Add numerical stability guard for fallback distribution - Add 8 fast synthetic unit tests - Add documentation section to assisted_decoding.md Reference: Wang & Kasa et al., "DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification", AISTATS 2026. https://arxiv.org/abs/2604.07622 Fixes #45865 * Address Cyril's review comments - Enforce 0.0 < assistant_ensemble_weight < 1.0 strictly in GenerationConfig.validate() (boundaries and out-of-range values now raise at config time). - Simplify runtime guards in _speculative_sampling and the greedy verification path to rely on the new validate() check instead of repeating the < 1.0 condition inline. - Tighten the prompt-lookup guard: any non-None weight now raises (previously only raised for w < 1.0). - Use modern type annotation 'float | None = None' on _speculative_sampling. - Remove '1.0 keeps decoding lossless' from the docstring; only None is lossless now. - Move tests from tests/generation/test_static_ensemble.py into existing files: - 4 sampling-behaviour tests into UtilsFunctionsTest in tests/generation/test_utils.py - 2 config tests into tests/generation/test_configuration_utils.py - Delete tests/generation/test_static_ensemble.py. Addresses review at https://github.com/huggingface/transformers/pull/45979#pullrequestreview-3490039503 * style: ruff format new test methods * Address Cyril's follow-up nits - Drop intermediate 'w = assistant_ensemble_weight' aliases in both _speculative_sampling and the greedy verification path; inline the parameter name directly for readability. - Scope the numerical-stability guard on 'p_prime_sum' to the ensemble path only (when assistant_ensemble_weight is not None). Standard (lossless) SD keeps its original fallback behavior. --------- Co-authored-by: Siva Rajesh Kasa <kasasiva@amazon.com> Co-authored-by: Siva Rajesh Kasa <kasakh@users.noreply.github.com>
Author
Parents
Loading