PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Pre-training of Large Autoregressive Image Models

Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, Armand Joulin

arXiv:2401.08541v1cs.CV

TL;DR

The paper asks whether autoregressive pre-training can produce visual representations that scale with model capacity and data, as in language modeling. It introduces AIM, adapting the approach with prefix attention and a prediction head for visual features. AIM’s features improve with capacity and training data, its objective correlates with downstream quality, and performance shows no sign of saturation in the reported scaling study.

  • Problem

    It is unclear whether autoregressive Transformer scaling in language modeling also yields competitive visual representations with similar scaling behavior.

    Source: §1 · pp. 1–2

  • Method

    AIM pre-trains Vision Transformers autoregressively on image patches, using prefix attention and a heavily parameterized token-level prediction head.

    Source: §1 · p. 2, §4.2 · p. 4

  • Results

    AIM features improve with model capacity and training data, and validation objective value correlates with downstream frozen-feature quality, with no observed saturation as image count increases.

    Source: §1 · p. 2

  • Takeaways & Limitations

    AIM scales to 7 billion parameters with a vanilla Transformer without stability-inducing techniques or extensive scale-specific hyperparameter adjustments.

    Source: §6 · p. 9

    Contrastive methods currently yield stronger representations at a given model size than generative methods such as AIM, but pose scalability and loss-tractability challenges.

    Source: §6 · p. 9

Abstract

from arXiv · show

This paper introduces AIM, a collection of vision models pre-trained with an autoregressive objective. These models are inspired by their textual counterparts, i.e., Large Language Models (LLMs), and exhibit similar scaling properties. Specifically, we highlight two key findings: (1) the performance of the visual features scale with both the model capacity and the quantity of data, (2) the value of the objective function correlates with the performance of the model on downstream tasks. We illustrate the practical implication of these findings by pre-training a 7 billion parameter AIM on 2 billion images, that achieves 84.0% on ImageNet-1k with a frozen trunk. Interestingly, even at this scale, we observe no sign of saturation in performance, suggesting that AIM potentially represents a new frontier for training large-scale vision models. The pre-training of AIM is similar to the pre-training of LLMs, and does not require any image-specific strategy to stabilize the training at scale.

1. Introduction

The paper asks whether autoregressive Transformers can scale for visual representation learning as they do for language. AIM applies this approach to images, and its features improve with model capacity and training data while objective value correlates with downstream performance.

1. Introduction

  • The paper tests whether autoregressive Transformer scaling, established in language modeling, also yields competitive visual representations.The motivation is that neither autoregressive objectives nor Transformers are inherently specific to language.Source: §1 · pp. 1–2
  • AIM uses prefix attention to permit bidirectional downstream attention and a heavily parameterized token-level prediction head that improves feature quality with little training overhead.The approach does not rely on stability-inducing techniques used by supervised or self-supervised methods.Source: §1 · p. 2
  • AIM features improve with model capacity and training data, while validation-objective value correlates with downstream feature quality.The authors report continued improvement with more images and no sign of saturation.Source: §1 · p. 2

2. Related Work

Prior image autoregression explored image-adapted architectures and Transformers, while newer work scales diverse self-supervised methods. The paper positions AIM against these approaches, emphasizing its observed benefit from scaling data and model size.

2. Related Work

  • Image autoregressive modeling has received less attention than language and speech autoregression [56] [57].The cited image work includes [68] [76].Source: §2 · p. 2
  • Earlier image autoregressive work found image-adapted architectures improved over generic ones, and later work adopted Transformers for better autoregressive models [76] [77] [79].Chen et al. reported continuous improvements with more compute; AIM builds on this line with more data and further architecture and training improvements.Source: §2 · p. 2
  • DINOv2 scales iBot [88] on 142M private images with a 460M-parameter model, while the cited MAE scaling study did not show strong self-supervised gains at billions of images [70].The authors contrast these reports with AIM’s observed feature-quality benefit from scaling to billions of parameters and images.Source: §2 · pp. 2–3

3. Pre-training Dataset

AIM is pre-trained on DFN-2B, a top-ranked subset of a larger image-text collection, with limited content-based curation. Training mixes this data with ImageNet-1k.

