Notesclub

created by hec & contributors

terms privacy

LLM を Elixir / Livebook で学ぶ 4

04_decoder_generation_with_bumblebee.livemd

Ryo Wakabayashi

@RyoWakabayashi

elixir-learning

Share to X

Share to Bluesky

More notebooks

LLM を Elixir / Livebook で学ぶ 4

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 系モデルで「次トークン予測」を体感する

前のノートブックでは、Transformerの主要部品とエンコーダー／デコーダー全体の順伝播を自分で組み立てました。

このノートブックでは、Bumblebee を使って GPT 系モデルを実際に動かしながら、

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

を確認します。

方針

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

仕組みを観察することが目的
長時間学習はしない
例文や解説はこの 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トークンになりやすい一方、長い語・派生語・珍しい綴りは、モデルが知っている小さな部品へ分割されることがあります。

prompt_input =
  Kino.Input.textarea("PROMPT",
    default: "Tokenization is unexpectedly interesting."
  )

prompt = Kino.Input.read(prompt_input)
prompt

tokenized = Bumblebee.apply_tokenizer(tokenizer, prompt)
token_ids = tokenized["input_ids"][[0]] |> Nx.to_flat_list()

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 ->
    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 ->
    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: 上位候補の確率に応じて抽選し、別の展開も選べるようにする

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
    }
  )

greedy_serving = Bumblebee.Text.generation(gpt2, tokenizer, greedy_config)
sampling_serving = Bumblebee.Text.generation(gpt2, tokenizer, sampling_config)

greedy_result =
  greedy_serving
  |> Nx.Serving.run(prompt)
  |> Map.get(:results)

sampling_result =
  sampling_serving
  |> Nx.Serving.run(prompt)
  |> Map.get(:results)

Kino.Layout.tabs([
  greedy: Kino.DataTable.new(greedy_result),
  sampling: Kino.DataTable.new(sampling_result)
])

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

greedy は同じ入力なら同じ経路を選びやすく、無難だが繰り返しも起きやすい
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)
    {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. 理論とのつながり

前回の内容とつなげると、生成の内部ではざっくり次のことが起きています。

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

7. まとめ

この notebook の要点

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

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

Other notebooks:

Dr. Christian Geuer-Pollmann
@chgeuer

livebook_on_azure

Christian's first LiveBook test

notebook1.livemd

tutorial advanced data-science axon exla nx

2022-8-18
@andyl

elix_util

MNIST

mnist.livemd

tutorial advanced data-science req axon exla nx

2022-8-18
@TomBers

livebookNotes

Attractors

attractors.livemd

advanced data-science decimal vega_lite kino

2022-8-18
@TomBers

livebookNotes

Trying Nx

NX.livemd

advanced data-science exla axon nx

2022-8-18
Ryo Wakabayashi
@RyoWakabayashi

elixir-learning

RESAS から人口構成を取得する

population.livemd

tutorial advanced data-science req explorer kino kino_vega_lite

2026-7-7
Ryo Wakabayashi
@RyoWakabayashi

elixir-learning

Text classification

text_classification.livemd

tutorial advanced data-science bumblebee nx exla kino

2026-7-7
Joel Heaps
@joelheaps

aoc24

Day 16

day16.livemd

advanced algorithms kino kino_aoc libgraph

2024-12-20

Back