世界を動かす技術を、日本語で。

最小の例を用いたトランスフォーマーの理解

概要

  • Transformer の内部動作を視覚化し、情報の流れや Attentionメカニズム の理解を支援
  • 最小限のデータセット と単純化したモデル設計で、プロセスを段階的に追跡
  • シンプルなトークナイゼーション と小規模な語彙リストの採用
  • 2層・2ヘッド構成のデコーダ専用Transformer による学習・可視化
  • Attentionの可視化 により、モデルがどのように文脈を捉えているかを明示

Transformerの内部可視化:最小構成で理解する仕組み

  • Transformer型Large Language Model(LLM) の内部情報の流れや Attentionメカニズム の動作は、数値が膨大なため人間には直感的に把握しづらい課題
  • 本記事では、 最小限のデータセット意図的に単純化したモデル を用いて、内部状態を視覚化しながら解説
  • 各層で情報がどのように変換されるか、および Attentionがどのトークンに注目するか を段階的に追跡可能
  • データセットとソースコード はMITライセンスで公開(https://github.com/rti/gptvis)

可視化手法

  • 食品アイテムのトークン埋め込みベクトルを 色付きボックスの積み重ね で表現
  • 各トークンの 20次元ベクトル を4次元ごとに1つのボックスとして5段重ねで可視化
  • 味覚トークン("juicy", "sour", "sweet", "spicy")の埋め込みには 共通性と個別性 が現れる

セットアップとデータセット

  • データ、トークナイゼーション、モデル構造 の3点を徹底的に単純化
  • 最小限かつ構造化されたデータセット :果物と味覚の関係に焦点を当てたシンプルな英文のみ
    • 繰り返しパターン と明確な意味的リンクが特徴
    • 検証用文 を1文だけ分離し、モデルが単なる暗記でなく 意味的連想を学習したか を評価

データセット例

  • 訓練データ:
    • "lemon tastes sour", "apple tastes sweet", "orange tastes juicy", "chili tastes spicy" など
    • "i like the spicy taste of chili", "i like the sweet taste of apple" など
    • "i like sour so i like lemon", "i like sweet so i like apple" など
  • 検証データ:
    • "i like spicy so i like chili"

トークナイゼーション

  • 単純な正規表現 で単語単位に分割
  • 語彙数19個 のリスト(例:"is", "the", "orange", "chili", "sour", "of", "taste", "apple", "sweet", "juicy", "a", "spicy", "so", "like", "tastes", "i", "lemon", "UNKNOWN", "PADDING")

簡易モデルアーキテクチャ

  • デコーダ専用Transformer、層数2、Attentionヘッド2
  • 20次元埋め込みパラメータ数約1万
  • 埋め込み行列の共有(Tied Embeddings) を採用し、入力・出力で同一ベクトル空間を利用
  • Google Gemma 等で使われる手法で、可視化・パラメータ削減に有効

学習・検証結果

  • 1万ステップ学習 で訓練・検証ともに低損失を達成
  • "i like spicy so i like" に対し "chili" を正しく予測
  • 未見データでの一般化能力 を確認

埋め込みと可視化

  • 各トークンは 20次元ベクトル で表現
  • 味覚トークン は個別性と共通性を持つ埋め込みを獲得
  • 入力トークン埋め込み は同じ単語なら同じベクトル
  • 最終層の出力ベクトル は、次トークンとして予測される"chili"の埋め込みと著しく類似
  • Tied Embeddings により、入力と出力で同一の埋め込み行列を利用

Transformer層内のAttentionメカニズム

  • 各層でトークンベクトルは 自己情報+他トークンの文脈情報 で更新
  • Attention機構 により、各トークンは前方トークンを重み付けして参照
    • Layer 1 では"spicy"が前の"i"へ強いAttentionを向ける
    • "so"は直前の"spicy"と最初の"i"両方に注目(Attentionヘッドごとに異なる色で表示)
    • 最後の"like"は両層で"spicy"に強く注目し、正しい予測に寄与

まとめ

  • 最小限構成のTransformer でも、 Attentionメカニズム により文脈理解と一般化が可能
  • 可視化 により、各トークンの埋め込みやAttentionの働きが直感的に把握可能
  • シンプルな構成でも本質的な動作は大型モデルと共通 であり、内部理解の足掛かりとして有用

Hackerたちの意見

すごくクールだね!

このトピックに関して数日前に話題になった別の記事もおすすめだよ。最後まで読んだら、全部理解できたしね。> 「AIはどうやって猫を識別するの?」

これをシェアしてくれて嬉しい!すごくわかりやすいし、内容が薄くなってないね。ありがとう!

正直な感想だけど、最初の部分を読んだときはすごくワクワクした。でも、結局のところ、以前よりも理解が深まったわけじゃなかった。参考までに、最初の理解度はかなり低かったんだ。要するに、a) 埋め込みが何かは基本的に知ってる b) トランスフォーマーは行列の掛け算で動く c) それは、事前に訓練された埋め込みの利点を持つマルチスレッドのマルコフ連鎖生成器みたいなもの。

