Transformers encontram Raciocinadores Algorítmicos Neurais

Transformers meet Neural Algorithmic Reasoners

June 13, 2024

Autores: Wilfried Bounsi, Borja Ibarz, Andrew Dudzik, Jessica B. Hamrick, Larisa Markeeva, Alex Vitvitskyi, Razvan Pascanu, Petar Veličković

cs.AI

Resumo

Os Transformers revolucionaram o aprendizado de máquina com sua arquitetura simples, porém eficaz. O pré-treinamento de Transformers em grandes conjuntos de dados textuais da Internet resultou em uma generalização sem precedentes para tarefas de compreensão de linguagem natural (NLU). No entanto, esses modelos de linguagem permanecem frágeis quando confrontados com formas algorítmicas de raciocínio, onde os cálculos precisam ser precisos e robustos. Para superar essa limitação, propomos uma abordagem inovadora que combina a compreensão de linguagem do Transformer com a robustez dos raciocinadores algorítmicos neurais (NARs) baseados em redes neurais de grafos (GNNs). Esses NARs se mostraram eficazes como solucionadores genéricos para tarefas algorítmicas, quando especificadas em forma de grafo. Para tornar seus embeddings acessíveis a um Transformer, propomos uma arquitetura híbrida com um procedimento de treinamento em duas fases, permitindo que os tokens no modelo de linguagem realizem atenção cruzada com os embeddings dos nós do NAR. Avaliamos nosso modelo TransNAR resultante no CLRS-Text, a versão textual do benchmark CLRS-30, e demonstramos ganhos significativos em relação a modelos baseados apenas em Transformers para raciocínio algorítmico, tanto dentro quanto fora da distribuição.

English

Transformers have revolutionized machine learning with their simple yet effective architecture. Pre-training Transformers on massive text datasets from the Internet has led to unmatched generalization for natural language understanding (NLU) tasks. However, such language models remain fragile when tasked with algorithmic forms of reasoning, where computations must be precise and robust. To address this limitation, we propose a novel approach that combines the Transformer's language understanding with the robustness of graph neural network (GNN)-based neural algorithmic reasoners (NARs). Such NARs proved effective as generic solvers for algorithmic tasks, when specified in graph form. To make their embeddings accessible to a Transformer, we propose a hybrid architecture with a two-phase training procedure, allowing the tokens in the language model to cross-attend to the node embeddings from the NAR. We evaluate our resulting TransNAR model on CLRS-Text, the text-based version of the CLRS-30 benchmark, and demonstrate significant gains over Transformer-only models for algorithmic reasoning, both in and out of distribution.