Discrepancies between scores shapes between doc/generate methods
Author: icannos
Pull request: no

### System Info


- `transformers` version: 4.26.1
- Platform: Linux-6.0.19-4-MANJARO-x86_64-with-glibc2.37
- Python version: 3.10.9
- Huggingface_hub version: 0.12.0
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Tensorflow version (GPU?): 2.8.0 (True)
- 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?

@gante 

### Information

- [ ] The official example scripts
- [X] 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

I use the generate method to retrieve the probability distributions at each step of the generation. I rely on `scores` to get them.

I noticed discrepancies between the shape of `scores` between the doc, the actual behavior and between different generation configuration (differences that should not change anything, I believe).


The following snippet highlight the main problem:

```python
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")

string_inputs = [
    "translate in french: I love cats",
    "Answer the question: what is 3+5 ?"
]

inputs = tokenizer(string_inputs, padding=True, truncation=True, return_tensors='pt')

outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, num_beams=7, num_return_sequences=5)

print(type(outputs))
print(outputs.scores[0].shape)

# Output:
# <class 'transformers.generation.utils.BeamSearchEncoderDecoderOutput'>
# torch.Size([14, 32128])
# This is consistent with the doc and is the expected behavior: 14 = (batch_size=2)*(num_beams=7)
# https://huggingface.co/docs/transformers/internal/generation_utils#transformers.generation.BeamSearchEncoderDecoderOutput


outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, num_beams=7, num_return_sequences=5, do_sample=True)
print(type(outputs))
print(outputs.scores[0].shape)

# Output:
# <class 'transformers.generation.utils.BeamSampleEncoderDecoderOutput'>
# torch.Size([70, 32128])
# This is not consistent with the documentation seems not to be the expected behavior:
# 70 = (batch_size=2)*(num_beams=7)*(num_return_sequences=5)
# https://huggingface.co/docs/transformers/internal/generation_utils#transformers.generation.BeamSampleEncoderDecoderOutput

```

It is worth pointing out that in the documentation, the shape `batch_size*num_beams*num_return_sequences` is to be expected with DecoderOnlyOutput; however I don't understand why there would be a difference, nor what the `num_return_sequences`  does there anyway.


### Expected behavior

The shape of an element of scores should always be `(batch_size*num_beam, vocab_size)`, and not `(batch_size*num_beam*num_return_sequences, vocab_size)`. 

If I missed something and I'm, in fact mistaken, at the very least the documentation is not consistent with the actual behavior.

Comments:
Hey @icannos 👋 Here the problem is the documentation, which is incomplete :) 

Have a look at my answer in our forum, for a similar question -- https://discuss.huggingface.co/t/t5-why-do-we-have-more-tokens-expressed-via-cross-attentions-than-the-decoded-sequence/31893

Updating the docs is in our TODO list!

> Hey @icannos wave Here the problem is the documentation, which is incomplete :)
> 
> Have a look at my answer in our forum, for a similar question -- https://discuss.huggingface.co/t/t5-why-do-we-have-more-tokens-expressed-via-cross-attentions-than-the-decoded-sequence/31893

Thanks for your quick reply but I'm not sure it answers my problem. It is not about the length of the generated sequences but about the number of beams reported.

I believe the behavior should be identical in both the examples I gave. I don't understand why `num_return_sequences` is at all involved in the shape of scores. Moreover, it is not involved when there is not sampling but is there when sampling is involved.

Hey @icannos -- my apologies. I was multitasking and did not read your issue with proper attention. You're right, my answer does not address your question.

The difference in behavior you describe stems from these two lines: [beam search](https://github.com/huggingface/transformers/blob/82e61f34451dbea2de8d2220d51b0609d605dfd7/src/transformers/generation/utils.py#L1470) // [beam sample](https://github.com/huggingface/transformers/blob/82e61f34451dbea2de8d2220d51b0609d605dfd7/src/transformers/generation/utils.py#L1507). According to the code specified there, the behavior you see is correct (even if not properly documented!).

This means that `beam_sample` runs `num_return_sequences` independent beam searches for a given input (with the next token being selected from sampling, rather than simply the most likely token) and keeps the top beam in each one, as opposed to simply drawing the top `num_return_sequences` beams from one beam search, resulting in `num_return_sequences` times more `scores`. This is done to increase diversity in the output: the top `num_return_sequences` of a given beam search tend to be very similar to each other.

Does this answer your questions? 🤗

Yes I had found those lines, but I did think it was a mistake or misleading behavior.  

I don't know how it could be done, but I feel like it would be good that the behaviors were the same for all generation methods.

So I guess, if I want the sequences of probability distribution for each sequence I should reshape the scores to be `batch_size, num_return_sequences, num_beam, vocab_size` and retrieve `scores[:, :, 0, :]` right ?

Thank for your help! 


@icannos 

> I don't know how it could be done, but I feel like it would be good that the behaviors were the same for all generation methods.

Yeah... it's hard to balance a consistent API with consistent results. In this particular case, they are at odds with each other :( We favor consistent results in an attempt to maximize usefulness, as only advanced users tend to fiddle with these details -- how can a beginner know that gathering multiple beams from a single stochastic beam search is a bad idea? 😛 

> So I guess, if I want the sequences of probability distribution for each sequence I should reshape the scores to be batch_size, num_return_sequences, num_beam, vocab_size and retrieve scores[:, :, 0, :] right ?

Either that or slicing, yes 

Thanks a lot for your help, I'll close the issue.

Yeah I get the results driven API. I hope at least this issue might help someone in the future stumbling upon these problems.

@gante 
I would like to reopen this issue, I'm getting it even with greedy decoding (which should not run beam search AFAIK). 

When calling `output = model.generate(**input, output_scores=True, renormalize_logits=True, max_new_tokens=3, min_new_tokens=1, do_sample=False, return_dict_in_generate=True)`, on accessing `output['scores']`, I get 3 tensors, each of which are of shape (batch_size, 32128). This does not make much sense because flan-t5 only has 32100 tokens in its vocab. 

It's not clear where these extra scores are coming from, and which ones I should ignore?

Hey @adivekar-utexas 👋 

Many models run computations with embedding sizes larger than the actual vocabulary size, for speed purposes (e.g. [see here why](https://twitter.com/karpathy/status/1621578354024677377)). Any time you see an embedding size larger than the vocabulary size, you can safely discard the tokens whose index is beyond the vocabulary size :)

Thanks for the confirmation @gante !

For reference for others finding this issue, the necessary fix is to truncate the vocab dimension to `len(tokenizer.get_vocab())`, i.e. `scores_at_timesteps: List[Tensor] = [scores_at_timestep[:, len(tokenizer.get_vocab())] for scores_at_timestep in output['scores']]`

`tokenizer.vocab_size` sometimes does not take into account special tokes.