Attention-only Transformerの性能検証──FFNなしでもパラメータ再配分で同等性能を実現

論文概要

Transformerのパラメータの約3分の2を占めるFeed-Forward Network (FFN) の必要性を検証するため、著者らはAttention-onlyのデコーダーモデル「Simple Attention Networks (SANs)」を提案し、標準的なTransformerと厳密な比較を行った。パラメータ数、計算量(FLOPs)、深さを制御した条件下で、最大105Bトークンの学習を実施。その結果、FFNを単純除去するだけでは性能が低下するものの、浮いたパラメータをAttention層の深さに再配分することで、損失差をわずか0.27%（0.006 nats）にまで抑えられることを実証した。

関連研究

従来、Transformerの性能はAttention層とFFN層の相互作用に依存すると考えられてきた。先行研究ではFFNの重要性が指摘されていたが、パラメータ数や計算量、深さを同時に制御した厳密な比較は不足していた。本研究は、これらの変数を固定した「制御実験」を行うことで、アーキテクチャの構成要素の役割を再定義した。

新規性と貢献

本研究の最大の貢献は、FFNの役割を「パラメータの格納場所」と「知識の保持」に分離・特定した点にある。Attention-onlyモデルはコンテキストに基づいた推論には優れるが、重みそのものに記憶された知識の再現性では標準モデルに劣る。また、QK-normalizationを導入することで、FFNなしでも48層という深層スタックの学習を安定させることに成功した。

提案手法の詳細

SANsはFFNを完全に削除し、その分のパラメータをAttention層の深さに割り当てたモデルである。なぜFFNなしで学習可能なのか、その鍵は重み行列のダイナミクスにある。分析の結果、ルーティング行列（Q/K）は学習初期に早期に決定（結晶化）し、コンテンツ行列が学習を通じてランクを蓄積する。FFNがない場合、このランク蓄積の役割がAttentionの出力プロジェクション層に転移する。この構造的適応により、モデルはFFNなしでも高い表現力を維持する。

評価・考察

実験では、105Bトークンの学習において、パラメータ数を一致させた場合、SANsと標準Transformerの損失差はわずか0.006 natsであった。性能差の分析（図2参照）では、この差は主に「知識密度が高いクエリ」に集中しており、コンテキストが豊富なタスクでは両者の差はほぼ消失する。つまり、Attentionはコンテキスト処理を担い、FFNはパラメータ記憶を担うという機能分担が明らかとなった。

応用例と今後の展望

本研究の結果は、推論コストの削減やモデルの軽量化を目指す日本のエンジニアにとって重要な指針となる。特に、コンテキスト処理が主体のRAG（検索拡張生成）システムや、特定のドメイン知識を外部データベースから取得するアプリケーションにおいて、FFNを削減したモデル構成を採用することで、メモリ効率と推論速度を最適化できる可能性がある。今後は、知識の保持を外部メモリに依存させる設計が、従来のFFNベースのモデルを完全に代替可能かどうかが焦点となる。

結論

Attention-onlyモデルは、パラメータを適切に再配分することで、FFNを搭載したモデルと同等の性能を達成できる。FFNの役割は主にパラメータの記憶に限定されており、コンテキスト処理能力についてはAttention層がその大半を担っていることが本研究により示された。

注釈

• FFN (Feed-forward Network): Transformerの各層に含まれる全結合層。モデルの知識を保持する主要な部位とされる。
• QK-normalization: QueryとKeyのベクトルを正規化する手法。深層モデルの学習を安定させる。
• nats: 自然対数を底とする情報量の単位。モデルの損失（perplexity）を測る指標。