TLDR: COGITAO is a new open-source visual reasoning framework and benchmark designed to systematically study how AI models compose learned concepts and generalize them to novel visual tasks. It generates rule-based tasks involving 28 composable transformations on objects in grid environments, allowing for millions of unique tasks and infinite samples. Baseline experiments with state-of-the-art vision models show consistent failure to generalize to new combinations or environmental variations, despite strong in-domain performance, highlighting current AI’s reliance on pattern recognition over true compositional understanding.
In the quest to develop truly intelligent machines, researchers are constantly striving to overcome a significant hurdle: the ability of artificial intelligence (AI) to compose learned concepts and apply them effectively in new, unfamiliar situations. This fundamental human capability, known as compositional generalization, remains a persistent challenge for even the most advanced machine learning models. To tackle this, a team of researchers from the Institute of Neuroinformatics at the University of Zurich and ETH Zurich, ETH Zurich, Centre for Artificial Intelligence at Zurich University of Applied Sciences, and RW AI AG have introduced a novel framework called COGITAO.
What is COGITAO?
COGITAO, which stands for COmpositional Generalization In Transformations And Objects, is a modular and extensible framework designed to systematically investigate how AI models handle compositionality and generalization in visual tasks. Inspired by the Abstract and Reasoning Corpus (ARC-AGI), COGITAO generates rule-based tasks where a series of transformations are applied to objects within grid-like environments. This framework offers unparalleled flexibility, supporting composition at various depths using a set of 28 interoperable transformations. It also provides extensive control over grid parameters and object properties, allowing for the creation of millions of unique task rules and virtually unlimited sample generation per rule.
How COGITAO Works
The core of COGITAO is its Python-based generator. It operates in a grid-like environment, similar to ARC-AGI, with pixels numbered from 0 to 9. The process involves three main components:
- Object Generation: Instead of fixed shapes, COGITAO pre-generates a vast library of 23,000 distinct objects. These objects vary in size, symmetry, connectivity, colors, color patterns, and footprints (like rectangles, disks, diamonds, or ellipses). This diversity is crucial for assessing generalization capabilities.
- Transformations: The framework includes 28 atomic object transformations. These transformations are designed to be composable with each other and to modify objects in unique ways that cannot be replicated by combining other transformations. Examples include translating objects (up, down, left, right), mirroring (horizontal, vertical), rotating by 90 degrees, cropping sides, changing colors, filling or emptying object interiors, extending contours, and padding or duplicating objects. The generator can combine these transformations into sequences of arbitrary depth, leading to an enormous number of possible tasks.
- Task Generation: The generator randomly samples objects and positions them on an input grid. Then, a chosen sequence of transformations is applied sequentially to each object to produce an output grid. The system ensures that objects remain within grid boundaries and do not collide during transformations, discarding invalid generations.
Studying Compositional and Environmental Generalization
The researchers used COGITAO to conduct two main types of studies:
- Compositional Generalization (CompGen): This study focuses on how well models can generalize to new combinations of transformations. For instance, models might be trained on individual transformations and simple two-step compositions, then tested on entirely new two-step compositions or even deeper sequences of transformations using the same atomic operations.
- Environmental Generalization (EnvGen): This study assesses a model’s ability to apply learned transformations in varied environments. This includes testing generalization to a greater number of objects, larger grid sizes, larger object dimensions, more complex object features (asymmetric, multi-colored), or a combination of all these challenges.
Baseline Experiments and Key Findings
The paper presents baseline experiments using state-of-the-art vision models, including a Vanilla Vision Transformer (ViT), a specialized Grid ViT (which incorporates biases for grid-based tasks), and LLaDA (a diffusion-based language model). The results consistently highlight a significant limitation: while these models often perform well on tasks similar to their training data (in-domain), their performance drops dramatically when faced with novel combinations of familiar elements or changes in the environment (out-of-domain). For example, in compositional tasks, accuracies plummeted to as low as 5.1% for Grid ViT and 6.4% for LLaDA when encountering unseen transformation sequences. This suggests that current vision models primarily rely on pattern recognition rather than developing a true compositional understanding.
Also Read:
- Evaluating AI: Bridging the Gap Between Benchmarks and Human Understanding
- Structuring Intelligence: Language Models Crafting Hierarchical Learning Environments for AI Agents
The Path Forward
COGITAO provides a controlled and flexible environment to diagnose these fundamental limitations in AI. By requiring models to generate output grids rather than just classify, it significantly raises the bar for evaluation. The framework is fully open-sourced, including all code and datasets, encouraging further research. Future directions include exploring in-context learning, curriculum learning, and analyzing internal model representations to understand how models develop object-centric or transformation-centric abstractions. This work is a crucial step towards building AI systems that can truly compose and generalize like humans. You can find the full research paper here: COGITAO: A Visual Reasoning Framework To Study Compositionality & Generalization.
