Notesclub

created by hec & contributors

terms privacy

LLM を Elixir / Livebook で学ぶ発展編

07_advanced_encoder_decoder_transformer.livemd

Ryo Wakabayashi

@RyoWakabayashi

elixir-learning

Share to X

Share to Bluesky

More notebooks

LLM を Elixir / Livebook で学ぶ発展編

# Transformerの数値計算と可視化に使うライブラリを準備する
Mix.install([
  {:nx, "~> 0.9"},
  {:kino, "~> 0.15"},
  {:kino_vega_lite, "~> 0.1"}
])

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

本編（第1〜6章）では、GPT 系のデコーダーのみの構成を扱いました。実は、Transformer が最初に発表された論文（Attention Is All You Need, 2017）は機械翻訳のためのもので、エンコーダーとデコーダーを組み合わせた構成でした。

この発展編では、その元祖の構成を組み立てます。翻訳のように「読む文」と「作る文」が別の系列であるタスクで、Transformer がどう働くかを見ます。

Kino.Mermaid.new("""
flowchart
  A["入力を数値にする"] --> B["次トークンの点数を出す"]
  B --> C["正解との差を損失にする"]
  C --> D["損失が下がるように重みを更新する"]
  D --> A
""")

学習の基本の流れは本編と変わりません。変わるのは、次を予測するときに使える情報 です。

モデル	次トークンを予測するために見るもの	得意なタスク
第4章のミニGPT（デコーダーのみ）	同じ系列の左側	文章の続きの生成
この章のTransformer（エンコーダー・デコーダー）	入力文全体 + 出力文の左側	翻訳・要約など、別の系列への変換

たとえばデコーダーが <BOS> わたしは まで読んだ時点を考えます。第2章の bigram に見える最後のトークンはどの例でも は でした。エンコーダー・デコーダー構成の Transformer は、エンコーダーが読んだ英語側を参照できるため、次を 緑茶、猫、本 などに選び分けられます。

この違いを作るために、第4章までに登場した部品（マルチヘッドアテンション・位置エンコーディング・フィードフォワード層・残差接続・レイヤー正規化）に、この章で初登場する

エンコーダー（入力文全体を読み、記憶を作る）
クロスアテンション（デコーダーからエンコーダーの記憶を参照する）

を加えて、Transformer 全体を 1 本の流れとして実装します。

この notebook の狙い

この notebook は2段階で進みます。

まずランダムな重みで各部品を動かし、データの流れと形状を見る
次に複数の翻訳例を使い、最後の出力層を学習して予測の変化を見る

全パラメータをまとめて学習する流れは第4章のミニGPTですでに体験しました。この章の主役はエンコーダーとクロスアテンションによるデータの流れなので、後半の小さな実験では、Transformer本体を文脈を作る装置として固定し、デコーダーの表現 -> 次トークン を結ぶ出力層だけを学習します。実用モデルはTransformer全体を学習しますが、損失・更新・学習前後の比較は同じ考え方です。

このnotebookで使う用語

この章ではTransformerの部品名がまとめて登場します。まず「どこで、何を見る部品か」を中心に整理します。

用語	英語・コード上の表記	この章での意味
Transformer	Transformer	アテンションなどの部品を積み重ね、文脈を扱うニューラルネットワーク
入力文／出力文	source / target	エンコーダーが読む文と、デコーダーが生成する文
エンコーダー	encoder	入力文全体を読み、デコーダーが参照する記憶を作る部分
デコーダー	decoder	ここまでの出力文とエンコーダーの記憶から、次トークンを予測する部分
自己アテンション	self-attention	同じ系列の中で、各位置がほかの位置を参照するアテンション
因果マスク付き自己アテンション	causal self-attention	出力文の未来を見ないよう、因果マスクを適用した自己アテンション
クロスアテンション	cross-attention	デコーダーからエンコーダーの記憶を参照するアテンション
マルチヘッドアテンション	multi-head attention	複数の異なる見方でアテンションを並列計算し、最後に結合する仕組み
位置エンコーディング	positional encoding	系列内の位置を埋め込みベクトルへ加える情報
フィードフォワード層	feed-forward layer	各位置のベクトルを、同じ小さなネットワークで個別に変換する層
残差接続	residual connection	層の入力を層の出力へ足し戻し、元の情報を流れやすくする接続
レイヤー正規化	layer normalization	各位置のベクトルの平均や分散を整え、計算を安定させる処理
モデル次元	`d_model`	各トークン位置を表すベクトルの要素数
ヘッド	attention head	マルチヘッドアテンションを構成する、1つの独立した見方
ロジット	logits	出力語彙の各トークンへ付ける、確率へ変換する前の点数

全体像

Kino.Mermaid.new("""
flowchart
  A["入力文のトークン ID"] --> B["入力文の埋め込み + 位置情報"]
  B --> C["エンコーダーブロック × N"]
  C --> D["エンコーダーの記憶"]
  E["出力文のトークン ID"] --> F["出力文の埋め込み + 位置情報"]
  F --> G["デコーダーブロック × N"]
  D --> G
  G --> H["出力語彙への線形変換"]
  H --> I["次トークン候補のスコア"]
""")

エンコーダーブロックとデコーダーブロック

Kino.Mermaid.new("""
flowchart TD
  subgraph encoder_block["エンコーダーブロック"]
    E1["マルチヘッド自己アテンション"] --> E2["残差接続 + レイヤー正規化"]
    E2 --> E3["フィードフォワード層"]
    E3 --> E4["残差接続 + レイヤー正規化"]
  end

  subgraph decoder_block["デコーダーブロック"]
    D1["因果マスク付きマルチヘッド自己アテンション"] --> D2["残差接続 + レイヤー正規化"]
    D2 --> D3["エンコーダーの記憶へのクロスアテンション"]
    D3 --> D4["残差接続 + レイヤー正規化"]
    D4 --> D5["フィードフォワード層"]
    D5 --> D6["残差接続 + レイヤー正規化"]
  end
""")

この 2 つは似ていますが、役割は少し違います。

エンコーダーは入力文を読んで、「この文は何を言っているか」を表す記憶を作る
デコーダーは出力文の途中経過を見ながら、必要に応じてエンコーダーの記憶を参照して次トークンを決める

部品	使用場所	参照範囲	役割
マルチヘッド自己アテンション	エンコーダー / デコーダー	同じ系列の各位置	文脈を読んで各位置の表現を更新する
因果マスク付き自己アテンション	デコーダーのみ	出力文の左側だけ	未来を見ずに次トークン予測を成立させる
クロスアテンション	デコーダーのみ	エンコーダーの記憶	入力文の情報を出力文の生成へ持ち込む
フィードフォワード層	エンコーダー / デコーダー	各位置を独立に変換	アテンション後の表現をさらに変換する
残差接続 + レイヤー正規化	エンコーダー / デコーダー	入力 + サブレイヤー出力	情報を保ちつつ学習を安定させる

マルチヘッドアテンションを直感で捉える（第4章の復習）

1本のアテンションだけだと、全部の関係を1種類の重み行列で表そうとします。

マルチヘッドアテンションでは、埋め込みを複数の小さな部分空間へ分けて、別々の見方 を並列に学べるようにします。

Kino.Mermaid.new("""
flowchart
  A["入力の埋め込み"] --> B["Query / Key / Value へ線形変換"]
  B --> C["複数のヘッドへ分割"]
  C --> D1["ヘッド 0 のアテンション"]
  C --> D2["ヘッド 1 のアテンション"]
  D1 --> E["各ヘッドを連結"]
  D2 --> E
  E --> F["出力用の線形変換"]
""")

初学者向けには、こう考えると分かりやすいです。

head 0 は語順や主語っぽさを見るかもしれない
head 1 は目的語っぽさや関連語を見るかもしれない

実際に何を見るかは学習で決まりますが、別の観点を並列に持てる のが本質です。

準備

