Грокнутые трансформеры являются неявными рассудителями: механистическое путешествие к грани обобщения

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.