注意メカニズムにバグがあり、ソフトマックスが犯人であり、すべてのトランスフォーマーに影響を与えている

「私は、8年間誰も発見できなかった注目度の式のバグを発見しました。GPTやLLaMAを含むすべてのTransformerモデルが影響を受けました。」

昨日、エヴァン・ミラーという統計エンジニアがAIコミュニティに波紋を呼ぶ発言をしました。

機械学習における注目の式は次のようになることがわかっています。

写真

この式は 2017 年に Transformer が登場して以来広く使用されてきましたが、現在、Evan Miller はこの式が間違っており、バグがあることを発見しました。

Evan Miller によるこのブログ投稿では、現在普及している AI モデルが重要な部分で失敗し、すべての Transformer モデルの圧縮と展開が困難になっている理由について説明しています。

要約すると、Evan Miller は、従来のソフトマックス関数に対する革新的な調整である、Softmax_1 とも呼ばれる新しい関数 Quiet Attention を導入しました。

あるネットユーザーは、このブログを「長すぎるので読まないでください」というセクションに要約した。ブログの著者は、アテンション メカニズムで使用されるソフトマックス式の分母に 1 を追加することを提案しています (最終出力ソフトマックスではありません)。アテンション ユニットのソフトマックスにより、キー/クエリの一致を確率として扱うことができます。これらの確率は、キー値ルックアップの連続値バージョンをサポートします (ルックアップの 1/0 出力として重みを取得するのではなく、重みが高いほど、目的のキー値ルックアップになります)。

分母に 1 を加えると、アテンション ユニットが変更され、実際の重み確率ベクトルが使用されなくなり、代わりに合計が 1 未満になる重みが使用されるようになります。理由は、ネットワークが高い重みを与えることを学習できるため、調整されたソフトマックスが確率ベクトルに非常に近くなるためです。すべて低い重みを提供する新しいオプションもあります (すべて低い出力重みを提供します)。これは、何に対しても高い信頼度を持たないことを選択できることを意味します。

「これが Microsoft RetNet がトランスフォーマーよりもパフォーマンスが優れている理由なのか?」と推測する人もいました。

一部のネットユーザーは、この研究はLLMの改善を促進し、それによって重みを大幅に圧縮し、より小さなモデルをより大きなモデルに匹敵させることができるとも述べている。

ミラー氏は次のように語っています。「以下に示すように、Softmax_1 関数は従来の softmax 関数と同じように使用できます。」

 import torch from softmax_one.softmax_one import softmax_one x = torch.randn(5)y = softmax_one(x, dim=0)

この変更に基づいて、ミラーは実験も行い、結果は次のとおりです。

次に、ミラーがどのようなエラーを見つけたかを見てみましょう。

外れ値

Evan Miller は量子化に関する論文を読んでいるときにこのバグを発見しました。現在、メモリとストレージは人工知能の発展を制限する重要な要因となっています。モデルを圧縮し、クラウドやエッジデバイスで大規模言語モデル (LLM) を実行しようとする取り組みが行われてきました。

コンピュータでは、情報はバイナリ データのストリームとして保存されます。データ ストリームが高度に予測可能である場合 (たとえば、常に制限された範囲内に含まれる場合)、比較的少ないビットを使用してそれらを保存できます。一方、数字の列が予測不可能で、めったに現れないほど大きな数字である場合は、それをエンコードして保存するために、より多くの 2 進数が必要になります。 Transformer モデルには、いくつかの外れ値の重みが含まれています。

Qualcomm AI Researchが6月に発表した論文「Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing」の中で、研究チームはこれらの外れ値の存在を注意メカニズムのソフトマックス関数にまで遡って追跡した。

Qualcomm 論文アドレス: https://arxiv.org/abs/2306.12929

これは驚くべきことのように聞こえますが、Evan Miller 氏はそれが正しいと考え、さらに softmax 関数にバグがあることを発見しました。

エヴァン・ミラーが、ソフトマックス関数がアテンションメカニズムに適したツールではないことをどのように示しているかを見てみましょう。

Softmaxの問題点

