How to generate text: using different decoding methods for language generation with Transformers
Patrick von Platen
Note: Edited on July 2023 with up-to-date references and examples.IntroductionIn recent years, there has been an increasing interest in open-endedlanguage generation thanks to the rise of large transformer-basedlanguage models trained on millions of webpages, including OpenAI's ChatGPTand Meta's LLaMA.The results on conditioned open-ended language generation are impressive, having shown togeneralize to new tasks,handle code,or take non-text data as input.Besides the improved transformer architecture and massive unsupervisedtraining data, better decoding methods have also played an importantrole.This blog post gives a brief overview of different decoding strategiesand more importantly shows how you can implement them with very littleeffort using the popular transformers library!All of the following functionalities can be used for auto-regressivelanguage generation (herea refresher). In short, auto-regressive language generation is basedon the assumption that the probability distribution of a word sequencecan be decomposed into the product of conditional next worddistributions:P(w1:T∣W0)=∏t=1TP(wt∣w1:t−1,W0) ,with w1:0=∅, P(w_{1:T} | W_0 ) = \prod_{t=1}^T P(w_{t} | w_{1: t-1}, W_0) \text{ ,with } w_{1: 0} = \emptyset, P(w1:T​∣W0​)=t=1∏T​P(wt​∣w1:t−1​,W0​) ,with w1:0​=∅,and W0W_0W0​ being the initial context word sequence. The length TTTof the word sequence is usually determined on-the-fly and correspondsto the timestep t=Tt=Tt=T the EOS token is generated from P(wt∣w1:t−1,W0)P(w_{t} | w_{1: t-1}, W_{0})P(wt​∣w1:t−1​,W0​).We will give a tour of the currently most prominent decoding methods,mainly Greedy search, Beam search, and Sampling.Let's quickly install transformers and load the model. We will use GPT2in PyTorch for demonstration, but the API is 1-to-1 the same forTensorFlow and JAX.!pip install -q transformersfrom transformers import AutoModelForCausalLM, AutoTokenizerimport torchtorch_device = "cuda" if torch.cuda.is_available() else "cpu"tokenizer = AutoTokenizer.from_pretrained("gpt2")# add the EOS token as PAD token to avoid warningsmodel = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id).to(torch_device)Greedy SearchGreedy search is the simplest decoding method.It selects the word with the highest probability asits next word: wt=argmaxwP(w∣w1:t−1)w_t = argmax_{w}P(w | w_{1:t-1})wt​=argmaxw​P(w∣w1:t−1​) at each timestepttt. The following sketch shows greedy search.Starting from the word "The",\text{"The"},"The", the algorithm greedily choosesthe next word of highest probability "nice"\text{"nice"}"nice" and so on, sothat the final generated word sequence is ("The","nice","woman")(\text{"The"}, \text{"nice"}, \text{"woman"})("The","nice","woman")having an overall probability of 0.5×0.4=0.20.5 \times 0.4 = 0.20.5×0.4=0.2 .In the following we will generate word sequences using GPT2 on thecontext ("I","enjoy","walking","with","my","cute","dog")(\text{"I"}, \text{"enjoy"}, \text{"walking"}, \text{"with"}, \text{"my"}, \text{"cute"}, \text{"dog"})("I","enjoy","walking","with","my","cute","dog"). Let'ssee how greedy search can be used in transformers:# encode context the generation is conditioned onmodel_inputs = tokenizer('I enjoy walking with my cute dog', return_tensors='pt').to(torch_device)# generate 40 new tokensgreedy_output = model.generate(**model_inputs, max_new_tokens=40)print("Output:" + 100 * '-')print(tokenizer.decode(greedy_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with my dog. I'm not sure if I'll ever be able to walk with my dog.I'm not sureAlright! We have generated our first short text with GPT2 😊. Thegenerated words following the context are reasonable, but the modelquickly starts repeating itself! This is a very common problem inlanguage generation in general and seems to be even more so in greedyand beam search - check out Vijayakumar etal., 2016 and Shao etal., 2017.The major drawback of greedy search though is that it misses highprobability words hidden behind a low probability word as can be seen inour sketch above:The word "has"\text{"has"}"has"with its high conditional probability of 0.90.90.9is hidden behind the word "dog"\text{"dog"}"dog", which has only thesecond-highest conditional probability, so that greedy search misses theword sequence "The","dog","has"\text{"The"}, \text{"dog"}, \text{"has"}"The","dog","has" .Thankfully, we have beam search to alleviate this problem!Beam searchBeam search reduces the risk of missing hidden high probability wordsequences by keeping the most likely num_beams of hypotheses at eachtime step and eventually choosing the hypothesis that has the overallhighest probability. Let's illustrate with num_beams=2:At time step 1, besides the most likely hypothesis ("The","nice")(\text{"The"}, \text{"nice"})("The","nice"),beam search also keeps track of the secondmost likely one ("The","dog")(\text{"The"}, \text{"dog"})("The","dog").At time step 2, beam search finds that the word sequence ("The","dog","has")(\text{"The"}, \text{"dog"}, \text{"has"})("The","dog","has"),has with 0.360.360.36a higher probability than ("The","nice","woman")(\text{"The"}, \text{"nice"}, \text{"woman"})("The","nice","woman"),which has 0.20.20.2 . Great, it has found the most likely word sequence inour toy example!Beam search will always find an output sequence with higher probabilitythan greedy search, but is not guaranteed to find the most likelyoutput.Let's see how beam search can be used in transformers. We setnum_beams > 1 and early_stopping=True so that generation is finishedwhen all beam hypotheses reached the EOS token.# activate beam search and early_stoppingbeam_output = model.generate(**model_inputs,max_new_tokens=40,num_beams=5,early_stopping=True)print("Output:" + 100 * '-')print(tokenizer.decode(beam_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I'm not sure if I'll ever be able to walk with him again. I'm not sureWhile the result is arguably more fluent, the output still includesrepetitions of the same word sequences.One of the available remedies is to introduce n-grams (a.k.a word sequences ofn words) penalties as introduced by Paulus et al.(2017) and Klein et al.(2017). The most common n-gramspenalty makes sure that no n-gram appears twice by manually settingthe probability of next words that could create an already seen n-gramto 0.Let's try it out by setting no_repeat_ngram_size=2 so that no 2-gramappears twice:# set no_repeat_ngram_size to 2beam_output = model.generate(**model_inputs,max_new_tokens=40,num_beams=5,no_repeat_ngram_size=2,early_stopping=True)print("Output:" + 100 * '-')print(tokenizer.decode(beam_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I've been thinking about this for a while now, and I think it's time for me toNice, that looks much better! We can see that the repetition does notappear anymore. Nevertheless, n-gram penalties have to be used withcare. An article generated about the city New York should not use a2-gram penalty or otherwise, the name of the city would only appearonce in the whole text!Another important feature about beam search is that we can compare thetop beams after generation and choose the generated beam that fits ourpurpose best.In transformers, we simply set the parameter num_return_sequences tothe number of highest scoring beams that should be returned. Make surethough that num_return_sequences <= num_beams!# set return_num_sequences > 1beam_outputs = model.generate(**model_inputs,max_new_tokens=40,num_beams=5,no_repeat_ngram_size=2,num_return_sequences=5,early_stopping=True)# now we have 3 output sequencesprint("Output:" + 100 * '-')for i, beam_output in enumerate(beam_outputs):print("{}: {}".format(i, tokenizer.decode(beam_output, skip_special_tokens=True)))Output:----------------------------------------------------------------------------------------------------0: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I've been thinking about this for a while now, and I think it's time for me to1: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with her again.I've been thinking about this for a while now, and I think it's time for me to2: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I've been thinking about this for a while now, and I think it's a good idea to3: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I've been thinking about this for a while now, and I think it's time to take a4: I enjoy walking with my cute dog, but I'm not sure if I'll ever be able to walk with him again.I've been thinking about this for a while now, and I think it's a good idea.As can be seen, the five beam hypotheses are only marginally differentto each other - which should not be too surprising when using only 5beams.In open-ended generation, a couple of reasons have been broughtforward why beam search might not be the best possible option:Beam search can work very well in tasks where the length of thedesired generation is more or less predictable as in machinetranslation or summarization - see Murray et al.(2018) and Yang et al.(2018). But this is not the casefor open-ended generation where the desired output length can varygreatly, e.g. dialog and story generation.We have seen that beam search heavily suffers from repetitivegeneration. This is especially hard to control with n-gram- orother penalties in story generation since finding a good trade-offbetween inhibiting repetition and repeating cycles of identicaln-grams requires a lot of finetuning.As argued in Ari Holtzman et al.(2019), high quality humanlanguage does not follow a distribution of high probability nextwords. In other words, as humans, we want generated text to surpriseus and not to be boring/predictable. The authors show this nicely byplotting the probability, a model would give to human text vs. whatbeam search does.So let's stop being boring and introduce some randomness 🤪.SamplingIn its most basic form, sampling means randomly picking the next word wtw_twt​ according to its conditional probability distribution:wt∼P(w∣w1:t−1) w_t \sim P(w|w_{1:t-1}) wt​∼P(w∣w1:t−1​)Taking the example from above, the following graphic visualizes languagegeneration when sampling.It becomes obvious that language generation using sampling is notdeterministic anymore. The word ("car")(\text{"car"})("car") is sampled from theconditioned probability distribution P(w∣"The")P(w | \text{"The"})P(w∣"The"), followedby sampling ("drives")(\text{"drives"})("drives") from P(w∣"The","car")P(w | \text{"The"}, \text{"car"})P(w∣"The","car") .In transformers, we set do_sample=True and deactivate Top-Ksampling (more on this later) via top_k=0. In the following, we willfix the random seed for illustration purposes. Feel free to change theset_seed argument to obtain different results, or to remove it for non-determinism.# set seed to reproduce results. Feel free to change the seed though to get different resultsfrom transformers import set_seedset_seed(42)# activate sampling and deactivate top_k by setting top_k sampling to 0sample_output = model.generate(**model_inputs,max_new_tokens=40,do_sample=True,top_k=0)print("Output:" + 100 * '-')print(tokenizer.decode(sample_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog for the rest of the day, but this had me staying in an unusual room and not going on nights out with friends (which will always be wondered for a mere minute or so at this point).Interesting! The text seems alright - but when taking a closer look, itis not very coherent and doesn't sound like it was written by ahuman. That is the big problem when sampling word sequences: The modelsoften generate incoherent gibberish, cf. Ari Holtzman et al.(2019).A trick is to make the distribution P(w∣w1:t−1)P(w|w_{1:t-1})P(w∣w1:t−1​) sharper(increasing the likelihood of high probability words and decreasing thelikelihood of low probability words) by lowering the so-calledtemperature of thesoftmax.An illustration of applying temperature to our example from above couldlook as follows.The conditional next word distribution of step t=1t=1t=1 becomes muchsharper leaving almost no chance for word ("car")(\text{"car"})("car") to beselected.Let's see how we can cool down the distribution in the library bysetting temperature=0.6:# set seed to reproduce results. Feel free to change the seed though to get different resultsset_seed(42)# use temperature to decrease the sensitivity to low probability candidatessample_output = model.generate(**model_inputs,max_new_tokens=40,do_sample=True,top_k=0,temperature=0.6,)print("Output:" + 100 * '-')print(tokenizer.decode(sample_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog, but I don't like to chew on it. I like to eat it and not chew on it. I like to be able to walk with my dog."So how did you decideOK. There are less weird n-grams and the output is a bit more coherentnow! While applying temperature can make a distribution less random, inits limit, when setting temperature →0\to 0→0, temperature scaledsampling becomes equal to greedy decoding and will suffer from the sameproblems as before.Top-K SamplingFan et. al (2018) introduced asimple, but very powerful sampling scheme, called Top-K sampling.In Top-K sampling, the K most likely next words are filtered and theprobability mass is redistributed among only those K next words. GPT2adopted this sampling scheme, which was one of the reasons for itssuccess in story generation.We extend the range of words used for both sampling steps in the exampleabove from 3 words to 10 words to better illustrate Top-K sampling.Having set K=6K = 6K=6, in both sampling steps we limit our sampling poolto 6 words. While the 6 most likely words, defined asVtop-KV_{\text{top-K}}Vtop-K​ encompass only ca. two-thirds of the wholeprobability mass in the first step, it includes almost all of theprobability mass in the second step. Nevertheless, we see that itsuccessfully eliminates the rather weird candidates (“not",“the",“small",“told")(\text{``not"}, \text{``the"}, \text{``small"}, \text{``told"})(“not",“the",“small",“told") in the second sampling step.Let's see how Top-K can be used in the library by setting top_k=50:# set seed to reproduce results. Feel free to change the seed though to get different resultsset_seed(42)# set top_k to 50sample_output = model.generate(**model_inputs,max_new_tokens=40,do_sample=True,top_k=50)print("Output:" + 100 * '-')print(tokenizer.decode(sample_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog for the rest of the day, but this time it was hard for me to figure out what to do with it. (One reason I asked this for a few months back is that I had aNot bad at all! The text is arguably the most human-sounding text sofar. One concern though with Top-K sampling is that it does notdynamically adapt the number of words that are filtered from the nextword probability distribution P(w∣w1:t−1)P(w|w_{1:t-1})P(w∣w1:t−1​). This can beproblematic as some words might be sampled from a very sharpdistribution (distribution on the right in the graph above), whereasothers from a much more flat distribution (distribution on the left inthe graph above).In step t=1t=1t=1, Top-K eliminates the possibility to sample("people","big","house","cat")(\text{"people"}, \text{"big"}, \text{"house"}, \text{"cat"})("people","big","house","cat"), which seem like reasonablecandidates. On the other hand, in step t=2t=2t=2 the method includes thearguably ill-fitted words ("down","a")(\text{"down"}, \text{"a"})("down","a") in the sample pool ofwords. Thus, limiting the sample pool to a fixed size K could endangerthe model to produce gibberish for sharp distributions and limit themodel's creativity for flat distribution. This intuition led AriHoltzman et al. (2019) to createTop-p- or nucleus-sampling.Top-p (nucleus) samplingInstead of sampling only from the most likely K words, in Top-psampling chooses from the smallest possible set of words whosecumulative probability exceeds the probability p. The probability massis then redistributed among this set of words. This way, the size of theset of words (a.k.a the number of words in the set) can dynamicallyincrease and decrease according to the next word's probabilitydistribution. Ok, that was very wordy, let's visualize.Having set p=0.92p=0.92p=0.92, Top-p sampling picks the minimum number ofwords to exceed together p=92%p=92\%p=92% of the probability mass, defined asVtop-pV_{\text{top-p}}Vtop-p​. In the first example, this included the 9 mostlikely words, whereas it only has to pick the top 3 words in the secondexample to exceed 92%. Quite simple actually! It can be seen that itkeeps a wide range of words where the next word is arguably lesspredictable, e.g. P(w∣"The”)P(w | \text{"The''})P(w∣"The”), and only a few words whenthe next word seems more predictable, e.g.P(w∣"The","car")P(w | \text{"The"}, \text{"car"})P(w∣"The","car").Alright, time to check it out in transformers! We activate Top-psampling by setting 0 < top_p < 1:# set seed to reproduce results. Feel free to change the seed though to get different resultsset_seed(42)# set top_k to 50sample_output = model.generate(**model_inputs,max_new_tokens=40,do_sample=True,top_p=0.92,top_k=0)print("Output:" + 100 * '-')print(tokenizer.decode(sample_output[0], skip_special_tokens=True))Output:----------------------------------------------------------------------------------------------------I enjoy walking with my cute dog for the rest of the day, but this had me staying in an unusual room and not going on nights out with friends (which will always be my yearning for such a spacious screen on my deskGreat, that sounds like it could have been written by a human. Well,maybe not quite yet.While in theory, Top-p seems more elegant than Top-K, both methodswork well in practice. Top-p can also be used in combination withTop-K, which can avoid very low ranked words while allowing for somedynamic selection.Finally, to get multiple independently sampled outputs, we can againset the parameter num_return_sequences > 1:# set seed to reproduce results. Feel free to change the seed though to get different resultsset_seed(42)# set top_k = 50 and set top_p = 0.95 and num_return_sequences = 3sample_outputs = model.generate(**model_inputs,max_new_tokens=40,do_sample=True,top_k=50,top_p=0.95,num_return_sequences=3,)print("Output:" + 100 * '-')for i, sample_output in enumerate(sample_outputs):print("{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))Output:----------------------------------------------------------------------------------------------------0: I enjoy walking with my cute dog for the rest of the day, but this time it was hard for me to figure out what to do with it. When I finally looked at this for a few moments, I immediately thought, "1: I enjoy walking with my cute dog. The only time I felt like walking was when I was working, so it was awesome for me. I didn't want to walk for days. I am really curious how she can walk with me2: I enjoy walking with my cute dog (Chama-I-I-I-I-I), and I really enjoy running. I play in a little game I play with my brother in which I take pictures of our houses.Cool, now you should have all the tools to let your model write yourstories with transformers!ConclusionAs ad-hoc decoding methods, top-p and top-K sampling seem toproduce more fluent text than traditional greedy - and beam searchon open-ended language generation. There isevidence that the apparent flaws of greedy and beam search -mainly generating repetitive word sequences - are caused by the model(especially the way the model is trained), rather than the decodingmethod, cf. Welleck et al.(2019). Also, as demonstrated inWelleck et al. (2020), it looks astop-K and top-p sampling also suffer from generating repetitive wordsequences.In Welleck et al. (2019), theauthors show that according to human evaluations, beam search cangenerate more fluent text than Top-p sampling, when adapting themodel's training objective.Open-ended language generation is a rapidly evolving field of researchand as it is often the case there is no one-size-fits-all method here,so one has to see what works best in one's specific use case.Fortunately, you can try out all the different decoding methods intransfomers 🤗 -- you can have an overview of the available methodshere.Thanks to everybody, who has contributed to the blog post: Alexander Rush, Julien Chaumand, Thomas Wolf, Victor Sanh, Sam Shleifer, Clément Delangue, Yacine Jernite, Oliver Åstrand and John de Wasseige.Appendixgenerate has evolved into a highly composable method, with flags to manipulate the resulting text in manydirections that were not covered in this blog post. Here are a few helpful pages to guide you:How to parameterize generateHow to stream the outputFull list of decoding optionsgenerate API referenceLLM score leaderboardIf you find that navigating our docs is challenging and you can't easily find what you're looking for, drop us a message in this GitHub issue. Your feedback is critical to set our future direction! 🤗