Notesclub

created by hec & contributors

terms privacy

LLM を Elixir / Livebook で学ぶ 3

03_attention_causal_mask_position.livemd

Ryo Wakabayashi

@RyoWakabayashi

elixir-learning

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 の限界から出発する

第2章の bigram モデルには、はっきりした限界がありました。

直前のトークン	bigram が出せる予測	本当は使いたい情報
は（ねこはの続き）	ひるね / さかな / さんぽ / … が混ざった分布	主語が `ねこ` であること
は（いぬはの続き）	上とまったく同じ分布	主語が `いぬ` であること

bigram は「直前の 1 トークン」しか見ないので、は の前に何があったかを予測へ活かせません。

この章では、この限界を解決する Transformer の中心の仕組みを 3 つ、実際に計算しながら学びます。

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

このnotebookで使う用語

用語	英語・コード上の表記	この章での意味
アテンション	attention	現在位置に必要な情報を、ほかの位置から重み付きで集める仕組み
Query／Key／Value	query / key / value	参照したい情報、参照先との相性、実際に取り出す情報を表す3種類のベクトル
softmax	softmax	複数の点数を、合計1になる確率や重みへ変換する関数
スケーリング	scaling / `sqrt(d_k)`	内積の点数を次元数の平方根で割り、softmaxが極端に尖るのを防ぐ処理
自己アテンション	self-attention	同じ系列の中で、各位置がほかの位置を参照するアテンション
因果マスク	causal mask	現在位置から未来のトークンを見えなくする制限
位置エンコーディング	positional encoding	トークンが何番目にあるかをベクトルへ加える仕組み
デコーダーのみのモデル	decoder-only model	左側の文脈から次トークンを予測するデコーダーを積み重ねたモデル

準備

defmodule LLMScratch.Visuals do
  # 行データを、カテゴリごとの大きさを比べる棒グラフへ変換する
  def bar_chart(rows, title, x_field, y_field, opts \\ []) do
    width = Keyword.get(opts, :width, 520)
    height = Keyword.get(opts, :height, 260)

    VegaLite.new(width: width, height: height, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:bar, tooltip: true)
    |> VegaLite.encode_field(
      :x, Atom.to_string(x_field),
      type: :nominal,
      title: Atom.to_string(x_field)
    )
    |> VegaLite.encode_field(
      :y, Atom.to_string(y_field),
      type: :quantitative,
      title: Atom.to_string(y_field)
    )
    |> Kino.VegaLite.new()
  end

  # 同じ形式のデータを、値の推移を追う折れ線グラフへ変換する
  def line_chart(rows, title, x_field, y_field, opts \\ []) do
    width = Keyword.get(opts, :width, 520)
    height = Keyword.get(opts, :height, 260)
    color_field = Keyword.get(opts, :color_field)

    chart =
      VegaLite.new(width: width, height: height, title: title)
      |> VegaLite.data_from_values(rows)
      |> VegaLite.mark(:line, point: true)
      |> VegaLite.encode_field(:x, Atom.to_string(x_field),
        type: :quantitative,
        title: Atom.to_string(x_field)
      )
      |> VegaLite.encode_field(:y, Atom.to_string(y_field),
        type: :quantitative,
        title: Atom.to_string(y_field)
      )

    # 系列を色分けしたい場合だけcolorエンコーディングを追加する
    chart =
      if color_field do
        VegaLite.encode_field(chart, :color, Atom.to_string(color_field), type: :nominal)
      else
        chart
      end

    Kino.VegaLite.new(chart)
  end

  # QueryとKeyの全組み合わせを、色の濃さで読めるヒートマップにする
  def heatmap(rows, title, opts \\ []) do
    width = Keyword.get(opts, :width, 520)
    height = Keyword.get(opts, :height, 260)

    VegaLite.new(width: width, height: height, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:rect, tooltip: true)
    |> VegaLite.encode_field(
      :x, "key", type: :nominal, title: "注目されるトークン",
      sort: [field: "index_k"]
    )
    |> VegaLite.encode_field(
      :y, "query", type: :nominal, title: "見ているトークン",
      sort: [field: "index_q"]
    )
    |> VegaLite.encode_field(
      :color, "score", type: :quantitative, scale: [scheme: "blues"])
    |> Kino.VegaLite.new()
  end
