Dynamic number of speculative tokens in order to accelerate speculative decoding
Author: jmamou
Pull request: yes
Labels: name: Generation

# What does this PR do?

This PR adds support to dynamic number of _speculative tokens_ in order to accelerate speculative decoding.
It is an unsupervised version of the dynamic speculation lookahead from [Dynamic Speculation Lookahead Accelerates Speculative Decoding of Large Language Models](https://arxiv.org/abs/2405.04304).
 
We add the argument `assistant_confidence_threshold` to the generation configuration of the model. It is a confidence threshold for the assistant model. If the assistant model's confidence in its prediction for the current token is lower than this threshold, the assistant model stops the current token generation iteration, even if the number of _speculative tokens_ (defined by `num_assistant_tokens`) is not yet reached.

## Before submitting
- [X ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ V] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ V] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ V] Did you write any new necessary tests?

## Who can review?

@gante @amyeroberts

Comments:
> Thank you for opening the PR @jmamou 🤗 The idea in the paper is really cool, much better than my scrappy heuristic! 🔥
> 
> I've added a few nits in the PR review that should be trivial to solve.
> 
> A proposal for a follow-up PR: Different flavors of speculative decoding/assisted generation are (mostly) invisible to the user other than execution speed. According to your paper, DISCO should be faster than the current default. Do you think you can find a single default value for `assistant_confidence_threshold` such that it beats the current default in most situations? If so, I'd be more that happy to make DISCO the default 🤗

Thanks @gante for your feedback. 

Concerning the follow-up PR, I have run experiments with vicuna-13b/vicuna-68m as target/assistant on the datasets Alpaca and CNN-DM with do_sample True and False. The values of `assistant_confidence_threshold=0.4` and `num_assistant_tokens=20` are consistently almost optimal. 

> Thank you for iterating! To make our CI happy, run `make fixup` and push the changes :)

sure!
I am done, except for the test "examples_tensorflow" that does not seem to be related to my code. 

> Thank you for iterating! To make our CI happy, run `make fixup` and push the changes :)

sure!
I am done, except for the test "examples_tensorflow" that does not seem to be related to my code. 

> The values of assistant_confidence_threshold=0.4 and num_assistant_tokens=20 are consistently almost optimal.

@jmamou feel free to open a PR to update the defaults, I'd be glad to accept it! 🤗