Los Transformadores Grokked son Razonadores Implícitos: Un Viaje Mecanicista al Límite de la Generalización

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

May 23, 2024

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

cs.AI

Resumen

Estudiamos si los transformadores pueden aprender a razonar implícitamente sobre conocimiento paramétrico, una habilidad con la que incluso los modelos de lenguaje más avanzados tienen dificultades. Centrándonos en dos tipos de razonamiento representativos, composición y comparación, encontramos consistentemente que los transformadores pueden aprender razonamiento implícito, pero solo a través del "grokking", es decir, entrenamiento extendido mucho más allá del sobreajuste. Los niveles de generalización también varían según los tipos de razonamiento: cuando se enfrentan a ejemplos fuera de distribución, los transformadores no logran generalizar sistemáticamente para la composición, pero sí lo hacen para la comparación. Profundizamos en los mecanismos internos del modelo durante el entrenamiento, realizando experimentos analíticos que revelan: 1) el mecanismo detrás del grokking, como la formación del circuito generalizador y su relación con la eficiencia relativa de los circuitos de generalización y memorización, y 2) la conexión entre la sistematicidad y la configuración del circuito generalizador. Nuestros hallazgos guían la configuración de datos y entrenamiento para inducir mejor el razonamiento implícito y sugieren posibles mejoras en la arquitectura del transformador, como fomentar el intercambio de conocimiento entre capas. Además, demostramos que para una tarea de razonamiento desafiante con un espacio de búsqueda grande, GPT-4-Turbo y Gemini-1.5-Pro, basados en memoria no paramétrica, fracasan rotundamente independientemente del estilo de "prompting" o la ampliación de recuperación, mientras que un transformador completamente "grokked" puede alcanzar una precisión casi perfecta, mostrando el poder de la memoria paramétrica para el razonamiento complejo.

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.