defmodule LLMScratch.Visuals do
  # 2つのカテゴリ軸と数値を、色の濃さで読めるヒートマップにする
  def heatmap(rows, title, x_field, y_field, color_field, opts \\ []) do
    width = Keyword.get(opts, :width, 620)
    height = Keyword.get(opts, :height, 320)
    x_title = Keyword.get(opts, :x_title, Atom.to_string(x_field))
    y_title = Keyword.get(opts, :y_title, Atom.to_string(y_field))

    VegaLite.new(width: width, height: height, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:rect, tooltip: true)
    |> VegaLite.encode_field(:x, Atom.to_string(x_field), type: :nominal, title: x_title)
    |> VegaLite.encode_field(:y, Atom.to_string(y_field), type: :nominal, title: y_title)
    |> VegaLite.encode_field(:color, Atom.to_string(color_field), type: :quantitative, scale: [scheme: "tealblues"])
    |> Kino.VegaLite.new()
  end

  # 学習回数に対する損失など、連続値の推移を折れ線で表示する
  def line_chart(rows, title, x_field, y_field, opts \\ []) do
    x_title = Keyword.get(opts, :x_title, Atom.to_string(x_field))
    y_title = Keyword.get(opts, :y_title, Atom.to_string(y_field))

    VegaLite.new(width: 620, 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: x_title
    )
    |> VegaLite.encode_field(:y, Atom.to_string(y_field),
      type: :quantitative,
      title: y_title
    )
    |> Kino.VegaLite.new()
  end
end

defmodule LLMScratch.Transformer do
  def linear(x, w, b) do
    # 最後の次元へ重み行列を掛け、バイアスを加える
    Nx.dot(x, [1], w, [0])
    |> Nx.add(b)
  end

  def layer_norm(x, gamma, beta, eps \\ 1.0e-5) do
    # 系列内の各位置について、特徴次元の平均と分散を求める
    mean = Nx.mean(x, axes: [1], keep_axes: true)
    variance =
      x
      |> Nx.subtract(mean)
      |> Nx.pow(2)
      |> Nx.mean(axes: [1], keep_axes: true)

    # 平均0・分散1付近へそろえた後、学習可能なgammaとbetaで調整する
    normalized =
      x
      |> Nx.subtract(mean)
      |> Nx.divide(Nx.sqrt(Nx.add(variance, eps)))

    normalized
    |> Nx.multiply(gamma)
    |> Nx.add(beta)
  end

  def positional_encoding(seq_len, d_model) do
    # 位置と次元の組ごとに、周期の異なるsin/cos信号を作る
    0..(seq_len - 1)
    |> Enum.map(fn pos ->
      0..(d_model - 1)
      |> Enum.map(fn dim ->
        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)
    end)
    |> Nx.tensor(type: {:f, 32})
  end

  def embed(token_ids, embedding_table) do
    # 各トークンIDに対応する行を埋め込み表から取り出す
    Nx.take(embedding_table, token_ids)
  end

  def split_heads(x, num_heads) do
    {seq_len, d_model} = Nx.shape(x)
    head_dim = div(d_model, num_heads)

    # {系列長, モデル次元}を{ヘッド数, 系列長, ヘッド次元}へ並べ替える
    x
    |> Nx.reshape({seq_len, num_heads, head_dim})
    |> Nx.transpose(axes: [1, 0, 2])
  end

  def combine_heads(x) do
    {num_heads, seq_len, head_dim} = Nx.shape(x)

    # 分割したヘッドを、元のモデル次元を持つ系列へ戻す
    x
    |> Nx.transpose(axes: [1, 0, 2])
    |> Nx.reshape({seq_len, num_heads * head_dim})
  end

  def row_softmax(scores) do
    # QueryごとにKeyへの点数を合計1のアテンション重みへ変換する
    exps =
      scores
      |> Nx.subtract(Nx.reduce_max(scores, axes: [1], keep_axes: true))
      |> Nx.exp()

    Nx.divide(exps, Nx.sum(exps, axes: [1], keep_axes: true))
  end

  def causal_mask(seq_len) do
    # j > iとなる右上部分だけを1にし、未来の位置を示す
    for i <- 0..(seq_len - 1) do
      for j <- 0..(seq_len - 1) do
        j > i
      end
    end
    |> Nx.tensor(type: {:u, 8})
  end

  def scaled_dot_product_attention(q, k, v, mask \\ nil) do
    head_dim = q |> Nx.shape() |> elem(1)

    # QueryとKeyの相性を内積で測り、値が極端にならないようsqrt(head_dim)で割る
    scores =
      Nx.dot(q, [1], Nx.transpose(k), [0])
      |> Nx.divide(:math.sqrt(head_dim))

    # マスク対象へ巨大な負値を置くと、softmax後の重みがほぼ0になる
    scores =
      if mask do
        Nx.select(mask, Nx.broadcast(-1.0e9, Nx.shape(scores)), scores)
      else
        scores
      end

    weights = row_softmax(scores)

    # Keyへの重みを使ってValueを混ぜ、各Query位置の出力を作る
    output = Nx.dot(weights, [1], v, [0])
    {output, weights}
  end

  def multi_head_attention(query, key, value, params, num_heads, mask \\ nil) do
    # 入力をそれぞれ別の線形変換へ通し、Query／Key／Valueを作る
    q = linear(query, params.wq, params.bq)
    k = linear(key, params.wk, params.bk)
    v = linear(value, params.wv, params.bv)

    # モデル次元を複数のヘッドへ均等に分割する
    q_heads = split_heads(q, num_heads)
    k_heads = split_heads(k, num_heads)
    v_heads = split_heads(v, num_heads)

    # 各ヘッドは独立したアテンション重みと出力を計算する
    head_results =
      for head <- 0..(num_heads - 1) do
        scaled_dot_product_attention(q_heads[head], k_heads[head], v_heads[head], mask)
      end

    head_outputs = Enum.map(head_results, &elem(&1, 0)) |> Nx.stack()
    head_weights = Enum.map(head_results, &elem(&1, 1)) |> Nx.stack()

    # 全ヘッドを連結し、最後の線形変換で情報を混ぜ合わせる
    output =
      head_outputs
      |> combine_heads()
      |> linear(params.wo, params.bo)

    {output, head_weights}
  end

  def feed_forward(x, params) do
    # 各位置を独立にd_model -> d_ff -> d_modelと変換する
    x
    |> linear(params.w1, params.b1)
    |> Nx.max(0.0)
    |> linear(params.w2, params.b2)
  end

  def residual_layer_norm(x, sublayer_output, ln_params) do
    # サブレイヤーへ渡した元の情報を足し戻してから正規化する
    x
    |> Nx.add(sublayer_output)
    |> layer_norm(ln_params.gamma, ln_params.beta)
  end

  def encoder_block(x, params, num_heads) do
    # エンコーダーの自己アテンションでは、Query／Key／Valueがすべて同じ系列
    {attn_out, attn_weights} =
      multi_head_attention(x, x, x, params.self_attention, num_heads)

    # 各サブレイヤーの後に、残差接続とレイヤー正規化を行う
    x1 = residual_layer_norm(x, attn_out, params.ln1)
    ff_out = feed_forward(x1, params.feed_forward)
    x2 = residual_layer_norm(x1, ff_out, params.ln2)

    {x2, %{self_weights: attn_weights, attention_output: attn_out, feed_forward_output: ff_out}}
  end

  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)

    # 次に、出力文側をQuery、エンコーダーの記憶をKey／Valueとして参照する
    {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,
       cross_attention_output: cross_attn_out,
       feed_forward_output: ff_out
     }}
  end

  def encoder(token_ids, params, config) do
    # トークン埋め込みへ位置情報を加え、最初のブロックへの入力を作る
    x =
      token_ids
      |> embed(params.src_embedding)
      |> Nx.add(positional_encoding(tuple_size_from_ids(token_ids), config.d_model))

    # 前のブロック出力を次のブロック入力として順番に渡す
    {infos, memory} =
      Enum.map_reduce(params.encoder_blocks, x, fn block_params, acc ->
        {output, info} = encoder_block(acc, block_params, config.num_heads)
        {info, output}
      end)

    {memory, infos}
  end

  def decoder(token_ids, memory, params, config) do
    # 出力文側にも、トークン埋め込みと位置情報を加える
    x =
      token_ids
      |> embed(params.tgt_embedding)
      |> Nx.add(positional_encoding(tuple_size_from_ids(token_ids), config.d_model))

    # 全デコーダーブロックで共通して使う因果マスク
    mask = causal_mask(tuple_size_from_ids(token_ids))

    # 各ブロックが同じエンコーダーの記憶を参照しながら表現を更新する
    {infos, decoder_output} =
      Enum.map_reduce(params.decoder_blocks, x, fn block_params, acc ->
        {output, info} = decoder_block(acc, memory, block_params, config.num_heads, mask)
        {info, output}
      end)

    {decoder_output, infos, mask}
  end

  def transformer(src_ids, tgt_ids, params, config) do
    # 入力文を記憶へ変換し、その記憶を参照して出力文側の表現を作る
    {memory, encoder_infos} = encoder(src_ids, params, config)
    {decoder_output, decoder_infos, tgt_mask} = decoder(tgt_ids, memory, params, config)
    # 各デコーダー位置のベクトルを、出力語彙ごとのロジットへ射影する
    logits = linear(decoder_output, params.output_projection.w, params.output_projection.b)

    %{
      encoder_memory: memory,
      decoder_output: decoder_output,
      logits: logits,
      encoder_infos: encoder_infos,
      decoder_infos: decoder_infos,
      tgt_mask: tgt_mask
    }
  end

  defp tuple_size_from_ids(token_ids) do
    token_ids
    |> Nx.shape()
    |> elem(0)
  end