なぜソフトマックスはアテンション メカニズムに適していないのでしょうか? アテンション メカニズムで何ができるのかから始める必要があります。

一般的に数値エラーはプログラムエラーによって発生することが多いのですが、プログラムにエラーがない場合は複雑な数式の修復から始める必要があり、時間がかかります。

エヴァン・ミラーは、あるアイデアを思いつくまでに約 50 件の arXiV 論文を読みました。ミラーは「入力埋め込み」から始めます。これは、入力文字列内の単語を表す浮動小数点ベクトルとして理解できます。

たとえば、Meta が最近リリースした LLaMA 2 モデルでは、通常 30,000 ～ 50,000 のエントリを含む語彙内の 1 つの単語を表すためだけに、半精度浮動小数点数として表現される長さ 3204 の埋め込みベクトルを使用します。つまり、単語の埋め込みベクトルは 6KB 以上のストレージ スペースを占有します。技術の発展に伴い、「入力埋め込み」の長さは徐々に長くなり、それに応じて占有されるストレージスペースも増加します。

ストレージの使用に非常に敏感な C プログラマーの場合、この数字を受け入れることができないかもしれません。2 バイトで保存できるものをなぜ 6 KB を使用するのでしょうか? 2 バイトに基づいて計算すると、語彙サイズが 2^16=65384 未満であれば、エントリを表すのに必要なのは 16 ビットだけです。

ただし、実際には、Transformer は次のように動作します。入力ベクトルを同じサイズの出力ベクトルに変換し、最終的な 6 KB の出力ベクトルを使用して次のトークンを予測します。動作中、Transformer ワークフローの各レイヤーは元の単語ベクトルに情報を追加します。残余接続もこれに使用されます。すべての注意メカニズムは、元の 2 バイトの情報に補足資料を追加し、LLM がより長いコンテキストを分析できるようにします。

Transformer の最後のステップは、この出力ベクトルを長方形行列で乗算し、結果の語彙長ベクトルをソフトマックス関数に圧縮し、これらの指数的にスケーリングされた出力を次のトークンの確率として扱うことです。これは理にかなっています。しかし、私たち全員が知っているように、これらの出力確率が正しいかどうかは確信が持てないため、完全に正しいわけではありません。代わりに、すべての Transformer 実装とその派生は、サンプリング メカニズムを使用して、ソフトマックスの過剰表現の確率が低いという事実を隠します。

次にミラー氏はソフトマックスの開発の歴史を紹介しました。ソフトマックスは、エネルギーレベルに基づいて状態分布を予測する方法として統計学に初めて登場しました。その形式は次のとおりです。

経済学者は後にこれを次のように修正した。

この変更により、softmax には複数のロジスティック関数が備わります。ミラー氏はソフトマックス関数について広範囲にわたる研究を行ってきたため、ソフトマックスが不適切に使用されている箇所を特定することができました。

ソフトマックスは広く使用されています。物理学では非常にうまく機能します。経済学ではそれほど正確ではないかもしれませんが、機械学習の分野に適用すると、離散的な選択が含まれる限り常に機能するようです。

ミラー氏はさらに、ソフトマックスの鍵は、一部の項目を保持したくない場合はソフトマックスを変更する必要があり、そうしないと結果が歪んでしまうことだと述べました。

たとえば、LLM のコンテキストでは、非意味的トークン (コンマなど) に重みを付けることによって歪みが生じます。重みが大きくなると、圧縮が困難な外れ値となり、調査が難しくなります。 Qualcomm の AI 研究者もこの現象を発見しました。LLM では、異常な活性化の 97% 以上がスペースと句読点の位置で発生しました。

次に、ミラー氏は、問題がどこにあるのかを見つけるために、アテンションでソフトマックスがどのように使用されるかを紹介しました。

上記の式を分解すると、デコーダーのみのモデルでは、𝑄、𝐾、𝑉は同じ入力シーケンスから生成されます。これらはまったく同じではありません。つまり、投影方法が異なります。しかし、各レイヤーでは、すべて同じ注釈埋め込みベクトルから始まります。

