Les Transformers rencontrent les Raisonneurs Algorithmiques Neuronaux

Transformers meet Neural Algorithmic Reasoners

June 13, 2024

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

cs.AI

papers.abstract

Les Transformers ont révolutionné l'apprentissage automatique grâce à leur architecture simple mais efficace. Le pré-entraînement des Transformers sur d'immenses ensembles de données textuelles provenant d'Internet a permis une généralisation inégalée pour les tâches de compréhension du langage naturel (NLU). Cependant, ces modèles de langage restent fragiles lorsqu'ils sont confrontés à des formes de raisonnement algorithmique, où les calculs doivent être précis et robustes. Pour pallier cette limitation, nous proposons une nouvelle approche qui combine la compréhension du langage du Transformer avec la robustesse des raisonneurs algorithmiques neuronaux (NAR) basés sur des réseaux de neurones graphiques (GNN). Ces NAR se sont avérés efficaces en tant que solveurs génériques pour les tâches algorithmiques, lorsqu'elles sont spécifiées sous forme de graphe. Pour rendre leurs embeddings accessibles à un Transformer, nous proposons une architecture hybride avec une procédure d'entraînement en deux phases, permettant aux tokens du modèle de langage de réaliser une attention croisée sur les embeddings de nœuds provenant du NAR. Nous évaluons notre modèle TransNAR résultant sur CLRS-Text, la version textuelle du benchmark CLRS-30, et démontrons des gains significatifs par rapport aux modèles basés uniquement sur les Transformers pour le raisonnement algorithmique, à la fois en distribution et hors distribution.

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.