arXivに投稿された新しい論文「On the Invariants of Softmax Attention」では、Softmax注意機構の内在構造が詳細に調査されています。著者のWonsuk Lee博士は「エネルギーフィールド」(energy field)を定義し、これは行中心化されたアテンションロジットであり、多様なモデル、アーキテクチャ、入力にわたって不変性を示すことを明らかにしました。
研究では2つのクラスの不変性が見出されました。
メカニズムレベルの不変性:これらはSoftmax注意の代数構造に由来し、行ごとのロジットの合計がゼロになる制約、ヘッド次元によって決まるランク境界、およびそれらに続くスペクトル特性を含みます。これらの性質はモデルに関係なく常に成立し、数学的な証明と実験によって確認されています。
モデルレベルの規則性:メカニズムによって強制されるものではありませんが、テストされたすべての自己回帰言語モデル(GPT、LLaMAなど複数のアーキテクチャファミリーにわたる)で一貫して見られます。エネルギーフィールドの分散はキー位置全体に分布し、少数の位置に集中することはありません。この非局在化は、キー行列の特性である「キーインコヒーレンス」(key incoherence)に起因します。この特性は、キー行列の行ベクトル間の相関が低いことを意味し、アテンション分布をより均一にします。
これらの不変性には実用的な意味があります。ランク境界はエネルギーフィールドを低次元部分空間に制限し、注意表現の有効次元を理解するのに役立ち、モデル圧縮に応用できる可能性があります。キーインコヒーレンスはヘッドごとの訓練モニターを提供し、訓練中の異常(例えば、特定のアテンションヘッドが退化するなど)を検出するために使用できます。すべての結果は、複数のコンテキスト長および入力テキスト(短いテキストから長いコンテキストまで)で検証されています。
この論文は注意機構の理論的理解を深めるだけでなく、深層学習の実践者に有用な分析ツールを提供します。将来的には、これらの不変性を活用した注意機構の設計改善や訓練手法の開発が期待されます。