Beam search calculates mean logprobs wrong?
Author: TomerRonen34
Pull request: no
Labels: name: WIP

### System Info

- `transformers` version: 4.33.3
- Platform: Linux-5.15.0-76-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- 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?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@gante (recommended for generate-related issues)
@patrickvonplaten (wrote the code according to git-blame)

### 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

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

device = "cuda"
model_name = "gpt2"
short_prompt = "Once upon a time"
long_prompt = """In a hole in the ground there lived a hobbit. Not a nasty, dirty, wet hole, filled with the ends of bworms and an oozy smell, nor yet a dry, bare, sandy hole with nothing in it to sit down on or to eat: it a was a hobbit-hole, and that means comfort.
It had a perfectly round door like a porthole, painted green, with a shiny yellow brass knob in the exact middle. The door opened on to a tube-shaped hall like a tunnel: a very comfortable tunnel without smoke, with panelled walls, and floors tiled and carpeted, provided with polished chairs, and lots and lots of pegs for hats and coats- the hobbit was fond of visitors. The tunnel wound on and on – going fairly but not quite straight into the side of the hill – The Hill, as all the people for many miles around called it – and many little round doors opened out of it, first on one side and then on another. No going upstairs for the hobbit: bedrooms, bathrooms, cellars, pantries (lots of these), wardrobes (he had whole rooms devoted to clothes), kitchens, dining-rooms, all were on the same floor, and indeed on the same passage. The best rooms were all on the lefthand side (going in), for these were the only ones to have windows, deep-set round windows looking over his garden, and meadows beyond, sloping down to the river.
This hobbit was a very well-to-do hobbit, and his name was Baggins. The Bagginses have lived in the neighbourhood of The Hill for time out of mind, and people considered them very respectable, not only because most of them were rich, but also because they never had any adventures or did anything unexpected: you could tell what a Baggins would say on any question without the bother of asking him. This is a story of how a Baggins had an adventure, and"""

beam_size = 5
max_new_tokens = 1
num_return_sequences = beam_size

with torch.device(device), torch.no_grad():
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    for prompt_name, prompt in [("short_prompt", short_prompt), ("long_prompt", long_prompt)]:
        batch = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**batch, num_beams=beam_size, num_return_sequences=num_return_sequences,
                                 max_new_tokens=max_new_tokens, output_scores=True, return_dict_in_generate=True)
        print(f"{prompt_name}: scores={outputs.sequences_scores.tolist()}")
```

Prints:
```
short_prompt: scores=[-0.17024032771587372, -0.5479289293289185, -0.6405749320983887, -0.6600505113601685, -0.7051623463630676]
long_prompt: scores=[-0.0049241515807807446, -0.006088315974920988, -0.006767737679183483, -0.006866625044494867, -0.006999899633228779]
```

### Expected behavior

When doing beam search, beam scores are normalized to represent the average token logprob. However, the current implementation divides the sum of **_generated token logprobs_** by the length of the **_entire sequence, including prompt_**. This creates inconsistencies between the scores of sequences of different lengths, and also prefers shorter generations. [Code here](https://github.com/huggingface/transformers/blob/75a33d60f25d99ff8cdd657d6ba685dc4336a0d1/src/transformers/generation/beam_search.py#L938).

My reproduction example shows that the absolute values of beam scores returned by generating a single token with a long prompt are orders of magnitude smaller than beam scores returned by a short prompt.

The main scenario where this behavior is problematic is for beams that terminate with an EOS before `max_new_tokens` is reached, since the denominator in their score calculation will be skewed. For example, if we have 2 candidates with lengths `l1` and `l2`, where all token logprobs are `s` and the prompt length is `p`, we'll have:
`score_i = l_i * s / (l_i + p) = s / (1 + p / l_i)`, showing a preference for shorter generations since `s < 0`.

Comments:
Hi @TomerRonen34 👋 

A few points that help explain what you're seeing, and how to disable certain default behaviors:
1. It's actually the other way around, there is a default preference for long generations :) Taking the first scores from your example, `-0.0049241515807807446` (long prompt) > `-0.17024032771587372` (short prompt). You can also see that the sequence scores are sorted in descending order.
2. This preference for long generations by default is intentional. The score of a sequence always goes down as we add words (the logits are always negative). Therefore, if a longer generation has a similar score to a shorter generation, then the longer generation's is composed of high-probability tokens, which usually means it is a better answer to the prompt.
3. The default for `length_penalty` is `1.0`. Counter-intuitively, summarization models usually set it to `2.0` -- these models are trained to output short sequences (i.e. quickly generate an EOS token), so when a high-probability sequence of tokens comes along, it is almost always relevant to the output.
4. You can disable this behavior by passing `length_penalty=0.0` to `generate` (or by setting it in the [generation config](https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration))
5. Since you're playing with the scores, [`model.compute_transition_scores`](https://huggingface.co/docs/transformers/v4.34.0/en/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores) may be relevant to you :)

I hope this helps explaining what's going on, and how to control it.

Hi @gante , I understand `length_penalty` and the effect of averaging logprobs instead of summing them. I was unaware of `compute_transition_scores`, it looks very convenient. Thanks!

I'm afraid my previous code example wasn't clear enough, as it didn't highlight the main issue which I addressed in the "Expected behavior" section. I hereby provide a better reproduction example which clearly shows the preference of the current beam search implementation to shorter generations, which I think is due to a bug in the way the beam score is calculated [here](https://github.com/huggingface/transformers/blob/75a33d60f25d99ff8cdd657d6ba685dc4336a0d1/src/transformers/generation/beam_search.py#L938), where the sum of the generated logprobs is divided by the entire sequence length, including the prompt.

In the following example, all generations should have the exact same score, since I force the logprobs of all generated tokens to be exactly the same. However, we see that the score is actually
`logprob * gen_length / (gen_length + prompt_length - 1)`.

Running this code prints:
```
    Beam search outputs 3 sequences, that should all have the same score
    (logprob=-10.825), but instead they have these scores, which prefer shorter generations:
    buggy_score = logprob * gen_length / (gen_length + prompt_length - 1)
    [-1.3531131744384766, -2.706226348876953, -3.4919049739837646]
