[Generation] Add static ensemble verification for lossy speculative decoding
Author: kasakh
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=45979)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45979)
<!-- ci-dashboard-badge:end -->

## What does this PR do?

Adds support for **static ensemble verification** in speculative decoding (assisted generation), a training-free method that increases draft token acceptance rates by relaxing the verification distribution.

Fixes #45865

### Motivation

Standard speculative decoding rejects many plausible tokens because it strictly enforces that the output distribution matches the target model. Static ensemble verification blends the target and draft distributions during verification:

```
v(x) = w * p_target(x) + (1 - w) * q_draft(x)
```

This provably achieves the Pareto-optimal tradeoff between acceptance rate and distributional bias (Proposition 1 in the paper). The acceptance probability increases from `1 - TV(q,p)` to `1 - w*TV(q,p)`.

### Changes

- Adds `assistant_ensemble_weight` parameter to `GenerationConfig` (float in (0, 1], default `None` = lossless)
- Modifies `_speculative_sampling` to use ensemble acceptance ratio: `v(x)/q(x) = 1 - w + w*(p(x)/q(x))`
- Supports greedy decoding (`do_sample=False`) by comparing against `argmax(v)` instead of `argmax(p)`
- Raises clear error when used with candidate generators that do not return logits (e.g., prompt lookup)
- Adds numerical stability guard for the fallback distribution
- Adds 8 fast synthetic unit tests
- Adds documentation section to `assisted_decoding.md`

### Usage

```python
# Sampling mode
outputs = model.generate(
    **inputs,
    assistant_model=assistant_model,
    do_sample=True,
    assistant_ensemble_weight=0.7,
)

# Greedy mode
outputs = model.generate(
    **inputs,
    assistant_model=assistant_model,
    do_sample=False,
    assistant_ensemble_weight=0.7,
)
```

### Design decisions

1. **Fallback distribution unchanged**: Since `[v-q]+ = w*[p-q]+` normalizes to the same distribution as `[p-q]+`, we compute the standard fallback directly for numerical stability.
2. **Acceptance ratio form**: Uses `1 - w + w*(p_i/q_i)` which is algebraically equivalent to `v_i/q_i` but avoids constructing the full ensemble distribution.
3. **Bonus token**: When all candidates are accepted, the bonus token uses `p` (target) rather than `v`, since `q` at position N+1 is unavailable. This is conservative and documented.
4. **Greedy ensemble**: For `do_sample=False`, compares draft tokens against `argmax(v)` — "greedy under the ensemble verifier."

### Tests

8 fast synthetic tests (all pass, ~0.15s total):
- `w=None`/`w=1.0` equivalence (backward compatibility)
- Deterministic acceptance ratio (w=0.7 accepts where w=1.0 rejects)
- Fallback distribution finite and normalized
- Numerical stability with near-zero residual
- Greedy ensemble: `argmax(v)` accepts where `argmax(p)` rejects
- ValueError for missing candidate logits
- GenerationConfig round-trip serialization
- Default value is None

### Results (from the paper)

On CNN/DailyMail with Llama-3.1-8B-Instruct (target) + Llama-3.2-1B-Instruct (draft), temperature=0:
- Standard SD (w=1.0): ~65% acceptance rate
- Static ensemble (w=0.7): ~78% acceptance rate, ROUGE-L within 0.5pt of target

### References

