Triplet Edge Attention for Algorithmic Reasoning
This work addresses the problem of improving generalization in neural algorithmic reasoning for researchers in graph neural networks and algorithmic learning.
The paper tackles the challenge of developing expressive graph neural networks for neural algorithmic reasoning by introducing Triplet Edge Attention (TEA), an edge-aware graph attention layer, achieving a 5% average improvement on the CLRS benchmark and a 30% improvement for string algorithms.
This work investigates neural algorithmic reasoning to develop neural networks capable of learning from classical algorithms. The main challenge is to develop graph neural networks that are expressive enough to predict the given algorithm outputs while generalizing well to out-of-distribution data. In this work, we introduce a new graph neural network layer called Triplet Edge Attention (TEA), an edge-aware graph attention layer. Our algorithm works by precisely computing edge latent, aggregating multiple triplet messages using edge-based attention. We empirically validate our TEA layer in the CLRS benchmark and demonstrate a $5%$ improvement on average. In particular, we achieve a $30%$ improvement for the string algorithms compared to the state-of-the-art model.