end

この章でも、第1章と同じ文と手作り埋め込みを使います。

tokens = ["ねこ", "は", "ひるね", "が", "すき"]

# 第1章と同じ、説明用の手作り2次元埋め込み
embedding_map = %{
  "ねこ" => [0.90, 0.15],
  "は" => [0.10, 0.05],
  "ひるね" => [0.95, 0.85],
  "が" => [0.12, 0.10],
  "すき" => [0.88, 0.70]
}

embedding_tensor =
  tokens
  |> Enum.map(&embedding_map[&1])
  |> Nx.tensor(type: {:f, 32})

embedding_tensor

1. 最初のアイデア: 前の全トークンを平均する

「直前の 1 トークンでは足りないなら、前にあるトークンを全部使えばよい」——これが出発点です。

いちばん簡単な方法は、前にある全トークンの埋め込みを同じ重みで平均することです。

# 文末「すき」の位置で、全トークンの埋め込みを均等な重みで混ぜる
uniform_weight = 1.0 / length(tokens)

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

uniform_rows =
  Enum.map(tokens, fn token ->
    %{token: token, weight: uniform_weight}
  end)

Kino.Layout.grid(
  [
    LLMScratch.Visuals.bar_chart(uniform_rows, "均等な重み: 全部 0.2", :token, :weight),
    Kino.DataTable.new([
      %{
        混ぜた結果: inspect(uniform_mixture |> Nx.to_flat_list() |> Enum.map(&Float.round(&1, 3)))
      }
    ])
  ],
  columns: 1
)

これで「文全体の情報」は手に入りましたが、新しい問題が生まれます。

ひるね のような大事な語も、は のような機能語も、同じ 0.2 の重みで混ざる
どの位置から見ても同じ混ぜ方になるので、「いま何を知りたいか」を反映できない

必要なのは、今の位置が何を知りたいかに応じて、参照先ごとに重みを変える仕組みです。それがアテンションです。

2. アテンション: 重み付きで文脈を集める

アテンションは、今の位置の判断に必要な情報を、ほかの位置からどの割合で集めるか を決める仕組みです。

たとえば ねこはひるねがすき の最後にある すき を考えます。すき というトークン単体だけでは、「何が好きなのか」は分かりません。そこで前にある各位置へ重みを付け、ベクトルを混ぜ合わせます。

参照先	重みの例	集めたい情報のイメージ
ねこ	0.20	誰の話か
は	0.03	文のつながり
ひるね	0.60	何が好きか
が	0.05	文のつながり
すき	0.12	現在位置自身の情報

この 0.20 や 0.60 がアテンション重みです。「参照する」とは、文字列を読み直すことではなく、各位置のValueベクトルを重みに応じて足し合わせ、今の位置の表現を更新する ことを指します。実際の重みは学習によって決まり、上の数値は説明用です。

では、その重みはどう決めるのでしょうか。第1章の最後で見た 内積で相性を測る 考え方をここで使います。

softmax_1d = fn scores ->
  # 最大値を先に引くと、expの値が大きくなりすぎるのを防げる
  exps = Nx.exp(Nx.subtract(scores, Nx.reduce_max(scores)))
  Nx.divide(exps, Nx.sum(exps))
end

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

  # Queryと各Keyの内積を次元数の平方根で割り、参照先ごとの点数を作る
  scores = Nx.divide(Nx.dot(keys, [1], query, [0]), :math.sqrt(d_model))
  weights = softmax_1d.(scores)

  # 合計1の重みでValueを混ぜ、現在位置へ集める情報を作る
  output = Nx.dot(weights, [0], values, [0])
  {output, weights}
end

# 今回は最後の「すき」をQueryとして、文中の全位置を参照させる
query_index = Enum.find_index(tokens, &(&1 == "すき"))
query = embedding_tensor[query_index]

