線形写像から自己注意へ：モデルの「メモリ・ストア」と想起能力

現代の大規模言語モデルは、世界のあらゆる事実を想起できると考えられています。友人とのボードゲームでルールを確認したいと Astra に頼んだ際、その記憶を読み上げてドライブ中に私の知識を復元してくれた事例はその好例です。しかし、我々はモデルに何でも聞いてしまうほど慣れ込み、彼らが膨大な情報を備えていることを当然のこととして受け止めています。では、彼らはどのようにしてこれほど多くの情報を覚えているのでしょうか？

本稿では、線形写像（Linear Maps）、MLP（多層パーセプトロン）、そして 自己注意機構（Self-Attention） をそれぞれ異なる種類の「メモリ・ストア（記憶庫）」として捉える直感的な視点を提示します。この視点から、各レイヤーが記憶できる量の目安となるヒューリスティクスを導き出します。

---

1. メモリ・ストアの定義と干渉の問題

メモリ・ストアとは

メモリ・ストアとは、$M$ つの記憶 $i = 1, \ldots, M$ を保持する関数 $f$ です。一つの「記憶」は、キー・ベクトル（検索に使用される）と バリュー・ベクトル（取り出したい値）のペアです：

• キー：$k_i \in \reals^{d_k}$
• バリュー：$v_i \in \reals^{d_v}$

関数 $f$ は入力 $x$ でクエリされ、出力 $f(x)$ を想起します。良いメモリ・ストアは、類似したキーでクエリされた際に、対応するバリューを正しく想起することを目指します（これを「結合的記憶」と呼びます）。

干渉対シグナル比

事実を無限に蓄積できても、必要な時に思い出しできないのであれば意味がありません。記憶数が増えるほど、他の記憶が相互に干渉し合い始めます。本稿は 「干渉対シグナル比（Interference-to-Signal Ratio）」 に焦点を当てます：

• シグナル: 目標となる記憶のバリュー ($v_j$)
• 干渉: クエリキーと類似した他の記憶からのノイズ

容量 (Capacity) とは、ストアが干渉がシグナルと同等の大きさ（比=1）まで達する前に保持できる記憶数のことです。

---

2. 外積型メモリ・ストア（線形写像）

最も基本的なストアは、単にキーとバリューの外積を通じて記憶を蓄積する単純な線形層です。

仕組み

重み行列 $W$ は、すべての記憶の外積の和として表現されます： $$ W = \sum_{i=1}^{M} v_i k_i\transpose = V\transpose K $$ ここで、新しい記憶 $(k, v)$ を保存するには、既存の重み $W$ に $vk\transpose$ を単純に足し合わせるだけです。

想起における干渉

入力を行列に掛けることで想起を行います： $$ f(x) = V\transpose K x = \sum_{i=1}^{M} v_i,(k_i\transpose x) $$ つまり、想起は保存されたキーに対する類似度によるルックアップです。

もしすべてのキーが直交していれば干渉はゼロになりますが、$\reals^{d_k}$ 空間には $d_k$ つしか直交方向がありません。記憶数 $M$ がキー次元 $d_k$ を超过した（$M > d_k$）時点では、完璧な想起は不可能になります。

容量の限界

干渉項はランダムウォークのように加算され、干渉対シグナル比は以下のようになります： $$ \frac{\norm{\text{interference}}}{\norm{\text{signal}}} \sim \sqrt{\frac{M}{d_k}} $$ この比を 1 と置くことで、容量は $M \approx d_k$ と見積もられます。つまり、キー次元数以上の記憶を正しく想起することは外積型ストアでは困難です。

改善策：最適な線形メモリ・ストア

外積型ストアは干渉を削減するためにはキー次元 $d_k$ を大きくする必要がありますが、これはハードウェア制約によって制限されます。より賢い方法は、最小二乗回帰問題として重み $W$ を最適化することです： $$ \operatorname*{minimize}W \sum{i=1}^M \norm{W k_i - v_i}_2^2 $$

これにより導かれる「最適な線形メモリ・ストア」では、キーの共分散を除去（白化）し、共有方向に対する過大評価を防ぎます。結果として：

• $M \leq d_k$ の場合：干渉ゼロで想起可能。
• $M > d_k$ の場合：容量が $2d_k$ まで拡張されます。

つまり、最適化を行うことで、外積型ストアの容量を約 2 倍 にできることが示唆されます（64 次元キーの場合、容量は 64 から 128 に）。

---

3. MLP による非線形メモリ・ストア

外積型ストアや最適化された線形ストアは類似度に対して線形的なため、小さい類似度と大きい類似度を同様に拡大縮小します。より良い性能を得るには、非線形性（Non-linearity） を用いて類似度が低いキーからの干渉を選択的に抑制する必要があります。

ReLU と二乗 ReLU

各類似度を非線形関数 $\sigma$ で処理することで干渉を削減できます： $$ f(x) = \sum_{i=1}^{M} v_i,\sigma(k_i\transpose x) $$

• ReLU: 負の類似度を持つキー（$k_i\transpose x < 0$）をゼロにします。容量が $2d_k$ に拡張されます。
• 二乗 ReLU ($\operatorname{ReLU}(s)^2$): 類似度が低いキーをさらに強く抑制します（例：類似度 0.1 は 0.01、類似度 1 は 1）。これにより容量は劇的に増大し、約 2,700 に達します。

SwiGLU と超置合

現代の言語モデルで見られる SwiGLU（Swish ゲート付き MLP）や GLU は、この非線形性を活用しています。これらの MLP は「超置合（Superposition）」と呼ばれる現象を通じて事実を記憶しており、一つの重みが複数の事実を表します。

---

4. より高いべき乗による干渉低減

