技術的概要：ラプラシアンヘッドはトークン表現を平滑化することでトランスフォーマーを改善する

問題提起

標準的なトランスフォーマーアーキテクチャは、マルチヘッドアテンションと残差接続を通じてトークン表現を更新する。この定式化において、各アテンションヘッド $i$ はソフトマックスアテンション行列 $P^{(i)}$ を計算し、重み付き和を介して入力 $X$ を更新する：$X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$。

効果的である一方で、このメカニズムはトークン表現のシーケンシャル内分散に対する明示的な制御を欠いている。グラフニューラルネットワーク（GNN）や深層トランスフォーマーに関する文献における支配的な見解は、「過剰平滑化（トークン表現が類似しすぎ、区別性を失うこと）」を有害な失敗モードとして扱っている。したがって、これまでの研究の大半は、正規化、エッジのドロップ、多様性損失を通じて平滑化を防止することに焦点を当ててきた。本論文は、トークンの過剰平滑化が本質的に有害であるという仮定に挑戦し、制御された平滑化は表現学習にとって有益であり得ると提案する。

手法

著者は、パラメータフリーのアーキテクチャ変更を提案する：標準的なアテンションヘッドの一部をラプラシアンヘッドに置き換えることである。

ラプラシアン演算子

標準的な更新において、ヘッドは $P^{(i)}X$ を計算する。これはトークンのアテンション重み付き平均である。提案されたラプラシアンヘッドは、アテンション行列 $P^{(i)}$ を、単位行列 $I$ を用いたグラフラプラシアン演算子 $I - P^{(i)}$ に置き換える。更新式は以下のようになる：
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
ここで、$A$ は標準アテンションヘッドの集合、$L$ はラプラシアンヘッドの集合である。

動機

本論文はこの変更に対する 2 つの理論的動機を提供する：

1. 平均と分散の直接制御：標準的なアテンションヘッドはトークン表現の平均を更新する。対照的に、項 $(I - P^{(i)})X$ は各トークンとそのアテンション重み付き平均との差を計算する。この差を加えることで、モデルは標準的なトランスフォーマーには存在しない能力である、トークン表現のシーケンシャル内分散を直接制御できるようになる。
2. グラフ拡散：トークンをエッジ重みが $P^{(i)}$ で定義されたグラフのノードと見なす場合、$I - P^{(i)}$ はランダムウォーク正規化グラフラプラシアンである。この演算子を適用することは、トークングラフ上での熱拡散の単一ステップに対応し、シーケンシャル全体にわたって信号（トークン表現）を漸進的に平滑化する。

主要な貢献

1. アーキテクチャの提案：ラプラシアンヘッドを標準的なアテンションヘッドと併用する、シンプルでパラメータフリーなトランスフォーマーアーキテクチャの変更。
2. 実証的検証：この変更が、教師あり分類、言語モデル、自己教師あり学習という 3 つの異なる学習パラダイムにおいて性能を向上させることの証明。
3. 幾何学的特徴付け：ラプラシアンヘッドが学習されたトークン表現の幾何学をどのように変化させるかについての詳細な分析。有害な過剰平滑化ではなく、有益な「崩壊」と平滑化を誘発することを示す。

実験結果

1. 画像分類（教師あり学習）

CIFAR-10、CIFAR-100、ImageNet-1k において DeiT-3 ファミリー（ViT-B）を使用：

• 性能：ラプラシアンヘッドの追加は、Top-1 精度を一貫して向上させた。向上度はラプラシアンヘッドの数とほぼ単調増加であった。
  • ImageNet において、12 個のラプラシアンヘッドを持つモデル（ViT-B-12L）は、ベースラインの**81.45%に対して82.32%**の精度を達成した。
  • CIFAR-100 において、向上は約 2%（75.12% 対 77.04%）であった。
• 幾何学：
  • 分散分析（ANOVA）：ラプラシアンヘッドはシーケンシャル内分散を有意に減少させ、クラス間分散を増加させた。これは、同一シーケンス内のトークンが共通の平均に向かって崩壊し、異なるシーケンスの平均がより効果的に分離することを示している。
  • ニューラル・クラッシュ：モデルはニューラル・クラッシュ幾何学への収束を強めた。トークン表現は各クラスに対して明確に分離された、単体のようなクラスター（Simplex ETF）を形成し、シーケンス平均はグローバルなクラス幾何学と整合した。

2. 言語モデル（自己回帰）

FineWeb-Edu で事前学習され、推論/コードベンチマークで微調整された GPT-2 スタイルのモデル（5 億 6100 万パラメータ）を使用：

• 性能：ラプラシアンヘッドを持つモデルは、ARC-Easy、ARC-Challenge、MMLU、GSM8K、HumanEval において平均してベースラインを上回った。
  • 5 個のラプラシアンヘッドを持つモデルは、最高の平均性能（30.06%）を達成し、ARC-Easy（+16%）と ARC-Challenge（+7%）で顕著な向上を示した。
• 幾何学：
  • スペクトル減衰：ラプラシアンヘッドは、より急速に減衰するスペクトルを持つトークン表現を生み出した。すべてのヘッドがラプラシアンである場合、分散の 99% を説明するために必要な主成分の数は、547.7 から 430.4 へと**21%**減少した。
  • 言語的クラッシュ：モデルは、同じ次のトークン予測を共有するトークン表現間の分離性を高め、同様の意味的役割を持つトークンがクラスター化する、より強力な形態の「言語的クラッシュ」を示した。

3. 自己教師あり学習（DINO）

ImageNet-1k における DINO 事前学習に ViT-S モデルを使用：

• 性能：ラプラシアンヘッドは、ImageNet 上の線形プロービング精度と ADE20K 上のセグメンテーション性能を向上させた。
  • ViT-S-5L は、ベースライン（29.81%）と比較して、最良の ADE20K mIoU（32.27%）を達成した。
• 幾何学：
  • セグメンテーション：PCA 可視化は、ラプラシアンヘッドがベースラインと比較して、より鮮明な物体境界を生成し、ヘルメットの形状やジャージの番号などのより微細な细节を保持することを示した。
  • スペクトル減衰：言語モデルと同様に、これらのモデルはより急速なスペクトル減衰を示し、画像内トークン表現を低次元部分空間に集中させた。

意義と主張

本論文は、トークンの過剰平滑化が普遍的に有害であるという支配的な見解は不完全であると論じる。代わりに、著者は平滑化は表現学習のメカニズムであり、制御された場合に有益であると主張する。

• タスク固有の利益：ラプラシアンヘッドによって誘発される「平滑化」は失敗モードではなく、タスクの目的と整合する機能である：
  • 教師あり分類において、シーケンス内のトークンを崩壊させ、シーケンス平均をクラス幾何学と整合させることでニューラル・クラッシュを促進する。
  • 言語モデルにおいて、言語的クラッシュを強化し、同一の次のトークンターゲットを持つトークンの分離性を向上させる。
  • 自己教師あり学習において、表現を低次元部分空間に集中させ、セグメンテーションなどの下流タスクを支援する。
• メカニズム：グラフ拡散（ラプラシアンを介して）をアーキテクチャに明示的に組み込むことで、モデルはシーケンシャル内分散を直接制御する能力を獲得し、より頑健で幾何学的に構造化された表現へと至る。

著者は、固定戦略の実装は最適ではない可能性が高いと結論づけているが、その結果は、平滑化を単に緩和すべき問題としてではなく、表現学習の基本的な構成要素としてさらに調査されるべきであることを示唆している。