# 最小例なので、同じ埋め込みtensorをKeyとValueにも使う
{attention_output, attention_weights} =
  scaled_dot_attention.(query, embedding_tensor, embedding_tensor)

%{
  query_token: Enum.at(tokens, query_index),
  output_vector: attention_output,
  weights: attention_weights
}

# tensorの重みを、トークン名付きのグラフ用データへ変換する
weight_rows =
  tokens
  |> Enum.with_index()
  |> Enum.map(fn {token, index} ->
    %{
      token: token,
      weight: Nx.to_number(attention_weights[index])
    }
  end)

LLMScratch.Visuals.bar_chart(weight_rows, "`すき` がどのトークンを見るか", :token, :weight)

先ほどの均等な 0.2 と違い、棒の高さがトークンごとに変わりました。棒が高いほど、その位置のベクトルが attention_output に強く混ざります。この例では、すき と意味的に近く置いた ひるね や ねこ への重みが大きくなりやすくなります。

もちろん実際の LLM では、埋め込みは人手ではなく学習で作られます。

Query／Key／Valueは何を分けているか

実際のアテンションの説明では Query、Key、Value という3つの名前が登場します。図書館で資料を探す場面に置き換えると、次のように読めます。

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

処理は Query と各 Key の相性を点数化 -> softmax で合計1の重みにする -> その重みで Value を混ぜる の順です。この最小例では理解しやすさを優先し、同じ埋め込みを Query / Key / Value として使っています。実際の Transformer では、それぞれ別の線形変換を通します。同じトークンでも「探すとき（Query）」と「見つけてもらうとき（Key）」で違う顔を持てるようにするためです。この線形変換は第4章で実装します。

3. なぜスケーリングが必要か

埋め込み次元が大きいと、内積の値が大きくなりやすく、softmax が極端に尖ることがあります。

# 比較を再現できるよう、乱数の初期状態を固定する
:rand.seed(:exsss, {11, 22, 33})

# 値のばらつきが大きい点数と、小さい点数を用意する
large_scores =
  1..12
  |> Enum.map(fn _ -> :rand.normal() * 8 end)

small_scores =
  1..12
  |> Enum.map(fn _ -> :rand.normal() end)

# 2種類の点数へ同じsoftmaxを適用し、確率の偏り方を比べる
large_softmax = softmax_1d.(Nx.tensor(large_scores))
small_softmax = softmax_1d.(Nx.tensor(small_scores))

score_rows =
  Enum.with_index(large_scores)
  |> Enum.map(fn {score, index} ->
    %{
      index: index,
      large_scale: score,
      large_softmax: Nx.to_number(large_softmax[index]),
      small_softmax: Nx.to_number(small_softmax[index])
    }
  end)

Kino.DataTable.new(score_rows)

# 2系列を同じグラフに重ねられる縦長の行データへ変換する
softmax_rows =
  score_rows
  |> Enum.flat_map(fn row ->
    [
      %{index: row.index, value: row.large_softmax, series: "値が大きいとき"},
      %{index: row.index, value: row.small_softmax, series: "値が小さいとき"}
    ]
  end)

LLMScratch.Visuals.line_chart(
  softmax_rows,
  "softmax の尖り方の比較",
  :index,
  :value,
  color_field: :series
)

値が大きいときは、ほぼ 1 か所へ重みが集中してしまいます。こうなると「複数の位置から少しずつ情報を集める」というアテンションの良さが失われ、学習も不安定になります。

この偏りを少しやわらげるために、Transformerでは通常 sqrt(d_k) で割るスケーリング付きドット積アテンション（scaled dot-product attention）を使います。先ほどの scaled_dot_attention にも、この割り算がすでに入っています。

4. 因果マスクを可視化する

第1章で見たように、GPTは ここまでのトークンから次を当てる デコーダーのみのモデルです。その学習では、完成した文章を次のような問題にずらして使います。

モデルに見せる範囲	当ててほしい正解
わたし	は
わたしは	きょう
わたしはきょう	本