3Blue1Brownのトランスフォーマーについての動画見た?

別のところに投稿したコメントを貼るね:私が好きなリソース:セバスチャン・ラシュカの本「ゼロから作るディープラーニング」 これらの動画 / プレイリスト: https://youtube.com/playlist?list=PLoROMvodv4rOY23Y0BoGoBGgQ... https://youtube.com/playlist?list=PLoROMvodv4rOwvldxftJTmoR3... https://youtube.com/playlist?list=PL7m7hLIqA0hoIUPhC26ASCVs_... https://www.youtube.com/live/uIsej_SIIQU?si=RHBetDNa7JXKjziD これが私がtinystoriesで訓練した基本的な実装だよ: https://gist.github.com/nikki93/f7eae83095f30374d7a3006fd5af... (これは私にとって新しい分野だったから、上記を手伝うためにClaudeのコードをたくさん使ったよ) (前はCとmlxでやったけど、結局Pythonに屈した笑)でも、全体的にはこうなる: - テキストをトークン化する - 簡単なNNでトークンを埋め込む(それぞれをベクトルにマッピング) - 各トークンがどこにいるかをエンコードするために位置情報を適用する - アテンションを行う。これが重要で、私にとってもすごく興味深い部分。3つのニューラルネットワーク:Q、K、Vが各トークンに適用される。新しい埋め込みのシーケンスを生成して、各位置にはその位置のQと他の位置のKのドット積で重み付けされた全トークンのVが加算される。新しい埋め込みは前の層に「加算」される(こういう加算を「残差」と呼ぶ) - アテンションなしで別のNNパスも行い、出力を再度加算する(残差)。実際には、異なるQ、K、Vを持つ「ヘッド」が複数あって、出力はその2回目のNNパスの前にまとめられる。各段階で数値を適切に保つための正規化も行われる。アテンションとフォワードブロックを何度も繰り返して、最終層の出力の最後の埋め込みがサンプリングの基準になる。これがどれだけ早く一貫した文法を生成し始めるかに驚いたよ。訓練ループが各訓練段階で例の出力を示す生成ステップも行うのが、出力が時間とともに質的に改善されるのを見るのに役立ったし、学ぶ様子を「見る」のがちょっとクールだった。これにはMoEやスパース対デンスアテンション、さらにこれらの上にRLがあること(人間のフィードバック用か「バックトラッキングとスパース報酬での検索」のためか)については触れてない。まだコーディングしてないけど、ちょっと読んだだけ… さて、これは何かのプロセスを学ぶためのセットアップで、重みの中にそれを実行するプロセスが広がっているけど、そのプロセスが何なのかはまだよくわからない。「機械的解釈可能性」はそれに取り組むための分野で、最近調べてるんだ。

同じように、HNで見出しを読んだときはすごくワクワクしたし、これは電気機器についての話だと思った。LLMの意味が、他の修飾語がない場合のデフォルトとして、電気の意味を超えてしまったのかな、こういうコミュニティでは。

私にとっては、もう少し説明があった方がいい気がする。簡潔だけど、文法やリズムがちょっとぎこちない。