Transformerアーキテクチャを連想記憶として理解する

MLPや自己注意(セルフアテンション)層を含む現代のTransformerアーキテクチャは、直感的に連想記憶ストアとして理解できます。これらの層をキーと値のペアを保存および取得するためのメカニズムと見なすことで、干渉対信号比(保存された記憶がターゲットとなる事実の取得をどれだけ妨げるかを定量化する指標)に基づいて、その容量制限を導き出すことができます。

線形メモリストア

最も基本的なメモリストアは線形層であり、外積を使用して記憶を蓄積します:$W = \sum_{i=1}^{M} v_i k_i^T$。このモデルでは、クエリ $x$ はクエリとすべての保存されたキーとの間の類似度を計算することによって値を取得し、これらの類似度によって重み付けされた値の合計を返します。

しかし、このアプローチは干渉の問題を抱えています。$M$ 個の記憶を保存する場合、ターゲットキー $k_j$ の想起は $f(k_j) = v_j + \sum_{i\neq j} v_i(k_i^T k_j)$ となります。$M$ が増加するにつれて、干渉項が蓄積されます。ランダムなキーの場合、容量(干渉が信号と一致する点)は、キーの次元を $d_k$ とすると、およそ $M \approx d_k$ となります。

最適な線形ストアとMLPによる容量の向上

メモリへの書き込みと読み出しの方法を洗練させることで、容量を向上させることができます。最適な線形ストアは、最小二乗法を用いてキーを「白色化(whiten)」し、干渉を効果的に打ち消すことで、容量を $M \approx 2d_k$ まで倍増させます。

MLPは、干渉をさらに抑制する非線形性を導入します。類似度を ReLU や二乗 ReLU などの関数 $\sigma$ に通すことで、モデルはノイズの原因となる低類似度のキーを抑制します。

  • ReLU MLP: 負の類似度をゼロにすることで干渉の分散が減少し、容量は $M \approx 2d_k$ まで押し上げられます。
  • 二乗 ReLU (GLU/SwiGLU): 低類似度のキーをさらに抑制することで、容量が大幅に向上し、$d_k^2$ に比例してスケーリングします。

これらの知見は、Transformerのフィードフォワード層が、認識可能な入力パターンに基づいてニューロンが活性化するキー・値ストアとして機能するという研究と一致しています。

非パラメトリックメモリとしての自己注意

自己注意は、ソフトマックス関数を使用して重み付き取得を行うことで、このアプローチの限界を体現しています。キーと値が固定パラメータであるMLPとは異なり、自己注意は入力トークンからこれらを計算するため、容量がコンテキスト長とともにスケーリングする非パラメトリックなストアとなります。

ソフトマックスに固有の指数関数を使用することで、自己注意は膨大な容量を実現できます。$\beta = \sqrt{d_k}$ でスケーリングすると、容量は多項式ではなく指数関数的に増加し、標準的な構成で100万を超える記憶に達します。長文コンテキストのシナリオでこのパフォーマンスを維持するために、モデルはしばしば $\log M$ ヒューリスティック(QwenやYaRNに見られるようなもの)を使用してアテンションロジットをスケーリングし、コンテキストの増大に伴って干渉対信号比を安定させています。

知見と議論

このフレームワークは、通信工学とディープラーニングの架け橋となります。コミュニティでの議論で指摘されているように、「干渉対信号」という視点は電気エンジニアには非常になじみ深いものであり、チャネル容量を最大化するための技術が、ニューラルネットワークが情報をどのように保存するかを理解するために直接適用できることを示唆しています。

数学的な抽象化は強力ですが、一部の実務家は、これらのシステムを理解するには具体的な例が必要であると強調しています。LLMの文脈では、「キー」はコンテキストウィンドウ内のトークンであり、「値」はモデルが後続のトークンを予測するために使用する対応する隠れ状態表現です。あるコメント投稿者が指摘したように、従来のホップフィールドネットワークと比較したTransformerの効率性は、反復的な収束を必要とせず、単一のフォワードパスで記憶を取得できる能力にあります。

最終的に、Transformerをメモリストアとして見ることは、これらのモデルにおける知能が、壊滅的な干渉なしに膨大な量の情報を圧縮および取得する能力と深く結びついていることを浮き彫りにしています。

Sources

関連

  • Dispatch
  • Dispatch
  • Dispatch
  • プロジェクト
  • Dispatch