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

2024-10-08

Trans. Mach. Learn. Res.

AI Summary

Vision Transformers (ViTs) exhibit poor performance on the Abstraction and Reasoning Corpus (ARC), a structured visual reasoning benchmark, due to their inability to learn implicit input-output image mappings—stemming from inherent representational limitations. To address this, we propose ViTARC: (1) a pixel-level input representation that preserves fine-grained spatial fidelity; (2) a spatially aware patching mechanism that respects geometric continuity; and (3) object-level positional encodings derived from automatic segmentation, explicitly injecting structural inductive biases into the architecture. Evaluated on 400 publicly available ARC tasks, ViTARC achieves near-perfect (≈100%) solution rates on over 50% of tasks—substantially outperforming standard ViTs even under million-sample training regimes. This is the first demonstration that explicit structural priors are critical for few-shot, compositional visual reasoning, establishing a new paradigm for integrating domain knowledge into vision foundation models.

Technology Category

• Computer Vision: Visual Reasoning & Symbolic Representations
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Machine Learning: Deep Neural Architectures and Foundation Models

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Evaluating Vision Transformers' ability to solve ARC visual reasoning tasks
• Addressing ViT's failure in learning structured mappings from ARC input-output pairs
• Proposing ViTARC with 2D-aware enhancements for improved abstract visual reasoning

Innovation

• Pixel-level input representation for ViT
• Spatially-aware tokenization scheme design
• Object-based positional encoding with segmentation