【論文瞬読】Transformer の限界を探る: 表現の崩壊と Over-squashing の影響

2024年6月17日 14:49

こんにちは、みなさん。株式会社AI Nestです。
今日は、自然言語処理の分野で大きな注目を集めている研究についてご紹介したいと思います。Federico Barbero 氏らによる "Transformers need glasses! Information over-squashing in language tasks" という論文です。この研究は、最先端の大規模言語モデル (LLMs) の基盤となっている decoder-only Transformer アーキテクチャの限界を、理論と実験の両面から明らかにしています。

タイトル：Transformers need glasses! Information over-squashing in language tasks
URL：https://arxiv.org/abs/2406.04267
所属：University of Oxford, Google DeepMind
著者：Federico Barbero, Andrea Banino, Steven Kapturowski, Dharshan Kumaran, João G.M. Araújo, Alex Vitvitskyi, Razvan Pascanu, Petar Veličković

研究の背景

LLMs は、GPT-3 や BERT など、Transformer アーキテクチャを基盤とした大規模な言語モデルの総称です。これらのモデルは、大量のテキストデータを用いて事前学習され、対話システム、マルチモーダル入力の理解、コード生成など、さまざまなタスクで驚くべき性能を示しています。特に、GPT-3 は、少ないプロンプトで複雑なタスクを遂行できることから、「Foundation Models」とも呼ばれ、AI 分野に大きなインパクトを与えています。

しかし、その一方で、LLMs はカウンティングやコピーといった一見シンプルなタスクで苦戦することが報告されています。たとえば、「1」が連続する系列の中に、最後に「0」が1つだけ現れるようなケースです。人間なら簡単に対処できるこうしたタスクに、なぜ LLMs は苦戦するのでしょうか？ Barbero 氏らは、こうした失敗例に着目し、その原因を探ることで、LLMs のさらなる改善につなげようとしています。

表現の崩壊と Over-squashing

この研究では、理論的な信号伝播分析により、2つの興味深い現象が見出されました。

1つ目は「表現の崩壊 (representational collapse)」です。これは、ある種の異なる入力系列に対して、Transformer の最終層の最後のトークンの表現が任意に近くなってしまい、モデルがこれらの系列に対して異なる応答ができなくなってしまう問題です。具体的には、系列の長さが増すにつれて、最後のトークンの表現が収束してしまうのです。この現象は、低精度の浮動小数点を使用することで悪化することも示されました。

2つ目は「over-squashing」です。これは、decoder-only Transformer が入力内の特定のトークンに対する感度を失ってしまう現象で、グラフニューラルネットワークの分野でよく知られていました。Transformer は、入力系列をグラフ構造とみなすことができるため、グラフニューラルネットワークの知見が応用できるのです。Over-squashing は、グラフ上の情報伝播が不十分になることで起こります。

著者らは、これらの現象が、Transformer の causal attention mechanism に起因することを理論的に示しました。つまり、現在の LLMs の設計自体に、こうした限界の原因があるというわけです。

実験による裏付け

理論的な知見は、実際の LLMs での実験によって裏付けられています。著者らは、最新の LLMs を使って、コピーとカウンティングのタスクを行い、系列の長さが増すにつれてパフォーマンスが急激に低下することを示しました。

たとえば、コピータスクでは、系列の最後の要素をコピーする方が、最初の要素をコピーするよりも難しいことが示されました。これは、表現の崩壊によって、系列の最後の方の情報が失われやすいためだと解釈できます。

カウンティングタスクでは、連続する「1」の数を数えるタスクで、系列が長くなるほど正解率が下がることが示されました。興味深いことに、LLMs は「100」のような区切りのよい数字を出力しがちで、実際に数えているというよりは、大雑把に推測している様子が伺えました。

カウンティングタスクにおいて LLMs が「100」のような
区切りのよい数字を出力しがちであることを示す図

これらの結果は、表現の崩壊と over-squashing が実際の LLMs の挙動に影響を与えていることを強く示唆しています。

今後の展望

この研究は、Transformer ベースの LLMs の限界を明確に示した重要な一歩だと言えます。カウンティングやコピーは、より複雑な推論タスクを解く上での基本的な構成要素です。したがって、こうした限界を理解し改善していくことは、LLMs の応用可能性を広げる上で不可欠でしょう。

著者らは、表現の崩壊を緩和するために、系列内に追加のトークンを導入するというシンプルな解決策を提案しています。たとえば、長い系列の中に時々記号を挿入することで、表現の収束を防ぐことができます。こうした知見は、今後の LLMs の設計に活かされていくことが期待されます。

また、この研究は、LLMs の解釈性 (interpretability) の重要性も浮き彫りにしています。LLMs が何を学習し、どのように推論を行っているのかを理解することは、より信頼性の高い AI システムを構築する上で欠かせません。この研究で用いられた理論的・実験的アプローチは、そうした解釈性の研究にも応用できるでしょう。

結論

Barbero 氏らの研究は、LLMs の限界を理論と実験の両面から明らかにした重要な成果です。Transformer アーキテクチャの根本的な問題点を指摘し、その改善の方向性を示唆しています。この分野の研究者や実践者にとって、ぜひ一読をおすすめしたい論文です。

LLMs は、自然言語処理だけでなく、さまざまな AI タスクにおいて大きな可能性を秘めています。しかし、その実用化のためには、こうした基礎研究の積み重ねが欠かせません。モデルの限界を見極め、それを克服していく地道な努力が、AI 技術の健全な発展につながるのです。

Barbero 氏らの研究は、そうした努力の一つの重要な成果だと言えるでしょう。この研究を出発点として、さらなる探求が進められていくことを期待したいと思います。LLMs の可能性を最大限に引き出し、より信頼性の高い AI システムを構築していくために、私たちにできることは何でしょうか。この論文を読んで、みなさんも一緒に考えてみてください。