Powered by AppSignal & Oban Pro
Notesclub
Sign in with GitHub
created by hec & contributors
terms privacy

LLM を Elixir / Livebook で学ぶ 3

03_transformer_from_building_blocks.livemd

avatar
Ryo Wakabayashi
avatar
avatar
Share to X
Share to Bluesky
More notebooks

LLM を Elixir / Livebook で学ぶ 3

Mix.install([
  {:nx, "~> 0.9"},
  {:kino, "~> 0.15"},
  {:kino_vega_lite, "~> 0.1"}
])

Bigram から Transformer へ: 部品・学習・予測をつなげる

第2章では、現在の1トークンだけ から次を予測する bigram モデルを学習しました。この章でも学習の基本は変わりません。

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

変わるのは、次を予測するときに使える情報 です。

モデル 次トークンを予測するために見るもの 苦手・得意
第2章の bigram 直前の1トークン 簡単だが、長い文脈を区別できない
この章のTransformer 入力文全体 + 出力文の左側 離れた情報や別の文を参照できる

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

この違いを作るために、

  • マルチヘッドアテンション
  • 位置エンコーディング
  • フィードフォワード層
  • 残差接続
  • レイヤー正規化
  • エンコーダー
  • デコーダー
  • Transformer 全体

を 1 本の流れとして実装します。

この notebook の狙い

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

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

すべての層を同時に学習すると、初学者には「どの部分を観察しているか」が見えにくくなります。そこで後半の小さな実験では、Transformer本体を文脈を作る装置として固定し、デコーダーの表現 -> 次トークン を結ぶ出力層だけを学習します。実用モデルはTransformer全体を学習しますが、損失・更新・学習前後の比較は同じ考え方です。

全体像

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 つは似ていますが、役割は少し違います。

  • エンコーダーは入力文を読んで、「この文は何を言っているか」を表す記憶を作る
  • デコーダーは出力文の途中経過を見ながら、必要に応じてエンコーダーの記憶を参照して次トークンを決める 
component_rows = [
  %{
    部品: "マルチヘッド自己アテンション",
    使用場所: "エンコーダー / デコーダー",
    参照範囲: "同じ系列の各位置",
    役割: "文脈を読んで各位置の表現を更新する"
  },
  %{
    部品: "因果マスク付き自己アテンション",
    使用場所: "デコーダーのみ",
    参照範囲: "出力文の左側だけ",
    役割: "未来を見ずに次トークン予測を成立させる"
  },
  %{
    部品: "クロスアテンション",
    使用場所: "デコーダーのみ",
    参照範囲: "エンコーダーの記憶",
    役割: "入力文の情報を出力文の生成へ持ち込む"
  },
  %{
    部品: "フィードフォワード層",
    使用場所: "エンコーダー / デコーダー",
    参照範囲: "各位置を独立に変換",
    役割: "アテンション後の表現をさらに変換する"
  },
  %{
    部品: "残差接続 + レイヤー正規化",
    使用場所: "エンコーダー / デコーダー",
    参照範囲: "入力 + サブレイヤー出力",
    役割: "情報を保ちつつ学習を安定させる"
  }
]

Kino.DataTable.new(component_rows, keys: [:部品, :使用場所, :参照範囲, :役割])