もし1行目の予測中に、正解である右側の は や、その先の きょう まで見えていたら、モデルは答えを写すだけで損失を下げられます。それでは、生成時のように「右側がまだ存在しない状況」で予測する力が身につきません。

アテンションは放っておくと系列の全位置を参照できてしまうので、GPTの学習では各位置から右側を見えなくします。これが因果マスクです。

# 因果マスクを観察するための、デコーダー入力と手作り埋め込み
decoder_tokens = ["わたし", "は", "きょう", "本", "を", "読む"]

decoder_embedding_map = %{
  "わたし" => [0.80, 0.10, 0.20],
  "は" => [0.10, 0.05, 0.05],
  "きょう" => [0.55, 0.85, 0.20],
  "本" => [0.92, 0.70, 0.15],
  "を" => [0.12, 0.12, 0.02],
  "読む" => [0.75, 0.88, 0.92]
}

# 各トークンの埋め込みを系列順に並べる
decoder_embeddings =
  decoder_tokens
  |> Enum.map(&decoder_embedding_map[&1])
  |> Nx.tensor(type: {:f, 32})

decoder_embeddings

# Query位置iより右にあるKey位置jを「参照禁止=1」として記録する
causal_mask_rows =
  for {query, i} <- Enum.with_index(decoder_tokens),
      {key, j} <- Enum.with_index(decoder_tokens) do
    %{
      index_q: i,
      index_k: j,
      query: query,
      key: key,
      score: if(j > i, do: 1.0, else: 0.0)
    }
  end

LLMScratch.Visuals.heatmap(causal_mask_rows, "1 の場所は未来なので見てはいけない")

row_softmax = fn matrix ->
  # 行ごとにsoftmaxを計算し、各QueryからKeyへの重みの合計を1にする
  exps =
    Nx.exp(
      Nx.subtract(matrix, Nx.reduce_max(matrix, axes: [1], keep_axes: true))
    )

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

# 全Queryと全Keyの内積を一度に計算し、{系列長, 系列長}の点数表を作る
score_matrix =
  Nx.divide(
    Nx.dot(decoder_embeddings, [1], Nx.transpose(decoder_embeddings), [0]),
    :math.sqrt(3)
  )

# j > i、つまり現在位置より未来だけが1になる因果マスク
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})

# 未来の点数を非常に小さくすると、softmax後の重みはほぼ0になる
masked_scores =
  Nx.select(mask_tensor, Nx.broadcast(-1.0e9, Nx.shape(score_matrix)), score_matrix)

masked_weights = row_softmax.(masked_scores)

# マスク適用後のtensorを、トークン名付きのヒートマップ用データへ戻す
masked_weight_rows =
  for {query, i} <- Enum.with_index(decoder_tokens),
      {key, j} <- Enum.with_index(decoder_tokens) do
    %{
      index_q: i,
      index_k: j,
      query: query,
      key: key,
      score: Nx.to_number(masked_weights[[i, j]])
    }
  end

LLMScratch.Visuals.heatmap(masked_weight_rows, "因果マスク適用後のアテンション重み")

ヒートマップは、縦軸の現在位置から横軸のどこを参照したかを表します。右上側の重みが 0 になり、各トークンが自分より右側を見ていないことを確認してください。これで学習時と生成時の条件がそろいます。

「マスクする」の実装は、見てはいけない位置の点数を softmax の前に -10億のような巨大な負の値へ置き換える だけです。softmax は指数関数を使うので、巨大な負の点数の重みは実質 0 になります。

5. 位置情報はどう入るのか

アテンションの計算をもう一度見直すと、内積と softmax と重み付き和だけでできています。つまり、何も補わなければ どのトークンが何番目にあったか を直接は持ちません。ところが、次の2文は同じトークンを含んでいても意味が違います。

ねこがいぬを追う
いぬがねこを追う

そこで各トークンの埋め込みに、座席番号の名札のような 位置ベクトル を足します。同じ ねこ の埋め込みでも、0番目用の位置ベクトルを足した結果と、2番目用を足した結果は別のベクトルになります。

