Gegrokte Transformers zijn Impliciete Redeneerders: Een Mechanistische Reis naar de Grens van Generalisatie

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

Samenvatting

We onderzoeken of transformers kunnen leren om impliciet te redeneren over parametrische kennis, een vaardigheid waar zelfs de meest capabele taalmodelen moeite mee hebben. Door ons te richten op twee representatieve redeneertypes, compositie en vergelijking, vinden we consistent dat transformers impliciet redeneren kunnen leren, maar alleen door middel van grokking, d.w.z. uitgebreide training ver voorbij overfitting. De niveaus van generalisatie variëren ook tussen redeneertypes: wanneer ze worden geconfronteerd met voorbeelden buiten de distributie, slagen transformers er niet in om systematisch te generaliseren voor compositie, maar wel voor vergelijking. We duiken in de interne werking van het model gedurende de training, waarbij we analytische experimenten uitvoeren die het volgende onthullen: 1) het mechanisme achter grokking, zoals de vorming van het generaliserende circuit en de relatie met de relatieve efficiëntie van generaliserende en memoriserende circuits, en 2) de verbinding tussen systematiciteit en de configuratie van het generaliserende circuit. Onze bevindingen begeleiden de opzet van data en training om impliciet redeneren beter te induceren en suggereren mogelijke verbeteringen aan de transformer-architectuur, zoals het stimuleren van kennisuitwisseling tussen lagen. Bovendien tonen we aan dat voor een uitdagende redeneertaak met een grote zoekruimte, GPT-4-Turbo en Gemini-1.5-Pro, gebaseerd op niet-parametrisch geheugen, slecht presteren ongeacht de promptstijlen of retrieval-augmentatie, terwijl een volledig gegrokked transformer bijna perfecte nauwkeurigheid kan bereiken, wat de kracht van parametrisch geheugen voor complex redeneren aantoont.

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.