- GPT-2に似たデコーダー専用Transformerを、学習なしで重みだけを直接設計し、
(aab)* パターンを予測させた実験
"aabaabaabaab..." の予測には直前の2トークンを見る必要があるため、単純な交互パターンよりもアテンションの動作を示すのに適している- モデルは
N_CTX=5、N_VOCAB=2、N_EMBED=8 の小さな構成で、a=0、b=1 のトークン化とワンホット埋め込みを使う
- 単一のアテンションヘッドは直近2トークンに0.5ずつ注意を配分し、
a=1、b=-1 エンコーディングの打ち消しを利用して次のトークンを計算する
- 曖昧でない文脈では100.0%（27/27）の精度を示したが、5トークン文脈で1回予測するのに約4,000 FLOPsが必要で、直接実装したルールよりはるかに非効率
学習なしで重みを直接指定したミニGPT-2
- 目的は、Transformerとアテンションの各構成要素が実際に何をしているのかを直感的に理解すること
- モデルは学習せず、事前学習済み重みも使わず、夕方のひとときで各重みを直接指定する方法で作られた
- 構造はGPT-2に似たデコーダー専用Transformerで、実装はjaymodyのpicoGPT実装を基に簡略化されている
- layer normを削除
- multi-head attentionの代わりに単一ヘッドを使用
- transformer blockの
mlp feed-forward layerを削除
(aab)* シーケンスを選んだ理由
- 最初は
"ababababab" のようなシーケンスを予測しようとしたが、Transformerはshifted sequenceを予測するため、簡単すぎるタスクになる
a なら b、そうでなければ a を予測すればよく、位置埋め込みを使う必要がない
- 最終的なタスクは
"aabaabaabaab..."、つまり (aab)* シーケンスの予測に決まった
- 直前の2トークンが
ab または ba なら、次のトークンは a
- 直前の2トークンが
aa なら、次のトークンは b
bb はタスク範囲外のケースとして扱う
- トークン化は2つの記号だけを扱う単純な方式
モデルの次元と計算フロー
- 選んだモデルパラメータは3つ
N_CTX = 5: モデルが一度に見る最大コンテキスト長N_VOCAB = 2: a、b の2トークンN_EMBED = 8: トークン・位置・計算用の空間を収める埋め込みサイズ
- 実際のタスクには直前の2トークンだけが必要だが、
N_CTX=5 として、関係ないトークンを無視しなければならない状況を含めている
gpt 関数は次の順序で動作する
wte[inputs] + wpe[range(len(inputs))] でトークン埋め込みと位置埋め込みを足す- 1つのtransformer blockを通過する
- 最後に
x @ wte.T で語彙空間のlogitsを作る
位置とトークンをワンホットで格納する埋め込み
wpe は位置埋め込みで、先頭5つの埋め込み次元を位置ワンホットとして使う
- position 0は
[1, 0, 0, 0, 0]
- position 4は
[0, 0, 0, 0, 1]
wte はトークン埋め込みで、次の2つの次元をトークンワンホットとして使う
- token
a は該当するトークン次元で [1, 0]
- token
b は [0, 1]
- 8番目の埋め込み位置は最初は使わず、transformer block内でscratch spaceとして使う
- 例えば
"aabaa" は 5 x 8 の埋め込み行列として表現され、各行は位置ワンホットとトークンワンホットをあわせて含む
アテンションヘッドが直近2トークンを選ぶ方法
- transformer blockは、1つのattention headと、attention結果を再び埋め込み空間に戻す
c_proj 線形層で構成される
c_attn は embed_size x (embed_size * 3) サイズの線形層で、入力埋め込みを qkv 行列に変換した後、q、k、v に分割するk は位置埋め込みを分離し、各トークンが持つ位置情報を表すq は各位置が探したい位置範囲を表し、q @ k.T によってattention score行列を作る- softmaxとcausal maskの後、attention行列は次の性質を持つ
- 最初の行は最初のトークンだけに100% attention
- 以降の行は、アクセス可能な最新2トークンにそれぞれ
0.5 ずつattention
- causal maskは未来のトークン位置に非常に小さい値、実際のコードでは
-1e10 のような値を足し、未来トークンを見られないようにする
- この手作りモデルは未来を覗き見るようには設計されていないが、GPT-2の構造に近く保つためにmaskを残している
np.sqrt(q.shape[-1]) で割るscalingは、実際の学習では勾配改善に役立つが、この手製モデルには影響しない
v エンコーディングと足し合わせの打ち消しで作る予測
v はトークンワンホットを a=1、b=-1 エンコーディングに変換する- attention結果が直近2トークンを0.5ずつ平均するため、このエンコーディングは次のルールを計算する
a, b → 0.5 * 1 + 0.5 * (-1) = 0b, a → 0.5 * (-1) + 0.5 * 1 = 0a, a → 0.5 * 1 + 0.5 * 1 = 1
- 結果として、行の7番目の位置には次の値が生じる
0 なら a を予測すべき場合1 なら b を予測すべき場合
- 入力
"aabaa" では最初の行は情報が足りず b 予測になる可能性があるが、それ以降の予測は (aab)* ルールと一致する
予測値を語彙空間へ戻す
c_proj はattention結果の7番目の位置の値を、トークンワンホット形式に戻す- 単に
[..., 1, 0, ...] または [..., 0, 1, ...] を作るのではなく、1024 でスケールしたワンホットを作る
embedding[row, 5] = 1024 + (-1024) * predictionembedding[row, 6] = 0 + 1024 * prediction
- transformer blockにはresidual connectionがあり、
x = x + causal_self_attention(...) によって元の埋め込みが足される
- residual signalが不要に残るため、
1024 スケールを使ってこの信号を圧倒する
- 最後に
x @ wte.T を計算してlogitsを作り、softmaxを適用する
"aabaa" 文脈では、最終予測行が b を指す- 学習時にはすべての行の予測が有用だが、推論では最後の行だけが必要
生成結果と精度
complete 関数は最後の最大5トークンをモデルに入れ、softmax結果の最後の行から argmax で次のトークンを選ぶ- 生成例は次のとおり
complete("a") → a :: baabaabaabcomplete("ba") → ba :: abaabaabaacomplete("abaab") → abaab :: aabaabaaba
- 範囲外の入力でも反復パターンに復帰する場合がある
complete("ababa") → ababa :: abaabaabaacomplete("bbbbb") → bbbbb :: aabaabaaba
"aab" * 10 テストで曖昧でない文脈だけを評価すると、精度は100.0%（27/27）
4,000 FLOPsと8命令の違い
- 5トークン文脈全体を使う場合、このモデルは単一トークン予測に約4,000 floating point operationsを必要とする
- 大半はattention計算で使われる
- context windowの縮小、fused multiply-add、kv cachingなどで削減できる
- それでも単一トークン予測には数百個のマシン命令が必要
- 同じ
(aab)* ルールを直接書いたx64アセンブリは、8命令で次のトークンを計算する
- 自然言語生成において、現在のモデルより1000倍効率のよい言語モデルを学習できるのかという問いが残る
1件のコメント
Hacker News のコメント
関連する研究として「Thinking Like Transformers」があります。
RASP というプリミティブなプログラミング言語を紹介しており、Transformer の構成要素でモデル化できる演算から成り、ヒストグラムやソートのようなプログラムを書けることを示しています。
Sasha Rush と Gail Weiss による優れたブログ記事もあり、後続研究では RASP 系のプログラムを、学習なしで実際のモデル重みにコンパイルできることが示されました。
[1] https://arxiv.org/abs/2106.06981
[2] https://srush.github.io/raspy/
[3] https://arxiv.org/abs/2301.05062
この分野が面白いなら、Transformer モデルの重みを自分で選び、人間が小学校で習う方法に近い形で筆算の足し算を行わせる私の HandCrafted Transformers の取り組みも、一度見る価値があると思います。
[1] https://colab.research.google.com/github/newhouseb/handcraft...
ニューラルネットワークからコードへ向かう方向も、説明可能性の観点から非常に興味深いと思います。
Transformer をよく理解していると思っていましたが、自分で実装したことはありませんでした。
ある日自分で実装してみると、標準の PyTorch Transformer ほどにはうまく動作せず、学習も進まず、結局 dropout を無視していたことが原因だと気づきました。
数字の足し算を学習させていて、同じペアを二度見せたことがなかったので過学習は不可能だと思っていたのですが、dropout の役割は想像以上に大きかったのです。
要するに、Transformer はとにかく自分で実装してみるのがよく、低レイヤーからやるほどなお良いということです。
やった人はみな予想外のことを学んでいて、トークン単位の学習の並列化から、バックプロパゲーションが実際にどう動くのかまで、人によって気づく点はさまざまでした。
Karpathy 系の資料も良いですが、私にとって Transformer がようやく理解できたのはこの動画でした: https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k
似たような考えをしばらく持っていました。
ドメイン専門家が手で調整して学習を加速できる、モデル重みに対する直感的なインターフェースを作れないでしょうか。
例えば視覚モデルでトラフィックコーンを検出する際に、「オレンジ色らしさ」に相当する重みのまとまりを上げるような形です。
そうすれば、「オレンジ色らしさ」を正しく補正するために何千・何百万もの追加例を要求する代わりに、人間が加速できます。
もちろん難しいのは、このインターフェースが異なる意味を持つ重みのまとまりにマッピングされなければならない点ですが、技術的に不可能な理由があるのか気になります。
[1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
短く言えば、AI に人間が手助けする方式は、より多くの計算資源で回す方式より、ほとんど常に費用対効果が低いということです。
人間がオレンジ色のトラフィックコーンを検出するよう重みの層を補正している間に、GPU クラスタはトラフィックコーン、信号機、木、他の車、少し違うオレンジ色のトラフィックコーンまで検出するよう AI を学習させてしまいます。
オレンジ色をよりよく見るように画像を調整したとしても、同時に他のすべての色の精度を監視できなければ、知らないうちに別の色で問題を作ってしまう可能性があります。
特定のニューロンのクラスタが特定の概念に対応するという形ではなく、すべてがおおむねすべての仕事を少しずつ行っています。
関与するパラメータが多すぎます。
Transformer 論文は技術的すぎて、いつも浅くでも理解したいと思いながら難しく感じていました。
この記事は動作の仕組みを理解するのに本当に役立ち、少なくとも例はとても明確でした。
おかげで大学時代に学んだ行列も思い出せました。
Turing マシンや正規表現をパースするマシンのような、一種の抽象機械ではないのか？
まずトークンの一覧を定義し、たとえば分かりやすく 24 文字だとしましょう
この機械はトークンの入力シーケンスを受け取り、決定論的な行列演算を行ったあと、すべてのトークンの確率一覧を出力します
「学習」とは、その演算に使われる行列内の数値の一部を設定する過程にすぎません
最終的なコードには
if文が 1 つしかなく、それも結果の正確さを評価するためのものだという点は注目に値しますすべての「論理」は行列演算の結果から生じます
ニューラルネットワークでは、すべてが概して少しぼんやりしていて、
if/elseのようなものはほとんどありませんが、Transformer の例のように値を 0 や -∞ で「マスク」する場合はあります出力もほぼ常にスコアや確率のまとまりなので、猫と犬の写真を区別するモデルが
dog:0.95 cat:0.05のような結果を出せば、犬のスコアのほうが高いので犬を予測したと言いますTransformer の中核であるアテンション機構は、一種のソフトな参照演算に基づいています
ぼんやりしていないシステムなら、シーケンス内の各トークンを巡回し、現在のトークンと関係があるか確認して、関係があれば何らかの動作をするでしょうが、Transformer では関連性は二値の判断ではありません
代わりに、シーケンス内のすべてのトークンのペアの間で連続的な関連性スコアを計算し、そのスコアを使って次の動作を行います
ただし、二値判断ベースのシステムから直接一般化しにくいものもあります
たとえば、そうした関連性スコアは語彙トークンに対する加重平均を計算する重みとして使われ、現在位置に対する「平均トークン」を得ます
これを分岐ロジックに基づくプロセスの拡張として簡単に解釈する方法はなさそうです
Linear Transformers が実は Fast Weight Programmers であることを説明しているこの論文を見るとよいです: https://arxiv.org/abs/2102.11174
重みを慎重に設定すれば、どんな計算でも実行させられます
ただし、近似に基づかないコンパイラがあるとよいのですが
「自分でモデルを作ってみたくなるかもしれない」というのは、好奇心を満たす学習練習以外に何に使おうとしているのか気になります
複雑な機械学習モデルは、家でブログを読む人が扱うには現実的ではないように感じ始めています
それらしく見える古い英単語をたくさん作り出し、英語文法の基礎や戯曲の形式などを学びます
あれほど短い時間でそこまで到達するのはかなり驚きでした
ローカルで複数のモデルを Shakespeare-from-Wish.com レベルの忠実度で学習してみれば、よいアーキテクチャを見つけたか、そしてスケールさせてみる時期が来たかを判断する助けになりそうです
Transformer とアテンションをもっとよく理解したくて、The Illustrated Transformer を読んだものの、アテンションのいろいろな部品が実際に何をしているのか直感的に腑に落ちなかった、という内容です
qとkの違い、ましてvについてはなおさらだった、という感じです単に好奇心を満たすだけでなく、理解を作り、深めるのに役立ちます
タイトルに neural network のような表現を入れられるとよさそうです
これは 2 つの回路を電磁的に結合するコイルの束ではなく、機械学習の「Transformer」アーキテクチャに関する内容です