Transformadores Grokked são Raciocinadores Implícitos: Uma Jornada Mecanicista até o Limite da Generalização

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

Resumo

Investigamos se os transformadores podem aprender a raciocinar implicitamente sobre conhecimento paramétrico, uma habilidade com a qual até os modelos de linguagem mais capazes lutam. Focando em dois tipos representativos de raciocínio, composição e comparação, consistentemente descobrimos que os transformadores podem aprender raciocínio implícito, mas apenas através do fenômeno de "grokking", ou seja, treinamento prolongado muito além do ponto de sobreajuste. Os níveis de generalização também variam entre os tipos de raciocínio: ao enfrentar exemplos fora da distribuição, os transformadores falham em generalizar sistematicamente para a composição, mas têm sucesso na comparação. Investigamos os mecanismos internos do modelo ao longo do treinamento, conduzindo experimentos analíticos que revelam: 1) o mecanismo por trás do grokking, como a formação do circuito generalizador e sua relação com a eficiência relativa dos circuitos de generalização e memorização, e 2) a conexão entre a sistematicidade e a configuração do circuito generalizador. Nossas descobertas orientam a configuração de dados e treinamento para melhor induzir o raciocínio implícito e sugerem melhorias potenciais na arquitetura do transformador, como o incentivo ao compartilhamento de conhecimento entre camadas. Além disso, demonstramos que, para uma tarefa desafiadora de raciocínio com um grande espaço de busca, o GPT-4-Turbo e o Gemini-1.5-Pro, baseados em memória não paramétrica, falham drasticamente, independentemente dos estilos de prompt ou da ampliação de recuperação, enquanto um transformador totalmente "grokked" pode alcançar precisão quase perfeita, destacando o poder da memória paramétrica para raciocínios complexos.

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.