end

defmodule LLMScratch.TrainableHead 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(w, b, features, target_one_hot) do
    # Transformerの出力表現を、学習対象のwとbで次トークン確率へ変換する
    probabilities =
      features
      |> Nx.dot(w)
      |> Nx.add(b)
      |> row_softmax()

    # 正解トークンの負の対数確率を平均する
    target_one_hot
    |> Nx.multiply(Nx.log(Nx.add(probabilities, 1.0e-9)))
    |> Nx.sum(axes: [1])
    |> Nx.negate()
    |> Nx.mean()
  end

  defn update(w, b, features, target_one_hot, learning_rate) do
    # 出力層のwとbだけについて勾配を求める
    {grad_w, grad_b} =
      grad({w, b}, fn {current_w, current_b} ->
        loss(current_w, current_b, features, target_one_hot)
      end)

    {
      Nx.subtract(w, Nx.multiply(grad_w, learning_rate)),
      Nx.subtract(b, Nx.multiply(grad_b, learning_rate))
    }
  end

  def probabilities(w, b, features) do
    # 学習前後の比較で使う、出力層の順伝播
    features
    |> Nx.dot(w)
    |> Nx.add(b)
    |> row_softmax()
  end
end

1. 複数の短い翻訳例を用意する

# 入力文側と出力文側で、それぞれ別の語彙を使う
src_vocab = [
  "<PAD>", "<BOS>", "<EOS>", "i", "really", "like",
  "green", "tea", "cats", "coffee", "books"
]

tgt_vocab = [
  "<PAD>", "<BOS>", "<EOS>", "わたし", "は", "緑茶",
  "が", "好き", "猫", "コーヒー", "本"
]

src_token_to_id = src_vocab |> Enum.with_index() |> Map.new()
tgt_token_to_id = tgt_vocab |> Enum.with_index() |> Map.new()

# decoder_inputを1位置左へずらしたexpectedが、各位置の正解になる
training_examples = [
  %{
    label: "緑茶の例",
    source: ["i", "really", "like", "green", "tea", "<EOS>"],
    decoder_input: ["<BOS>", "わたし", "は", "緑茶", "が", "好き"],
    expected: ["わたし", "は", "緑茶", "が", "好き", "<EOS>"]
  },
  %{
    label: "猫の例",
    source: ["i", "really", "like", "cats", "<PAD>", "<EOS>"],
    decoder_input: ["<BOS>", "わたし", "は", "猫", "が", "好き"],
    expected: ["わたし", "は", "猫", "が", "好き", "<EOS>"]
  },
  %{
    label: "コーヒーの例",
    source: ["i", "really", "like", "coffee", "<PAD>", "<EOS>"],
    decoder_input: ["<BOS>", "わたし", "は", "コーヒー", "が", "好き"],
    expected: ["わたし", "は", "コーヒー", "が", "好き", "<EOS>"]
  },
  %{
    label: "本の例",
    source: ["i", "really", "like", "books", "<PAD>", "<EOS>"],
    decoder_input: ["<BOS>", "わたし", "は", "本", "が", "好き"],
    expected: ["わたし", "は", "本", "が", "好き", "<EOS>"]
  }
]

# 複数例の入力と正解の対応を確認しやすい表示形式へ整える
example_rows =
  Enum.map(training_examples, fn example ->
    %{
      例: example.label,
      エンコーダー入力: Enum.join(example.source, " "),
      デコーダー入力: Enum.join(example.decoder_input, " "),
      正解の次トークン列: Enum.join(example.expected, " ")
    }
  end)

Kino.DataTable.new(
  example_rows,
  keys: [:例, :エンコーダー入力, :デコーダー入力, :正解の次トークン列]
)

4例とも出力文側の最初は <BOS> わたしは で同じです。次に来る目的語を選ぶには、デコーダー側だけでなく、エンコーダーが読む green tea、cats、coffee、books を見分ける必要があります。

以下ではまず green tea の1例を使って、Transformer内部を順番に観察します。後半で4例すべてを学習に使います。

# 最初の1例を、Transformer内部を順に観察する代表例として使う
walkthrough_example = hd(training_examples)

src_tokens = walkthrough_example.source
tgt_tokens = walkthrough_example.decoder_input
tgt_expected_next_tokens = walkthrough_example.expected

# 文字列の系列を、埋め込み表から参照できる整数IDのtensorへ変換する
src_ids = src_tokens |> Enum.map(&src_token_to_id[&1]) |> Nx.tensor(type: :s64)
tgt_ids = tgt_tokens |> Enum.map(&tgt_token_to_id[&1]) |> Nx.tensor(type: :s64)

%{
  src_tokens: src_tokens,
  tgt_tokens: tgt_tokens,
  tgt_expected_next_tokens: tgt_expected_next_tokens,
  src_ids: src_ids,
  tgt_ids: tgt_ids
}

この例では、入力文側に

i really like green tea

が入り、出力文側では

<BOS> わたしは緑茶が好き

という途中状態から、各位置で次トークンを予測していきます。

# 入力文、デコーダー入力、正解を位置ごとに対応付ける
sequence_rows =
  0..(length(src_tokens) - 1)
  |> Enum.map(fn position ->
    %{
      位置: position,
      入力文のトークン: Enum.at(src_tokens, position),
      デコーダーに見せるトークン: Enum.at(tgt_tokens, position),
      正解の次トークン: Enum.at(tgt_expected_next_tokens, position)
    }
  end)

Kino.DataTable.new(
  sequence_rows,
  keys: [:位置, :入力文のトークン, :デコーダーに見せるトークン, :正解の次トークン]
)

この表の見方は次の通りです。

入力文のトークン はエンコーダーが読む単語
デコーダーに見せるトークン はデコーダーがその位置までに見えている単語
正解の次トークン は、その位置で本来高く出てほしい次トークン

2. Transformer の設定

ここでは、これから作るTransformerの大きさと形を先に決めます。

次のconfigに入れる数値は、学習によって変わる重みそのものではありません。「ベクトルを何次元にするか」「同じ処理を何段重ねるか」といった、モデルの設計を決める値です。この設計に合わせて、次の節で重みを初期化します。

# 学習される重みではなく、モデル内部の幅・分割数・深さをここで決める
config = %{
  d_model: 8,
  num_heads: 2,
  d_ff: 16,
  num_encoder_blocks: 2,
  num_decoder_blocks: 2
}

