Why is assisted decoding slower than the paper reported peformance?
Author: cliangyu
Pull request: no

### System Info

- `transformers` version: 4.34.0.dev0
- Platform: Linux-6.2.0-26-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - deepspeed_config: {'gradient_accumulation_steps': 1, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': False, 'zero_stage': 2}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@gante

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1.We are not sure why the speculative decoding boosting performance as reported in the [blog](https://huggingface.co/blog/assisted-generation) falls largely behind the reported values in the DeepMind and Google papers.
2. The HF implementation uses the argmax token of the large model probability. However, in the papers the authors sampled (instead of argmax). Why is it so? Because argmax is faster than sampling?

### Expected behavior

The speculative decoding should be as fast as the reported in papers.

Comments:
I'm also confused about this, specially the following piece:

```python
# 4. Compare the argmax from the original model logits with the assistant forecasted tokens. We can keep
# the assistant forecasted tokens until the first mismatch, or until the max length is reached.
candidate_new_tokens = candidate_input_ids[:, -candidate_length:]
n_matches = ((~(candidate_new_tokens == selected_tokens[:, :-1])).cumsum(dim=-1) < 1).sum()
```

While it determines `n_matches` based on **the equality between tokens**, the speculative decoding paper determines it based on **the value of logits**.

To highlight the difference, we can imagine that we choose the original model itself as the assistant model, so the logits from original model and assistant model are always the same. In this case, The output from assistant model will always be accepted if we implement speculative decoding following the paper, and will be likely rejected because of the randomness in sampling if we implement it following the hf way.

Hey @cliangyu @daquexian 👋 

First of all, let me start the discussion by highlighting that assisted generation was originally designed without the knowledge of speculative decoding -- our internal development started way before speculative decoding got onto our radar :) The two techniques rest on the same principle, that GPU memory bandwidth is the bottleneck, and using a smaller model to generate candidate sequences can alleviate that. 

We decided to follow this path because it was the simplest to bring the concept of a smaller model assisting a larger one to life. We also dynamically adapt the query size to the assistant model, where speculative decoding does not -- this is because we noticed high variability of the optimal speculation size across tasks. These are the two main differences between the methods. As a result, speculative decoding is superior for tasks with sampling (because of the better sampling strategy), assisted generation is better for greedy approaches (because of the dynamic size of the assistant query).

Because speculative decoding is indeed better with sampling (and thus with most LLM tasks), we want to bring it to our assisted generation + sampling code path. I'm out of bandwidth for the next weeks, but we will gladly accept contributions! 🤗 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.