# Towards Better Out-of-Distribution Generalization of Neural Algorithmic Reasoning Tasks

arXiv: https://arxiv.org/abs/2211.00692 · HTML: https://mlsanity.com/paper/2211.00692/

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Builds on 59 of its 64 references that are in the map.
- First posted on arXiv 1 Nov 2022.

## Builds on (most cited first, 30 of 59)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Can You Trust Your Model's Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift](https://mlsanity.com/paper/1906.02530/) (2019)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://mlsanity.com/paper/1810.02244/) (2018)
- [DropEdge: Towards Deep Graph Convolutional Networks on Node Classification](https://mlsanity.com/paper/1907.10903/) (2019)
- [Provably Powerful Graph Networks](https://mlsanity.com/paper/1905.11136/) (2019)
- [Improving Graph Neural Network Expressivity via Subgraph Isomorphism Counting](https://mlsanity.com/paper/2006.09252/) (2020)
- [The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2110.06296/) (2021)
- [Accuracy on the Line: On the Strong Correlation Between Out-of-Distribution and In-Distribution Generalization](https://mlsanity.com/paper/2107.04649/) (2021)
- [Can Graph Neural Networks Count Substructures?](https://mlsanity.com/paper/2002.04025/) (2020)
- [Weisfeiler and Lehman Go Topological: Message Passing Simplicial Networks](https://mlsanity.com/paper/2103.03212/) (2021)
- [Weisfeiler and Lehman Go Cellular: CW Networks](https://mlsanity.com/paper/2106.12575/) (2021)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Exploring Length Generalization in Large Language Models](https://mlsanity.com/paper/2207.04901/) (2022)

## Built on by (most cited first, 30 of 30)

- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Generalization on the Unseen, Logic Reasoning and Degree Curriculum](https://mlsanity.com/paper/2301.13105/) (2023)
- [Does Invariant Graph Learning via Environment Augmentation Learn Invariance?](https://mlsanity.com/paper/2310.19035/) (2023)
- [How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad](https://mlsanity.com/paper/2406.06467/) (2024)
- [Neural Algorithmic Reasoning with Causal Regularisation](https://mlsanity.com/paper/2302.10258/) (2023)
- [SwinGNN: Rethinking Permutation Invariance in Diffusion Models for Graph Generation](https://mlsanity.com/paper/2307.01646/) (2023)
- [Over-Squashing in Graph Neural Networks: A Comprehensive survey](https://mlsanity.com/paper/2308.15568/) (2023)
- [Neural Algorithmic Reasoning for Combinatorial Optimisation](https://mlsanity.com/paper/2306.06064/) (2023)
- [Neural Algorithmic Reasoning Without Intermediate Supervision](https://mlsanity.com/paper/2306.13411/) (2023)
- [How Interpretable Are Interpretable Graph Neural Networks?](https://mlsanity.com/paper/2406.07955/) (2024)
- [SALSA-CLRS: A Sparse and Scalable Benchmark for Algorithmic Reasoning](https://mlsanity.com/paper/2309.12253/) (2023)
- [On the Markov Property of Neural Algorithmic Reasoning: Analyses and Methods](https://mlsanity.com/paper/2403.04929/) (2024)
- [Explicitly Encoding Structural Symmetry is Key to Length Generalization in Arithmetic Tasks](https://mlsanity.com/paper/2406.01895/) (2024)
- [Discrete Neural Algorithmic Reasoning](https://mlsanity.com/paper/2402.11628/) (2024)
- [Parallel Algorithms Align with Neural Execution](https://mlsanity.com/paper/2307.04049/) (2023)
- [Chameleon: A Data-Efficient Generalist for Dense Visual Prediction in the Wild](https://mlsanity.com/paper/2404.18459/) (2024)
- [Recurrent Aggregators in Neural Algorithmic Reasoning](https://mlsanity.com/paper/2409.07154/) (2024)
- [Measuring Diversity: Axioms and Challenges](https://mlsanity.com/paper/2410.14556/) (2024)
- [Deep Equilibrium Algorithmic Reasoning](https://mlsanity.com/paper/2410.15059/) (2024)
- [Open-Book Neural Algorithmic Reasoning](https://mlsanity.com/paper/2501.00072/) (2024)
- [NAR-*ICP: Neural Execution of Classical ICP-based Pointcloud Registration Algorithms](https://mlsanity.com/paper/2410.11031/) (2024)
- [A Hybrid System for Systematic Generalization in Simple Arithmetic Problems](https://mlsanity.com/paper/2306.17249/) (2023)
- [On the Minimal Degree Bias in Generalization on the Unseen for non-Boolean Functions](https://mlsanity.com/paper/2406.06354/) (2024)
- [Beyond Interpolation: Extrapolative Reasoning with Reinforcement Learning and Graph Neural Networks](https://mlsanity.com/paper/2502.04402/) (2025)
- [KNARsack: Teaching Neural Algorithmic Reasoners to Solve Pseudo-Polynomial Problems](https://mlsanity.com/paper/2509.15239/) (2025)
- [FloydNet: A Learning Paradigm for Global Relational Reasoning](https://mlsanity.com/paper/2601.19094/) (2026)
- [Algorithm-Informed Graph Neural Networks for Leakage Detection and Localization in Water Distribution Networks](https://mlsanity.com/paper/2408.02797/) (2024)
- [Tackling GNARLy Problems: Graph Neural Algorithmic Reasoning Reimagined through Reinforcement Learning](https://mlsanity.com/paper/2509.18930/) (2025)
- [Is Code Better Than Language for Algorithmic Reasoning](https://mlsanity.com/paper/2606.15589/) (2026)
- [Fitting Every Hint Is Not Enough: Premature Floods in Recurrent Algorithmic Reasoners](https://mlsanity.com/paper/2610.14982/) (2026)

## Abstract (arXiv)

In this paper, we study the OOD generalization of neural algorithmic reasoning tasks, where the goal is to learn an algorithm (e.g., sorting, breadth-first search, and depth-first search) from input-output pairs using deep neural networks. First, we argue that OOD generalization in this setting is significantly different than common OOD settings. For example, some phenomena in OOD generalization of image classifications such as accuracy on the line are not observed here, and techniques such as data augmentation methods do not help as assumptions underlying many augmentation techniques are often violated. Second, we analyze the main challenges (e.g., input distribution shift, non-representative data generation, and uninformative validation metrics) of the current leading benchmark, i.e., CLRS citepdeepmind2021clrs, which contains 30 algorithmic reasoning tasks. We propose several solutions, including a simple-yet-effective fix to the input distribution shift and improved data generation. Finally, we propose an attention-based 2WL-graph neural network (GNN) processor which complements message-passing GNNs so their combination outperforms the state-of-the-art model by a 3% margin averaged over all algorithms. Our code is available at: urlhttps://github.com/smahdavi4/clrs.