트랜스포머가 신경 알고리즘 추론기와 만나다

Transformers meet Neural Algorithmic Reasoners

June 13, 2024

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

cs.AI

초록

트랜스포머는 단순하면서도 효과적인 아키텍처로 머신러닝 분야에 혁명을 일으켰습니다. 인터넷에서 수집한 방대한 텍스트 데이터셋에 대한 트랜스포머의 사전 학습은 자연어 이해(NLU) 작업에서 뛰어난 일반화 능력을 이끌어냈습니다. 그러나 이러한 언어 모델은 계산이 정확하고 견고해야 하는 알고리즘적 추론 형태의 작업에서는 여전히 취약한 모습을 보입니다. 이러한 한계를 해결하기 위해, 우리는 트랜스포머의 언어 이해 능력과 그래프 신경망(GNN) 기반의 신경 알고리즘 추론기(NAR)의 견고성을 결합한 새로운 접근 방식을 제안합니다. 이러한 NAR은 그래프 형태로 명시된 알고리즘 작업에 대해 일반적인 솔버로서 효과적인 것으로 입증되었습니다. 트랜스포머가 이러한 임베딩에 접근할 수 있도록, 우리는 언어 모델의 토큰들이 NAR의 노드 임베딩에 교차 주의(cross-attend)할 수 있는 2단계 학습 절차를 가진 하이브리드 아키텍처를 제안합니다. 우리는 결과적으로 얻은 TransNAR 모델을 CLRS-30 벤치마크의 텍스트 기반 버전인 CLRS-Text에서 평가하고, 분포 내외에서 알고리즘적 추론에 대해 트랜스포머 단독 모델 대비 상당한 성능 향상을 입증합니다.

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.