𝑄𝐾^𝑇 項は、異なる位置にあるトークン ベクトル間の相関関係を見つけるために使用され、基本的に各列と行がトークンの位置に対応する相関行列 (ドット積がスケーリングされる) を構築します。次に、この正方行列の各行に対してソフトマックス演算を実行し、結果の確率を𝑉行列の中央値ベクトルの混合関数として使用します。確率混合 𝑉 は入力ベクトルに追加され、その合計がニューラル ネットワークに渡されてさらに処理されます。

マルチヘッドアテンションは、上記のプロセスをレイヤーごとに複数回並列に実行します。本質的に、このアプローチは埋め込みベクトルを分割し、各ヘッドはベクトル全体の情報を使用して、出力ベクトルの（重複しない）セグメントに注釈を付けます。これは、オリジナルの Transformer 論文の連結操作です。

ソフトマックスを使用する場合の問題は、出力ベクトルに追加する情報がない場合でも、各アテンション ヘッドに注釈を付けることを強制することです。

Softmax_1 と QuietAttention

さあ、Softmax Super-Mod が LLM チャネルに点火する様子をご覧ください。

ちょっと残念ですよね？ミラーがしたのは分母に1を足しただけだった。これにより、必要に応じてベクトル全体をゼロに近づけることができます。それ以外の場合、値はわずかに縮小され、縮小された値は、注目後に行われる正規化プロセスで補正されます。

主な違いは、𝑥 のエントリがゼロより大幅に小さい場合の負の値に対する制限であり、モデルは注釈を完全に回避しようとします。元のソフトマックスの次の制限動作

新しく改良された softmax_1 と比較してください。

通常の softmax は常に同じ合計重みを出力します。softmax_1 はほぼ同じに見えますが、負の象限に「脱出ハッチ」があります。明確に言えば、ここでの核心的な問題は数値的なものではなく、数学的なものである。追加の精度ではソフトマックスを保存できず、すべてのトランスフォーマーが影響を受けます。

softmax_1 については、他にもいくつか注目すべき点があります。導関数は正なので、常にゼロ以外の勾配があり、その合計は 0 から 1 の間なので、出力は制御不能になりません。この関数は以下の特性を維持する

つまり、出力ベクトル内の相対値は変更されません。

もともとミラーは、この関数を ghostmax と呼ぶつもりでした。これは、この関数が、結果を弱める V 行列内に追加のゼロ値のエントリとゼロ ベクトルを持つものと考えられるためです。

softmax_1 は表面的には退屈に思えるかもしれませんが、ミラー氏はそれが外れ値のフィードバック ループを解決し、量子化をカスケード研究のトピックにすることを 99.44% 確信しています。ミラー氏は、彼の意見が正しいことを証明するために実験を行いたい場合は、彼に連絡を取ることができると述べている。彼は論文を書くだろう。

改良されたメカニズムは QuietAttention と呼ばれ、これにより注意ヘッドが「静か」な状態を維持できるようになります。

ミラー氏は、すぐにテストを統合できるようになると考えています。各入力コンテキストの先頭にゼロのベクトルを追加し、選択したニューラル ネットワークがバイアス (位置エンコーディングを含む) を追加しないようにすると、ゼロは通過しても変更されず、後続の各ソフトマックス分母に 1 を追加する効果があります。こうすることで、グラデーション コードを扱う際に正気を失うことがなくなります。ミラー氏は、固定埋め込みと特別なプレフィックス トークンを備えた LLaMA モデルを使用してこれを実行できると示唆しました。

モデルを再トレーニングする必要があるため、まだ Raspberry Pi でこれを試さないでください。しかし、ミラー氏は、これらの重み尖度と活性化無限大ノルムを数回実行した後、それがどのように見えるかを知りたいと考えました。彼は、これがクアルコムの AI 研究チームの論文であろうと、ビブラテックスを計算する LLM チャンネルの誰かであろうと、影響力のある研究になると信じていますが、最初にそれを発見したのは彼です。

• プロジェクトアドレス: https://github.com/kyegomez/AttentionIsOFFByOne

• ブログリンク: https://www.evanmiller.org/attention-is-off-by-one.html?cnotallow=5d0e431f4edf1d8cccea47871e82fbc4