I Transformers incontrano i Neural Algorithmic Reasoners

Transformers meet Neural Algorithmic Reasoners

June 13, 2024

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

cs.AI

Abstract

I Transformer hanno rivoluzionato l'apprendimento automatico con la loro architettura semplice ma efficace. Il pre-addestramento dei Transformer su enormi dataset testuali provenienti da Internet ha portato a una generalizzazione senza pari per i compiti di comprensione del linguaggio naturale (NLU). Tuttavia, tali modelli linguistici rimangono fragili quando vengono incaricati di forme algoritmiche di ragionamento, dove i calcoli devono essere precisi e robusti. Per affrontare questa limitazione, proponiamo un approccio innovativo che combina la comprensione del linguaggio del Transformer con la robustezza dei risolutori neurali algoritmici (NAR) basati su reti neurali a grafo (GNN). Tali NAR si sono dimostrati efficaci come risolutori generici per compiti algoritmici, quando specificati in forma di grafo. Per rendere i loro embedding accessibili a un Transformer, proponiamo un'architettura ibrida con una procedura di addestramento in due fasi, che consente ai token nel modello linguistico di eseguire cross-attention sugli embedding dei nodi provenienti dal NAR. Valutiamo il nostro modello TransNAR risultante su CLRS-Text, la versione testuale del benchmark CLRS-30, e dimostriamo miglioramenti significativi rispetto ai modelli basati esclusivamente su Transformer per il ragionamento algoritmico, sia all'interno che al di fuori della distribuzione.

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.