Infusing Lattice Symmetry Priors in Attention Mechanisms for Sample-Efficient Abstract Geometric Reasoning

TL;DR

This paper introduces LatFormer, a model that embeds lattice symmetry priors into attention mechanisms, significantly improving sample efficiency and geometric reasoning capabilities in AI tasks like ARC and LARC.

Contribution

The paper proposes a novel attention modification that incorporates lattice symmetry priors, enabling models to better generalize in geometric reasoning tasks.

Findings

LatFormer requires 100 times less data than standard transformers.

Incorporating geometric priors improves performance on synthetic geometric reasoning.

Preliminary results suggest deep learning models can handle complex geometric datasets.

Abstract

The Abstraction and Reasoning Corpus (ARC) (Chollet, 2019) and its most recent language-complete instantiation (LARC) has been postulated as an important step towards general AI. Yet, even state-of-the-art machine learning models struggle to achieve meaningful performance on these problems, falling behind non-learning based approaches. We argue that solving these tasks requires extreme generalization that can only be achieved by proper accounting for core knowledge priors. As a step towards this goal, we focus on geometry priors and introduce LatFormer, a model that incorporates lattice symmetry priors in attention masks. We show that, for any transformation of the hypercubic lattice, there exists a binary attention mask that implements that group action. Hence, our study motivates a modification to the standard attention mechanism, where attention weights are scaled using soft masks generated by a convolutional network. Experiments on synthetic geometric reasoning show that LatFormer requires 2 orders of magnitude fewer data than standard attention and transformers. Moreover, our results on ARC and LARC tasks that incorporate geometric priors provide preliminary evidence that these complex datasets do not lie out of the reach of deep learning models.