トランスフォーマーとニューラルアルゴリズム推論器の融合

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

要旨

Transformerは、そのシンプルでありながら効果的なアーキテクチャにより、機械学習に革命をもたらしました。インターネット上の大規模なテキストデータセットでTransformerを事前学習させることで、自然言語理解（NLU）タスクにおいて比類のない汎化性能が実現されました。しかし、計算が正確かつロバストでなければならないアルゴリズム的な推論を必要とするタスクにおいては、このような言語モデルは依然として脆弱です。この制約に対処するため、我々はTransformerの言語理解能力と、グラフニューラルネットワーク（GNN）ベースのニューラルアルゴリズム推論器（NAR）のロバスト性を組み合わせた新しいアプローチを提案します。このようなNARは、グラフ形式で指定されたアルゴリズムタスクに対する汎用ソルバーとして有効であることが証明されています。Transformerがこれらの埋め込みにアクセスできるようにするため、2段階のトレーニング手順を備えたハイブリッドアーキテクチャを提案し、言語モデルのトークンがNARのノード埋め込みにクロスアテンションを適用できるようにします。我々は、提案するTransNARモデルをテキストベースのCLRS-30ベンチマークであるCLRS-Textで評価し、アルゴリズム推論において、分布内および分布外の両方で、Transformerのみのモデルを大幅に上回る性能を示しました。

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.