まずは5つの設定をざっくり捉える

設定	このnotebookでの値	何を決めるか	直感的な捉え方
モデル次元 `d_model`	8	1つのトークンを何個の数値で表すか	各トークンを記録する欄の幅
ヘッド数 `num_heads`	2	アテンションを何個の見方に分けるか	文を同時に見る視点の数
フィードフォワード層の中間次元 `d_ff`	16	各位置の情報を一時的に何次元まで広げるか	情報を加工する作業スペースの広さ
エンコーダーブロック数 `num_encoder_blocks`	2	入力文を読む処理を何段重ねるか	入力文を読み直して表現を更新する回数
デコーダーブロック数 `num_decoder_blocks`	2	出力文を作る処理を何段重ねるか	次トークンを考える表現を更新する回数

このnotebookでは内部の値を追いやすくするため、どれも非常に小さくしています。値を大きくするとモデルが表現できる関係も増えますが、必要な重み、メモリ、計算時間も増えます。

`d_model: 8`は「1トークンを8個の数値で表す」

今回の入力文は、<PAD>や<EOS>を含めて6トークンです。各トークンを8次元のベクトルに変換すると、系列全体の形状は{6, 8}になります。

6個のトークン × 各トークンを表す8個の数値
                 ↓
              {6, 8}

{6, 8}の6はトークンの位置数、8は各位置が持つ数値の数です。8個の数値に最初から「主語」「動詞」のような決まった意味があるわけではありません。学習を通じて、予測に役立つ情報の持ち方が作られます。

`num_heads: 2`は「8次元を2つの見方に分ける」

マルチヘッドアテンションでは、8次元の表現を2つのヘッドに分けます。

モデル次元 8 ÷ ヘッド数 2 = 1ヘッドあたり4次元

{6, 8}
   ↓ 2つのヘッドへ分割
{2, 6, 4}
   ↓ 各ヘッドでアテンションを計算して再結合
{6, 8}

{2, 6, 4}は、2ヘッド × 6位置 × 4次元という意味です。2つのヘッドは別々の重みを計算できるため、たとえば一方は近くの語、もう一方は離れた関連語を重視する、といった異なる見方を学べます。ただし、実際に何を見るかは学習によって決まります。

この実装では均等に分割するため、d_modelはnum_headsで割り切れる必要があります。今回は8 ÷ 2 = 4なので分割できます。

`d_ff: 16`は「各位置を8次元から16次元へ広げて戻す」

フィードフォワード層は、アテンションで集めた情報を各位置で個別に加工します。

各位置の8次元
    ↓ 1回目の線形変換
各位置の16次元
    ↓ 活性化関数
各位置の16次元
    ↓ 2回目の線形変換
各位置の8次元

つまり、系列全体では{6, 8} → {6, 16} → {6, 8}と変化します。最後に8次元へ戻すのは、元の入力と足し合わせる残差接続を行えるようにするためです。

ブロック数の2は「同じ形の処理を2段重ねる」

エンコーダーブロックは、主に「自己アテンション」と「フィードフォワード層」を1組にした処理です。デコーダーブロックには、さらに「因果マスク付き自己アテンション」と「クロスアテンション」があります。

Kino.Mermaid.new("""
flowchart
  S["入力文"] --> E1["エンコーダーブロック 1"]
  E1 --> E2["エンコーダーブロック 2"]
  E2 --> M["エンコーダーの記憶"]

  T["ここまでの出力文"] --> D1["デコーダーブロック 1"]
  D1 --> D2["デコーダーブロック 2"]
  M --> D1
  M --> D2
  D2 --> P["次トークン候補のスコア"]
""")

ブロックを通るたびに各位置の表現は更新されますが、外側の形状は基本的に{系列長, d_model}のままです。つまり、ブロック数はベクトルの幅ではなく、処理の深さを決めます。

今回の設定を形状でまとめる

項目	形状	意味
1トークンの埋め込み	`{d_model} = {8}`	1つの位置の表現
系列全体の埋め込み	`{系列長, d_model} = {6, 8}`	系列全体の表現
2つのヘッドへ分割した後	`{ヘッド数, 系列長, ヘッド次元} = {2, 6, 4}`	8次元を4次元 × 2に分ける
ヘッドごとのアテンション重み	`{系列長, 系列長} = {6, 6}`	各位置がどこを見るか
フィードフォワード層の中間	`{系列長, d_ff} = {6, 16}`	各位置を一時的に16次元へ広げた表現

ここで、系列長の6や語彙数はconfigには入っていないことにも注意してください。

系列長は、モデルへ渡す文章のトークン数で決まる
語彙数は、この章で用意した語彙の種類数で決まる
configの5項目は、文章を処理するモデル内部の幅・分割数・深さを決める

この区別を押さえると、後続コードに出てくるテンソルの形状を追いやすくなります。

3. パラメータを初期化する

# 同じ構成の実験を再現できるよう、乱数の初期状態を固定する
:rand.seed(:exsss, {101, 202, 303})

# 指定された形状を持つ、小さなランダム値のtensorを作る
random_tensor = fn shape ->
  shape
  |> Tuple.to_list()
  |> Enum.product()
  |> then(fn total ->
    1..total
    |> Enum.map(fn _ -> :rand.normal() * 0.12 end)
  end)
  |> Nx.tensor(type: {:f, 32})
  |> Nx.reshape(shape)
end

zeros = fn size ->
  Nx.broadcast(0.0, {size})
end

ones = fn size ->
  Nx.broadcast(1.0, {size})
end

# 1つのマルチヘッドアテンションで使う4組の線形変換
attention_params = fn d_model ->
  %{
    wq: random_tensor.({d_model, d_model}),
    bq: zeros.(d_model),
    wk: random_tensor.({d_model, d_model}),
    bk: zeros.(d_model),
    wv: random_tensor.({d_model, d_model}),
    bv: zeros.(d_model),
    wo: random_tensor.({d_model, d_model}),
    bo: zeros.(d_model)
  }
end

# d_modelからd_ffへ広げ、再びd_modelへ戻す2層分の重み
feed_forward_params = fn d_model, d_ff ->
  %{
    w1: random_tensor.({d_model, d_ff}),
    b1: zeros.(d_ff),
    w2: random_tensor.({d_ff, d_model}),
    b2: zeros.(d_model)
  }
end

# 初期状態では値をそのまま正規化できるgamma=1、beta=0を置く
layer_norm_params = fn d_model ->
  %{
    gamma: ones.(d_model),
    beta: zeros.(d_model)
  }
end

# エンコーダーブロック1段分のパラメータをまとめる
encoder_block_params = fn ->
  %{
    self_attention: attention_params.(config.d_model),
    feed_forward: feed_forward_params.(config.d_model, config.d_ff),
    ln1: layer_norm_params.(config.d_model),
    ln2: layer_norm_params.(config.d_model)
  }
end

# デコーダーには自己アテンションとクロスアテンションの両方が必要
decoder_block_params = fn ->
  %{
    self_attention: attention_params.(config.d_model),
    cross_attention: attention_params.(config.d_model),
    feed_forward: feed_forward_params.(config.d_model, config.d_ff),
    ln1: layer_norm_params.(config.d_model),
    ln2: layer_norm_params.(config.d_model),
    ln3: layer_norm_params.(config.d_model)
  }
end

# 語彙の埋め込み、各ブロック、最後の語彙射影をモデル全体へまとめる
params = %{
  src_embedding: random_tensor.({length(src_vocab), config.d_model}),
  tgt_embedding: random_tensor.({length(tgt_vocab), config.d_model}),
  encoder_blocks: Enum.map(1..config.num_encoder_blocks, fn _ -> encoder_block_params.() end),
  decoder_blocks: Enum.map(1..config.num_decoder_blocks, fn _ -> decoder_block_params.() end),
  output_projection: %{
    w: random_tensor.({config.d_model, length(tgt_vocab)}),
    b: zeros.(length(tgt_vocab))
  }
}

