この記事は、Elixir Advent Calendar 2025その3 の7日目です
昨日は、@t-yamanashi さんで 「AtomVM Raspberry Pi Picoのファームをビルドしてブートを高速化する」 でした
piacere です、ご覧いただいてありがとございます 
生成AIの基本的な概念と各パーツについて解説しつつ、Elixir生成AIエンジン「Bumblebee」のコードハックをしていきます
初回は、生成AIの全体像を見た後、生成AIの中核となる 「Attention」 について触れていきます
なお、Bumblebeeは各種モデル固有の処理を持っていますが、このコラムシリーズ中では下記の gemma2-2b を対象として解説します
Elixirアドベントカレンダー、応援お願いします
今年もやっています
ａ）生成AIとは？
生成AIは、2017年にGoogleが公開したディープラーニングモデル 「Transformer」 を基盤とすることで、ソレ以前のモデルでは難しかった文脈や長文の解釈力と表現力、そして大規模データを効率良く学習できる並列計算を叶え、2022年頃から広まり始めました
Transformerは、入力された文章（≒プロンプト）を 「Attention」 によって単語同士の関係から文脈の特徴を捉え、「FFN（Feed-Forward Network）」 でその特徴を強調します
この2つを繰り返すことで、文章全体の文脈を理解し、文脈に応じた質問応答や要約、画像生成などの高度なタスクを実現します
詳しくは、下記を2分ほどご覧いただくと、生成AIが Attention と FFN（動画では Multilayer Perceptron と呼んでいる）を繰り返す点が理解しやすいと思います
現在主流のLLMは、このTransformerをそのまま使うのでは無く、下記のような構造変更を行うことで大規模モデルを成立させています
-
Encoderを無くしDecoderのみで構成することで生成タスクに特化 -
LayerNormをブロックの前に置くことで大規模モデルを安定化 - その他、大規模化／効率化／表現力向上のための各種改良
ｂ）Attentionの前に行われていること
ユーザーが入力したプロンプトは、下記の流れによってAttentionが扱える入力である QKV へと変換されます
| 用語 | 説明 |
|---|---|
| ①Tokenizer | Attention前にプロンプトをトークンに分解／数値化 |
| ②Embedding | Attention前にQKVの元となるベクトルを生成するニューラルネットワーク層 |
| ③Positional Encoding | Attention前に単語の位置情報を埋め込みベクトルに加える手法 |
| ④QKV | Attentionへの入力としてEmbeddingから作られる下記3つ ・Query：単語自身がどのような他単語を問いたいか？ ・Key：単語自分がどんな役割を提供できるか？ ・Value：単語自身の意味そのもの |
Q（Query） は「単語自身がどの他単語を参照したいか？」ですが、たとえば「食べた」という単語なら、「誰が？」「何を？」といった情報をプロンプト中から参照しようとする問いを保持します
K（Key） は「単語自分がどんな役割を提供できるか？」ですが、上記の「食べた」に対して、「Aさん」という単語であれば「主語としての情報を提供できる役割」を保持し、「カレー」という単語であれば「目的語としての情報を提供できる役割」を保持します
V（Value） は「単語自身の意味そのもの」ですが、上記の「カレー」なら食べ物としての意味、「Aさん」なら人物としての意味を保持します
ｃ）Attentionで行われていること
Attentionが、単語同士の関係から文脈の特徴を捉える仕組みは、下記の流れで実現されます
| 処理 | 説明 |
|---|---|
| ①スコア行列 | プロンプト内の全単語の組み合わせに対して類似度スコアを出したもので、全単語同士の「文脈上の必要性」を網羅した行列を作成 類似度スコアは、Qの「単語自身がどのような他単語を問いたいか？」と、Kの「単語自分がどんな役割を提供できるか？」のマッチ度（≒QとKの内積で算出） |
| ②Softmax | スコア行列に対し、マッチ度の高い単語を強調し、マッチ度の低い単語の影響を小さくする「重み」を作成 |
| ③加重和 | Vの「単語自身の意味そのもの」に重みを掛けることで、Qの「問いたい」に応えるVが強調され、それを足し合わせることで複数のVを繋いだ文脈が作られる |
なお、「①スコア行列」と「③加重和」の計算の際、過去トークンの K ないしは V をキャッシュすることで計算を効率良くするための K/Vキャッシュ も使っています
それから、生成AIにおけるAttentionは、正確には 「Self-Attention」 であり、下記のような違いがあります
| 処理 | 説明 |
|---|---|
| Attention | プロンプト内の各単語が、自身と他単語をどれだけ重要とみなすかを読み取るニューラルネットワーク層 |
| Self-Attention | プロンプト内の各単語が、自身と他単語が遠く離れていても関係性を文脈全体から読み取れるAttentionで、この特徴により長い文章を理解できるようになり、LLM（≒大規模言語モデル）が構築可能となる |
Attention内で行われている代表的なアルゴリズムには、下記があります
| 処理 | 説明 |
|---|---|
| Scaling | Attention内でスコアの勾配消失を防ぎ、安定化させる |
| Masking | Attention内で後ろのトークンが前のトークンに影響を及ぼさないようSoftmax時に0とするための手法（Softmax前に-∞を入れる） |
| Dropout | Attention内で「重み」に対してランダムでゼロにすることで過学習を防ぐ（≒汎化と呼ばれる処理） |
Attentionのメカニズムは、下記動画がとても分かりやすいのでご参考ください
BumblebeeのAttention
QKVに対し、Multi-Head Attentionを実行する下記の Bumblebee.Layers.attention が、Attentionの中核処理となります
引数のうち opts の dropout_rate 以外は、後続の関数に渡されているだけのため、ここでは割愛し、後続関数ないしは処理内で解説します
defmodule Bumblebee.Layers do
@moduledoc false
…
def attention(query, key, value, key_mask, head_mask, bias, offset, opts \\ []) do
opts = Keyword.validate!(opts, [:window_size, causal: false, scale: true, dropout_rate: 0.0])
weights =
Axon.layer(
&attention_weights_impl/7,
[
query,
key,
Axon.optional(key_mask),
Axon.optional(head_mask),
Axon.optional(bias),
Axon.optional(offset)
],
causal: opts[:causal],
window_size: opts[:window_size],
scale: opts[:scale]
)
|> Axon.dropout(rate: opts[:dropout_rate])
output = Axon.layer(&attention_output_impl/3, [weights, value], opts)
{output, weights}
end
最初の Keyword.validate! は、opts で指定されたオプション引数が関数仕様通りかバリデーションしますが、うち window_size と causal はMaskingで使用、scale はScalingで使用、dropout_rate はDropoutで使用されます
Axon.layer は、ニューラルネットワーク層を構築するもので、層として定義されている attention_weights_impl は、上記ｃ）で解説した下記を順に行い、QKV から 重み を作成します
| 処理 | 説明 |
|---|---|
| スコア行列 | プロンプト内の全単語の組み合わせに対して類似度スコアを出したもので、全単語同士の「文脈上の必要性」を網羅した行列を作成 類似度スコアは、Qの「単語自身がどのような他単語を問いたいか？」と、Kの「単語自分がどんな役割を提供できるか？」のマッチ度（≒QとKの内積で算出） |
| Scaling | Attention内でスコアの勾配消失を防ぎ、安定化させる |
| Masking | Attention内で後ろのトークンが前のトークンに影響を及ぼさないようSoftmax時に0とするための手法（Softmax前に-∞を入れる） |
| Softmax | スコア行列に対し、マッチ度の高い単語を強調し、マッチ度の低い単語の影響を小さくする「重み」を作成 |
その後、Axon.dropout で 重み に対して、上記ｃ）で解説した下記を行い、 重み を安定させます
| 処理 | 説明 |
|---|---|
| Dropout | Attention内で「重み」に対してランダムでゼロにすることで過学習を防ぎ、安定させる（≒汎化と呼ばれる処理） |
最後に attention_output_impl では、Attentionの最終的な出力として、上記ｃ）で解説した下記を行い、V と 重み を掛け合わせることで、文脈情報を返します
| 処理 | 説明 |
|---|---|
| ③加重和 | Vの「単語自身の意味そのもの」に重みを掛けることで、Qの「問いたい」に応えるVが強調され、それを足し合わせることで複数のVを繋いだ文脈が作られる |
スコア行列 → Scaling → Masking → Softmax
上述した attention_weights_impl は、同モジュール中の下記Nx関数で定義されています
なお引数の key_mask は、入力データの長さが不均一なのを埋め合わせるための、部脈上は無意味な「パディングトークン」を無効化するために使われ、head_mask はMulti-Head Attentionにおいて一部のヘッドをDropoutするために使われます
bias は、T5などのEncoderモデルでは使われますが、Self-AttentionのようなDecoderモデルでは空（≒%Axon.None{}）が指定されます
opts の causal オプションは、未来トークン参照を許可する（≒false）か禁止する（≒true）かの指定で、未来トークン参照を禁止するSelf-Attentionであれば必ずtrueになります
window_size オプションは、近接するトークンのみAttention処理を行うための指定ですが、Self-Attentionは無制限で処理するので、必ず未設定（≒nil）になります
defmodule Bumblebee.Layers do
@moduledoc false
…
defnp attention_weights_impl(query, key, key_mask, head_mask, bias, offset, opts \\ []) do
opts = keyword!(opts, [:window_size, mode: :inference, scale: true, causal: false])
query = Nx.transpose(query, axes: [0, 2, 1, 3])
key = Nx.transpose(key, axes: [0, 2, 1, 3])
weights = Nx.dot(query, [3], [0, 1], key, [3], [0, 1])
weights =
if opts[:scale] do
depth = Nx.axis_size(query, -1)
weights / Nx.as_type(Nx.sqrt(depth), Nx.type(query))
else
weights
end
key_mask =
case key_mask do
%Axon.None{} ->
Nx.broadcast(1, {1, 1, 1, 1})
key_mask ->
case Nx.rank(key_mask) do
2 -> key_mask |> Nx.new_axis(1) |> Nx.new_axis(1)
4 -> key_mask
end
end
query_sequence_length = Nx.axis_size(query, 2)
key_sequence_length = Nx.axis_size(key, 2)
offset = ensure_offset(offset)
causal_and_window_mask =
case {opts[:causal], opts[:window_size]} do
{false, nil} ->
Nx.broadcast(1, {1, 1})
{false, {left_size, right_size}} ->
window_mask(query_sequence_length, key_sequence_length, offset, left_size, right_size)
{true, nil} ->
causal_mask(query_sequence_length, key_sequence_length, offset)
{true, {left_size, _right_size}} ->
window_mask(query_sequence_length, key_sequence_length, offset, left_size, 0)
end
|> Nx.new_axis(0)
|> Nx.new_axis(0)
mask = key_mask and causal_and_window_mask
bias =
case bias do
%Axon.None{} ->
Nx.select(
mask,
Nx.tensor(0.0, type: Nx.type(query)),
Nx.Constants.min_finite(Nx.type(query))
)
bias ->
Nx.select(
Nx.broadcast(mask, max_shape(mask, bias)),
bias,
Nx.Constants.min_finite(Nx.type(query))
)
end
weights = weights + bias
weights = Axon.Activations.softmax(weights, axis: -1)
case head_mask do
%Axon.None{} ->
weights
head_mask ->
head_mask = Nx.reshape(head_mask, {1, :auto, 1, 1})
Nx.multiply(weights, head_mask)
end
end
最初に、Q と K の内積を取ることで類似度を計算し、スコア行列を作成し、weights に保持し、scale オプションがtrueであればScalingも行います
それから、key_mask が未設定（≒%Axon.None{}）であればパディングの無効化を行わず、設定されていれば key_mask を次元に合わせて次元変換を行い、causal オプションと window_size オプションの組み合わせ毎に下記4パターンのいずれかで causal_and_window_mask に保持し、最終的に key_mask でフィルタし、mask に保持します
なお、Self-Attentionでは3番目のパターン使用が確定です
| causal | window_size | 説明 | Attention動作 |
|---|---|---|---|
| false | nil | マスク無 | 制限無Attention |
| false | {left_size, right_size} | 因果性無＋ウィンドウサイズあり → window_mask を呼び出す |
近接トークン限定Attention |
| true | nil | 因果性あり＋ウィンドウサイズ無 → causal_mask を呼び出す |
因果マスク |
| true | {left_size, _right_size} | 因果性あり＋ウィンドウサイズあり → window_mask を呼び出す（右側 right_size は0に強制設定） |
制限付因果マスク |
参考までに、各Attention動作がどのようなタスクに使われるかも挙げておきます
| Attention動作 | 使われるタスク例 |
|---|---|
| 制限無Attention | BERTのような未来トークンを参照するEncoderモデル |
| 近接トークン限定Attention | ViT（Vision Transformer）のような画像処理、長文解析 |
| 因果マスク | GPTシリーズやGeminiのような現在主流のDecoderモデル |
| 制限付因果マスク | 大規模言語モデルの推論、直近のみ参照のストリーミング |
その後、後ろのトークンが前のトークンに影響を及ぼさないよう、Softmax時に0とするために-∞（≒Nx.Constants.min_finite）を bias に保持し、Maskingを行い、Softmaxを行った結果を weights に更新します
最後に、head_mask が未設定（≒%Axon.None{}）であればパディングの無効化を行わず、設定されていれば head_mask を weights と同次元に調整し、weights と乗算することで head_mask 中の値がゼロの場合は weights が無効化されDropputを実現します
因果マスク
上述した causal_mask は、同モジュール中の下記Nx関数で定義されています
引数の query_sequence_length は Q の長さ、key_sequence_length は K の長さ、offset は「参照可能な過去のKVの長さ」、つまり「K/Vキャッシュ の長さ」が指定されます
defmodule Bumblebee.Layers do
@moduledoc false
…
defnp causal_mask(query_sequence_length, key_sequence_length, offset) do
Nx.greater_equal(
Nx.iota({query_sequence_length, 1}) + offset,
Nx.iota({1, key_sequence_length})
)
end
結果として、Q の長さ x K の長さのマスク行列を返却します
加重和
上述した attention_output_impl は、同モジュール中の下記Nx関数で定義されています
defmodule Bumblebee.Layers do
@moduledoc false
…
defnp attention_output_impl(weights, value, _opts \\ []) do
value = Nx.transpose(value, axes: [0, 2, 1, 3])
out = Nx.dot(weights, [3], [0, 1], value, [2], [0, 1])
Nx.transpose(out, axes: [0, 2, 1, 3])
end
weights と V の内積を計算するために、Transformerブロックの標準的な形式である V の軸順を転置してから内積を行います
最後に、Transformerブロックの標準的な形式に戻し、Attention後のFFNに渡せるようにして返却します
AttentionとFFNはどのように呼び出される？
最後に、Attentionがどのように呼ばれ、Attentionの後続のFFNに渡されるかを見ます
まずBumblebeeで生成AIを使う際、モデルロードから実行までは下記のようなコードで行われます
repo = {:hf, "google/gemma-2-2b"}
{:ok, model} = Bumblebee.load_model(repo)
{:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
generation = Bumblebee.Text.TextGeneration.generation(model, tokenizer)
{:ok, result} = generation.("「先端ぴあちゃん」から連想されるストーリーを3つ挙げて")
Bumblebee.load_model でのモデルロード時に Bumblebee.【モデル名】.decoder が呼ばれ、その中でTransformerのニューラルネットワーク層を定義する Bumblebee.Layers.Transformer.blocks でAttentionとFFNを含むニューラルネットワーク層として設定されます
なお decoder は、Bumblebee.load_model → Bumblebee.build_model → Bumblebee.【各モデル】.model → Bumblebee.【各モデル】.core → Bumblebee.【各モデル】.decoder という順で呼び出されています
その後の Bumblebee.Text.TextGeneration.generation 内で、Nx.Serving.new にてAttentionとFFNを含むニューラルネットワーク層が実行されます
gemma2-2b であれば、Bumblebee.【モデル名】.decoder の部分は下記 Bumblebee.Text.Gemma が選択され、その中の blocks でAttentionとFFNを含むニューラルネットワーク層が作られています
defmodule Bumblebee.Text.Gemma do
alias Bumblebee.Shared
…
defp decoder(
hidden_state,
position_ids,
attention_mask,
attention_head_mask,
cache,
spec,
opts
) do
name = opts[:name]
+ Layers.Transformer.blocks(hidden_state,
+ attention_mask: attention_mask,
+ attention_head_mask: attention_head_mask,
+ attention_head_size: spec.attention_head_size,
+ cache: cache,
+ num_blocks: spec.num_blocks,
+ num_attention_heads: spec.num_attention_heads,
+ num_key_value_heads: spec.num_key_value_heads,
hidden_size: spec.hidden_size,
kernel_initializer: kernel_initializer(spec),
layer_norm:
&Layers.rms_norm(&1, shift: 1.0, name: &2, epsilon: spec.layer_norm_epsilon, upcast: :all),
+ ffn:
+ &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
+ name: &2,
+ activation: spec.activation
+ ),
block_type: :norm_first,
+ causal: true,
rotary_embedding: [
position_ids: position_ids,
max_positions: spec.max_positions,
base: spec.rotary_embedding_base,
scaling_strategy: spec.rotary_embedding_scaling_strategy
],
+ query_use_bias: spec.use_attention_bias,
+ key_use_bias: spec.use_attention_bias,
+ value_use_bias: spec.use_attention_bias,
+ output_use_bias: spec.use_attention_bias,
name: join(name, "blocks")
)
end
blocks は下記の通りで、AttentionやFFNの層を含む様々なニューラルネットワーク層が設定されます
Attentionは、attentions層として Layers.append されます
FFNは、opts の ffn で指定され、block_opts_keys → block_opts → hidden_state と経由され、hidden_state層として追加されます
defmodule Bumblebee.Layers.Transformer do
@moduledoc false
…
def blocks(hidden_state, opts) do
validate_required_keys!(opts, [:num_blocks, :num_attention_heads, :hidden_size, :ffn])
block_opts_keys = [
+ :num_attention_heads,
+ :num_key_value_heads,
+ :causal,
:hidden_size,
+ :ffn,
:kernel_initializer,
+ :attention_head_size,
+ :dropout_rate,
+ :attention_dropout_rate,
+ :query_use_bias,
+ :key_use_bias,
+ :value_use_bias,
+ :output_use_bias,
:layer_norm,
:block_type,
+ :attention_window_size,
+ :scale_attention_weights
]
opts =
Keyword.validate!(
opts,
block_opts_keys ++
[
:name,
:num_blocks,
:rotary_embedding,
attention_mask: Layers.none(),
attention_head_mask: Layers.none(),
attention_relative_bias: nil,
share_attention_relative_bias: false,
cross_hidden_state: nil,
cross_attention_mask: Layers.none(),
cross_attention_head_mask: Layers.none(),
cache: Layers.none()
]
)
name = opts[:name]
num_blocks = opts[:num_blocks]
attention_mask = opts[:attention_mask]
attention_head_mask = opts[:attention_head_mask]
cross_hidden_state = opts[:cross_hidden_state]
cross_attention_mask = opts[:cross_attention_mask]
cross_attention_head_mask = opts[:cross_attention_head_mask]
cache = opts[:cache]
rotary_embedding = opts[:rotary_embedding]
+ block_opts = Keyword.take(opts, block_opts_keys)
{attention_mask, cache} = Layers.Decoder.cached_attention_mask(attention_mask, cache)
offset = Layers.Decoder.get_cache_offset(cache)
state = %{
hidden_state: hidden_state,
hidden_states: Axon.container({hidden_state}),
attentions: Axon.container({}),
cross_attentions: Axon.container({}),
cache: cache,
attention_relative_bias: Layers.none()
}
outputs =
for idx <- 0..(num_blocks - 1), reduce: state do
state ->
block_attention_head_mask = Axon.nx(attention_head_mask, & &1[idx])
block_cross_attention_head_mask = Axon.nx(cross_attention_head_mask, & &1[idx])
block_cache = Layers.Decoder.get_block_cache(state.cache, idx)
attention_relative_bias =
if opts[:share_attention_relative_bias] and idx > 0 do
state.attention_relative_bias
else
opts[:attention_relative_bias] || Layers.none()
end
block_rotary_embedding =
case rotary_embedding do
nil -> nil
fun when is_function(fun, 1) -> fun.(idx)
config when is_list(config) -> config
end
{hidden_state, attention, cross_attention, block_cache, attention_relative_bias} =
block(
state.hidden_state,
[
attention_mask: attention_mask,
attention_head_mask: block_attention_head_mask,
attention_relative_bias: attention_relative_bias,
cross_hidden_state: cross_hidden_state,
cross_attention_mask: cross_attention_mask,
cross_attention_head_mask: block_cross_attention_head_mask,
block_cache: block_cache,
offset: offset,
rotary_embedding: block_rotary_embedding,
name: join(name, idx)
+ ] ++ block_opts
)
cache = Layers.Decoder.put_block_cache(state.cache, idx, block_cache)
%{
+ hidden_state: hidden_state,
hidden_states: Layers.append(state.hidden_states, hidden_state),
+ attentions: Layers.append(state.attentions, attention),
cross_attentions: Layers.append(state.cross_attentions, cross_attention),
attention_relative_bias: attention_relative_bias,
cache: cache
}
end
update_in(outputs.cache, &Layers.Decoder.update_cache_offset(&1, hidden_state))
end
終わりに
生成AIとAttentionの基礎を解説した後、Elixir生成AIエンジン「Bumblebee」のAttentionコードの中身をハックすることで、概念と実装の両面から、生成AIの中核であるAttentionを解説してみました
生成AIは使ったことあるけど、中身のメカニズムの理解や、どのように精度高い推論が出来ているのかをご存知無い方は多いと思いますので、この機会にBumblebeeコードと共に深淵を覗いてみてはいかがでしょう？
次回は、Attentionと双璧を成す「FFN」について、Bumblebeeコードハックと共に解説します
なお、今回の解説では一部を除き、ほぼ「Single-Head Attention」の範囲で完結する解説が多かったのですが、GPTシリーズやGemini、Claude等の現代型モデルや、Deepseek／GPT-4以降の最新モデルで導入された下記についても、機会あればコラム化したいと思います
| 用語 | 説明 |
|---|---|
| Single-Head Attention | 単一のヘッドを持つAttentionで、Transformer以前のSeq2SeqモデルやRNN／CNNベースのモデルで使われていた |
| Multi-Head Attention | Transformerで採用された、複数のヘッドを持つAttention（≒異なる視点での解析を複数回同時に行うことで多角的な情報を効率良く抽出）で、GPTシリーズやGemini、Claude等、現代的なモデルはこのAttentionを利用している |
| Grouped Query Attention | Deepseekで実用化され、GPT-4以降にも実装されたMulti-Head Attentionの改良版で、複数のQで同じKとVを共有し、計算速度と効率を向上 |
| MoE | DeepSeekが先駆者で、GPT-4以降にも実装された機構（「Mixture of Experts」の略）で、FFN内で各専門領域に特化した複数モデルを使い分けることで、計算コスト削減／モデル大規模化を実現 |
| router | FFN内でどのMoEを使うかを振り分けるニューラルネットワーク層（学習時は入力トークンからExpertへの振り分けを学ぶ） |
p.s.このコラムが、面白かったり、役に立ったら…
にて、どうぞ応援よろしくお願いします
明日は、 @t-yamanashi さんで 「ExAtomVMでElixirの関数の互換性を確認してみる」 です