Elixir と Livebook で学ぶ LLM / Transformer 入門 〜 トークン化からミニGPTの自作学習まで

はじめに

LLM（大規模言語モデル）や Transformer の仕組みは、記事や動画で断片的に説明されていることが多く、「アテンションだけ分かった」「学習の流れだけ分かった」という状態で終わってしまいがちです

この記事では、Elixir と Livebook を使って、LLM の中心にある1つの問いを少しずつ広げながら理解していく教材シリーズを紹介します

モデルは、何を手がかりに次のトークンを予測するのか？

対象読者は次のような方です

• Elixir はある程度書けるが、LLM や Transformer の内部構造は詳しく知らない
• 数式だけの説明より、実際にコードを動かして表やグラフで確認したい

ノートブックは GitHub で公開しています

事前準備

Livebook のインストール方法は、公式サイトを参照してください
デスクトップアプリ・Docker・Fly.io など複数の方法があります

以下のバッジから、各ノートブックを直接 Livebook で開くこともできます（Livebook が起動していれば、そのままクリックして開けます）

シリーズの入口はこちらです

シリーズの全体構成

このシリーズは本編6章と発展編1章の、合計7つのノートブックで構成しています

章	ノートブック	内容
1	01_llm_overview_tokens_embeddings.livemd	LLMの全体像・トークン化・one-hotの限界・埋め込み
2	02_tiny_bigram_language_model_training.livemd	極小の言語モデル（bigram）を実際に学習する
3	03_attention_causal_mask_position.livemd	アテンション・因果マスク・位置エンコーディング
4	04_mini_gpt_from_scratch.livemd	ミニGPTを組み立てて、埋め込みからアテンションまで全部学習する
5	05_decoder_generation_with_bumblebee.livemd	事前学習済み GPT-2 で次トークン予測と生成を体感する
6	06_alignment_data_and_preferences.livemd	SFT / DPO によるアラインメントの直感を学ぶ
発展編	07_advanced_encoder_decoder_transformer.livemd	エンコーダー・デコーダー構成とクロスアテンション

第1章: LLMの全体像・トークン化・埋め込み

最初の章では、細かい部品に入る前に「GPT系モデルは結局何をしているのか」を確認します

GPT系モデルは、ここまでのトークンを読み、次の1トークンを予測する ことを繰り返しているだけです
この章では、図の a と b、つまり「文章をトークンに分けて、ベクトルにする」ところまでを扱います

トークン化と one-hot の限界

まずは空白区切りの簡単な文をトークンに分けます

sentence = "ねこ は ひるね が すき"
tokens = String.split(sentence, " ")

トークンには整数の ID を振りますが、ID をそのまま計算に使うのは無理があります

そこで最初に思いつくのが one-hot 表現です
ただし、one-hot には大きな弱点があります
どの2語を比べても「まったく似ていない」ことになるのです

語彙の全トークン同士で内積（ベクトルの近さを測る計算）を取ってみると、それがよく分かります

# 語彙の全トークン分のone-hotを作り、全組み合わせの内積を求める
vocab_one_hot =
  Nx.equal(
    Nx.new_axis(Nx.iota({vocab_size}), 1),
    Nx.iota({vocab_size})
  )
  |> Nx.as_type(:f32)

one_hot_similarity = Nx.dot(vocab_one_hot, [1], vocab_one_hot, [1])

one_hot_similarity_rows =
  for {token_a, i} <- Enum.with_index(vocab),
      {token_b, j} <- Enum.with_index(vocab) do
    %{
      token_a: token_a,
      token_b: token_b,
      dot: Nx.to_number(one_hot_similarity[[i, j]])
    }
  end

LLMScratch.Visuals.heatmap(
  one_hot_similarity_rows,
  "one-hot 表現の内積（全組み合わせ）",
  :token_b,
  :token_a,
  :dot
)

実行結果

自分自身との内積は 1.0、それ以外はすべて 0.0 になります

つまり one-hot の世界では、「ねこ」と「ひるね」の関係も、「ねこ」と「が」の関係も、まったく同じ「無関係」として扱われてしまいます

埋め込みなら「近さ」を表現できる

そこで LLM では、one-hot の代わりに密な実数ベクトルである埋め込み（embedding）を使います
この章では説明用に、人手で2次元の埋め込みを割り当てます

embedding_map = %{
  "ねこ" => [0.90, 0.15],
  "は" => [0.10, 0.05],
  "ひるね" => [0.95, 0.85],
  "が" => [0.12, 0.10],
  "すき" => [0.88, 0.70]
}