%{
  src_embedding_shape: Nx.shape(params.src_embedding),
  tgt_embedding_shape: Nx.shape(params.tgt_embedding),
  projection_shape: Nx.shape(params.output_projection.w)
}

4. 位置エンコーディングを可視化する

# {位置, 次元}ごとの位置エンコーディングをヒートマップ用に展開する
positional_encoding = LLMScratch.Transformer.positional_encoding(length(src_tokens), config.d_model)

position_rows =
  for position <- 0..(length(src_tokens) - 1),
      dim <- 0..(config.d_model - 1) do
    %{
      位置: "位置_#{position}",
      次元: "次元_#{dim}",
      値: Nx.to_number(positional_encoding[[position, dim]])
    }
  end

LLMScratch.Visuals.heatmap(
  position_rows,
  "位置エンコーディング",
  :次元,
  :位置,
  :値,
  x_title: "埋め込み次元",
  y_title: "位置"
)

ここで見たいのは、トークンそのもの とは別に、何番目にいるか の信号が加わることです。

# トークンの情報と位置の情報を、同じ形状のtensor同士で足し合わせる
src_token_embeddings = LLMScratch.Transformer.embed(src_ids, params.src_embedding)
src_position_embeddings = LLMScratch.Transformer.positional_encoding(length(src_tokens), config.d_model)
src_embeddings_with_position = Nx.add(src_token_embeddings, src_position_embeddings)

# 各ベクトルは長いため、先頭4要素だけを表示して加算前後を比べる
position_detail_rows =
  for position <- 0..(length(src_tokens) - 1) do
    token_preview =
      src_token_embeddings[position][0..3]
      |> Nx.to_flat_list()
      |> Enum.map(&Float.round(&1, 3))
      |> inspect()

    pe_preview =
      src_position_embeddings[position][0..3]
      |> Nx.to_flat_list()
      |> Enum.map(&Float.round(&1, 3))
      |> inspect()

    sum_preview =
      src_embeddings_with_position[position][0..3]
      |> Nx.to_flat_list()
      |> Enum.map(&Float.round(&1, 3))
      |> inspect()

    %{
      位置: position,
      トークン: Enum.at(src_tokens, position),
      トークン埋め込みの先頭4要素: token_preview,
      位置情報の先頭4要素: pe_preview,
      加算後の先頭4要素: sum_preview
    }
  end

Kino.DataTable.new(
  position_detail_rows,
  keys: [:位置, :トークン, :トークン埋め込みの先頭4要素, :位置情報の先頭4要素, :加算後の先頭4要素]
)

この表では、同じ d_model の中に

トークン自身の意味っぽい成分
位置の成分

が足し合わされる様子を見ています。

5. マルチヘッドアテンションを単独で動かす

まずはエンコーダー側の入力埋め込みに対して、マルチヘッド自己アテンションを1回だけかけてみます。

# 自己アテンションなので、同じ入力をQuery／Key／Valueへ渡す
src_embeddings = src_embeddings_with_position

{mh_output, mh_weights} =
  LLMScratch.Transformer.multi_head_attention(
    src_embeddings,
    src_embeddings,
    src_embeddings,
    hd(params.encoder_blocks).self_attention,
    config.num_heads
  )

%{
  input_shape: Nx.shape(src_embeddings),
  output_shape: Nx.shape(mh_output),
  attention_weight_shape: Nx.shape(mh_weights)
}

attention_weight_shape = {head, query_position, key_position} という形状になっていれば、ヘッドごとの参照先を分けて見られます。

# ヘッド分割前のQueryを再現し、形状の変化を単独で確認する
q_projection =
  LLMScratch.Transformer.linear(
    src_embeddings,
    hd(params.encoder_blocks).self_attention.wq,
    hd(params.encoder_blocks).self_attention.bq
  )

q_heads = LLMScratch.Transformer.split_heads(q_projection, config.num_heads)

head_shape_rows = [
  %{段階: "入力の埋め込み", 形状: inspect(Tuple.to_list(Nx.shape(src_embeddings))), 説明: "各位置8次元"},
  %{段階: "Queryへの線形変換", 形状: inspect(Tuple.to_list(Nx.shape(q_projection))), 説明: "まだ8次元のまま"},
  %{段階: "ヘッドへ分割", 形状: inspect(Tuple.to_list(Nx.shape(q_heads))), 説明: "2ヘッド × 4次元へ分割"}
]

Kino.DataTable.new(head_shape_rows, keys: [:段階, :形状, :説明])

# ヘッドごとのQueryベクトルを別々のタブへ表示する
q_head_tabs =
  0..(config.num_heads - 1)
  |> Enum.map(fn head ->
    rows =
      for position <- 0..(length(src_tokens) - 1),
          dim <- 0..(div(config.d_model, config.num_heads) - 1) do
        %{
          位置とトークン: "#{position}:#{Enum.at(src_tokens, position)}",
          次元: "次元_#{dim}",
          値: Float.round(Nx.to_number(q_heads[[head, position, dim]]), 4)
        }
      end

    {"Queryヘッド #{head}",
     LLMScratch.Visuals.heatmap(
       rows,
       "ヘッド #{head} のQueryベクトル",
       :次元,
       :位置とトークン,
       :値,
       x_title: "ヘッド内の次元",
       y_title: "位置"
     )}
  end)

Kino.Layout.tabs(q_head_tabs)

ここで見えるのは、1本の8次元ベクトル が 4次元のヘッド0 と 4次元のヘッド1 に分かれてから、それぞれ独立にアテンションへ入っていくことです。

# 各ヘッドの{参照元, 参照先}ごとのアテンション重みを表示する
head_tabs =
  0..(config.num_heads - 1)
  |> Enum.map(fn head ->
    rows =
      for q <- 0..(length(src_tokens) - 1),
          k <- 0..(length(src_tokens) - 1) do
        %{
          参照元: "#{q}:#{Enum.at(src_tokens, q)}",
          参照先: "#{k}:#{Enum.at(src_tokens, k)}",
          重み: Float.round(Nx.to_number(mh_weights[[head, q, k]]), 4)
        }
      end

    {"ヘッド #{head}",
     LLMScratch.Visuals.heatmap(
       rows,
       "エンコーダーの自己アテンション: ヘッド #{head}",
       :参照先,
       :参照元,
       :重み,
       x_title: "参照先",
       y_title: "見ている位置"
     )}
  end)

Kino.Layout.tabs(head_tabs)

重みはランダム初期化なので意味解釈はまだ難しいですが、ヘッドごとに別の重み行列を持つ という構造はここで確認できます。

6. フィードフォワード層・残差接続・レイヤー正規化

エンコーダーブロックの前半と後半で何が起きるかを、形状と一緒に見ます。

アテンションは「他の位置を読む」層ですが、フィードフォワード層は違います。

フィードフォワード層は 各位置を独立に 非線形変換します。つまり、位置0と位置1を混ぜるのではなく、各位置のベクトルを個別に押し広げて変換しています。

Kino.Mermaid.new("""
flowchart
  A["位置 0 のベクトル"] --> B0["共通のフィードフォワード層"]
  C["位置 1 のベクトル"] --> B1["共通のフィードフォワード層"]
  D["位置 2 のベクトル"] --> B2["共通のフィードフォワード層"]
  B0 --> E["更新後の位置 0"]
  B1 --> F["更新後の位置 1"]
  B2 --> G["更新後の位置 2"]
""")

# 1段目だけを単独実行し、各サブレイヤーの入出力を観察する
first_encoder_block = hd(params.encoder_blocks)

{encoder_block_output, encoder_block_info} =
  LLMScratch.Transformer.encoder_block(src_embeddings, first_encoder_block, config.num_heads)

