Unexpected usage of `next_token_scores` and `beam_scores` in `beam_sample()`
Author: jmzhao
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.2
- Platform: Linux
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu111 (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?: No
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- T5, BART, Marian, Pegasus, EncoderDecoder: @patrickvonplaten
- Blenderbot, MBART: @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet, BigBird: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder, UniSpeech, UniSpeechSAT, SEW, SEW-D, Speech2Text: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten @narsil
- Tokenizers: @LysandreJik
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


@patrickvonplaten 

## Information

Not sure if the `beam_sample()` function called by `generate()` is using `next_token_scores` and `beam_scores` as intended for sampling and score accumulation. 

Model I am using (Bert, XLNet ...): distilgpt2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

```python3
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "distilgpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "Cats meow, dogs"
input_ids = tokenizer.encode(text, return_tensors="pt")
input_len = input_ids.size()[-1]
max_new_tokens = 3
temperature = 0.2

sample_outputs = model.generate(
    input_ids,
    max_length=input_len + max_new_tokens,
    do_sample=True,
    num_beams=5,
    top_k=0,
    temperature=temperature,
    return_dict_in_generate=True,
    output_scores=True,
)

# Expect the generated sequence to be "Cats meow, dogs, and cats".
tokenizer.batch_decode(sample_outputs.sequences)  # => ['Cats meow, dogs, and cats']
# The generated tokens are picked as the ones with the top `beam_scores`.
top_score_ids = [score.argmax().item() for score in sample_outputs.scores]
generated_token_ids = sample_outputs.sequences[0][input_len:].tolist()
top_score_ids == generated_token_ids  # => True

beam_score = [score.max().item() for score in sample_outputs.scores]
beam_score  # => [-6.53633451461792, -41.86376953125, -221.3215789794922]

log_probs = []
warped_log_probs = []
for i in range(max_new_tokens):
    outputs = model(sample_outputs.sequences[:, : input_len + i])
    sampled_id = sample_outputs.sequences[0][input_len + i]
    next_token_logits = outputs.logits[0][-1].detach()
    log_probs.append(next_token_logits.log_softmax(-1)[sampled_id].item())
    warped_log_probs.append(
        (next_token_logits / temperature).log_softmax(-1)[sampled_id].item()
    )
log_probs  # => [-1.3072705268859863, -1.8364229202270508, -2.400545835494995]
warped_log_probs  # => [-0.005173391196876764, -0.08735305070877075, -0.016386810690164566]

def approx_eq_list(l1, l2):
    return all(abs(x / y - 1) < 1e-4 for x, y in zip(beam_score_1[1:], beam_score))

# The expected `beam_score` based on Algorithm (1).
beam_score_1 = [0]
for p in log_probs:
    beam_score_1.append((beam_score_1[-1] + p) / temperature)
beam_score_1  # => [0, -6.536352634429932, -41.86387777328491, -221.32211804389954]
approx_eq_list(beam_score_1[1:], beam_score)  # => True
# This is fine if Algorithm (1) is intended.

# The expected `beam_score` based on Algorithm (2).
beam_score_2 = [0]
for p in warped_log_probs:
    beam_score_2.append(beam_score_2[-1] + p)
beam_score_2  # => [0, -0.005173391196876764, -0.09252644190564752, -0.10891325259581208]
approx_eq_list(beam_score_2[1:], beam_score)  # => False
# Expect this to be True if Algorithm (2) is intended.

sample_outputs.sequences_scores.tolist()  # => [-24.591285705566406]
# (?) Seems unrelated to none of the values above.
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

I found that `logits_warper`[ is applied after adding ](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_utils.py#L2227)`beam_scores`, within which [the temperature is handled](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_logits_process.py#L127). This way, we sample the next tokens using the beam_scores *plus* the current next_token_scores. The temperature term will telescope along the sequence, instead of applying to just the next token distribution. 

That is, let `next_token_logits[i]` be the model's output logits at the `i`-th new token, `beam_score[i]` be the beam score after the `i`-th generated token, `beam_score[0] = 0`, and `sequence[i]` be the id of the `i`-th generated token. The current implementation seems to be doing `beam_scores[i] = (next_token_logits[i].log_softmax(-1) + beam_score[i-1]) / temperature`, and use `beam_scores[i]` to [sample next tokens](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_utils.py#L2251) to get `sequence[i]` and `beam_score[i] = beam_scores[i][sequence[i]]`. The resulting `beam_score[num_new_tokens - 1]` would be equal to `next_token_logits[i].log_softmax(-1)[sequence[i]] / temperature ** (num_new_tokens - i)` summing over `i = 0 ... num_new_tokens - 1`. Let's call this Algorithm (1).

I'd like to know if this is the intended way of using temperature term for beam search. Based on what I understand from reading posts and papers, it seems like the temperature term should be applied to just the probability of each of the next tokens. That is, to use `(next_token_logits[i] / temperature).log_softmax(-1)` to sample next tokens to get `sequence[i]`, and set `beam_scores[i] = (next_token_logits[i] / temperature).log_softmax(-1) + beam_score[i-1]` and `beam_score[i] = beam_scores[i][sequence[i]]`. The resulting `beam_score[num_new_tokens - 1]` would be equal to `(next_token_logits[i] / temperature).log_softmax(-1)[sequence[i]]` summing over `i = 0 ... num_new_tokens - 1`. Let's call this Algorithm (2).

In code, this would mean to apply `log_softmax` after the [`logits_warper`](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_utils.py#L2227) and delay [the addition of `beam_scores`](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_utils.py#L2226) to [right before passing it to `beam_scorer`](https://github.com/huggingface/transformers/blob/27b3031de2fb8195dec9bc2093e3e70bdb1c4bff/src/transformers/generation_utils.py#L2263). 

Algorithm (1) is also likely the cause for [issue#11267](https://github.com/huggingface/transformers/issues/11267), where telescopic `/ temperature` leads to all `-inf` `beam_scores` at some point, which leads to the all `-inf` `next_token_scores`, which leads to all `nan` `probs` tenor, which leads to the reported "RuntimeError: probability tensor contains either `inf`, `nan` or element < 0".

Comments:
Hey @jmzhao,

Thanks for the in-detail description. What you are saying is correct and we have chose to stick to this algorithm simply because that's how we implemented it the first time and we didn't want to break backwards compatibility anymore afterwards. The problem with `beam_sample()` IMO is that there is no clear consensus on how to exactly implement this algorithm. Only very few people seem to use it in the first place because the other sampling algorithms like `beam_search`, `sample`, and `greedy_search` are very well-defined and mathematically sound algorithms. 

In case you know of a paper that clearly defines beam sampling or even better the paper that created beam sampling I think we could change the algorithm here. IMO, there is no clear advantage why Algorithm 1 should be better than Algorithm 2. E.g., if we follow Algorithm 2, it would mean that previous beam scores are not taken into account when shaping the probability distribution of the next word, *e.g.* topk and topp would not take into account any previous beam scores in which case the algorithm would be very similar to just the `sample()` algorithm. 

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/master/CONTRIBUTING.md) are likely to be ignored.