```
The code:
```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import LogitsProcessor, LogitsProcessorList


class UniformScoreFixedLengthsLogitsProcessor(LogitsProcessor):
    """
    This logits processor forces the beam search to output sequences with predetermined lengths
    (e.g. [3,7,10]), where all the tokens have the exact same logprob. This logprob represents a
    uniform distribution over the vocabulary.
    """

    def __init__(self, generation_lengths: list[int], eos_token_id: int):
        assert len(set(generation_lengths)) == len(generation_lengths), \
            "This example is designed for unique lengths"
        self.generation_lengths = sorted(generation_lengths)
        self.eos_token_id = eos_token_id
        self.i_seq = 0
        self.times_called = 0

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        self.times_called += 1
        uniform_logprob = _calculate_uniform_logprob(vocab_size=scores.shape[1])
        scores[:] = uniform_logprob
        scores[:self.i_seq, :] = -float("inf")
        if (self.i_seq < len(self.generation_lengths) and
                self.times_called == self.generation_lengths[self.i_seq]):
            scores[self.i_seq, :] = -float("inf")
            scores[self.i_seq, self.eos_token_id] = uniform_logprob
            self.i_seq += 1
        return scores


def _calculate_uniform_logprob(vocab_size: int) -> float:
    return torch.zeros(vocab_size).log_softmax(-1)[0].item()


def run_example():
    device = "cuda"
    model_name = "gpt2"
    prompt = "a a a a a a a a a a a a a a a a a a a a a a"

    generation_lengths = [3, 7, 10]
    beam_size = len(generation_lengths)
    num_return_sequences = beam_size
    max_new_tokens = max(generation_lengths) + 100  # won't be reached anyway

    with torch.device(device), torch.no_grad():
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)

        logits_processor = UniformScoreFixedLengthsLogitsProcessor(generation_lengths=generation_lengths,
                                                                   eos_token_id=tokenizer.eos_token_id)

        batch = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(
            **batch, num_beams=beam_size, num_return_sequences=num_return_sequences,
            max_new_tokens=max_new_tokens, output_scores=True, return_dict_in_generate=True,
            logits_processor=LogitsProcessorList([logits_processor]))

    prompt_length = len(batch["input_ids"][0])
    uniform_logprob = _calculate_uniform_logprob(vocab_size=tokenizer.vocab_size)
    buggy_scores = [uniform_logprob * gen_length / (gen_length + prompt_length - 1)
                    for gen_length in sorted(generation_lengths)]
    assert torch.allclose(torch.tensor(buggy_scores), outputs.sequences_scores.cpu())

    print(f"""
    Beam search outputs {beam_size} sequences, that should all have the same score
    (logprob={uniform_logprob:.3f}), but instead they have these scores, which prefer shorter generations:
    buggy_score = logprob * gen_length / (gen_length + prompt_length - 1)
    {outputs.sequences_scores.tolist()}
    """)


run_example()
```

I see, accounting for the prompt length in the division (when `length_penalty != 0.0`) does change the beam search behavior. For context, beam search is mostly used with encoder-decoder models, whose prompt length is always `1`, and thus it wasn't a problem in the vast majority of use cases, nor it manifested as a preference for shorter sequences.

I agree that it should be fixed, but the priority is not high -- we want to refactor beam search first.


I had the same issue during my research experimentation with decoder-only LLMs, and found the exactly same cause as @TomerRonen34  already demoed. I have fixed the issue locally by making some minimal changes to the source codes, but I haven't submitted the PR yet. Actually I think this issue should be put into a higher priority because most recent LLMs are decoder-only, and the current problematic beam search with these LLMs have already been used in recently published research works without notice, e.g., https://github.com/lorenzkuhn/semantic_uncertainty. Since the `transformers` library is actively used by many AI researchers, many incorrect experimental results or misleading conclusions may be produced with the current buggy implementation. We should try to fix it ASAP. If @gante is busy with other more important refactoring works, I can submit a PR with a quick simple fix first.

It's great to hear that you have a fix @VsonicV! I wrote my own local fix as well, but mine relies on `len(beam_indices)` to find the number of generated tokens so far, so it only works when `.generate` is called with `return_dict_in_generate=True, output_scores=True`, otherwise `beam_indices` is None.

@VsonicV I'd love to merge the PR, assuming the fix is not large 🤗 

@gante I have created the PR #27351 for your review.

@TomerRonen34 Feel free to check and test the PR #27351