shape_rows = [
  %{データ: "ブロックへの入力", 形状: inspect(Tuple.to_list(Nx.shape(src_embeddings)))},
  %{データ: "自己アテンションの出力", 形状: inspect(Tuple.to_list(Nx.shape(encoder_block_info.attention_output)))},
  %{データ: "フィードフォワード層の出力", 形状: inspect(Tuple.to_list(Nx.shape(encoder_block_info.feed_forward_output)))},
  %{データ: "ブロック全体の出力", 形状: inspect(Tuple.to_list(Nx.shape(encoder_block_output)))}
]

Kino.DataTable.new(shape_rows, keys: [:データ, :形状])

# フィードフォワード層が各位置の値をどう変えたか、先頭4要素で比べる
ff_preview_rows =
  for position <- 0..(length(src_tokens) - 1) do
    input_preview =
      src_embeddings[position][0..3]
      |> Nx.to_flat_list()
      |> Enum.map(&Float.round(&1, 3))
      |> inspect()

    output_preview =
      encoder_block_info.feed_forward_output[position][0..3]
      |> Nx.to_flat_list()
      |> Enum.map(&Float.round(&1, 3))
      |> inspect()

    %{
      位置とトークン: "#{position}:#{Enum.at(src_tokens, position)}",
      変換前の先頭4要素: input_preview,
      変換後の先頭4要素: output_preview
    }
  end

Kino.DataTable.new(ff_preview_rows, keys: [:位置とトークン, :変換前の先頭4要素, :変換後の先頭4要素])

このブロックでは次の順で処理しています。

マルチヘッド自己アテンション
入力との足し算である残差接続
レイヤー正規化
フィードフォワード層
もう一度、残差接続とレイヤー正規化

残差接続は、深いネットワークでも元の情報が流れやすくなるのを助けます。レイヤー正規化は、各位置の表現を安定させる役割を持ちます。

# 自己アテンションの出力へ、ブロック入力を足し戻す
residual_before_norm =
  src_embeddings
  |> Nx.add(encoder_block_info.attention_output)

# その結果を位置ごとにレイヤー正規化する
residual_after_norm =
  LLMScratch.Transformer.layer_norm(
    residual_before_norm,
    first_encoder_block.ln1.gamma,
    first_encoder_block.ln1.beta
  )

# 正規化の前後で、各位置の平均と分散を比較する
norm_rows =
  for position <- 0..(length(src_tokens) - 1) do
    before = residual_before_norm[position]
    normalized = residual_after_norm[position]

    %{
      位置とトークン: "#{position}:#{Enum.at(src_tokens, position)}",
      正規化前の平均: Float.round(Nx.to_number(Nx.mean(before)), 4),
      正規化前の分散: Float.round(Nx.to_number(Nx.mean(Nx.pow(Nx.subtract(before, Nx.mean(before)), 2))), 4),
      正規化後の平均: Float.round(Nx.to_number(Nx.mean(normalized)), 4),
      正規化後の分散: Float.round(Nx.to_number(Nx.mean(Nx.pow(Nx.subtract(normalized, Nx.mean(normalized)), 2))), 4)
    }
  end

Kino.DataTable.new(
  norm_rows,
  keys: [:位置とトークン, :正規化前の平均, :正規化前の分散, :正規化後の平均, :正規化後の分散]
)

この表では、レイヤー正規化の後に各位置ベクトルの平均が0付近、分散が1付近へ寄ることが見えます。

つまり、残差接続で情報を足し戻しつつ、レイヤー正規化で値のスケールを整える という役割分担です。

7. エンコーダーを実行する

エンコーダーとデコーダーは混同しやすいので、先に役割を表で整理しておきます。

部分	入力	主なアテンション	未来の参照	出力	役割
エンコーダー	入力文のトークン	自己アテンション	入力文全体を見てよい	エンコーダーの記憶	入力文を理解した表現を作る
デコーダー	ここまでの出力文トークン	因果マスク付き自己アテンション + クロスアテンション	出力文の未来は見ない	隠れ状態 / 次トークン候補のスコア	入力文の記憶を参照しながら出力文を生成する

# 全エンコーダーブロックを通し、デコーダーが参照する記憶を得る
{encoder_memory, encoder_infos} =
  LLMScratch.Transformer.encoder(src_ids, params, config)

%{
  encoder_blocks: length(encoder_infos),
  encoder_memory_shape: Nx.shape(encoder_memory)
}

エンコーダーは入力文の系列を読み取り、後でデコーダーが参照する記憶を作ります。

この例では、エンコーダーは単に tea だけを見るのではなく、

really
like
green
tea

の並び全体を 6 個の位置表現へ変換します。
デコーダーはこの記憶を見ることで、「何を好きなのか」「どの語がどこにあったか」の手がかりを受け取ります。

# エンコーダーの記憶を{入力位置, 次元}のヒートマップへ展開する
encoder_memory_rows =
  for position <- 0..(length(src_tokens) - 1),
      dim <- 0..(config.d_model - 1) do
    %{
      位置とトークン: "#{position}:#{Enum.at(src_tokens, position)}",
      次元: "次元_#{dim}",
      値: Float.round(Nx.to_number(encoder_memory[[position, dim]]), 4)
    }
  end

LLMScratch.Visuals.heatmap(
  encoder_memory_rows,
  "エンコーダーの記憶",
  :次元,
  :位置とトークン,
  :値,
  x_title: "埋め込み次元",
  y_title: "入力文の位置"
)

8. デコーダーを実行する

# 出力文のID列とエンコーダーの記憶を、全デコーダーブロックへ渡す
{decoder_output, decoder_infos, tgt_mask} =
  LLMScratch.Transformer.decoder(tgt_ids, encoder_memory, params, config)

%{
  decoder_blocks: length(decoder_infos),
  decoder_output_shape: Nx.shape(decoder_output),
  target_mask_shape: Nx.shape(tgt_mask)
}

デコーダーでは2種類のアテンションが使われます。

因果マスク付き自己アテンション: 出力文の左側だけを見る
クロスアテンション: エンコーダーの記憶を参照する

今回の例なら、デコーダーの最後の位置 好き の直後では、本来 <EOS> が来そうです。

その1つ前の位置 が では、次に 好き が来るかどうかを、出力文の左側とエンコーダーの記憶の両方から判断します。

# 各位置で因果マスクにより見える接頭辞と、その位置の正解を並べる
decoder_step_rows =
  0..(length(tgt_tokens) - 1)
  |> Enum.map(fn position ->
    visible_prefix =
      tgt_tokens
      |> Enum.take(position + 1)
      |> Enum.join(" ")

    %{
      位置: position,
      デコーダーから見える範囲: visible_prefix,
      予測すべき次トークン: Enum.at(tgt_expected_next_tokens, position)
    }
  end)

Kino.DataTable.new(
  decoder_step_rows,
  keys: [:位置, :デコーダーから見える範囲, :予測すべき次トークン]
)

Kino.Mermaid.new("""
flowchart
  A["ここまでの出力文"] --> B["因果マスク付き自己アテンション"]
  B --> C["デコーダーの中間表現"]
  D["エンコーダーの記憶"] --> E["クロスアテンション"]
  C --> E
  E --> F["次トークン予測へ"]
""")

# 因果マスクの1を「参照禁止」、0を「参照可能」として可視化する
mask_rows =
  for q <- 0..(length(tgt_tokens) - 1),
      k <- 0..(length(tgt_tokens) - 1) do
    %{
      現在位置: "#{q}:#{Enum.at(tgt_tokens, q)}",
      参照先: "#{k}:#{Enum.at(tgt_tokens, k)}",
      マスク値: if(Nx.to_number(tgt_mask[[q, k]]) == 1, do: 1.0, else: 0.0)
    }
  end

LLMScratch.Visuals.heatmap(
  mask_rows,
  "デコーダーの因果マスク",
  :参照先,
  :現在位置,
  :マスク値,
  x_title: "未来を含む参照先",
  y_title: "現在位置"
)

9. デコーダーの因果マスク付き自己アテンションとクロスアテンションを見る

# 1段目のデコーダーから、2種類のアテンション重みを取り出す
first_decoder_info = hd(decoder_infos)

