Showing individual token and corresponding score during beam search


I am using beam search with a pre-trained T5 model for summarization. I would like to visualize the beam search process by showing the tokens with the highest scores, and eventually the chosen beam like this diagram:

(Taken from How to generate text: using different decoding methods for language generation with Transformers)

I am unsure how I can show the tokens and their corresponding scores.

I followed the discussion [Announcement] GenerationOutputs: Scores, Attentions and Hidden States now available as outputs to generate and

Following the docs, upon calling generate, I have set return_dict_in_generate=True, output_scores=True

generated_outputs = model_t5summary.generate(,
  max_length=input_ids.shape[-1] + 2,

Now I have an instance of BeamSearchEncoderDecoderOutput.

If I understand the docs (Utilities for Generation — transformers 4.2.0 documentation) correctly, scores will provide me with what I want but I am unsure on how to use the scores.

Any help/pointers from the community would be greatly appreciated, thank you :pray:

1 Like

Tagging @patrickvonplaten reposted my question from Github, thanks for directing me to the forum

Hey @monmanuela,

Good question!

So in the case of beam_search, the scores correspond to the log probability of all words + the log probability all previous scores in your beam.

So regarding the image this means that scores[0][0] will correspond to the log probabilities of all possible words in the vocabulary, so assuming your vocab would only consist of dog, nice, car and the probs are the same as in the diagram, the values would correspond to log(0.4), log(0.5), log(0.1)

Scores[1][0] then corresponds to the chosen word of time step one (e.g. dog) and all possible values again, se:
log(0.4) + log(0.05), log(0.4) + log(0.05), log(0.4) + log(0.9) using the diagram above again.

1 Like

Hi @patrickvonplaten, thanks for your detailed answer! I have managed to manually visualize the beam search process thanks to your help :blush:

Hi, @patrickvonplaten
can you add more details? based on your example,
yes, I verified that

torch.sum(torch.exp(aa.scores[0][0])) == 1

but for the second token:

torch.sum(torch.exp(aa.scores[1][0] - aa.scores[0][0].max())) != 1

which is wrong, right?