Scores returned by beam search are not useful
Author: felix-schneider
Pull request: no

## Environment info

- `transformers` version: 4.12.3
- Platform: Linux-3.10.0-1160.45.1.el7.x86_64-x86_64-with-centos-7.9.2009-Core
- Python version: 3.7.11
- PyTorch version (GPU?): 1.9.0 (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?: yes
- Using distributed or parallel set-up in script?: yes
- 
### Who can help
- Text generation: @patrickvonplaten

## Information

The typical use-case why one might specify `output_scores` when using `model.generate` is to get the scores of each of the tokens returned in `sequences` (see #14086) . However, that is not what is being returned.

@patrickvonplaten gives an example [here](https://discuss.huggingface.co/t/generation-probabilities-how-to-compute-probabilities-of-output-scores-for-gpt2/3175) about how to calculate the scores, but this is not correct.

In beam search, it is false to assume that the token `sequences[i, j]` was chosen from the probability distribution `scores[j-1][i]`. It could have come from any of the other beams. Therefore, to get the actual per-token probabilities, one would have to completely reimplement beam search, similar to #14065 (although as I point out below, that solution is also not correct).

For sampling search, finding the per-token probabilities after the fact is completely impossible.

The returned scores right now are only useful for visualizing the beam search algorithm and not for the use-case that most users would want, i.e. to know the probabilities for each token in `sequences`.

## To reproduce

Adapted from [Patrick's example](https://discuss.huggingface.co/t/generation-probabilities-how-to-compute-probabilities-of-output-scores-for-gpt2/3175)

```python
import transformers
import datasets
import torch

model = transformers.AutoModelForSeq2SeqLM.from_pretrained("lidiya/bart-base-samsum")
tok = transformers.AutoTokenizer.from_pretrained("lidiya/bart-base-samsum")
ds = datasets.load_dataset("samsum")["validation"]
bsz = 10
num_beams = 4
batch = tok(ds[:bsz]["dialogue"], padding=True, truncation=True, return_tensors="pt")
output = model.generate(
    **batch, num_beams=num_beams, return_dict_in_generate=True, output_scores=True
)
tokens = output.sequences
seq_len = tokens.size(1)
assert list(tokens.size()) == [bsz, seq_len]
# No softmax is necessary, the scores are already logprobs (if the doc is correct)
scores = torch.stack(
    output.scores, dim=1
)  # Not dim=-1 so the dimensions are in the same order as in tokens
vocab_size = scores.size(2)
assert list(scores.size()) == [bsz * num_beams, seq_len - 1, vocab_size]

per_token_scores = torch.gather(
    # Step size of num_beams, so at least the gather
    # fetches probabilities from the correct batch index
    input=scores[::num_beams],
    # beam_search also returns the start token even though it is not an output from the model
    # that token has no probability, so we need to start at sequence index 1
    index=tokens[:, 1:, None],
    dim=2,
).squeeze(2)

# This result has the correct shape: (bsz, seq_len-1),
# but it is *not* actually the cumulative probabilities for each token in tokens,
# because tokens[0, i] did not necessarily come from scores[0, i-1],
# it could have come from another beam that assigned it a higher probability.
# Without reimplementing beam search, we cannot get the per-token probabilities.


# To actually get the per-token probabilities (without reimplementing beam search):
logits = model.forward(
    **batch,
    decoder_input_ids=tokens,
    decoder_attention_mask=tokens.ne(tok.pad_token_id)
).logits
actual_per_token_scores = torch.gather(
    input=logits.log_softmax(-1), index=tokens[:, 1:, None], dim=2
).squeeze(2)

assert not torch.allclose(per_token_scores, actual_per_token_scores.cumsum(1))


# Two naive possibilities to get the probabilities which will work often, but not always:
# The reason is that the same token may be chosen by two different beams, so without
# more history (and effectively a full re-implementation of beam search), we cannot get an accurate
# result in that case.

# To actually get the per-token probabilities
max_scores = scores.view(bsz, num_beams, seq_len - 1, vocab_size).max(1).values
maybe_token_scores = torch.gather(
    input=max_scores, index=tokens[:, 1:, None], dim=2
).squeeze(2)

# Alternative
maybe_token_scores = []
for i in range(1, seq_len):
    values, indices = torch.topk(
        scores[:, i - 1].view(bsz, num_beams * vocab_size), k=num_beams, dim=1
    )
    selected_tokens = indices % vocab_size
    assert list(selected_tokens.size()) == list(values.size()) == [bsz, num_beams]
    # Multiply instead of mask select because mask select flattens the array
    maybe_token_scores.append((values * selected_tokens.eq(tokens[:, i])).max(1).values)
maybe_token_scores = torch.stack(maybe_token_scores, dim=1)


# Right now the only definitive way to get per-token probabilities
logits = model.forward(
    **batch,
    decoder_input_ids=tokens,
    decoder_attention_mask=tokens.ne(tok.pad_token_id)
).logits
actual_per_token_scores = torch.gather(
    input=logits.log_softmax(-1), index=tokens[:, 1:, None], dim=2
).squeeze(2)

assert not torch.allclose(per_token_scores, actual_per_token_scores.cumsum(1))
```

## Expected behavior

Either `scores` or a new field in the beam search output should return the probability for each token in `sequences`. Because of padding, these sequences would be of different length.

Comments:
For anyone else looking for the same thing: I edited the original post to add some different ways to get the actual probabilities. I think the forward pass (the last option) is the only one guaranteed to be correct, but also the slowest.

It would be much easier if we just got this output out of beam search.

@felix-schneider thanks for your write-up here. Could you give more detail regarding the statement

> "but this is not correct"

In the example of the forum post, I do the following:

1. Return all output logits at every step (which can be seen as `softmax^(-1)(prob)`)
2. Compute the softmax to get the probability that of token i to be sampled at time-step t
3. Gather the probabilities of the generated tokens
4. Take the product of the probabilities to get the overall probability 

=> which step is not correct here? And how?

Note that the example is just for sampling - not for beam search

For beam search it is true that this operation is not that simple because it's hard to know from which beam the output tokens are coming from at each time step - one could compute this by taking the `topk` values of the `scores` output.
However in order to facilitate the computation for the beam scores I'd be happy to also return the `beam_idx` to make it easier to map the generated tokens to the corresponding score.

See: https://discuss.huggingface.co/t/generation-probabilities-how-to-compute-probabilities-of-output-scores-for-gpt2/3175/15?u=patrickvonplaten

I'll improve the behavior for beam search

https://github.com/huggingface/transformers/pull/14654

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.