# 出力文内の因果マスク付き自己アテンションをヘッド別に表示する
self_tabs =
  0..(config.num_heads - 1)
  |> Enum.map(fn head ->
    rows =
      for q <- 0..(length(tgt_tokens) - 1),
          k <- 0..(length(tgt_tokens) - 1) do
        %{
          参照元: "#{q}:#{Enum.at(tgt_tokens, q)}",
          参照先: "#{k}:#{Enum.at(tgt_tokens, k)}",
          重み: Float.round(Nx.to_number(first_decoder_info.self_weights[[head, q, k]]), 4)
        }
      end

    {"自己アテンション #{head}",
     LLMScratch.Visuals.heatmap(
       rows,
       "デコーダーの因果マスク付き自己アテンション: ヘッド #{head}",
       :参照先,
       :参照元,
       :重み,
       x_title: "出力文内の参照先",
       y_title: "出力文の現在位置"
     )}
  end)

# 出力文から入力文の記憶を見るクロスアテンションをヘッド別に表示する
cross_tabs =
  0..(config.num_heads - 1)
  |> Enum.map(fn head ->
    rows =
      for q <- 0..(length(tgt_tokens) - 1),
          k <- 0..(length(src_tokens) - 1) do
        %{
          出力文の位置: "#{q}:#{Enum.at(tgt_tokens, q)}",
          入力文の位置: "#{k}:#{Enum.at(src_tokens, k)}",
          重み: Float.round(Nx.to_number(first_decoder_info.cross_weights[[head, q, k]]), 4)
        }
      end

    {"クロスアテンション #{head}",
     LLMScratch.Visuals.heatmap(
       rows,
       "デコーダーのクロスアテンション: ヘッド #{head}",
       :入力文の位置,
       :出力文の位置,
       :重み,
       x_title: "入力文側の記憶",
       y_title: "出力文の現在位置"
     )}
  end)

Kino.Layout.tabs(self_tabs ++ cross_tabs)

ここで、デコーダーが 出力文の左側 と エンコーダーの記憶 を別のアテンションで読むことが見えます。

10. Transformer 全体を実行する

# エンコーダーから語彙ロジットまでを、一度の呼び出しで実行する
transformer_run =
  LLMScratch.Transformer.transformer(src_ids, tgt_ids, params, config)

summary_rows = [
  %{データ: "エンコーダーの記憶", 形状: inspect(Tuple.to_list(Nx.shape(transformer_run.encoder_memory)))},
  %{データ: "デコーダーの出力", 形状: inspect(Tuple.to_list(Nx.shape(transformer_run.decoder_output)))},
  %{データ: "次トークン候補のロジット", 形状: inspect(Tuple.to_list(Nx.shape(transformer_run.logits)))}
]

Kino.DataTable.new(summary_rows, keys: [:データ, :形状])

logits（ロジット）の形状が {出力文の系列長, 出力語彙数} になっていれば、出力文の各位置で「次に出しそうなトークン」をスコアとして出していることが分かります。

ここは少し抽象的なので、今回の例に当てはめるとこう読めます。

位置0のロジット: <BOS> を見たあとに何が来そうか
位置1のロジット: <BOS> わたし のあとに何が来そうか
位置2のロジット: <BOS> わたしは のあとに何が来そうか

つまり、ロジットは「文全体を一度に完成させる出力」ではなく、各位置での次トークン候補表 です。

# Transformer内部の各段階について、tensorの形状と役割を並べる
pipeline_rows = [
  %{
    段階: "入力文の埋め込み + 位置情報",
    形状: inspect(Tuple.to_list(Nx.shape(src_embeddings))),
    意味: "入力文の各位置を数値ベクトル化したもの"
  },
  %{
    段階: "エンコーダーの記憶",
    形状: inspect(Tuple.to_list(Nx.shape(transformer_run.encoder_memory))),
    意味: "入力文全体を読んだ後の文脈表現"
  },
  %{
    段階: "デコーダーの出力",
    形状: inspect(Tuple.to_list(Nx.shape(transformer_run.decoder_output))),
    意味: "出力文の各位置に対応する隠れ状態"
  },
  %{
    段階: "次トークン候補のロジット",
    形状: inspect(Tuple.to_list(Nx.shape(transformer_run.logits))),
    意味: "各位置で出力語彙へ変換したスコア"
  }
]

Kino.DataTable.new(pipeline_rows, keys: [:段階, :形状, :意味])

11. 学習前の出力語彙への射影を確認する

# 各デコーダー位置で、ロジットが高い次トークン候補を3件取り出す
top_predictions =
  0..(length(tgt_tokens) - 1)
  |> Enum.map(fn position ->
    top3 =
      transformer_run.logits[position]
      |> Nx.to_flat_list()
      |> Enum.zip(tgt_vocab)
      |> Enum.sort_by(fn {logit, _token} -> logit end, :desc)
      |> Enum.take(3)
      |> Enum.map(fn {logit, token} -> "#{token}(#{Float.round(logit, 3)})" end)
      |> Enum.join(" / ")

    %{
      位置: position,
      デコーダー入力: Enum.at(tgt_tokens, position),
      正解の次トークン: Enum.at(tgt_expected_next_tokens, position),
      スコア上位3候補: top3
    }
  end)

Kino.DataTable.new(top_predictions, keys: [:位置, :デコーダー入力, :正解の次トークン, :スコア上位3候補])

まだランダム重みなので 正解の次トークン が上位に来るとは限りません。ここではまず、

デコーダーの各位置ごとに
出力語彙全体へのスコアが出て
その中から次トークンを選べる

という Transformer 出力の形そのものです。

12. 4つの例で出力層を学習する

ここから第2章と同じ 学習前 -> 損失を下げる -> 学習後 の流れに戻ります。

各翻訳例をTransformerへ通すと、デコーダーの各位置に8次元の表現ができます。この実験では、その表現を作るTransformer本体は固定し、最後の 8次元 -> 出力語彙の確率 だけを学習します。

# 4つの翻訳例を同じTransformerへ通し、学習に使う特徴量を作る
training_runs =
  Enum.map(training_examples, fn example ->
    example_src_ids =
      example.source
      |> Enum.map(&src_token_to_id[&1])
      |> Nx.tensor(type: :s64)

    example_tgt_ids =
      example.decoder_input
      |> Enum.map(&tgt_token_to_id[&1])
      |> Nx.tensor(type: :s64)

    run =
      LLMScratch.Transformer.transformer(
        example_src_ids,
        example_tgt_ids,
        params,
        config
      )

    %{
      label: example.label,
      decoder_output: run.decoder_output,
      expected: example.expected
    }
  end)

sequence_length = length(tgt_tokens)
training_example_count = length(training_examples)

# {例, 位置, d_model}を{全24位置, d_model}へまとめる
training_features =
  training_runs
  |> Enum.map(& &1.decoder_output)
  |> Nx.stack()
  |> Nx.reshape({training_example_count * sequence_length, config.d_model})

# 4例分の正解トークンも、特徴量と同じ並び順で1列へまとめる
training_target_ids =
  training_examples
  |> Enum.flat_map(& &1.expected)
  |> Enum.map(&tgt_token_to_id[&1])
  |> Nx.tensor(type: :s64)

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

%{
  examples: training_example_count,
  positions_per_example: sequence_length,
  training_rows: elem(Nx.shape(training_features), 0),
  feature_shape: Nx.shape(training_features),
  target_shape: Nx.shape(training_target_one_hot)
}

training_rows = 4例 x 6位置 = 24 です。24個の次トークン問題をまとめて解き、正解トークンの確率が高くなるよう出力層を更新します。

# Transformer本体は固定し、既存の出力射影だけを学習対象として取り出す
initial_head_w = params.output_projection.w
initial_head_b = params.output_projection.b

head_epochs = 3000
head_learning_rate = 0.5

