Re Reference AI

技術

ビームサーチとは

Beam Search

複数の候補系列を並行して保持・評価しながら、より高い確率の系列全体を探索するデコーディング手法

デコーディング探索アルゴリズム

ひとことで言うと

複数の答えの候補を同時に比べながら、一番よさそうな答えを探す方法。

概要

ビームサーチ(Beam Search)とは、テキスト生成の各ステップで確率の高い上位1つのトークンだけを選ぶ貪欲法とは異なり、一定数(ビーム幅)の候補系列を並行して保持し、それぞれを1トークンずつ拡張しながら系列全体としての確率が高い候補を残していくデコーディング手法。 目先のステップだけでなく系列全体の確率を考慮するため、貪欲法より質の高い出力系列を見つけやすい一方、複数の候補を同時に計算・保持するためビーム幅に比例して計算コストが増加する。 ビーム幅を大きくするほど探索の網羅性は高まるが計算コストも比例して増加するため、機械翻訳や要約等、出力の正確性が重視されるタスクで適度なビーム幅とともに使われることが多い。

背景

貪欲法は各ステップで最も確率の高いトークンを選ぶが、その選択が必ずしも系列全体として最も確率の高い出力につながるとは限らないという問題があった。 ビームサーチは、複数の候補を並行して保持することで局所的に最適な選択に固執せず、系列全体としてより確率の高い出力を見つけやすくするために考案された。

アーキテクチャ

各ステップで、現在保持しているビーム幅個の候補系列それぞれに対し次のトークン候補を生成し、それら全ての拡張候補の中から系列全体の対数確率が高い上位ビーム幅個を新たな候補として残す。 これを終了条件(文の終わりを示すトークンの生成や最大長への到達等)を満たすまで繰り返し、最終的に最も確率の高い系列を出力とする。

ワークフロー

現在のビーム幅個の候補系列それぞれに次のトークン候補を生成 → 全拡張候補から系列全体の対数確率が高い上位ビーム幅個を残す → 終了条件を満たすまで繰り返す → 最終的に最も確率の高い系列を出力する。

コード例

Transformersライブラリでビームサーチを実行する

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer("要約すると、", return_tensors="pt")

outputs = model.generate(
    **inputs, max_new_tokens=30, num_beams=5, early_stopping=True
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

利点

  • 系列全体としての確率を考慮するため、貪欲法より質の高い出力を得やすい
  • 機械翻訳や要約等、唯一の正解に近い出力を求めるタスクに適する
  • num_beamsのようなパラメータ1つで、貪欲法から切り替えて導入できる手軽さがある

欠点

  • ビーム幅に比例して計算コストとメモリ使用量が増加する
  • 確率の高い無難な出力に偏りやすく、創造的・多様な文章生成には不向きな場合がある
  • 長い系列を生成するタスクでは、計算コストの増加が特に大きくなりやすい

比較

  • デコーディング — ビームサーチは、複数の候補系列を並行評価する代表的なデコーディング戦略
  • Temperature — ビームサーチが確率の高い系列を追求するのに対し、Temperatureによるサンプリングは出力に多様性を持たせる方向のパラメータ
  • 探索木 — ビームサーチは、探索木の各深さで有望な候補を一定数に絞り込みながら進める、探索木の効率化手法の1つ

関連用語

デコーディングTemperatureTop-k探索木

よくある質問

ビームサーチはLLMのチャット用途でも使われる？

機械翻訳・要約等、唯一の正解に近い出力を求めるタスクでは有効だが、対話生成では単調な出力になりやすく、サンプリングベースのデコーディングの方が好まれやすい。

ビーム幅はいくつに設定すればいい？

タスクや許容できる計算コストによるが、機械翻訳等では4〜10程度がよく使われ、大きくしすぎても品質向上は頭打ちになりやすい。

参考文献

関連Zenn記事