3. Pre-training Dataset

  • DFN-2B contains the top 15% of samples from DataComp’s 12.8B image-text pairs, after privacy and safety filters and deduplication against evaluation sets.The selection adds no curation based on image content; because AIM does not require text, the authors note it could use larger unpaired or low-alignment image collections.Source: §3 · p. 3
  • The DFN-2B+ training mixture samples DFN-2B with probability 0.8 and ImageNet-1k with probability 0.2.The authors motivate this mix by the LLM-pretraining practice of oversampling high-quality data sources [75].Source: §3 · p. 3

4. Approach

AIM trains on raster-ordered image patches by predicting each patch from preceding context, using a pixel-level objective. Its architecture adapts causal pre-training for downstream transfer while keeping the trunk fixed during evaluation.

4.1. Training Objective

  • AIM factorizes each fixed-length image patch sequence into conditional predictions, using raster order and preceding patches as context.The default training loss is normalized pixel-level squared error between predicted and ground-truth patches.Source: §4.1 · p. 3
  • Discrete-token cross-entropy works, but produces weaker features than pixel-wise loss.The discrete patches are created with an offline tokenizer.Source: §4.1 · p. 3

4.2. Architecture

  • Prefix attention makes the initial patches bidirectional context and computes loss only on the remaining patches, enabling mask removal at downstream adaptation.The prefix length is sampled uniformly during pre-training; removing the causal mask improves downstream performance.Source: Figure 3 · p. 4, §4.2 · p. 4
  • For AIM variants with at least 1B parameters, pre-training uses 1.2M iterations, corresponding to 1.2 trillion patches or 5B images.Table 1 also reports model dimensions and pre-training learning rate and batch size.Source: Table 1 · p. 4
  • AIM uses a ViT trunk with independently applied MLP prediction blocks, and scales capacity primarily by increasing width.The MLP head is discarded for transfer; the approach uses shared optimization hyperparameters across model sizes without requiring listed stability mechanisms.Source: §4.2 · p. 4
  • Downstream evaluation freezes all model weights and trains only a classification head, using attention pooling to form a global feature from patch features.The attentive probe retains a low additional parameter count and reduced overfitting risk relative to linear probing’s stated advantages.Source: §4.2 · p. 5

5. Results

AIM scales with data and model capacity, and its validation objective tracks downstream feature quality. Across architecture, objective, and evaluation comparisons, AIM shows strong frozen-trunk performance, with further gains from LoRA fine-tuning.

5.1. Impact of scaling

  • DFN-2B+—mostly DFN-2B with a small amount of IN-1k—achieves the best downstream performance, while IN-1k alone overfits and DFN-2B alone avoids overfitting but converges similarly amid distribution shift.The dataset comparison covers downstream performance across 15 benchmarks.Source: §5.1 · p. 5, Table 2 · p. 6, Figure 5 · p. 6
  • Extending pre-training from 2B to 5B images lowers validation loss, and lower-capacity, longer-trained models can match higher-capacity, shorter-trained models at similar FLOPs.The authors defer further investigation of the suggested scaling laws.Source: §5.1 · p. 6

5.2. Architecture and Design

  • Normalized-pixel MSE and raster traversal yield the strongest results among tested targets and traversal patterns; more uniform patch-prediction difficulty accompanies better models.The authors attribute the pattern result to predicting patches throughout the sequence requiring the model to retain more image information.Source: §5.2 · p. 6
  • Prefix attention supports effective transfer with either causal or bidirectional attention, with the best reported performance using bidirectional attention downstream.Causal-attention pre-training is effective for transfer only when the causal mask is retained downstream.Source: §5.2 · p. 7

5.3. Pre-training objective

  • With other architectural and optimization components fixed, AIM’s autoregressive objective outperforms a masking objective on frozen-trunk IN-1k performance.The masking baseline replaces randomly selected patches with learnable mask tokens.Source: Table 5 · p. 7, §5.3 · p. 8

5.2. Architecture and Design

  • An MLP prediction head improves over no head, and greater MLP depth or width improves downstream performance; attentive pooling also boosts performance across ablations.A transformer head yields only marginal improvement over the MLP at substantially higher computational cost.Source: §5.2 · pp. 7–8, Table 4 · p. 7

5.4. Comparison with other methods

  • AIM performs strongly across 15 frozen-trunk attentive-probing benchmarks, outperforming other generative and joint-embedding methods except DINOv2, which uses higher-resolution inputs.AIM also outperforms DINO and iBOT on average across benchmarks.Source: Table 6 · p. 8, §5.4 · p. 9

