Task embeddings improve rule discovery in abstraction reasoning tests
Implicit Rule Induction with Test-Time Task Embeddings in ARC-like Tasks
Artificial IntelligenceMachine Learning
Summary
Some AI tests check if computers can figure out rules behind puzzles, but it’s hard to tell if the AI really understands the rule or just guesses shortcuts. The authors studied a method that fine-tunes a small, special part of the AI (called a task embedding) first, then adjusts the main model. This method helps the AI better capture the real rules behind the puzzles, making it easier to match new tasks to known rules and improving its ability to solve new problems. They also found that the AI learns to combine and interpolate rules well, though it struggles to extend rules beyond what it has seen.
What this means in practice
- •For machine learning engineers: Improve AI models’ ability to discover underlying rules in new visual reasoning tasks by fine-tuning lightweight task embeddings before full model adjustment.
- •For automated reasoning developers: Enhance system capabilities to match and solve structured reasoning problems through embedding-based retrieval aligned with task semantics.
Authors
Adrien Deliège, Claas Beger, Marc Van Droogenbroeck, Melanie Mitchell
Abstract
The Abstraction and Reasoning Corpus and related benchmarks evaluate whether AI models can solve novel reasoning tasks, but often leave unclear whether success reflects inference of the intended underlying rule or reliance on shortcuts. We address this gap by studying test-time task embeddings in Vision ARC (VARC), a model in which a pre-trained backbone is complemented by a trainable embedding representing the transformation rule. In the original VARC, test-time training (TTT) is jointly applied to the backbone and task embedding. Here we introduce a novel two-step TTT protocol: first finetune only the task embedding (Embed-TTT), then freeze it and finetune the backbone. Across ARC-AGI-1, ConceptARC, and two controlled datasets with known rules, Embed-TTT consistently yields improved task embeddings, ones that align better with underlying task rules, improve embedding-based retrieval, and enable accurate linear probing of known rules. Qualitatively, Embed-TTT identifies more semantically meaningful relations between test and train tasks on ARC-AGI-1. We also show that optimizing only task embeddings (less than 0.01% of model parameters) already solves a non-trivial fraction of ARC-AGI-1, ConceptARC, and Mini-ARC tasks, while the full two-step pipeline improves final performance. Finally, we show that Embed-TTT recovers the underlying geometric structure of parametric rules and learns compositional capabilities that enable rule-wise interpolation, but not extrapolation. These findings support a clearer separation between rule induction and rule execution in ARC-like evaluations, motivating benchmarks that better distinguish in-distribution from out-of-distribution rules.