内容語（ねこ・ひるね・すき）を右上寄りに、助詞（は・が）を左下寄りに置いてあります
同じように内積を計算すると、今度は結果が変わります

embedding_similarity = Nx.dot(embedding_tensor, [1], embedding_tensor, [1])

実行結果（「すき」から見た各トークンとの内積）

相手のトークン	内積
ねこ	0.897
は	0.123
ひるね	1.431
が	0.176
すき	1.264

one-hot ではどの組み合わせも 0 でしたが、埋め込みでは すき と ひるね の内積が大きく、すき と は の内積は小さくなります
ベクトルの向きが近いほど内積が大きい、というこの感覚が、次章以降のアテンションの土台になります

第2章: 極小の言語モデル（bigram）を学習する

第1章で「文章を数値に直す」ところまで進んだので、第2章では実際に学習をしてみます
ただし、いきなり Transformer 全体は作らず、1個前のトークンだけを見る bigram 言語モデル を使います

学習用のコーパスには、この後の章でも使い回す6つの短い文を用意しました

corpus_sentences = [
  "ねこ は ひるね が すき",
  "ねこ は さかな が すき",
  "ねこ は まどべ で ひるね",
  "いぬ は さんぽ が すき",
  "いぬ は にわ を かける",
  "とり は そら を とぶ"
]

現在のトークン -> 次のトークン のペアに分解し、Nx.Defn.grad で自動微分しながら重み行列を更新していきます

defn loss(weights, input_ids, target_one_hot) do
  logits = Nx.take(weights, input_ids)
  probs = row_softmax(logits)

  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

学習が進むにつれて損失が下がり、学習前はランダムだった予測が、コーパスの傾向に沿ったものへ変わっていきます

bigram の限界（この章でいちばん大事な観察）

この章の主役は、次の表です

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章: アテンション・因果マスク・位置エンコーディング

bigram の限界を解決するために、Transformer の中心となる3つの仕組みを実際に計算しながら学びます

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

まずは「均等に混ぜる」を試して、限界を見る

「前のトークンを全部使えばいい」という発想の、いちばん単純な形は均等な重みで平均することです

uniform_weight = 1.0 / length(tokens)
uniform_weights = Nx.broadcast(uniform_weight, {length(tokens)})
uniform_mixture = Nx.dot(uniform_weights, [0], embedding_tensor, [0])

実行結果: [0.59, 0.37]

文全体の情報は手に入りますが、ひるね のような大事な語も は のような機能語も同じ重み（0.2）で混ざってしまいます
必要なのは、今の位置が何を知りたいかに応じて、参照先ごとに重みを変える仕組みです

アテンション: Query / Key / Value

scaled_dot_attention = fn query, keys, values ->
  d_model = keys |> Nx.shape() |> elem(1)

  scores = Nx.divide(Nx.dot(keys, [1], query, [0]), :math.sqrt(d_model))
  weights = softmax_1d.(scores)

  output = Nx.dot(weights, [0], values, [0])
  {output, weights}
end

すき を Query として全トークンを参照させると、均等な 0.2 とは違い、意味的に近い ひるね や ねこ への重みが大きくなります

名前	図書館でのたとえ	アテンションでの役割
Query	今探しているテーマ	現在位置が「どんな情報がほしいか」を表す
Key	各資料の見出し	Query と比べ、参照先としてどれくらい合うかを測る
Value	資料の中身	Key が合った位置から、実際に持ち帰って混ぜる情報

因果マスク: 未来を見せない

GPT の学習では、完成した文章を 現在まで -> 次の正解 にずらして使います
もし右側の答えまで見えていたら、モデルは答えを写すだけで損失を下げられてしまいます

mask_tensor =
  for i <- 0..(length(decoder_tokens) - 1) do
    for j <- 0..(length(decoder_tokens) - 1) do
      j > i
    end
  end
  |> Nx.tensor(type: {:u, 8})

masked_scores =
  Nx.select(mask_tensor, Nx.broadcast(-1.0e9, Nx.shape(score_matrix)), score_matrix)

ヒートマップの右上（未来への参照）が 0 になっていれば、各トークンが自分より右側を見ていないことが確認できます

位置エンコーディング: 語順を補う

アテンションの計算だけでは、どのトークンが何番目にあったか を直接は持ちません
そこで sin / cos を使った位置エンコーディングを、埋め込みに足し込みます

