Latent Space Representations of Neural Algorithmic Reasoners

TL;DR

This paper analyzes the structure of latent spaces in neural algorithmic reasoning using GNNs, identifies failure modes, and proposes solutions that improve performance on the CLRS-30 benchmark.

Contribution

It provides a detailed analysis of GNN latent spaces in neural algorithmic reasoning and introduces methods to mitigate identified failure modes.

Findings

Softmax aggregator improves resolution in latent space

Latent space decay helps handle out-of-range values

Proposed methods enhance algorithm performance on CLRS-30

Abstract

Neural Algorithmic Reasoning (NAR) is a research area focused on designing neural architectures that can reliably capture classical computation, usually by learning to execute algorithms. A typical approach is to rely on Graph Neural Network (GNN) architectures, which encode inputs in high-dimensional latent spaces that are repeatedly transformed during the execution of the algorithm. In this work we perform a detailed analysis of the structure of the latent space induced by the GNN when executing algorithms. We identify two possible failure modes: (i) loss of resolution, making it hard to distinguish similar values; (ii) inability to deal with values outside the range observed during training. We propose to solve the first issue by relying on a softmax aggregator, and propose to decay the latent space in order to deal with out-of-range values. We show that these changes lead to improvements on the majority of algorithms in the standard CLRS-30 benchmark when using the state-of-the-art Triplet-GMPNN processor. Our code is available at https://github.com/mirjanic/nar-latent-spaces