# COGITAO: A Visual Reasoning Framework To Study Compositionality & Generalization

arXiv: https://arxiv.org/abs/2509.05249 · HTML: https://mlsanity.com/paper/2509.05249/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Multimodal reasoning, one of 17 directions in LLM reasoning.
- Builds on 25 of its 72 references that are in the map.
- First posted on arXiv 5 Sep 2025.

## Placed in

- LLM reasoning / [Multimodal reasoning](https://mlsanity.com/lr/multimodal-reasoning/)

## Builds on (most cited first, 25 of 25)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [MONet: Unsupervised Scene Decomposition and Representation](https://mlsanity.com/paper/1901.11390/) (2019)
- [Learning Deep Transformer Models for Machine Translation](https://mlsanity.com/paper/1906.01787/) (2019)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [On the Measure of Intelligence](https://mlsanity.com/paper/1911.01547/) (2019)
- [CATER: A diagnostic dataset for Compositional Actions and TEmporal Reasoning](https://mlsanity.com/paper/1910.04744/) (2019)
- [The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers](https://mlsanity.com/paper/2108.12284/) (2021)
- [Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?](https://mlsanity.com/paper/2410.08292/) (2024)
- [Deep Learning using Rectified Linear Units (ReLU)](https://mlsanity.com/paper/1803.08375/) (2018)
- [Provable Compositional Generalization for Object-Centric Learning](https://mlsanity.com/paper/2310.05327/) (2023)
- [On the generalization of language models from in-context learning and finetuning: a controlled study](https://mlsanity.com/paper/2505.00661/) (2025)
- [Addressing the Abstraction and Reasoning Corpus via Procedural Example Generation](https://mlsanity.com/paper/2404.07353/) (2024)
- [Tackling the Abstraction and Reasoning Corpus with Vision Transformers: the Importance of 2D Representation, Positions, and Objects](https://mlsanity.com/paper/2410.06405/) (2024)
- [Scaling can lead to compositional generalization](https://mlsanity.com/paper/2507.07207/) (2025)
- [Compositional-ARC: Assessing Systematic Generalization in Abstract Spatial Reasoning](https://mlsanity.com/paper/2504.01445/) (2025)
- [Systematic Generalization: What Is Required and Can It Be Learned?](https://mlsanity.com/paper/1811.12889/) (2018)

## Abstract (arXiv)

The ability to compose learned concepts and apply them in novel settings is key to human intelligence, but remains a persistent limitation in state-of-the-art machine learning models. To address this issue, we introduce COGITAO, a modular and extensible data generation framework and benchmark designed to systematically study compositionality and generalization in visual domains. Drawing inspiration from ARC-AGI's problem-setting, COGITAO constructs rule-based tasks which apply a set of transformations to objects in grid-like environments. It supports composition, at adjustable depth, over a set of 28 interoperable transformations, along with extensive control over grid parametrization and object properties. This flexibility enables the creation of millions of unique task rules -- surpassing concurrent datasets by several orders of magnitude -- across a wide range of difficulties, while allowing virtually unlimited sample generation per rule. We provide baseline experiments using state-of-the-art vision models, highlighting their consistent failures to generalize to novel combinations of familiar elements, despite strong in-domain performance. COGITAO is fully open-sourced, including all code and datasets, to support continued research in this field.