positional_encoding = fn max_position, d_model ->
  for pos <- 0..(max_position - 1) do
    for dim <- 0..(d_model - 1) do
      angle = pos / :math.pow(10_000, (2 * div(dim, 2)) / d_model)
      if rem(dim, 2) == 0, do: :math.sin(angle), else: :math.cos(angle)
    end
  end
  |> Nx.tensor(type: {:f, 32})
end

これで、Transformer を組み立てるための部品がそろいました

第4章: ミニGPTを組み立てて、端から端まで学習する

この章がシリーズの山場です
第1〜3章の部品を全部つなげて、デコーダーのみのミニGPTを実装し、第2章とまったく同じコーパスで、埋め込みからアテンションまで全パラメータを学習します

第3章のアテンションに、この章で初めて次の4つが加わります

部品	役割のイメージ
マルチヘッドアテンション	アテンションを複数の「見方」に分けて並列に計算し、最後に結合する
フィードフォワード層	アテンションで集めた情報を、各位置ごとに個別に加工する
残差接続	層の入力を出力へ足し戻し、元の情報が消えないようにする
レイヤー正規化	各位置のベクトルのスケールを整え、学習を安定させる

マルチヘッド自己アテンション

def self_attention(x, params, num_heads, mask) do
  q = x |> linear(params.wq, params.bq) |> split_heads(num_heads)
  k = x |> linear(params.wk, params.bk) |> split_heads(num_heads)
  v = x |> linear(params.wv, params.bv) |> split_heads(num_heads)

  head_dim = q |> Nx.shape() |> elem(3)

  scores =
    Nx.dot(q, [3], [0, 1], k, [3], [0, 1])
    |> Nx.divide(:math.sqrt(head_dim))

  masked_scores =
    Nx.select(
      Nx.broadcast(mask, Nx.shape(scores)),
      Nx.broadcast(-1.0e9, Nx.shape(scores)),
      scores
    )

  weights = softmax(masked_scores)

  output =
    Nx.dot(weights, [3], [0, 1], v, [2], [0, 1])
    |> combine_heads()
    |> linear(params.wo, params.bo)

  {output, weights}
end

GPTブロックは、このアテンションとフィードフォワード層を、残差接続とレイヤー正規化で包んだものです

def gpt_block(x, params, num_heads, mask) do
  {attention_output, attention_weights} =
    self_attention(x, params.attention, num_heads, mask)

  x1 = layer_norm(Nx.add(x, attention_output), params.ln1.gamma, params.ln1.beta)
  ff_output = feed_forward(x1, params.feed_forward)
  x2 = layer_norm(Nx.add(x1, ff_output), params.ln2.gamma, params.ln2.beta)

  {x2, attention_weights}
end

パラメータは木構造、学習の流れは第2章と同じ

第2章では重みが1枚の行列でしたが、ミニGPTのパラメータは map と tuple を入れ子にした木になります
Nx.Defn.value_and_grad は、その木と同じ形をした勾配をまとめて返してくれるので、木の葉（tensor）1つ1つに勾配降下を適用する関数を用意します

defmodule LLMScratch.SGD do
  def apply_gradients(%Nx.Tensor{} = param, %Nx.Tensor{} = gradient, learning_rate) do
    Nx.subtract(param, Nx.multiply(gradient, learning_rate))
  end

  def apply_gradients(params, gradients, learning_rate) when is_map(params) do
    Map.new(params, fn {key, value} ->
      {key, apply_gradients(value, Map.fetch!(gradients, key), learning_rate)}
    end)
  end

  def apply_gradients(params, gradients, learning_rate) when is_tuple(params) do
    params
    |> Tuple.to_list()
    |> Enum.zip(Tuple.to_list(gradients))
    |> Enum.map(fn {param, gradient} -> apply_gradients(param, gradient, learning_rate) end)
    |> List.to_tuple()
  end
end

学習ループ自体は、第2章の weights - grads * learning_rate を全パラメータに広げただけです

value_and_grad =
  Nx.Defn.value_and_grad(fn {params, ids, targets} ->
    LLMScratch.MiniGPT.loss(params, ids, targets, config.num_heads)
  end)

{trained_params, loss_rows} =
  Enum.reduce(0..epochs, {initial_params, []}, fn epoch, {params, rows} ->
    {loss, {gradients, _ids_grad, _targets_grad}} =
      value_and_grad.({params, input_ids, target_one_hot})

    next_params = LLMScratch.SGD.apply_gradients(params, gradients, learning_rate)
    {next_params, [%{epoch: epoch, loss: Nx.to_number(loss)} | rows]}
  end)

