Tackling the Abstraction and Reasoning Corpus with Vision Transformers: the Importance of 2D Representation, Positions, and Objects (2410.06405v1)

Abstract: The Abstraction and Reasoning Corpus (ARC) is a popular benchmark focused on visual reasoning in the evaluation of Artificial Intelligence systems. In its original framing, an ARC task requires solving a program synthesis problem over small 2D images using a few input-output training pairs. In this work, we adopt the recently popular data-driven approach to the ARC and ask whether a Vision Transformer (ViT) can learn the implicit mapping, from input image to output image, that underlies the task. We show that a ViT -- otherwise a state-of-the-art model for images -- fails dramatically on most ARC tasks even when trained on one million examples per task. This points to an inherent representational deficiency of the ViT architecture that makes it incapable of uncovering the simple structured mappings underlying the ARC tasks. Building on these insights, we propose ViTARC, a ViT-style architecture that unlocks some of the visual reasoning capabilities required by the ARC. Specifically, we use a pixel-level input representation, design a spatially-aware tokenization scheme, and introduce a novel object-based positional encoding that leverages automatic segmentation, among other enhancements. Our task-specific ViTARC models achieve a test solve rate close to 100% on more than half of the 400 public ARC tasks strictly through supervised learning from input-output grids. This calls attention to the importance of imbuing the powerful (Vision) Transformer with the correct inductive biases for abstract visual reasoning that are critical even when the training data is plentiful and the mapping is noise-free. Hence, ViTARC provides a strong foundation for future research in visual reasoning using transformer-based architectures.

• The paper introduces ViTARC to boost ARC performance from 18% to 75% test accuracy by integrating advanced 2D, positional, and object-based encodings.
• The methodology employs enhanced pixel-level tokenization and a novel Positional Encoding Mixer to overcome limitations of standard Vision Transformers in visual reasoning tasks.
• Experimental results reveal that ViTARC solves nearly 100% of test instances for over 50% of ARC tasks, marking a significant breakthrough in abstract visual reasoning.

Overview of "Tackling the Abstraction and Reasoning Corpus with Vision Transformers"

The paper "Tackling the Abstraction and Reasoning Corpus with Vision Transformers: the Importance of 2D Representation, Positions, and Objects" presents a novel approach using Vision Transformers (ViT) for solving the Abstraction and Reasoning Corpus (ARC) tasks. The ARC benchmark is designed to assess artificial intelligence's capability in abstract visual reasoning, requiring systems to transform small 2D images based on few-shot input-output pairs without the use of textual or heuristic knowledge.

Core Contributions

The authors initially address the drawbacks faced by a vanilla Vision Transformer (ViT) architecture when applied to ARC tasks. Despite the simplicity of the image structures involved, the ViT model fails to effectively learn the implicit mapping between input and output images due to inherent architectural limitations, demonstrating a mere 18% test accuracy. This suggests that ViT's standard methods do not adequately capture the necessary spatial and structural mappings required for complex visual reasoning tasks.

To overcome these limitations, the paper introduces ViTARC, a modified ViT-style architecture that incorporates specific design adaptations enhancing its ability to reason visually:

1. 2D Visual Representation: The authors introduce a pixel-level input representation and augment it with a spatially-aware tokenization scheme involving visual tokens, including new 2D padding techniques and border tokens. This significantly improves the test solve rate from 18% to approximately 66%.
2. Enhanced Positional Encoding: Building on the insights regarding spatial deficiencies, the paper implements a combination of absolute, relative, and object positional encoding strategies, which are further refined with a novel Positional Encoding Mixer (PEmixer). The enhancements boost the test accuracy to 75%.
3. Object-based Encoding: This involves leveraging external segmentation techniques to assign object indices, which increases spatial awareness crucial for complex object reasoning tasks.

Experimental Evaluation

The authors conduct comprehensive experiments on the entire ARC benchmark, generating one million samples per task using the RE-ARC generator. They robustly analyze failure cases of the vanilla ViT and demonstrate how the proposed enhancements within ViTARC lead to improved performance. The final model successfully solves nearly 100% of test instances for over 50% of ARC tasks, indicating a substantial advancement over existing architectures.

Implications and Future Directions

This work underscores the critical need for incorporating effective inductive biases in Transformer architectures to succeed in abstract visual reasoning tasks. The emphasis on 2D spatial configurations and positional encoding suggests a broader paradigm where structural information can be further exploited for tasks beyond ARC.

For future developments, this research paves the way for creating more generalized AI systems capable of adaptive reasoning across various problem domains. Enhancing generalization capabilities to solve ARC tasks under fewer data conditions or different abstractions remains an open research area. Moreover, applying similar methods to other vision-based tasks could inform ongoing developments in AI reliability and flexibility.

In conclusion, the paper presents a rigorous exploration into the application of Transformer models in visual reasoning, offering valuable insights and a strong foundation for further research in AI-driven abstract reasoning tasks.