[論文レビュー] Unveiling Transformers with LEGO: a synthetic reasoning task

Yi Zhang, Artūrs Bačkurs+4 more|arXiv (Cornell University)|Jun 9, 2022

Explainable Artificial Intelligence (XAI)被引用数 4

ひとこと要約

本論文は、構造化されたアテンションメカニズムを通じて論理的推論の連鎖を追跡する能力を評価する合成的推論タスクである LEGO を導入する。著者らは同一のトークンに注目する新たな「関連付け」アテンションパターンを同定し、構造化されたアテンションが事前学習によって性能向上をもたらす仕組みを示し、大規模事前学習において FLOPs を削減しながらも精度を維持または向上させる LEGO アテンションモジュールを提案する。

ABSTRACT

We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel \emph{association} pattern that globally attends only to identical tokens. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regime the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose the LEGO attention module, a drop-in replacement for vanilla attention heads. This architectural change significantly reduces Flops and maintains or even \emph{improves} the model's performance at large-scale pretraining.

研究の動機と目的

• トランスフォーマーの逐次的推論能力を隔離して評価できる合成的推論タスクを設計すること。
• 関連しない NLP タスクでの事前学習およびデータセット構成が、推論タスクにおける一般化性能に与える影響を調査すること。
• 注目ヘッドの行動を同定・分析すること、特に「関連付け」として特徴づけられる新しい構造化アテンションパターンを対象とすること。
• 事前学習が推論タスクの性能向上に寄与する理由を理解し、実験的にその仮説を検証すること。
• 標準的なマルチヘッドアテンションの即時置き換えが可能で、計算コストを低減する効率的なアテンション機構「LEGO アテンション」を提案すること。

提案手法

• LEGO タスクは、同一のトークンの連鎖を追跡・推論する必要がある合成的推論ベンチマークとして設計されている。
• 著者らは、異なる連鎖長を想定して LEGO データでトランスフォーマーを学習し、活性化およびアテンションパターン分析を通じて注目ヘッドの挙動を分析した。
• 同一のトークンにのみグローバルに注目する新たな「関連付け」アテンションパターンを導入し、繰り返し現れる要素の正確な追跡を可能にした。
• 事前学習が構造化されたアテンションパターンを誘発することで性能向上をもたらすという仮説を提示し、制御されたアブレーション実験を通じて検証した。
• 標準的なマルチヘッドアテンションの代替として、同一トークンへの注目に制限することで FLOPs を削減する LEGO アテンションモジュールを提案した。
• 複数のデータレジーム、特にトレーニング時とテスト時の異なる連鎖長を想定した実験を通じて、耐久性およびショートカット学習の有無を評価した。

実験結果

リサーチクエスチョン

• RQ1 関連しない NLP タスクでの事前学習およびデータセット構成は、推論タスクにおける未学習の連鎖長への一般化能力にどのように影響するか？
• RQ2 LEGO タスクの学習過程でどのような特定のアテンションパターンが出現し、それらがトークン連鎖の推論をどのように支援するか？
• RQ3 なぜ事前学習が合成的推論タスクの性能向上に寄与するのか。これは、構造化されたアテンションパターンの出現に起因するのか？
• RQ4 どのようなデータレジームでモデルは「ショートカット」解を学習し、分布シフトに対して耐性を失うのか。このような行動はどのように是正できるか？
• RQ5 標準アテンションと比較して、効率的かつ高性能を達成できるように改良されたアテンション機構「LEGO アテンション」は、大規模事前学習において実現可能か？

主な発見

• 特定のヘッドが同一トークンにのみグローバルに注目する新たな「関連付け」アテンションパターンが同定され、シーケンス内での繰り返し要素の正確な追跡が可能となった。
• 関連しない NLP タスクでの事前学習は LEGO タスクの性能向上に寄与し、その向上は構造化されたアテンションパターンの出現と強く相関していた。
• 特定のデータレジームでは、真の推論ではなく表面的な統計的ヒントに依存するショートカット解が学習され、分布シフトに対する耐性が低下した。
• 提案された LEGO アテンションモジュールは、大規模事前学習ベンチマークで FLOPs を顕著に削減しながらも、性能を維持または向上させた。
• 研究結果は、アテンションパターンがランダムではなく構造的であることを確認し、重みの共有や畳み込み部品といったアーキテクチャ的選択が推論行動に影響を及ぼすことを示した。
• これらの発見は、事前学習が推論タスクに有益である理由が、単にデータの多様性に起因するのではなく、タスクに依存しないアテンションのインダクティブバイアスを誘発するからである可能性を示唆している。