AIM: predict patches to learn features

An interactive field guide

Continuous autoregression can be a pretraining task even when image generation is not the main evaluation.

Before this article: Attention and autoregression · ViT · Semantic representations

What was already possible

ViT normally reads all image patches to classify an image. Language models learn useful features by predicting the next token. AIM applies next-patch prediction to visual pretraining, using continuous pixels rather than a learned discrete image vocabulary.

What changes, and why

Order patches in raster order and use causal attention. A per-patch MLP predicts normalized pixels of the next patch with squared error. Prefix attention lets an initial set of observed patches interact bidirectionally; only later patches carry the autoregressive target. This exposes the backbone to bidirectional context before downstream use.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

1. Raster patch sequence[B,N,P²C] → [B,N,D]2. Prefix-causal backbone[B,N,D]3. Next-patch headh at i → pixels at i+1

Raster patch sequence

Continuous pixels supply the target, without VQ.

tokens = project(patchify(image))+sinusoidal_position

Prefix-causal backbone

The prefix is known context and is excluded from the AR target.

h = transformer(tokens,mask=prefix_causal_mask)

Next-patch head

Shift the target by one patch and normalize target pixels per patch.

prediction = mlp(h + head_position)
loss = mse(prediction[:,:-1],normalized_patches[:,1:],target_mask)

Inspect the pinned source

$$\mathcal{L}=\sum_{i\notin\mathrm{prefix}}\|f_\theta(x_{<i})-\operatorname{norm}(x_i)\|^2$$

Work through the mechanism

Select causal mode in the demo. Row i can use keys up to i, and its output predicts the next patch. Switch to prefix mode: the first three positions can see one another, while later positions keep their causal boundary. Known prefix pixels are inputs, not targets that leaked backward from an unknown future.

With a fixed Gaussian output variance, minimizing squared error is equivalent up to constants to maximizing Gaussian conditional likelihood. When several futures are plausible, the best squared-error estimate can be their mean. This may still create useful features even if the mean patch looks poor.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Query position i

3

0 causal · 1 prefix · 2 masked

0

Reset

Rows are queries; columns are keys. In masked mode, only even-indexed positions contain observed values.

Computed values for the current example

Quantity	Value
Selected row	3
Visible keys	4

Training is not sampling

After pretraining, discard the prediction head and evaluate the visual backbone. AIM uses an attentive probe with trainable attention pooling, so it is not a strict linear probe. The paper uses sinusoidal absolute positions at the input and before the prediction head.

What this does not establish

The contribution concerns scalable representation learning. It does not show that direct L2 prediction is the best distribution model for photorealistic generation. MAR finding poor generation from direct L2 does not contradict AIM learning useful features with that objective.

Open the derivation

A Gaussian conditional with fixed variance has a log-likelihood whose mean-dependent term is squared error. Under squared error, the optimal predictor is the conditional mean. That explains why a point prediction may average incompatible futures.

$$-\log\mathcal{N}(x;\mu(c),\sigma^2I)=\frac{\|x-\mu(c)\|^2}{2\sigma^2}+\mathrm{const}$$

Check your understanding

Why is an attentive probe different from a linear probe?

Its trainable attention pooling changes how patch features are combined; a strict linear probe has only a linear readout over the specified fixed features.

Sources and next steps

• AIM · 2401.08541v1
• Method / pinned implementation source

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.