COGITAO: A Visual Reasoning Framework To Study Compositionality & Generalization

TL;DR

COGITAO is an ARC-AGI-inspired data generator and benchmark that isolates compositional generalization in the visual domain from confounding visual complexity. It applies sequences of 28 object-centric atomic transformations (translate, rotate, mirror, crop, recolor, fill, extend contour, pad, duplicate) at arbitrary composition depth to procedurally generated objects on grids, yielding millions of unique task rules with effectively unlimited samples per rule. The authors define two study axes — CompGen (unseen compositions of familiar atomics) and EnvGen (unseen object counts, grid sizes, object sizes, object complexity) — and benchmark Vanilla ViT, a Grid ViT with ARC-tuned inductive biases, and the LLaDA diffusion language model, all at ~10M parameters. All three models score well in-domain but degrade sharply out-of-domain on compositional splits, reinforcing that current vision architectures still don’t generalize to novel combinations of learned primitives even when the primitives themselves are individually mastered.

Key claims

COGITAO’s 28 atomic transformations are designed to be mutually composable and mutually non-redundant, giving a theoretical task count of ~4.9×10⁹ at depth 6 [§3.1.2].
Tasks are input-output grid pairs conditioned on a token-level transformation embedding (e.g. [M, R] for mirror-then-rotate), so the CompGen study can present novel sequences of tokens the model has each seen individually [§4.1].
The CompGen study defines three splits: C1 (train atomics + composites → test unseen composites), C2 (train restricted composites only → test unseen composites), C3 (train depth-2 → test depth-3) [§4.1.1].
The EnvGen study defines five splits varying object count (G1), grid size (G2), object size (G3), object complexity/symmetry/color (G4), and all axes combined (G5), all with fixed transformation sequences [§4.1.2].
Grid ViT — a ViT with object positional encoding, PEMixer modules, register tokens, and modified positional encodings from ViTARC — achieves the strongest in-domain grid accuracy across nearly all settings [§4.4, Table 1].
LLaDA (diffusion LM) achieves the highest OOD accuracy on several compositional splits (C1, G4, G5), suggesting the diffusion-LM substrate has an OOD-generalization edge over both ViTs on this benchmark [§4.4, Table 1].
Out-of-domain grid accuracy collapses on the hardest splits — C3 drops to 6.4% for Vanilla ViT, G3 to 27.2%, G5 to 14.0% — despite strong in-domain scores, so “novel-combination generalization” remains an open problem even with ARC-tuned inductive biases [§4.4, Table 1].
The benchmark is deliberately trained from scratch at ~10M parameters and excludes foundation models — the authors frame COGITAO as measuring whether architectures compose, not whether large-scale pretraining papers over the failure [§2].
The paper’s design axis — decoupling compositional structure from visual complexity by using handcrafted objects on grids — is the same methodological bet as the ARC-AGI series and CVR, and explicitly avoids the conflation that CATER-style real-world benchmarks introduce [§2 Related Works].

Method

Object generation pre-samples 23,000 distinct objects by iterating over size, symmetry, connectivity, color palette, color pattern (stripes, uniform, hollow, random), and footprint (rectangle, disk, diamond, ellipse). A SetInitialGrid() routine positions objects on a grid subject to each transformation’s per-object constraints (some transformations require asymmetric or hole-bearing objects). sampleTransformations() builds the composition sequence, transform_and_position() applies it to each object and re-lays them into the grid, discarding samples where the output collides or leaves bounds. The 28 atomics are Translate (4 dirs), Mirror (2 axes), Rotate90, Crop (5 variants), ChangeColor, Filling (same/different color), Emptying, ExtendContour (same/different color), Padding (5 variants), Duplicate (5 variants). Models: (a) Vanilla ViT with learned absolute PE, (b) Grid ViT — a ViTARC-derived ViT with object-positional encoding, PEMixer, register tokens (à la Darcet), and modified RoPE-style PE, (c) LLaDA, a diffusion-based language model. All ~10M params, trained from scratch with AdamW + 200-step warmup + cosine anneal, and a 2-layer MLP grid-token classification head.

Results

Vanilla ViT: 71.1% ID on C1 → OOD 26.2%; 69.0% → 18.6% on C2; 78.9% → 6.4% on C3 [Table 1].
Grid ViT beats Vanilla ViT on nearly every ID setting and most OOD settings, quantifying the value of ARC-style inductive biases (OPE, register tokens, PEMixer) — though it still collapses on the hardest OOD splits [Table 1].
LLaDA (diffusion LM) is on par with or beats Grid ViT on OOD for C1, G4, and G5 despite worse ID performance, suggesting a substrate-level rather than inductive-bias-level effect [Table 1, §4.4].
Environmental splits are also hard: G3 (larger objects) drops to 27.2%, G5 (all axes combined) to 14.0% [Table 1].
Best in-domain grid accuracy of ~99% (G1, G2 with Grid ViT) coexists with OOD collapses in the same experiments — a direct empirical restatement of the “high ID, low OOD” failure mode compositional-generalization benchmarks are designed to expose [Table 1].

Why it’s interesting

This is the visual-domain analogue of what ARC-AGI-3: A New Challenge for Frontier Agentic Intelligence does at the agentic-interaction level: a benchmark constructed specifically so that dense LRM-style task-space coverage can’t shortcut compositional generalization. Both papers make the same critique — that ARC-AGI-1/2’s public sets were substantially memorization-attackable — but COGITAO addresses it by procedurally generating millions of composable tasks rather than by holding out a private set. Complements Less is More: Recursive Reasoning with Tiny Networks (TRM / recursive reasoning on ARC) and The Illusion of Thinking: Understanding the Strengths and Limitations of Reasoning Models via the Lens of Problem Complexity on the “reasoning collapses beyond a compositional-complexity threshold” axis — COGITAO’s Vanilla-ViT-6.4% on depth-3 composites is a small-model, from-scratch confirmation of that pattern in a different data regime. The LLaDA-vs-ViT OOD result is worth noting because it’s a rare head-to-head where a diffusion-LM substrate wins on generalization at fixed parameter budget — a small datapoint in the same direction as Autoregressive-to-Diffusion Vision Language Models (A2D-VL) and the broader diffusion-LM cluster in Diffusion Language Models.