LLM を Elixir / Livebook で学ぶ 5

# 事前学習済みモデルの推論と可視化に使うライブラリを準備する
Mix.install(
  [
    {:bumblebee, "~> 0.5"},
    {:nx, "~> 0.9", override: true},
    {:exla, "~> 0.9"},
    {:kino, "~> 0.15"},
    {:kino_vega_lite, "~> 0.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

事前学習済み GPT 系モデルで「次トークン予測」を体感する

前のノートブックでは、デコーダーのみのミニ GPT を自分で組み立て、小さなコーパスで端から端まで学習しました。

このノートブックでは、同じ構造を大規模化した実物のモデルを動かします。Bumblebee を使って GPT 系モデルを実行しながら、

• 入力がどうトークンに分かれるか
• 生成設定を変えると何が変わるか
• 生成トークン数が増えると時間がどう変わるか

を確認します。

方針

ここでは教材として扱いやすい gpt2 を使います。

• 仕組みを観察することが目的
• 長時間学習はしない
• 例文や解説はこの notebook 用の独自内容

このnotebookで使う用語

この章では、学習済みモデルへ文章を入力して生成するときの用語を使います。

準備

defmodule LLMScratch.Visuals do
  # プロンプトごとのトークン数など、カテゴリ間の値を比較する
  def bar_chart(rows, title, x_field, y_field, opts \\ []) do
    width = Keyword.get(opts, :width, 560)
    height = Keyword.get(opts, :height, 280)

    VegaLite.new(width: width, height: height, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:bar, tooltip: true)
    |> VegaLite.encode_field(:x, Atom.to_string(x_field), type: :nominal, title: Atom.to_string(x_field))
    |> VegaLite.encode_field(:y, Atom.to_string(y_field), type: :quantitative, title: Atom.to_string(y_field))
    |> Kino.VegaLite.new()
  end

  # 生成トークン数に対する実行時間など、連続値の変化を表示する
  def line_chart(rows, title, x_field, y_field) do
    VegaLite.new(width: 560, height: 280, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:line, point: true, tooltip: true)
    |> VegaLite.encode_field(:x, Atom.to_string(x_field), type: :quantitative, title: Atom.to_string(x_field))
    |> VegaLite.encode_field(:y, Atom.to_string(y_field), type: :quantitative, title: Atom.to_string(y_field))
    |> Kino.VegaLite.new()
  end
end

GPT 系モデルの生成ループ

Kino.Mermaid.new("""
flowchart
  A["プロンプト"] --> B["トークン化"]
  B --> C["モデルへ入力"]
  C --> D["次トークンの確率"]
  D --> E["1 トークン選ぶ"]
  E --> F["末尾に追加"]
  F --> C
""")

ポイントは、モデルが一度に文章全体を完成させているわけではなく、1 トークンずつ続きとして足している ことです。

モデルのダウンロード

# ダウンロード済みファイルを再利用できるよう、保存先を固定する
cache_dir = "/tmp/bumblebee_cache"
repo = {:hf, "gpt2", cache_dir: cache_dir}

# 同じリポジトリから、重み、トークナイザー、標準の生成設定を読み込む
{:ok, gpt2} = Bumblebee.load_model(repo)
{:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
{:ok, generation_config} = Bumblebee.load_generation_config(repo)

1. プロンプトをトークンとして見る

gpt2 は英語向けのモデルなので、英語で典型的に「1単語が複数トークンへ分かれる」例を観察します。

人が空白で数える単語と、モデルのトークンは別物です。頻出する文字列は1トークンになりやすい一方、長い語・派生語・珍しい綴りは、モデルが知っている小さな部品へ分割されることがあります。

# 入力を変えて再実行できるLivebookのテキスト欄を作る
prompt_input =
  Kino.Input.textarea("PROMPT",
    default: "Tokenization is unexpectedly interesting."
  )

# テキスト欄の現在値を、この後のトークン化と生成に使う
prompt = Kino.Input.read(prompt_input)
prompt

# 文章をモデルが扱うトークンID列へ変換する
tokenized = Bumblebee.apply_tokenizer(tokenizer, prompt)
token_ids = tokenized["input_ids"][[0]] |> Nx.to_flat_list()

# 各IDを1トークンずつ復号し、元の文字列のどの部品かを確認する
token_rows =
  token_ids
  |> Enum.with_index()
  |> Enum.map(fn {token_id, position} ->
    %{
      position: position,
      token_id: token_id,
      piece: Bumblebee.Tokenizer.decode(tokenizer, [token_id])
    }
  end)

Kino.DataTable.new(
  token_rows,
  keys: [:position, :token_id, :piece]
)

表の piece を上から読むと、たとえば Tokenization が Token と ization に分かれるはずです。先頭の空白が後ろの piece に含まれることもあります。

# 単語より小さなサブワードへ分かれやすい英単語を比べる
tokenization_examples = [
  %{english: "Tokenization", japanese: "トークン化"},
  %{english: "tokenizer", japanese: "トークナイザー"},
  %{english: "reuses", japanese: "再利用する"},
  %{english: "unexpectedly", japanese: "意外にも"}
]

token_piece_rows =
  tokenization_examples
  |> Enum.flat_map(fn example ->
    # 1単語をトークン化し、複数pieceなら複数行へ展開する
    ids =
      Bumblebee.apply_tokenizer(tokenizer, example.english)["input_ids"][[0]]
      |> Nx.to_flat_list()

    ids
    |> Enum.with_index()
    |> Enum.map(fn {token_id, piece_index} ->
      %{
        word: example.english,
        japanese_meaning: example.japanese,
        piece_index: piece_index,
        piece: Bumblebee.Tokenizer.decode(tokenizer, [token_id]),
        token_id: token_id
      }
    end)
  end)

Kino.DataTable.new(
  token_piece_rows,
  keys: [:word, :japanese_meaning, :piece_index, :piece, :token_id]
)

同じ word が複数行にまたがっていれば、1単語 = 1トークン ではありません。モデルは単語辞書だけでなく、再利用しやすい文字列の部品を組み合わせて文章を扱っています。これがサブワード分割です。

2. プロンプトごとの長さを比べる

トークン数は、そのまま計算量やコンテキスト長の感覚につながります。

# 文字数や空白区切りの単語数と、モデルのトークン数が違う例
prompt_examples = [
  "cat",
  "tokenizer",
  "Tokenization is interesting.",
  "The tokenizer reuses reusable pieces."
]

prompt_length_rows =
  prompt_examples
  |> Enum.map(fn text ->
    # 各文章を実際にトークン化し、ID列の長さを数える
    input_ids = Bumblebee.apply_tokenizer(tokenizer, text)["input_ids"][[0]] |> Nx.to_flat_list()

    %{
      prompt: text,
      token_count: length(input_ids)
    }
  end)

Kino.DataTable.new(prompt_length_rows)

LLMScratch.Visuals.bar_chart(prompt_length_rows, "プロンプトごとのトークン数", :prompt, :token_count)

cat は短く頻出なので1トークンになりやすい一方、tokenizer は複数 piece へ分かれます。文字数や空白区切りの単語数だけでは、モデルにとっての長さを決められません。計算量やコンテキスト上限を考えるときは、トークン数を見ます。

3. 生成設定を変えてみる

ここでは 2 つの設定を比べます。

• greedy: 毎回もっとも確率の高い候補を選ぶ
• sampling: 上位候補の確率に応じて抽選し、別の展開も選べるようにする

違いを観察しやすくするため、同じプロンプトをそれぞれ3回ずつ生成します。samplingでは試行ごとに異なる乱数seedを使います。seedをコード内で固定することで、3種類の結果を作りながら、notebookを再実行したときには同じ比較を再現できます。

# 毎回1位を選ぶ、再現性の高い貪欲探索
greedy_config =
  Bumblebee.configure(generation_config,
    max_new_tokens: 40,
    strategy: %{type: :greedy_search}
  )

# 上位候補から確率的に選ぶ、多様性のあるサンプリング
sampling_config =
  Bumblebee.configure(generation_config,
    max_new_tokens: 40,
    temperature: 1.1,
    strategy: %{
      type: :multinomial_sampling,
      top_k: 40,
      top_p: 0.9
    }
  )

# モデル、トークナイザー、生成設定を実行可能なServingへまとめる
greedy_serving = Bumblebee.Text.generation(gpt2, tokenizer, greedy_config)
sampling_serving = Bumblebee.Text.generation(gpt2, tokenizer, sampling_config)

# 3つのseedを使い、同じプロンプトで独立した3試行を行う
trial_seeds = [101, 202, 303]

run_trials = fn serving ->
  trial_seeds
  |> Enum.with_index(1)
  |> Enum.flat_map(fn {seed, trial} ->
    serving
    |> Nx.Serving.run(%{text: prompt, seed: seed})
    |> Map.fetch!(:results)
    |> Enum.map(&Map.put(&1, :trial, trial))
  end)
end

# greedyは乱数を使わないため、seedが違っても同じ結果になる
greedy_results = run_trials.(greedy_serving)

# samplingは確率的に抽選するため、seedごとに異なる結果になりやすい
sampling_results = run_trials.(sampling_serving)

# 生成文は先頭が改行になる場合があるため、表ではなく全文表示できるKino.Textを使う
generation_text = fn results ->
  results
  |> Enum.map(fn result ->
    case String.trim_leading(result.text) do
      "" -> "試行 #{result.trial}\n表示できる文字列は生成されませんでした"
      text -> "試行 #{result.trial}\n#{text}"
    end
  end)
  |> Enum.join("\n\n---\n\n")
  |> Kino.Text.new()
end

generation_tabs =
  Kino.Layout.tabs([
    {"greedy（3回）", generation_text.(greedy_results)},
    {"sampling（3回）", generation_text.(sampling_results)}
  ])

# 空白の違いを除いて、3回のうち何種類の文章が生成されたかを数える
unique_generation_count = fn results ->
  results
  |> Enum.map(&String.trim(&1.text))
  |> Enum.uniq()
  |> length()
end

diversity_rows = [
  %{
    方式: "greedy",
    試行数: length(greedy_results),
    異なる生成結果数: unique_generation_count.(greedy_results)
  },
  %{
    方式: "sampling",
    試行数: length(sampling_results),
    異なる生成結果数: unique_generation_count.(sampling_results)
  }
]

# 生成文とは別に、入力・出力トークン数も比較できるよう残す
generation_summary_rows =
  [
    {"greedy", greedy_results},
    {"sampling", sampling_results}
  ]
  |> Enum.flat_map(fn {method, results} ->
    Enum.map(results, fn result ->
      %{
        方式: method,
        試行: result.trial,
        入力トークン数: result.token_summary.input,
        出力トークン数: result.token_summary.output,
        パディング数: result.token_summary.padding
      }
    end)
  end)

Kino.Layout.grid(
  [
    Kino.DataTable.new(
      diversity_rows,
      keys: [:方式, :試行数, :異なる生成結果数]
    ),
    generation_tabs,
    Kino.DataTable.new(
      generation_summary_rows,
      keys: [:方式, :試行, :入力トークン数, :出力トークン数, :パディング数]
    )
  ],
  columns: 1
)

生成文が改行から始まると、1行の表では先頭の空行だけが表示され、:textが空に見えることがあります。ここでは先頭の空白を表示用に取り除き、Kino.Textで改行を含む全文を表示しています。

最初の表では、異なる生成結果数を確認してください。greedyは通常1、samplingは2または3になり、同じ入力でもsamplingの方が複数の続きを選びやすいことが分かります。その下のタブでは、実際にどこから文章が分かれていったかを3試行分の全文で比較できます。

初学者向けの観察ポイントは次の通りです。

• greedy は乱数を使わず毎回1位を選ぶので、同じ入力なら3回とも同じ結果になる
• sampling は確率に応じて抽選するので、同じ入力でも実行ごとに続きが変わり得る
• temperature を上げると確率差がなだらかになり、低確率だった候補も選ばれやすくなる
• top_k=40 と top_p=0.9 は、可能性が極端に低い候補まで抽選対象にしないための制限

ここで「創造性」という特別な能力を追加したわけではありません。次トークン候補から常に1位を取るか、上位候補を抽選するか を変えただけで、その小さな分岐が後続トークンへ連鎖します。

4. 同じ入力でも生成の感じが変わる

同じ中心場面 A small robot arrived at the station に、文体を示す短い書き出しを加えて比べます。英語生成ですが、比較の意図は日本語ラベルで確認できるようにします。

# 中心となる出来事はそろえ、書き出しだけで文体を指定する
prompt_variants = [
  %{
    style: "ニュース調",
    prompt: "News report: A small robot arrived at the station. According to witnesses,"
  },
  %{
    style: "子ども向け物語",
    prompt: "Children's story: A small robot arrived at the station. It smiled and said,"
  },
  %{
    style: "ミステリー",
    prompt: "Mystery story: A small robot arrived at the station. Nobody knew why,"
  }
]

variant_results =
  prompt_variants
  |> Enum.map(fn example ->
    # 確率的な生成なので、再実行すると結果が変わる場合がある
    output =
      sampling_serving
      |> Nx.Serving.run(example.prompt)
      |> Map.get(:results)

    %{
      style: example.style,
      prompt: example.prompt,
      result: inspect(output)
    }
  end)

Kino.DataTable.new(variant_results, keys: [:style, :prompt, :result])

3行は中心となる出来事を共通にしています。観察するときは、単語を逐語訳するより次を比べてください。

• ニュース調では、人・場所・目撃情報のような説明へ寄るか
• 子ども向けでは、会話や明るい展開へ寄るか
• ミステリーでは、理由・秘密・不穏な出来事へ寄るか

GPT-2 は指示に従うよう調整されたチャットモデルではないため、毎回きれいに文体が分かれるとは限りません。その不安定さも含めて、入力の少しの違いが次トークン確率を変え、生成全体の分岐になる ことを観察します。

5. 生成トークン数と時間の関係

トークンを 1 つずつ足すので、max_new_tokens を増やすと時間も伸びやすくなります。

# 生成上限だけを変え、処理時間との関係を測る
benchmark_prompt = "Elixir notebooks are useful because"

timing_rows =
  [10, 20, 40, 60]
  |> Enum.map(fn max_new_tokens ->
    # 比較対象ごとに最大生成トークン数を設定し直す
    config =
      Bumblebee.configure(generation_config,
        max_new_tokens: max_new_tokens,
        temperature: 0.9
      )

    serving = Bumblebee.Text.generation(gpt2, tokenizer, config)

    # :timer.tcは実行時間をマイクロ秒で返す
    {microseconds, _result} = :timer.tc(Nx.Serving, :run, [serving, benchmark_prompt])

    %{
      max_new_tokens: max_new_tokens,
      seconds: microseconds / 1_000_000
    }
  end)

Kino.DataTable.new(timing_rows)

LLMScratch.Visuals.line_chart(timing_rows, "生成トークン数と実行時間", :max_new_tokens, :seconds)

環境によって秒数はかなり変わりますが、傾向としては 出力を長くするほど待ち時間も増える と考えておくと理解しやすいです。

6. 理論とのつながり

第4章のミニ GPT とつなげると、生成の内部ではざっくり次のことが起きています。

• 入力文がトークン列になる
• 各トークンが埋め込みベクトルになる
• 因果マスク付き自己アテンションで左側の文脈だけを見る
• 最後の位置から「次トークン候補」の確率を出す
• 1 つ選んで末尾に足し、また同じ処理をする

これは第4章の next_token_probs と generate がやっていたことと同じです。違いは、語彙が 18 個から約 5 万個に、パラメータが約 5 千個から約 1.2 億個に増えていることです。

7. まとめ

この notebook の要点

• GPT 系モデルは、次トークン予測を繰り返して文章を伸ばす
• トークン数は、計算量とコンテキスト長の感覚に直結する
• temperature や top_k で生成の性格をある程度調整できる
• プロンプトの書き方も出力品質に大きく影響する

次のノートブックでは、モデルを「賢く、役に立つ応答に寄せる」ためのアラインメントの考え方を、軽量な例で学びます。