グルークド・トランスフォーマーは暗黙の推論者である：一般化の限界への機構論的探求

Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization

May 23, 2024

著者: Boshi Wang, Xiang Yue, Yu Su, Huan Sun

cs.AI

要旨

我々は、トランスフォーマーがパラメトリック知識を暗黙的に推論する能力を学習できるかどうかを研究する。これは、最も優れた言語モデルでさえも苦手とするスキルである。代表的な推論タイプである合成と比較に焦点を当て、トランスフォーマーが暗黙的推論を学習できるが、それはグロッキング（過学習をはるかに超えた長期訓練）を通じてのみ可能であることを一貫して見出した。また、推論タイプによって一般化のレベルも異なることが明らかとなった：分布外の例に直面した場合、トランスフォーマーは合成に対して体系的に一般化することに失敗するが、比較では成功する。我々は、訓練全体を通じてモデルの内部を詳細に調査し、以下のことを明らかにする分析実験を行った：1）グロッキングのメカニズム、例えば一般化回路の形成と、一般化回路と記憶回路の相対的な効率の関係、2）体系性と一般化回路の構成との関連性。我々の発見は、暗黙的推論をより良く誘導するためのデータと訓練設定の指針となり、トランスフォーマーアーキテクチャの潜在的な改善、例えば層間の知識共有を促進することを示唆する。さらに、大規模な探索空間を伴う難しい推論タスクにおいて、非パラメトリックメモリに基づくGPT-4-TurboとGemini-1.5-Proは、プロンプトスタイルや検索拡張に関わらず大きく失敗するのに対し、完全にグロッキングされたトランスフォーマーはほぼ完璧な精度を達成できることを示し、複雑な推論におけるパラメトリックメモリの力を実証した。

English

We study whether transformers can learn to implicitly reason over parametric knowledge, a skill that even the most capable language models struggle with. Focusing on two representative reasoning types, composition and comparison, we consistently find that transformers can learn implicit reasoning, but only through grokking, i.e., extended training far beyond overfitting. The levels of generalization also vary across reasoning types: when faced with out-of-distribution examples, transformers fail to systematically generalize for composition but succeed for comparison. We delve into the model's internals throughout training, conducting analytical experiments that reveal: 1) the mechanism behind grokking, such as the formation of the generalizing circuit and its relation to the relative efficiency of generalizing and memorizing circuits, and 2) the connection between systematicity and the configuration of the generalizing circuit. Our findings guide data and training setup to better induce implicit reasoning and suggest potential improvements to the transformer architecture, such as encouraging cross-layer knowledge sharing. Furthermore, we demonstrate that for a challenging reasoning task with a large search space, GPT-4-Turbo and Gemini-1.5-Pro based on non-parametric memory fail badly regardless of prompting styles or retrieval augmentation, while a fully grokked transformer can achieve near-perfect accuracy, showcasing the power of parametric memory for complex reasoning.