Transformers se encuentran con razonadores algorítmicos neuronales

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

Resumen

Los Transformers han revolucionado el aprendizaje automático con su arquitectura simple pero efectiva. El preentrenamiento de Transformers en grandes conjuntos de datos textuales de Internet ha llevado a una generalización sin precedentes para tareas de comprensión del lenguaje natural (NLU). Sin embargo, estos modelos de lenguaje siguen siendo frágiles cuando se enfrentan a formas algorítmicas de razonamiento, donde los cálculos deben ser precisos y robustos. Para abordar esta limitación, proponemos un enfoque novedoso que combina la comprensión del lenguaje del Transformer con la robustez de los razonadores algorítmicos neuronales (NARs) basados en redes neuronales de grafos (GNNs). Dichos NARs han demostrado ser efectivos como solucionadores genéricos para tareas algorítmicas, cuando se especifican en forma de grafo. Para hacer que sus incrustaciones sean accesibles a un Transformer, proponemos una arquitectura híbrida con un procedimiento de entrenamiento en dos fases, permitiendo que los tokens en el modelo de lenguaje atiendan de manera cruzada a las incrustaciones de nodos del NAR. Evaluamos nuestro modelo resultante, TransNAR, en CLRS-Text, la versión basada en texto del benchmark CLRS-30, y demostramos mejoras significativas sobre los modelos basados únicamente en Transformers para el razonamiento algorítmico, tanto dentro como fuera de la distribución.

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.