d_model: 16、num_heads: 2、d_ff: 32、ブロック数2 という設定で、パラメータ数は約5,000個です

GPT-2（第5章で使用）が約1.2億個、近年の大規模LLMが1兆個規模なのと比べると、けた違いに小さいモデルですが、構造はそのまま同じです

bigram が解けなかった問題が解けた瞬間

実際に学習を実行した結果がこちらです（EXLA バックエンドで学習時間は約6.5秒でした）

実行結果（損失の推移）

損失は 0 まで下がらず、0.30 付近で止まります
これは失敗ではありません
このコーパスには <BOS> の次や ねこ は の次など、文脈を完全に読んでも一意に決まらない箇所が残っているためで、損失の下限はモデルの能力ではなくデータの曖昧さで決まることがここで確認できます

そして、第2章で「は問題」として残した3つのプレフィックスに、学習後のモデルがどう答えるかがこちらです

実行結果（学習後の次トークン予測、上位3件）

<BOS> ねこ は: [
  {"ひるね", 0.380}, {"まどべ", 0.308}, {"さかな", 0.307}
]
<BOS> いぬ は: [
  {"さんぽ", 0.498}, {"にわ", 0.496}, {"そら", 0.002}
]
<BOS> とり は: [
  {"そら", 0.991}, {"にわ", 0.002}, {"さかな", 0.002}
]

3行とも直前のトークンは同じ は ですが、bigram では出せなかったそれぞれ違う予測が出ています
しかも、確率の値は、コーパス内の出現割合（ねこ は の続きは3候補で約1/3ずつ、いぬ は の続きは2候補で約1/2ずつ、とり は の続きは そら のみ）とほぼ一致しています
因果マスク付き自己アテンションが、2つ前の主語トークンを参照できているということです

最後に、学習済みモデルで貪欲法（greedy）による文章生成も試します

generate = fn params, max_steps ->
  Enum.reduce_while(1..max_steps, ["<BOS>"], fn _, acc ->
    [{next, _probability} | _] = next_token_probs.(params, acc)

    if next == "<EOS>" or length(acc) >= max_steps do
      {:halt, acc ++ [next]}
    else
      {:cont, acc ++ [next]}
    end
  end)
end

実行結果: ["<BOS>", "ねこ", "は", "ひるね", "が", "すき", "<EOS>"]

コーパスに含まれる自然な文がそのまま生成されました

第5章: 事前学習済み GPT-2 で生成を体感する

第4章で作ったミニGPTは、語彙18個・パラメータ約5,000個のおもちゃのようなモデルでした
第5章では、Bumblebee を使って本物の GPT-2（パラメータ約1.2億個）を動かし、同じ仕組みが実物になるとどう見えるかを確認します

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章では空白区切りの簡易トークナイザーを使いましたが、実物の GPT-2 は単語よりも細かい「サブワード」という単位を使います

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

Tokenization のような単語が Token と ization のように分割される様子が確認できます

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 の場合

試行 1
The first thing to note is that the API is not a singleton. It is a collection of objects. Each object is a collection of objects.

The first thing to note is

---

試行 2
The first thing to note is that the API is not a singleton. It is a collection of objects. Each object is a collection of objects.

The first thing to note is

---

試行 3
The first thing to note is that the API is not a singleton. It is a collection of objects. Each object is a collection of objects.

The first thing to note is

sampling の場合

試行 1
In the past, it has been suggested that it was just a minor nuisance or an obvious mistake and is now considered a bug in Ethereum and potentially even a fatal error. The only other reason Ethereum is

---

試行 2
For example, you may be able to figure out why the user has been turned off for a short period of time in a given set of contexts. In the case of a single example, the fact

---

試行 3
It may cause the system to shut down immediately after a critical operation to prevent that application from having enough data to do its job.

With all this thought, let's try something new. This

同じ入力でも、greedy は毎回同じ結果になり、sampling は実行のたびに違う続きが選ばれます
「創造性」という特別な能力が足されたわけではなく、次トークン候補から常に1位を取るか、上位候補を抽選するか という小さな違いが、生成全体の分岐につながっているだけだと分かります

第6章: SFT / DPO によるアラインメント

