LLM を Elixir / Livebook で学ぶ 2

# Nxの計算をEXLAで実行し、表とグラフも表示できるようにする
Mix.install(
  [
    {:nx, "~> 0.9"},
    {:exla, "~> 0.9"},
    {:kino, "~> 0.15"},
    {:kino_vega_lite, "~> 0.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

極小の言語モデルを学習して「学習」を理解する

第1章では、文章をトークン ID や埋め込みへ変換するところまでを見ました。この章では一歩進めて、next-token prediction を実際に学習してみます。

ただし、いきなり Transformer 全体を作るのではなく、まずは 1 個前のトークンだけを見る bigram 言語モデル を使います。

この notebook で分かること:

• 学習データはどんなペアに分解されるか
• 損失が下がると、次トークン予測がどう変わるか
• 小さなモデルでも「学習前」と「学習後」で生成が変わること
• それでも bigram には大きな限界があること

この notebook の位置づけ

これは Transformer の代わり ではありません。むしろ逆で、

• 「学習とは何か」を手で追いやすくする
• アテンションを入れる前の最小の次トークン学習を理解する

ための導入用モデルです。

入力を数値にする -> 次トークンの点数を出す -> 正解との差を損失にする -> 損失が下がるように重みを更新する というこの章の流れは、第4章のミニ GPT でも、実物の LLM でも、そのまま同じです。この章では「見る文脈」を最小の 1 トークンに絞ることで、学習そのものへ集中します。

なお、この章のモデルは埋め込みを使いません。トークン ID で重み行列の行を直接選ぶ、いちばん単純な構成です。

このnotebookで使う用語

この章では「モデルを学習する」ときに頻出する言葉が増えます。コードを読む前に、用語同士の関係だけ確認しておきます。

準備

defmodule LLMScratch.Visuals do
  # epochごとの損失など、連続した値の変化を表示する
  def line_chart(rows, title, x_field, y_field, opts \\ []) do
    width = Keyword.get(opts, :width, 620)
    height = Keyword.get(opts, :height, 280)

    VegaLite.new(width: width, height: height, 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

  # 現在トークンと次トークンの全組み合わせを色の濃さで表示する
  def heatmap(rows, title, opts \\ []) do
    width = Keyword.get(opts, :width, 620)
    height = Keyword.get(opts, :height, 360)

    VegaLite.new(width: width, height: height, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:rect, tooltip: true)
    |> VegaLite.encode_field(
      :x, "target", type: :nominal, title: "次トークン"
    )
    |> VegaLite.encode_field(
      :y, "source", type: :nominal, title: "現在トークン"
    )
    |> VegaLite.encode_field(
      :color, "score", type: :quantitative, scale: [scheme: "tealblues"]
    )
    |> Kino.VegaLite.new()
  end
end

defmodule TinyBigram do
  import Nx.Defn

  defn row_softmax(logits) do
    # 行ごとの最大値を引き、指数関数のオーバーフローを避ける
    shifted = Nx.subtract(logits, Nx.reduce_max(logits, axes: [1], keep_axes: true))
    exps = Nx.exp(shifted)
    Nx.divide(exps, Nx.sum(exps, axes: [1], keep_axes: true))
  end

  defn loss(weights, input_ids, target_one_hot) do
    # 各入力IDに対応する重み行を、次トークン候補のロジットとして取り出す
    logits = Nx.take(weights, input_ids)
    probs = row_softmax(logits)

    # 正解トークンの-log(確率)を全学習ペアで平均する
    target_one_hot
    |> Nx.multiply(Nx.log(Nx.add(probs, 1.0e-9)))
    |> Nx.sum(axes: [1])
    |> Nx.negate()
    |> Nx.mean()
  end

  defn update(weights, input_ids, target_one_hot, learning_rate) do
    # 自動微分で勾配を求め、損失が下がる向きへ重みを動かす
    grads = grad(weights, &loss(&1, input_ids, target_one_hot))
    weights - grads * learning_rate
  end

  defn probabilities(weights) do
    row_softmax(weights)
  end
end

1. 小さな学習コーパスを作る

今回は説明しやすいように、空白区切りの短い日本語文を小さく用意します。

# モデルに覚えさせる、ごく小さな学習用コーパス
corpus_sentences = [
  "ねこ は ひるね が すき",
  "ねこ は さかな が すき",
  "ねこ は まどべ で ひるね",
  "いぬ は さんぽ が すき",
  "いぬ は にわ を かける",
  "とり は そら を とぶ"
]

Kino.DataTable.new(Enum.map(corpus_sentences, &%{sentence: &1}))

2. 語彙とトークン ID

文のはじめと終わりが分かるように、<BOS> と <EOS> も入れます。

# 文の境界も予測対象にするため、文頭と文末の特殊トークンを加える
special_tokens = ["<BOS>", "<EOS>"]

# 全文をトークンへ分け、特殊トークンと合わせて重複を除く
vocab =
  corpus_sentences
  |> Enum.flat_map(&String.split(&1, " "))
  |> Kernel.++(special_tokens)
  |> Enum.uniq()
  |> Enum.sort()

# トークン文字列から整数IDを引くための対応表
token_to_id =
  vocab
  |> Enum.with_index()
  |> Map.new()

# 生成したIDをトークン文字列へ戻すための逆引き表
id_to_token =
  token_to_id
  |> Enum.map(fn {token, id} -> {id, token} end)
  |> Map.new()

%{
  vocab_size: length(vocab),
  vocab: vocab
}

# 語彙とIDの対応を確認する表示用データ
token_rows =
  vocab
  |> Enum.with_index()
  |> Enum.map(fn {token, id} ->
    %{
      token: token,
      token_id: id
    }
  end)

Kino.DataTable.new(token_rows, num_rows: 20)

3. next-token 学習用のペアへ分解する

bigram モデルでは、現在のトークン -> 次のトークン のペアを学習します。

# 各文の前後へ<BOS>と<EOS>を付け、文の境界も系列に含める
token_sequences =
  corpus_sentences
  |> Enum.map(fn sentence ->
    ["<BOS>"] ++ String.split(sentence, " ") ++ ["<EOS>"]
  end)

Kino.DataTable.new(Enum.map(token_sequences, &%{sequence: Enum.join(&1, " ")}))

# 各系列を1トークンずつずらし、「現在 -> 次」の学習問題へ分解する
training_pairs =
  token_sequences
  |> Enum.flat_map(fn sequence ->
    sequence
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.map(fn [source, target] ->
      %{
        source: source,
        target: target,
        source_id: token_to_id[source],
        target_id: token_to_id[target]
      }
    end)
  end)

Kino.DataTable.new(training_pairs, num_rows: 40)

4. 観察用に、出現回数の bigram 行列も作る

学習済みモデルの重みが、どのくらいこの分布に近づくかを見るための比較材料です。

このヒートマップは、コーパスを 現在トークン -> 次トークン の対応表として数えたものです。

• 縦軸で「現在のトークン」を1つ選ぶ
• その行を横へたどる
• 色が濃い列ほど、そのトークンが次に多く現れた

たとえば ねこ の行で は の列が濃ければ、学習データ中で ねこ -> は が何度も現れたことを表します。空白に見えるセルは、その組み合わせを一度も観察していないという意味です。

# 同じbigramがコーパス内に何回現れたかを数える
pair_counts =
  training_pairs
  |> Enum.frequencies_by(fn pair -> {pair.source, pair.target} end)

# 語彙の全組み合わせを作り、未出現の組み合わせは0回として埋める
count_rows =
  for source <- vocab, target <- vocab do
    %{
      source: source,
      target: target,
      score: Map.get(pair_counts, {source, target}, 0)
    }
  end

LLMScratch.Visuals.heatmap(count_rows, "学習コーパスの bigram 出現回数")

ヒートマップ全体を一度に読むのが難しい場合は、1行だけ表にすると分かりやすくなります。

# 候補が複数ある「は」の行だけを抜き出して確認する
focus_source = "は"

focus_count_rows =
  count_rows
  |> Enum.filter(&(&1.source == focus_source))
  |> Enum.reject(&(&1.score == 0))
  |> Enum.sort_by(& &1.score, :desc)

Kino.DataTable.new(focus_count_rows, keys: [:source, :target, :score])

は の次には複数の語が現れています。bigram モデルは は より前に ねこ があったのか いぬ があったのかを見られないので、この行にある候補のどれかを確率的に選ぶしかありません。この制約が、第3章でアテンションを学び、第4章でミニ GPT を作る理由につながります。この表の「は」の行は後の章でも再登場するので、覚えておいてください。

5. モデルに渡す tensor を作る

ここまでは、人が読める %{source: "ねこ", target: "は"}、つまり 現在トークン: "ねこ", 次トークン: "は" という形でした。しかし数値計算ライブラリ Nx に渡すには、同じ長さ・同じ型の数値の箱であるtensorにまとめます。

ここでは1件の学習例を、次の3段階で表現します。

# 人が読める学習ペアから、現在トークンと正解トークンのID列を作る
input_ids =
  training_pairs
  |> Enum.map(& &1.source_id)
  |> Nx.tensor(type: :s64)

target_ids =
  training_pairs
  |> Enum.map(& &1.target_id)
  |> Nx.tensor(type: :s64)

vocab_size = length(vocab)

# 各正解IDを、語彙数と同じ幅のone-hot表現へ変換する
target_one_hot =
  Nx.equal(Nx.new_axis(target_ids, 1), Nx.iota({vocab_size}))
  |> Nx.as_type(:f32)

%{
  input_shape: Nx.shape(input_ids),
  target_shape: Nx.shape(target_ids),
  target_one_hot_shape: Nx.shape(target_one_hot),
  vocab_size: vocab_size
}

形状の先頭の数は学習ペア数です。input_ids と target_ids は1ペアにつき1個のID、target_one_hot は1ペアにつき語彙数個の0/1を持ちます。

# 先頭5件について、文字列、ID、one-hotの対応を1行ずつ確認する
tensor_preview_rows =
  training_pairs
  |> Enum.take(5)
  |> Enum.with_index()
  |> Enum.map(fn {pair, row_index} ->
    %{
      row: row_index,
      human_pair: "#{pair.source} -> #{pair.target}",
      input_id: Nx.to_number(input_ids[row_index]),
      target_id: Nx.to_number(target_ids[row_index]),
      one_hot_1_position: pair.target_id
    }
  end)

Kino.DataTable.new(
  tensor_preview_rows,
  keys: [:row, :human_pair, :input_id, :target_id, :one_hot_1_position]
)

たとえば input_id=7 なら、モデルは重み行列の7行目を次トークン候補の点数として取り出します。one_hot_1_position は「その行の正解は何列目か」を示します。

6. 学習前の状態を見る

bigram モデルのパラメータは、[語彙数 x 語彙数] の行列です。

行は「今のトークン」、列は「次トークン候補」を表します。

# 同じ初期値から実験を再現できるよう、乱数の初期状態を固定する
:rand.seed(:exsss, {123, 456, 789})

# {現在トークン, 次トークン候補}の全組み合わせに小さな乱数を置く
initial_weights =
  1..(vocab_size * vocab_size)
  |> Enum.map(fn _ -> :rand.normal() * 0.01 end)
  |> Nx.tensor(type: {:f, 32})
  |> Nx.reshape({vocab_size, vocab_size})

Nx.shape(initial_weights)

# 学習前の重みで、コーパス全体に対する損失を測る
initial_loss =
  TinyBigram.loss(initial_weights, input_ids, target_one_hot)
  |> Nx.to_number()

%{initial_loss: initial_loss}

summarize_predictions = fn probabilities ->
  # 各現在トークンについて、確率が高い次トークンを3件ずつ選ぶ
  vocab
  |> Enum.with_index()
  |> Enum.map(fn {source, source_id} ->
    top_predictions =
      probabilities[source_id]
      |> Nx.to_flat_list()
      |> Enum.zip(vocab)
      |> Enum.map(fn {prob, token} -> {token, prob} end)
      |> Enum.sort_by(fn {_token, prob} -> prob end, :desc)
      |> Enum.take(3)
      |> Enum.map(fn {token, prob} -> "#{token}(#{Float.round(prob, 3)})" end)
      |> Enum.join(" / ")

    %{
      source: source,
      top3: top_predictions
    }
  end)
end

initial_probabilities = TinyBigram.probabilities(initial_weights)
Kino.DataTable.new(summarize_predictions.(initial_probabilities), num_rows: 20)

学習前はほぼランダムなので、ねこ の次に何が来るかもまだ分かっていません。

7. 学習を回す

今回は小さな全データをまとめて使う、フルバッチ学習にします。

1 epoch では、次の4段階を1回ずつ行います。

Kino.Mermaid.new("""
flowchart
  A["a. 全 training pairs を入力"] --> B["b. 次トークン確率を計算"]
  B --> C["c. 正解とのずれを損失にする"]
  C --> D["d. 損失が下がる方向へ重みを少し更新"]
  D --> A
""")

損失（loss）は、正解トークンに十分な確率を出せなかった度合いです。値そのものを暗記する必要はなく、まずは 学習が進むにつれて下がっているか を見ます。learning_rate は1回に重みをどれだけ動かすか、epochs はこの練習を何回繰り返すかです。

epochs = 400
learning_rate = 1.2e-1

# {現在の重み, 損失の履歴}をepochごとに更新する
{trained_weights, loss_rows} =
  Enum.reduce(0..epochs, {initial_weights, []}, fn epoch, {weights, acc} ->
    # 更新前の損失を記録してから、勾配降下を1回行う
    loss_value =
      TinyBigram.loss(weights, input_ids, target_one_hot)
      |> Nx.to_number()

    next_weights =
      TinyBigram.update(weights, input_ids, target_one_hot, learning_rate)

    row = %{epoch: epoch, loss: loss_value}

    # 先頭への追加は高速なので、最後に履歴全体を反転する
    {next_weights, [row | acc]}
  end)

loss_rows = Enum.reverse(loss_rows)

%{
  first_loss: hd(loss_rows).loss,
  last_loss: List.last(loss_rows).loss
}

LLMScratch.Visuals.line_chart(loss_rows, "学習中の損失", :epoch, :loss)

グラフの左から右へ進むにつれて損失が下がり、やがて変化が小さくなれば、モデルはこのコーパスから学べる規則をだいたい取り込んだと読めます。

8. 学習後の予測を見る

学習後は、重み行列の各行を確率へ変換して読みます。次の表では、各現在トークンのあとにモデルが高い確率を付けた上位3候補を表示します。

# 学習後の各重み行を、合計1の次トークン確率へ変換する
trained_probabilities = TinyBigram.probabilities(trained_weights)
Kino.DataTable.new(summarize_predictions.(trained_probabilities), num_rows: 20)

# {現在トークン, 次トークン}ごとの学習後確率をヒートマップ用に並べる
trained_rows =
  for {source, i} <- Enum.with_index(vocab),
      {target, j} <- Enum.with_index(vocab) do
    %{
      source: source,
      target: target,
      score: Float.round(Nx.to_number(trained_probabilities[[i, j]]), 4)
    }
  end

LLMScratch.Visuals.heatmap(trained_rows, "学習後の next-token 確率")

このヒートマップは先ほどの「出現回数」ではなく、モデルが学習後に出している 確率 です。

• <BOS> の行では、文頭に来やすい ねこ、いぬ、とり へ確率が集まる
• ねこ の行では、次に頻出した は が強くなる
• は の行では複数候補へ確率が分かれ、bigram の曖昧さが残る

出現回数ヒートマップと似た濃淡になれば、モデルがデータの傾向を重みに写し取ったと判断できます。

# 特徴の分かりやすい3トークンに絞って、学習前後の1位候補を比べる
comparison_sources = ["<BOS>", "ねこ", "は"]

before_after_rows =
  for source <- comparison_sources do
    source_id = token_to_id[source]

    # 1行分の確率から、最も確率が高い{確率, トークン}を返す
    best = fn probabilities ->
      probabilities[source_id]
      |> Nx.to_flat_list()
      |> Enum.zip(vocab)
      |> Enum.max_by(fn {probability, _token} -> probability end)
    end

    {before_probability, before_token} = best.(initial_probabilities)
    {after_probability, after_token} = best.(trained_probabilities)

    %{
      source: source,
      before_top: "#{before_token} (#{Float.round(before_probability, 3)})",
      after_top: "#{after_token} (#{Float.round(after_probability, 3)})"
    }
  end

Kino.DataTable.new(before_after_rows, keys: [:source, :before_top, :after_top])

この比較では、ランダムだった候補が、コーパスで実際によく続いた候補へ変わるところを確認してください。

9. 学習前後で生成を比べる

bigram なので、毎回「今の 1 語」だけを見て次トークンを選びます。

greedy_generate = fn probabilities, max_steps ->
  # 現在トークンの行から、確率最大の次トークンを1つ選ぶ
  next_token = fn source ->
    source_id = token_to_id[source]

    probabilities[source_id]
    |> Nx.to_flat_list()
    |> Enum.zip(vocab)
    |> Enum.map(fn {prob, token} -> {token, prob} end)
    |> Enum.reject(fn {token, _prob} -> token == "<BOS>" end)
    |> Enum.max_by(fn {_token, prob} -> prob end)
    |> elem(0)
  end

  # <BOS>から始め、<EOS>が出るか上限へ達するまで1語ずつ追加する
  1..max_steps
  |> Enum.reduce_while(["<BOS>"], fn _, acc ->
    source = List.last(acc)
    target = next_token.(source)

    if target == "<EOS>" do
      {:halt, acc ++ [target]}
    else
      {:cont, acc ++ [target]}
    end
  end)
end

format_generated = fn tokens ->
  # 人が読みやすいよう、特殊トークンを除いて空白で連結する
  tokens
  |> Enum.reject(&(&1 in ["<BOS>", "<EOS>"]))
  |> Enum.join(" ")
end

generation_rows = [
  %{
    stage: "before training",
    generated: format_generated.(greedy_generate.(initial_probabilities, 8))
  },
  %{
    stage: "after training",
    generated: format_generated.(greedy_generate.(trained_probabilities, 8))
  }
]

Kino.DataTable.new(
  generation_rows,
  keys: [:stage, :generated]
)

学習前は不自然な並びになりやすく、学習後はコーパスに寄った語順へ近づきます。

10. それでも bigram には限界がある

ここが大事です。

このモデルは「直前の 1 語」しか見ません。だから、

• ねこ は のあとに何が来るか（は しか見えないので、ねこ の情報を使えない）
• もっと前に出てきた主語や話題
• 長い依存関係

をほとんど扱えません。

たとえばこのコーパスでは、ねこ は の続きと いぬ は の続きは明らかに違います。しかし bigram モデルにとってはどちらも「現在トークン = は」なので、まったく同じ確率分布しか出せません。

LLM が Transformer を使うのは、こうした 長い文脈を読む必要 があるからです。この「は問題」が解けるようになる瞬間を、第4章で実際に確認します。

11. まとめ

この notebook の要点

• next-token 学習は、現在トークン -> 次トークン の予測問題として書ける
• 損失が下がると、次トークン確率の分布がデータに沿って整っていく
• ごく小さいモデルでも、学習前後の生成差ははっきり見える
• ただし bigram は 1 語しか見ないので、LLM の表現力には遠い

次のノートブックでは、bigram の限界を出発点に、

• 前にある全トークンを「重み付き」で参照するアテンション
• 未来のトークンを見ないための因果マスク
• 語順を扱うための位置エンコーディング

という、Transformer の中心となる 3 つの仕組みを可視化しながら学びます。

そして第4章では、この章とまったく同じコーパスでミニ GPT を学習し、bigram では出せなかった「ねこ は と いぬ は で違う予測」が実際に出せるようになることを確認します。