類似度を二乗することは効果的ですが、さらに高次の多項式を使えば容量は幾何級数的に増加します。$\sigma(s) = \operatorname{ReLU}(s)^p$ を用いると： $$ \frac{\norm{\text{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k^p}} $$ つまり、容量は $M \approx d_k^p$ となり、べき乗 $p$ を上げることで容量は爆発的に増加します。

• 二乗 ReLU ($p=2$): 容量は約 2,700
• 三乗 ReLU ($p=3$): 容量は約 35,000

これにより、MLP は潜在的に数百万もの記憶を保持する能力を持っています。

---

5. 非パラメトリックなメモリ・ストアとしての自己注意

重みのべき乗を増やすだけでは限界があります。ここで登場するのが 指数関数 です。指数関数はすべてのべき乗の加重和であり、これにより無限の表現力を持ちます： $$ e^{\beta s} = \sum_{p=0}^{\infty} \frac{(\beta s)^p}{p!} $$

自己注意の容量計算

指数関数（シャープネス $\beta$）を使用することで、干渉対シグナル比は以下の式で表されます： $$ \frac{\norm{\text{interference}}}{\norm{\text{signal}}} = \sqrt{M}, e^{\beta^2/d_k - \beta} $$ この比を 1 と置くと、容量は $M \approx e^{2\beta(1 - \beta/d_k)}$ となります。

ここで重要な発見があります：

• $\beta = \sqrt{d_k}$（通常 $d_k=64, \beta=8$）の場合、容量は指数関数的に増大します。
• 例えば $\beta = 8$ では、容量は約 120 万 ($e^{14} \approx 1209,790$) に達します。

これは線形ストアの約 20,000 倍の能力です。実際のトランスフォーマーで使用されるのは softmax です： $$ f(x) = V\transpose \operatorname{softmax}(\beta K x) $$ MLP の場合、キー・バリューは学習された重みで固定されていますが、自己注意の場合はコンテキスト内のトークンからその値を動的に計算するため、非パラメトリックなメモリ・ストアとなります。

---

6. 長いコンテキストにおける多数のペアの処理

モデルの記憶能力（コンテキスト長 $M$）が増加しても、シャープネス $\beta$ を固定すると干渉が蓄積してしまいます。$\beta \to \infty$ にするともう訓練できません（飽和現象）。これを克服するためには、コンテキスト長の増加に合わせて $\beta$ をスケールアップする必要があります。

スケーリング手法の比較

実装では以下のアプローチが存在します：

手法	ロジック概要	性能・特徴
固定 β	$\beta$ を一定に維持	コンテキストが長くなるにつれて干渉が蓄積し、想起失敗する。
Su (2021) / Scalable-Softmax	注意のエントロピーを一定にするため $\log M$ に比例して $\beta$ を上げる	数百万トークンのコンテキストでも干渉が抑制され、正しく想起可能になる。
LogN-Scaling (Qwen)	コンテキストが基準長 $L$ より長い場合にのみ $\log M$ スケールを適用	非常に長いコンテキストで安定した性能を発揮する。
YaRN (Qwen2.5-1M)	ロジットに $(0.1 \log(M/L) + 1)^2$ を掛ける	より急激にスケールするため、数百万規模でも比が $10^{-4}$ のレベルまで低下し、極めて安定する。

実験結果では、Su (2021) や Qwen LogN-Scaling、YaRN のいずれも、数百万トークンのコンテキスト（$M \approx 2,097,152$）においても干渉対シグナル比はほぼゼロに保たれ、想起が可能であることを示しています。

---

7. まとめ：メモリ・ストアの進化と展望

本研究では、線形な外積型ストアから始まり、MLP の非線形性を経由し、最終的に自己注意機構へと至る「メモリの進化」を理論的に示しました。

容量比較表 ($d_k = 64$ の場合)

以下は、キー次元 $d_k=64$ の場合の各アーキテクチャの推定容量です。

メモリ・ストア	関数形式 $f(x)$	推定容量
外積型	$V\transpose K x$	64
最適線形	$V\transpose K (K\transpose K)\inverse x$	128
ReLU MLP	$V\transpose \operatorname{ReLU}(Kx)$	128
二乗 ReLU (GLU)	$V\transpose \operatorname{ReLU}(Kx)^2$	約 2,700
SwiGLU MLP	$V\transpose (\operatorname{Swish}(Kx) \odot Kx)$	約 2,900
三乗 ReLU	$V\transpose \operatorname{ReLU}(Kx)^3$	約 35,000
自己注意 ($\beta = \sqrt{d_k}$)	$V\transpose \operatorname{softmax}(\beta K x)$	約 120 万
自己注意 ($\beta \propto \log M$)	スケーリングあり	200 万以上 (到達範囲外)

考察

• この容量見積もりは、単一のレイヤーが干渉を許容する限界を超えて何個のランダムベクトルを記憶できるかをカウントしています。
• トランスフォーマーモデルにおいて、自己注意機構（Attention） が、非常に高い想起能力と潜在的な無限に近い記憶容量を持つ「メモリ・ストア」として機能していることが理論的に裏付けられました。

今後の方向性

この研究は、以下の領域にもつながります：

• ホップフィールドネットワーク: 指数関数的容量を持つ密結合記憶。
• 推論中の更新: ファストウェイトや Mamba など、動的にメモリを更新する手法。
• パラメトリック/非パラメトリックのハイブリッド: プロダクトキー記憶や RETRO などの外部記憶利用。
• 記憶の編集と発見: ROME や MEMIT など、特定の情報を書き換える技術。

我々は機械がどのように記憶するかを理解することで、自分の記憶の限界についても学べるかもしれません。