ここでは固定の sin / cos 位置エンコーディングを小さく作ってみます。

positional_encoding = fn max_position, d_model ->
  # 位置ごと、次元ごとに周期の異なるsin/cosの値を計算する
  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)

      # 偶数次元にはsin、奇数次元にはcosを使う
      if rem(dim, 2) == 0 do
        :math.sin(angle)
      else
        :math.cos(angle)
      end
    end
  end
end

pe_matrix = positional_encoding.(12, 8)
pe_matrix |> Enum.take(4)

# 最初の4次元だけを取り出し、位置による波形の違いを可視化する
pe_rows =
  pe_matrix
  |> Enum.with_index()
  |> Enum.flat_map(fn {row, position} ->
    row
    |> Enum.take(4)
    |> Enum.with_index()
    |> Enum.map(fn {value, dim} ->
      %{
        position: position,
        value: value,
        dimension: "dim_#{dim}"
      }
    end)
  end)

LLMScratch.Visuals.line_chart(
  pe_rows,
  "位置エンコーディングの波形",
  :position,
  :value,
  color_field: :dimension
)

# 同じトークン埋め込みを3つの異なる位置へ置く
same_token_embedding = Nx.tensor([0.8, 0.2, -0.1, 0.5], type: {:f, 32})
small_pe = positional_encoding.(3, 4) |> Nx.tensor(type: {:f, 32})

same_token_position_rows =
  0..2
  |> Enum.map(fn position ->
    # トークン埋め込みへ位置ベクトルを足すと、位置ごとに異なる入力になる
    combined = Nx.add(same_token_embedding, small_pe[position])

    %{
      position: position,
      same_token_embedding: inspect(Nx.to_flat_list(same_token_embedding) |> Enum.map(&Float.round(&1, 3))),
      position_signal: inspect(Nx.to_flat_list(small_pe[position]) |> Enum.map(&Float.round(&1, 3))),
      model_input: inspect(Nx.to_flat_list(combined) |> Enum.map(&Float.round(&1, 3)))
    }
  end)

Kino.DataTable.new(
  same_token_position_rows,
  keys: [:position, :same_token_embedding, :position_signal, :model_input]
)

この表では same_token_embedding は3行とも同じですが、position_signal が違うため、足し算後の model_input は位置ごとに変わります。

sin / cos の波そのものに「主語」や「目的語」という意味が入っているわけではありません。異なる速さで動く複数の針を組み合わせて、各位置に重なりにくい目印を作っている、と捉えると十分です。モデルは学習を通じて、その目印と語順の関係を利用します。

6. ここまでのまとめ

この章の要点

前の全トークンを均等に混ぜるだけでは、大事な語と機能語を区別できない
アテンションは、Query と Key の内積で「参照先ごとの重み」を作り、その重みで Value を混ぜる
内積が大きくなりすぎて softmax が尖らないよう、sqrt(d_k) でスケーリングする
GPT では因果マスクによって未来のトークンを見ない
アテンションだけでは語順が分からないので、位置エンコーディングで補う

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

次のノートブックでは、この章のアテンション・因果マスク・位置エンコーディングに、マルチヘッド化・フィードフォワード層・残差接続・レイヤー正規化を加えてミニ GPT を組み立て、第2章と同じコーパスで端から端まで学習します。bigram では出せなかった「ねこは と いぬは で違う予測」が出せるようになるかを確かめます。

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
Nicolò G.
@nickgnd

programming-machine-learn...

SONAR, mines vs. rocks

sonar_classifier.livemd

tutorial advanced data-science exla nx vega_lite kino kino_vega_lite

2026-7-7
Chrisna Adhi Pranoto
@nethricia

curriculum

PicChat: Authentication

pic_chat_authentication.livemd

tutorial advanced testing debugging jason kino youtube hidden_cell

2026-7-12
@DeSchoel

Elixir_Curriculum

Rock Paper Scissors Lizard Spock

rock_paper_scissors_lizard_spock.livemd

tutorial intermediate jason kino youtube hidden_cell

2026-7-13

Back