Les Transformers Grokkés sont des Raisonneurs Implicites : Un Voyage Mécaniste aux Confins de la Généralisation

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

May 23, 2024

Auteurs: Boshi Wang, Xiang Yue, Yu Su, Huan Sun

cs.AI

Résumé

Nous étudions si les transformateurs peuvent apprendre à raisonner implicitement sur des connaissances paramétriques, une compétence avec laquelle même les modèles de langage les plus performants peinent. En nous concentrant sur deux types de raisonnement représentatifs, la composition et la comparaison, nous constatons systématiquement que les transformateurs peuvent apprendre le raisonnement implicite, mais uniquement à travers le phénomène de "grokking", c'est-à-dire un entraînement prolongé bien au-delà du surapprentissage. Les niveaux de généralisation varient également selon les types de raisonnement : face à des exemples hors distribution, les transformateurs échouent à généraliser systématiquement pour la composition, mais y parviennent pour la comparaison. Nous explorons les mécanismes internes du modèle tout au long de l'entraînement, en menant des expériences analytiques qui révèlent : 1) le mécanisme derrière le grokking, comme la formation du circuit de généralisation et sa relation avec l'efficacité relative des circuits de généralisation et de mémorisation, et 2) le lien entre la systématicité et la configuration du circuit de généralisation. Nos résultats guident la conception des données et des configurations d'entraînement pour mieux induire le raisonnement implicite et suggèrent des améliorations potentielles de l'architecture des transformateurs, comme l'encouragement du partage de connaissances inter-couches. De plus, nous démontrons que pour une tâche de raisonnement complexe avec un grand espace de recherche, GPT-4-Turbo et Gemini-1.5-Pro, basés sur une mémoire non paramétrique, échouent de manière significative, quels que soient les styles d'invite ou l'augmentation par récupération, tandis qu'un transformateur entièrement grokké peut atteindre une précision quasi parfaite, mettant en lumière la puissance de la mémoire paramétrique pour le raisonnement complexe.

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.