マルチヘッドアテンションを直感で捉える

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

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

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

    scores =
      if mask do
        Nx.select(mask, Nx.broadcast(-1.0e9, Nx.shape(scores)), scores)
      else
        scores
      end

    weights = row_softmax(scores)
    output = Nx.dot(weights, [1], v, [0])
    {output, weights}
  end

  def multi_head_attention(query, key, value, params, num_heads, mask \\ nil) do
    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
    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
    {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)

    {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
    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
    {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()

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例すべてを学習に使います。

walkthrough_example = hd(training_examples)

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

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 の設定

config = %{
  d_model: 8,
  num_heads: 2,
  d_ff: 16,
  num_encoder_blocks: 2,
  num_decoder_blocks: 2
}

d_model=8 はかなり小さいですが、形状を見るには十分です。num_heads=2 にすると、マルチヘッドアテンションの分割も追いやすくなります。

shape_concept_rows = [
  %{
    項目: "1トークンの埋め込み",
    形状: "{d_model} = {8}",
    意味: "1つの位置の表現"
  },
  %{
    項目: "系列全体の埋め込み",
    形状: "{系列長, d_model} = {6, 8}",
    意味: "系列全体の表現"
  },
  %{
    項目: "2つのヘッドへ分割した後",
    形状: "{ヘッド数, 系列長, ヘッド次元} = {2, 6, 4}",
    意味: "8次元を4次元 × 2に分ける"
  },
  %{
    項目: "ヘッドごとのアテンション重み",
    形状: "{系列長, 系列長} = {6, 6}",
    意味: "各位置がどこを見るか"
  }
]

Kino.DataTable.new(shape_concept_rows, keys: [:項目, :形状, :意味])

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

:rand.seed(:exsss, {101, 202, 303})

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

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

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

layer_norm_params = fn d_model ->
  %{
    gamma: ones.(d_model),
    beta: zeros.(d_model)
  }
end

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: "位置"
)

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

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)

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回だけかけてみます。

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} という形状になっていれば、ヘッドごとの参照先を分けて見られます。

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: [:段階, :形状, :説明])
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"]
""")
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: [:データ, :形状])
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要素])

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

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

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

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_decoder_rows = [
  %{
    部分: "エンコーダー",
    入力: "入力文のトークン",
    主なアテンション: "自己アテンション",
    未来の参照: "入力文全体を見てよい",
    出力: "エンコーダーの記憶",
    役割: "入力文を理解した表現を作る"
  },
  %{
    部分: "デコーダー",
    入力: "ここまでの出力文トークン",
    主なアテンション: "因果マスク付き自己アテンション + クロスアテンション",
    未来の参照: "出力文の未来は見ない",
    出力: "隠れ状態 / 次トークン候補のスコア",
    役割: "入力文の記憶を参照しながら出力文を生成する"
  }
]

Kino.DataTable.new(
  encoder_decoder_rows,
  keys: [:部分, :入力, :主なアテンション, :未来の参照, :出力, :役割]
)
{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. デコーダーを実行する

{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["次トークン予測へ"]
""")
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. デコーダーの因果マスク付き自己アテンションとクロスアテンションを見る

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> わたし は のあとに何が来そうか

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

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. 学習前の出力語彙への射影を確認する

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次元 -> 出力語彙の確率 だけを学習します。

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)

training_features =
  training_runs
  |> Enum.map(& &1.decoder_output)
  |> Nx.stack()
  |> Nx.reshape({training_example_count * sequence_length, config.d_model})

training_target_ids =
  training_examples
  |> Enum.flat_map(& &1.expected)
  |> Enum.map(&tgt_token_to_id[&1])
  |> Nx.tensor(type: :s64)

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個の次トークン問題をまとめて解き、正解トークンの確率が高くなるよう出力層を更新します。

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} ->
      loss =
        LLMScratch.TrainableHead.loss(
          w,
          b,
          training_features,
          training_target_one_hot
        )
        |> Nx.to_number()

      {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 ->
  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: [
    :入力文,
    :共通のデコーダー入力,
    :正解の次トークン,
    :学習前の予測,
    :学習後の予測
  ]
)

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

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 ->
  predicted_ids =
    probabilities
    |> Nx.argmax(axis: 1)
    |> Nx.to_flat_list()

  correct_ids = Nx.to_flat_list(training_target_ids)

  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は入力文全体と出力文の左側を使って次を予測できる
  • 同じデコーダー入力でも、クロスアテンションによって入力文に応じた次トークンを選び分けられる

次に進むなら、

  • 出力層だけでなく、埋め込み・アテンション・フィードフォワード層もまとめて学習する
  • パディングマスクを足す
  • デコーダーのみのGPTへ寄せて、因果言語モデル（causal LM）を作る

という順が自然です。

Other notebooks:

Back