事前学習だけを終えたモデルは、文章の続きは予測できても、必ずしも人の指示に分かりやすく答えるとは限りません
第6章では、その調整方法である SFT（教師あり微調整）と DPO（選好最適化）を、小さなデータで体験します

SFT では、指示 -> 望ましい応答 のペアを正解として学習します

defn sft_loss(policy_logits, prompt_ids, chosen_one_hot) do
  log_probs =
    policy_logits
    |> Nx.take(prompt_ids)
    |> log_softmax()

  log_probs
  |> Nx.multiply(chosen_one_hot)
  |> Nx.sum(axes: [1])
  |> Nx.negate()
  |> Nx.mean()
end

DPO では、同じ指示に対する chosen（望ましい応答）と rejected（望ましくない応答）を比べ、基準モデルから離れすぎないようにしながら、chosen を相対的に選びやすくします

defn dpo_loss(policy_logits, reference_logits, prompt_ids, chosen_one_hot, rejected_one_hot, beta) do
  policy_log_probs = policy_logits |> Nx.take(prompt_ids) |> log_softmax()
  reference_log_probs = reference_logits |> Nx.take(prompt_ids) |> log_softmax()

  policy_chosen = Nx.sum(policy_log_probs * chosen_one_hot, axes: [1])
  policy_rejected = Nx.sum(policy_log_probs * rejected_one_hot, axes: [1])
  reference_chosen = Nx.sum(reference_log_probs * chosen_one_hot, axes: [1])
  reference_rejected = Nx.sum(reference_log_probs * rejected_one_hot, axes: [1])

  relative_margin =
    Nx.subtract(policy_chosen, policy_rejected) -
      Nx.subtract(reference_chosen, reference_rejected)

  preferred_probability =
    Nx.divide(1.0, Nx.add(1.0, Nx.exp(Nx.negate(beta * relative_margin))))

  preferred_probability |> Nx.add(1.0e-9) |> Nx.log() |> Nx.negate() |> Nx.mean()
end

学習前は chosen と rejected が同確率、SFT 後は chosen の確率が上がり、DPO 後は chosen と rejected の対数確率マージンがさらに広がる、という3段階の変化を確認できます

発展編: エンコーダー・デコーダー構成とクロスアテンション

本編はここまでの GPT 系（デコーダーのみ）で完結していますが、Transformer が最初に発表された論文は機械翻訳のための、エンコーダーとデコーダーを組み合わせた構成でした
発展編では、その元祖の構成を組み立てます

デコーダーブロックには、第4章のミニGPTにはなかったクロスアテンションが追加されています
デコーダー側をQuery、エンコーダーの記憶をKey/Valueとして参照する仕組みです

def decoder_block(x, memory, params, num_heads, causal_mask) do
  {self_attn_out, self_weights} =
    multi_head_attention(x, x, x, params.self_attention, num_heads, causal_mask)

  x1 = residual_layer_norm(x, self_attn_out, params.ln1)

  {cross_attn_out, cross_weights} =
    multi_head_attention(x1, memory, memory, params.cross_attention, num_heads)

  x2 = residual_layer_norm(x1, cross_attn_out, params.ln2)
  ff_out = feed_forward(x2, params.feed_forward)
  x3 = residual_layer_norm(x2, ff_out, params.ln3)

  {x3, %{self_weights: self_weights, cross_weights: cross_weights}}
end

英日翻訳の小さな例（i really like green tea -> わたし は 緑茶 が 好き など4例）を使い、出力層を学習すると、デコーダー入力が同じ <BOS> わたし は でも、クロスアテンションが参照するエンコーダー側の記憶（green tea / cats / coffee / books）によって、続く目的語を選び分けられるようになります

まとめ

• LLM の中心は「ここまでのトークンから次の1トークンを予測する」というシンプルな問い
• one-hot では語同士の近さを表せず、埋め込みなら内積で近さを測れる
• bigram モデルには「直前の1語しか見られない」という原理的な限界がある
• アテンション・因果マスク・位置エンコーディングを組み合わせると、その限界を超えられる
• ミニGPTを自分の手で端から端まで学習すると、bigram には出せなかった予測が実際に出せる
• 実物の GPT-2 も、この教材で作ったミニGPTと同じ構造をスケールアップしたもの
• SFT / DPO は、事前学習済みモデルを人の指示や好みに合わせる調整の入り口
• エンコーダー・デコーダー構成とクロスアテンションを知ると、翻訳のような別系列間のタスクにも応用が効く

気になった章から、ぜひ手元の Livebook で実行してみてください