[Paper Review] A Generalist Neural Algorithmic Learner

Borja Ibarz, Vitaly Kurin+13 more|arXiv (Cornell University)|Sep 22, 2022

Machine Learning and Algorithms7 citations

TL;DR

This paper introduces a generalist neural algorithmic learner—a single graph neural network with shared parameters—that simultaneously learns to execute diverse classical algorithms such as sorting, shortest path, dynamic programming, and geometric computations. By improving input representation, training regime, and GNN architecture, the model achieves a 20%+ absolute improvement in average performance on the CLRS-30 benchmark and matches or exceeds specialist models in out-of-distribution generalization, demonstrating effective multi-task knowledge incorporation across heterogeneous control flows.

ABSTRACT

The cornerstone of neural algorithmic reasoning is the ability to solve algorithmic tasks, especially in a way that generalises out of distribution. While recent years have seen a surge in methodological improvements in this area, they mostly focused on building specialist models. Specialist models are capable of learning to neurally execute either only one algorithm or a collection of algorithms with identical control-flow backbone. Here, instead, we focus on constructing a generalist neural algorithmic learner -- a single graph neural network processor capable of learning to execute a wide range of algorithms, such as sorting, searching, dynamic programming, path-finding and geometry. We leverage the CLRS benchmark to empirically show that, much like recent successes in the domain of perception, generalist algorithmic learners can be built by "incorporating" knowledge. That is, it is possible to effectively learn algorithms in a multi-task manner, so long as we can learn to execute them well in a single-task regime. Motivated by this, we present a series of improvements to the input representation, training regime and processor architecture over CLRS, improving average single-task performance by over 20% from prior art. We then conduct a thorough ablation of multi-task learners leveraging these improvements. Our results demonstrate a generalist learner that effectively incorporates knowledge captured by specialist models.

Motivation & Objective

• To develop a single neural network capable of learning and executing a wide range of classical algorithmic tasks with shared parameters.
• To overcome the limitations of specialist models that are restricted to similar control-flow structures or single algorithms.
• To improve multi-task learning stability and performance in neural algorithmic reasoning by addressing numerical instabilities amplified in joint training.
• To demonstrate that knowledge from specialist models can be effectively incorporated into a generalist learner through improved training and architectural components.
• To establish a scalable, general-purpose framework for neural algorithmic reasoning applicable across diverse algorithmic domains.

Proposed method

• The model uses a unified graph neural network processor with shared weights to handle multiple algorithmic tasks via task-specific encoder-decoder heads.
• Input representations are enhanced with structured, task-agnostic features to improve model generalization and training stability.
• A chunked training regime is introduced to reduce gradient instability during multi-task learning, improving convergence and performance.
• Teacher forcing with noise is avoided to prevent distribution shift between training and inference, enhancing out-of-distribution generalization.
• Architectural improvements include a triplet reasoning module, gated message passing, and a Sinkhorn-based operator for differentiable routing and attention.
• Gradient clipping and efficiency optimizations are applied to stabilize training and accelerate convergence on the CLRS-30 benchmark.

Experimental results

Research questions

• RQ1 Can a single neural network effectively learn and execute a wide range of algorithmic tasks with fundamentally different control flows?
• RQ2 To what extent can multi-task training improve generalization across diverse algorithmic tasks compared to single-task specialist models?
• RQ3 How do improvements in input representation, training regime, and architecture collectively enhance multi-task algorithmic reasoning performance?
• RQ4 Does avoiding teacher forcing noise during training lead to better out-of-distribution generalization in algorithmic execution?
• RQ5 Can a generalist model match or exceed the performance of specialist models, especially on larger, unseen instances?

Key findings

• The proposed generalist model achieves a 20.3% absolute improvement in average performance on the CLRS-30 benchmark over prior state-of-the-art methods.
• The model matches or exceeds the out-of-distribution generalization performance of corresponding single-task specialist models on multiple tasks, including sorting and shortest path.
• The ablation study shows that teacher forcing noise significantly harms generalization, with performance dropping when ground-truth hints are used during training.
• The chunked training regime leads to a substantial performance gain over vanilla multi-task training, indicating improved training stability.
• The combination of architectural improvements—such as the triplet reasoning module and Sinkhorn operator—contributes significantly to the model's robustness and performance.
• The results demonstrate that effective multi-task learning in neural algorithmic reasoning is possible when numerical stability and representation quality are carefully addressed.