テキスト生成における decoding テクニック: Greedy search, Beam search, Top-K, Top-p
Transformer ベースの言語モデルが普及しているのは承知の通りだと思います。
中でも有名なのは BERT で、これは Transformer の Encoder のみを使う Autoencoding models と呼ばれるカテゴリのモデルです。入力の一部を隠してそれを復元するというタスクを大量に解かせることで事前学習を行うことになります。それゆえ、このタイプのモデルが最も向いているタスクは Token の分類（固有表現抽出等）や文章の分類など、入力の文章自体に興味がある場合です。
一方、Transformer の Decoder のみを使う Autoregressive models (GPT-2 など) や、 Encoder と Decoder の双方を使う Sequence-to-sequence models (BART や T5 など) は、翻訳や要約などの テキストを出力することを志向したモデル です。
このあたりのことは以下の記事にも書いてあります。
こうした Decoder を持つモデルでテキスト生成を実際に行う際には、auto-regressive なテキスト生成を行うことになります。これは「入力 + 今まで出力した単語列」から「次に出力される単語は何か」という条件付き確率を（EOSトークンが出力されるまで）各ステップで考えて、最終的にその積が出力された単語列の確率になるという単純な考え方です。
ただ、愚直にこれを考えると「語彙数」の「出力回数」乗 の計算が必要になり、とても計算できる量ではないので、何らかのヒューリスティックスを導入することになります。
この記事ではその際に使われるテクニックをいくつか説明します。
Greedy Search
日本語では貪欲法と呼ばれます。毎回最大の条件付き確率の単語のみを取ってくるというシンプルなやり方です。
画像は元記事より
計算が軽いというメリットの一方で、
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.
というように、同じ文章を繰り返し生成してしまう問題があることが知られています。
- Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence Models (AAAI 2018)
- Generating High-Quality and Informative Conversation Responses with Sequence-to-Sequence Models (EMNLP 2017)
Beam Search
ビーム探索では、毎ステップで上位 N (> 1) 本の確率のルートを保持します。これにより、Greedy Search に比べて無駄な文の繰り返しを低減できることが知られています。
画像は元記事より
さらに確実に繰り返しを減らしたい場合は、n-grams penalty を導入します。例えば、huggingface/transformers ライブラリでは no_repeat_ngram_size=2 などというパラメータを指定してあげれば、同じ二単語の連続が2度以上繰り返し出てくることはありません。
もちろん、このパラメータは注意して使う必要があります。例えば、no_repeat_ngram_size=2 によって "New York" というような単語列も2度以上出てこなくなってしまうからです（ニューヨークで発生した事件のニュース記事を翻訳する場合などでは明らかに不適当でしょう）。
また、num_return_sequences=5 などというパラメータを指定してあげることにより、最終的に確率の高かった上位 5 種類の文章を取得することができます。
ところが、Beam Search の上位候補を単に取得するだけでは上手くいかない場合もあることが知られています。それは 出力に多様性を持たせたい場合 です。
実際、 Beam Search の上位候補は、単語がわずかに異なるだけでほとんど同じような文になってしまうことが分かっています。
Top-K sampling
先ほどの話を言い換えると、尤度を最大化しようとするだけでは似通った文章になってしまう、ということになります。なので、最大化するのではなく、確率が大きめな候補からサンプリングしてランダム性を導入する、という発想が生まれます。
この方向性の手法でシンプルなものとして Top-K sampling があります。
画像は元記事より
上の画像は K=6 とした場合で、毎回単語を予測する際に上位6つのどれかをランダムに選びます（例えば "The" の次は "nice" が最も確率が高いですが、ここでは3番目に確率の高い "car" を選択しています）。
Top-K sampling は
Hierarchical Neural Story Generation (ACL 2018)
で用いられ、物語生成のようなケースで有効であることが示されました。
この手法はシンプルである反面、上位候補の分布については全く考慮していないという欠点がありあます。
例えば上の例で、"The" の次の単語は上位6個の候補のどれでも確率的に良さそうですが、"The car" の次は、上位6個の候補のうち下2つの "down" や "a" を選ぶのは明らかにマズそうです。
Top-p (nucleus) sampling
Top-k sampling の欠点を解消するために、Top-p sampling が提案されました。
これは、上位K個の候補を選ぶのではなく、確率の合計が p を超えるような最小の個数の候補を動的に決めるものです。
画像は元記事より
上の画像は top_p=0.92 とした場合です。これにより "The car" の次に "down" や "a" などが選択されることが防げます。
なお、Top-K と Top-p は対立する概念ではなく、実用上はどちらも上手くいくやり方であり、組み合わせて使うこともできます。huggingface/transformers ライブラリでは top_k=50, top_p=0.95 というようにパラメータを指定して使うことができます。
その他
最新の研究により、単純な Beam Search や Greedy Search が同じ単語列の繰り返しを発生させてしまうのは、decoding に問題があるのではなくモデルの学習自体に問題があるとされています。また、Top-K や Top-p のようなサンプリングによる decoding であってもそうした単語列の繰り返しは発生しうるそうです。
huggingface/transformers ライブラリでは、今回紹介したもの以外にも様々な decoding テクニックを generate() 関数のパラメータとして実装しています。色々試してみて、どれが良いのかを実験的に調べてみるのが良いでしょう。
Discussion