- Paper: [DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification](https://arxiv.org/abs/2604.07622) (AISTATS 2026)
- Code: https://github.com/comeusr/diversed
- Section 3.1 describes the static ensemble method

cc @Cyrilvallez

Comments:
Looks like AI code slop

In any case, I am quite reluctant to accept new feats in generation unless it is requested by community, so I will close this issue. If you are a real human, I recommend to host this as remote code via `custom generation strategies` (https://huggingface.co/docs/transformers/en/generation_strategies#custom-generation-methods)

That way we can see how much usage it has in the community, and consider shipping in core library

Hi @zucchini-nlp, I am a real human — I am one of the first authors (Siva Rajesh Kasa) of the DIVERSED paper (AISTATS 2026, https://arxiv.org/abs/2604.07622). I used AI tools to assist with drafting, which is permitted per the contributing guide, but I designed the approach and understand every line.

For context on how this PR came about: I opened issue #45865  first, @Rocketknight1 cc'd the generation team (https://github.com/huggingface/transformers/issues/45865#issuecomment-4421470190), and @Cyrilvallez explicitly encouraged me to submit a PR and pointed to the exact function to modify (https://github.com/huggingface/transformers/issues/45865#issuecomment-4436909551). @stevhliu also reviewed the docs here and left constructive inline feedback. I followed the process as intended.

Regarding community demand: speculative decoding is one of the most actively used features in transformers for inference acceleration, and lossy/relaxed verification has been a known gap. The feature request received positive engagement from multiple HF team members before this PR was submitted.

Would you be open to reconsidering, or could @Cyrilvallez weigh in given he approved the feature? I am happy to address any technical concerns or incorporate @stevhliu's docs suggestions.


Ah oke, I didn't see the context here, sorry. The PR looked like a AI agent one, so I closed it

I will re-open and lets wait for Cyril. TBH the idea is quite similar to the lenience factor from the very first paper from Leviathan and I do agree it might be of use for some ppl. My concern is only to keep the code clean, unless requested by several users, As you might know, maintaining cost grows with every new feature, and the field of speculative decoding currently has much more happening which we don't yet shipped in transformers. Anyway, yep, let's wait for more input from Cyril :)

Updated the PR to address @stevhliu's docs feedback (simplified docstring, nested heading, added recommended value, inlined paper reference) and fixed a bug where the error guard could fire incorrectly when candidate_logits is temporarily unavailable mid-generation. The ValueError now only triggers at config time for incompatible generators (prompt lookup).

Tested end-to-end on 6 model pairs across 4 families (Qwen, SmolLM, OPT, Pythia) — all show expected speedups with w=0.7.

Hi @Cyrilvallez — apologies for the bump. Whenever you have the bandwidth, would you mind taking a look at this PR? CI is green, and @stevhliu's docs feedback along with the candidate_logits bug fix from the @zucchini-nlp thread have all been addressed.

Also cc'ing @Rocketknight1 and @ArthurZucker who were originally tagged on the linked issue (#45865), in case any of you can point me to a more appropriate reviewer.

Thanks for all the work you do maintaining this library — happy to make any further changes whenever you get a chance to review.

Thanks for the review @Cyrilvallez! Addressed everything in 67ec46dfbf:

1. **Strict `0 < w < 1` validation moved into `GenerationConfig.validate()`** — `None` is the only lossless sentinel now; the boundary values `0.0` and `1.0` raise `ValueError` at config time. Updated the docstring to drop the "or `1.0` keeps decoding lossless" wording.
2. **Simplified runtime guards** — `_speculative_sampling` and the greedy verification path now just check `assistant_ensemble_weight is not None` (and `candidate_logits is not None` where relevant) since validate already enforces the range.
3. **Prompt-lookup guard now raises for any non-None weight** with the exact message you suggested.
4. **`float | None = None` annotation** on `_speculative_sampling`.
5. **Tests folded into existing generation test files** — `tests/generation/test_static_ensemble.py` is deleted. Four ensemble-sampling tests joined `UtilsFunctionsTest` in `tests/generation/test_utils.py` next to the existing `test_speculative_sampling*` methods, and two config-level tests (validate-rejects-invalid + default/round-trip) went into `tests/generation/test_configuration_utils.py`.

CI will rerun. Let me know if anything else needs tweaking.


Thanks for the fast turnaround on the second pass, @Cyrilvallez! All three suggestions were spot-on — the code reads much more cleanly now.

All 3 suggestions addressed in `6c41d1fcb4`:

1. **Numerical-stability guard scoped to the ensemble path** — `if assistant_ensemble_weight is not None and p_prime_sum <= torch.finfo(p_prime.dtype).tiny:`. Standard (lossless) SD keeps its original fallback behavior; the guard only fires on the new path when `p ≈ q` and the residual is numerically zero.
2. **Dropped `w = assistant_ensemble_weight` in `_speculative_sampling`** — the parameter name is used directly in the ratio expression.
3. **Same in the greedy verification path** — `nu_probs = assistant_ensemble_weight * p_probs + (1.0 - assistant_ensemble_weight) * q_probs`.

CI running now. Thanks again for the review time — more than happy to keep iterating if there's anything else.


<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=45979)
**Latest run:** [28978591668:2](https://github.com/huggingface/transformers/actions/runs/28978591668)
**Result:** `failure` | **Jobs:** 15 | **Tests:** 170,119 | **Failures:** 1 | **Duration:** 19h 57m

<!-- ci-dashboard-recap:end -->