5.5. Low-Rank Adaptation

  • LoRA fine-tuning improves AIM-7B by 3.9% over its last-layer frozen-trunk performance while tuning only 0.1% of trunk parameters.The paper reports that LoRA is compatible with AIM and provides large gains over frozen-trunk evaluation.Source: §5.5 · p. 9, Table 8 · p. 9

6. Discussion

AIM scales to large vision models and performs strongly across recognition benchmarks, while the authors note trade-offs against alternatives. They report no clear saturation with increased model or data scale.

6. Discussion

  • AIM performs strongly across 15 recognition benchmarks, outperforming prior state-of-the-art methods such as MAE and narrowing the gap with joint-embedding pre-training.The authors report no clear signs of saturation as model parameters or data scale.Source: §6 · p. 9
  • AIM scales to 7 billion parameters with a vanilla transformer, without stability-inducing techniques or extensive scale-specific hyperparameter adjustments.Source: §6 · p. 9
  • AIM’s pre-training-task performance strongly correlates with downstream performance.Source: §6 · p. 9
  • The authors see no clear saturation as parameters or data increase, suggesting potential for further gains with larger models and longer training schedules.Source: §6 · p. 9
  • MAE can learn good representations with little pre-training data, while contrastive methods yield stronger representations at a given model size but are harder to scale.The authors also note that MAE’s sample efficiency reduces overfitting risk relative to AIM.Source: §6 · p. 9

A. Datasets

AIM’s representations are evaluated on 15 recognition benchmarks spanning varied image domains and tasks. The benchmark table provides references, train and test image counts, and category counts.

A. Datasets

  • AIM is evaluated on 15 recognition benchmarks covering fine-grained recognition, medical imaging, satellite imagery, natural environments, and infographic images.Source: A. Datasets · p. 13
  • The benchmark table lists each dataset’s references, train and test image counts, and number of categories.Source: Table 9 · p. 13

B. Autoregression Patterns

The authors investigate alternative patterns for traversing images during pre-training and point to a figure illustrating those patterns.

B. Autoregression Patterns

C.1. Raster pattern validation loss

The raster-pattern validation loss dips at the second chunk before rising again, a behavior attributed to the ImageNet-1k validation set. The authors observe lower loss in most top rows for both regular and vertically flipped images.

C.1. Raster pattern validation loss

  • The second-chunk dip in raster-pattern validation loss is a side effect of using the ImageNet-1k validation set.Source: §C.1 · p. 13
  • Most top rows except the first have lower loss, whether computed on regular images or their vertically flipped counterparts.Source: §C.1 · p. 13

C.2. Downstream performance across layers

AIM’s strongest downstream features come from intermediate rather than final layers, peaking roughly two-thirds into model depth. Deeper layers still transfer strongly without a sharp performance decline.

C.2. Downstream performance across layers

  • Tables 7 and 6 report downstream gains from probing shallower layers rather than only the last layer.Source: §C.2 · p. 13
  • Peak downstream transfer features occur at roughly two-thirds of model depth across all AIM variants.The strongest features are extracted from layers other than the last.Source: §C.2 · p. 13
  • Deeper layers retain strong downstream performance, with no steep decline.Source: §C.2 · p. 13

D. Hyperparameters

AIM capacities share pre-training hyperparameters, though AIM-0.6B used a shorter schedule. The section also specifies attentive-probe settings and LoRA adaptation choices.

D. Hyperparameters

  • Pre-training uses AdamW with β1 = 0.9, β2 = 0.95, and a peak learning rate of 1e−3.Source: D. Hyperparameters · p. 14
  • The listed image preprocessing uses 224px sizing, scale [0.4, 1.0], ratio [0.75, 1.33], bicubic interpolation, and random horizontal flips.Source: D. Hyperparameters · p. 14
  • AIM variants use the same pre-training hyperparameters, except AIM-0.6B was trained for 500k iterations; scaling required no extra optimization tuning.No instability was observed while scaling model capacity.Source: D. Hyperparameters · pp. 13–14, Table 10 · p. 14
  • Attentive probing searches learning rates for AIM and baselines and averages the last six layers for generative models.The best learning-rate performance is reported for each method.Source: D. Hyperparameters · p. 14, Table 11 · p. 14
  • LoRA adapts rank-8 attention query, value, and output-projection matrices, with mixup alpha=0.8.Source: D. Hyperparameters · p. 14
Loading 2401.08541v1…