Sample- and Parameter-Efficient Auto-Regressive Image Models

TL;DR

XTRA introduces a block-based auto-regressive vision model that significantly improves sample and parameter efficiency, achieving state-of-the-art performance on multiple image recognition benchmarks with fewer samples and parameters.

Contribution

The paper proposes a novel block causal masking approach in auto-regressive vision models, enhancing their efficiency and scalability over previous methods.

Findings

XTRA surpasses previous models on 15 image benchmarks.

XTRA uses 152× fewer samples than prior models.

XTRA achieves comparable or better performance with 7-16× fewer parameters.

Abstract

We introduce XTRA, a vision model pre-trained with a novel auto-regressive objective that significantly enhances both sample and parameter efficiency compared to previous auto-regressive image models. Unlike contrastive or masked image modeling methods, which have not been demonstrated as having consistent scaling behavior on unbalanced internet data, auto-regressive vision models exhibit scalable and promising performance as model and dataset size increase. In contrast to standard auto-regressive models, XTRA employs a Block Causal Mask, where each Block represents k $\times$ k tokens rather than relying on a standard causal mask. By reconstructing pixel values block by block, XTRA captures higher-level structural patterns over larger image regions. Predicting on blocks allows the model to learn relationships across broader areas of pixels, enabling more abstract and semantically meaningful representations than traditional next-token prediction. This simple modification yields two key results. First, XTRA is sample-efficient. Despite being trained on 152$\times$ fewer samples (13.1M vs. 2B), XTRA ViT-H/14 surpasses the top-1 average accuracy of the previous state-of-the-art auto-regressive model across 15 diverse image recognition benchmarks. Second, XTRA is parameter-efficient. Compared to auto-regressive models trained on ImageNet-1k, XTRA ViT-B/16 outperforms in linear and attentive probing tasks, using 7-16$\times$ fewer parameters (85M vs. 1.36B/0.63B).