{{trained_head_w, trained_head_b}, head_loss_rows} =
  Enum.reduce(
    0..head_epochs,
    {{initial_head_w, initial_head_b}, []},
    fn epoch, {{w, b}, rows} ->
      # 現在の出力層で24問を予測したときの損失を記録する
      loss =
        LLMScratch.TrainableHead.loss(
          w,
          b,
          training_features,
          training_target_one_hot
        )
        |> Nx.to_number()

      # 自動微分した勾配を使い、wとbを1回更新する
      {next_w, next_b} =
        LLMScratch.TrainableHead.update(
          w,
          b,
          training_features,
          training_target_one_hot,
          head_learning_rate
        )

      # 先頭への追加で履歴を蓄え、ループ後に時系列順へ戻す
      {{next_w, next_b}, [%{epoch: epoch, loss: loss} | rows]}
    end
  )

head_loss_rows = Enum.reverse(head_loss_rows)

%{
  loss_before: hd(head_loss_rows).loss,
  loss_after: List.last(head_loss_rows).loss
}

# グラフの軸を日本語で表示するため、損失履歴のキーを変換する
head_loss_chart_rows =
  Enum.map(head_loss_rows, fn row ->
    %{学習回数: row.epoch, 損失: row.loss}
  end)

LLMScratch.Visuals.line_chart(
  head_loss_chart_rows,
  "Transformer の出力層を学習したときの損失",
  :学習回数,
  :損失,
  x_title: "学習回数（epoch）",
  y_title: "損失（loss）"
)

第2章と同じく、左から右へ損失が下がれば、出力層が24個の正解へ近づいています。

同じ `は` の次を、入力文に応じて変えられるか

4例すべてで、位置2のデコーダー入力は は です。bigramなら入力が同じなので同じ確率しか出せません。Transformerのデコーダー表現にはエンコーダーの情報が混ざるため、例ごとに別の目的語を予測できる可能性があります。

top_prediction = fn w, b, feature ->
  # 1位置の8次元表現を、出力語彙全体の確率へ変換する
  probabilities =
    LLMScratch.TrainableHead.probabilities(
      w,
      b,
      Nx.new_axis(feature, 0)
    )
    |> Nx.squeeze(axes: [0])
    |> Nx.to_flat_list()

  # 最も確率が高いトークンと、その確率だけを返す
  {probability, token} =
    probabilities
    |> Enum.zip(tgt_vocab)
    |> Enum.max_by(fn {probability, _token} -> probability end)

  "#{token} (#{Float.round(probability, 3)})"
end

object_position = 2

# デコーダー入力が同じ位置について、入力文ごとの予測差を比べる
object_prediction_rows =
  training_runs
  |> Enum.zip(training_examples)
  |> Enum.map(fn {run, example} ->
    feature = run.decoder_output[object_position]

    %{
      入力文: Enum.join(example.source, " "),
      共通のデコーダー入力: "<BOS> わたし は",
      正解の次トークン: Enum.at(example.expected, object_position),
      学習前の予測: top_prediction.(initial_head_w, initial_head_b, feature),
      学習後の予測: top_prediction.(trained_head_w, trained_head_b, feature)
    }
  end)

Kino.DataTable.new(
  object_prediction_rows,
  keys: [
    :入力文,
    :共通のデコーダー入力,
    :正解の次トークン,
    :学習前の予測,
    :学習後の予測
  ]
)

この表の中心は 共通のデコーダー入力 が全行で同じなのに、学習後の予測が入力文に応じて変わる点です。これはデコーダーのクロスアテンションがエンコーダーの記憶を読み、入力文の違いをデコーダーの表現へ持ち込めるためです。

# 24個の学習問題すべてについて、学習前後の確率を計算する
before_probabilities =
  LLMScratch.TrainableHead.probabilities(
    initial_head_w,
    initial_head_b,
    training_features
  )

after_probabilities =
  LLMScratch.TrainableHead.probabilities(
    trained_head_w,
    trained_head_b,
    training_features
  )

accuracy = fn probabilities ->
  # 各行で確率最大の語彙IDを予測結果とする
  predicted_ids =
    probabilities
    |> Nx.argmax(axis: 1)
    |> Nx.to_flat_list()

  correct_ids = Nx.to_flat_list(training_target_ids)

  # 正解IDと一致した割合を正答率として返す
  Enum.zip(predicted_ids, correct_ids)
  |> Enum.count(fn {predicted, correct} -> predicted == correct end)
  |> Kernel./(length(correct_ids))
end

Kino.DataTable.new([
  %{段階: "学習前", 正答率: accuracy.(before_probabilities)},
  %{段階: "学習後", 正答率: accuracy.(after_probabilities)}
])

正答率（accuracy）は24位置のうち正しい次トークンを1位にできた割合です。実用的な翻訳性能を測る値ではなく、この小さな学習データに対して出力層が変化したことを確認する指標です。

13. まとめ

この notebook の要点

マルチヘッドアテンションは、複数のアテンションを並列に走らせて結合する
位置エンコーディングは、トークンの並び順を表現へ加える
フィードフォワード層は各位置ごとに非線形変換をかける
残差接続とレイヤー正規化は深い構造を安定させる
エンコーダーは入力文を文脈表現へ変換し、デコーダーはそれを参照して出力文を生成する
Transformer全体は、これらのブロックを積み重ねた構造として順伝播を実装できる
学習の基本はbigramと同じだが、Transformerは入力文全体と出力文の左側を使って次を予測できる
同じデコーダー入力でも、クロスアテンションによって入力文に応じた次トークンを選び分けられる

次に進むなら、

第4章の学習ループを参考に、出力層だけでなく埋め込み・アテンション・フィードフォワード層もまとめて学習する
<PAD> を参照しないためのパディングマスクを足す
第4章のミニGPTとこの章のデコーダーを見比べ、クロスアテンションの有無が構成のどこに現れるかを確認する

という順が自然です。

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

Distributed Image Processing Main

distributed_main_en.livemd

tutorial advanced data-science gen-server nx evision flow req kino benchee

2026-7-7
@DeSchoel

Elixir_Curriculum

Concurrent Word Count

concurrent_word_count.livemd

advanced gen-server jason kino youtube hidden_cell

2026-7-12
Ryo Wakabayashi
@RyoWakabayashi

elixir-learning

Tellus から台風情報を取得する

typhoon.livemd

advanced apis req kino kino_maplibre

2026-7-12

Back

LLM を Elixir / Livebook で学ぶ 発展編

07_advanced_encoder_decoder_transformer.livemd

LLM を Elixir / Livebook で学ぶ 発展編

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

この notebook の狙い

このnotebookで使う用語

全体像

エンコーダーブロックとデコーダーブロック

マルチヘッドアテンションを直感で捉える（第4章の復習）

準備

1. 複数の短い翻訳例を用意する

2. Transformer の設定

まずは5つの設定をざっくり捉える

d_model: 8は「1トークンを8個の数値で表す」

num_heads: 2は「8次元を2つの見方に分ける」

d_ff: 16は「各位置を8次元から16次元へ広げて戻す」

ブロック数の2は「同じ形の処理を2段重ねる」

今回の設定を形状でまとめる

3. パラメータを初期化する

4. 位置エンコーディングを可視化する

5. マルチヘッドアテンションを単独で動かす

6. フィードフォワード層・残差接続・レイヤー正規化

7. エンコーダーを実行する

8. デコーダーを実行する

9. デコーダーの因果マスク付き自己アテンションとクロスアテンションを見る

10. Transformer 全体を実行する

11. 学習前の出力語彙への射影を確認する

12. 4つの例で出力層を学習する

同じ は の次を、入力文に応じて変えられるか

13. まとめ

この notebook の要点

Other notebooks:

LLM を Elixir / Livebook で学ぶ発展編

LLM を Elixir / Livebook で学ぶ発展編

`d_model: 8`は「1トークンを8個の数値で表す」

`num_heads: 2`は「8次元を2つの見方に分ける」

`d_ff: 16`は「各位置を8次元から16次元へ広げて戻す」

同じ `は` の次を、入力文に応じて変えられるか