Skip to content

← Flax step 45 of 100

Hard Primitives

Vision Transformer (Mean-Pool Variant)

Jump to the task

From image patches to a class score

A vision transformer processes an image as a sequence of patch vectors. A shared projection maps each patch to D features. Attention then lets patches exchange information, so each output vector can depend on the whole image rather than only its original patch.

This exercise builds a small model that averages the final patch vectors before classification. It has no class token. You initialize the model and run one forward pass; there is no training or pretrained checkpoint. The returned numbers are raw, untrained logits, not probabilities.

Patches and positions

For an image (H, W, C), use square patches of side P with no overlap. A convolution with kernel and stride (P, P) applies the same affine projection to every patch, including a bias. Its output has shape (H // P, W // P, D). Flatten the two spatial axes in row-major order to get T = (H // P) * (W // P) tokens of width D.

Add a learned position table of shape (T, D). Row 0 describes the first patch position, row 1 the next, and so on. Without this table, permuting the patch tokens would permute the encoder outputs, leaving their mean unchanged. With positions, moving patch content between fixed locations can change the prediction.

The task assumes both image dimensions are divisible by P. With VALID padding, incomplete edge patches would otherwise be omitted. Do not resize the input inside the function.

Average the patch rows

After the transformer blocks, normalize each token across its feature axis. Then average over the sequence axis, leaving one vector of width D. A final Dense layer maps that vector to num_classes logits.

Here is the pooling and head operation with chosen values:

import jax.numpy as jnp

# Two already-normalized token rows.
tokens = jnp.array([[1., -1.], [-1., 1.]])
kernel = jnp.array([[2., 0., 1.], [0., 3., -1.]])
bias = jnp.array([0.1, 0.2, 0.3])
pooled = tokens.mean(axis=0)  # [0., 0.]
logits = pooled @ kernel + bias  # [0.1, 0.2, 0.3]

For a (4, 4, 3) image with P=2 and D=8, patch projection produces (2, 2, 8), then reshaping gives (4, 8). The blocks and final LayerNorm keep (4, 8). Averaging axis 0 gives (8,), and a four-class head gives (4,). Averaging the last axis instead would produce one scalar per patch.

Transformer block

Each compact ViTBlock has two residual branches:

h = LayerNorm(x)
x = x + SelfAttention(h)
h = LayerNorm(x)
x = x + Dense(D)(ReLU(Dense(d_ff)(h)))

Use an unnamed nn.MultiHeadDotProductAttention with num_heads=num_heads and qkv_features=D, passing the normalized sequence as its sole input. Every token may attend to every token. Use no attention mask or dropout. Retain other Linen defaults, including biases and LayerNorm settings. Each block has its own parameters.

Within a block, the submodule names are LayerNorm_0, MultiHeadDotProductAttention_0, LayerNorm_1, Dense_0 and Dense_1. Their nesting and creation order matter for seeded initialization.

Problem

Implement vit_forward(seed, image, patch_size, d_model, num_heads, d_ff, num_layers, num_classes) using Flax Linen. Return a one-dimensional array of num_classes logits.

image is a floating-point array (H, W, C) with no batch axis. Dimension arguments are positive integers, except num_layers, which may be zero. H and W are divisible by patch_size, and d_model is divisible by num_heads. Treat the seed and dimension arguments as fixed Python integers when compiling a function over images.

Inside a compact MiniViT module:

  1. Call an unnamed nn.Conv(d_model, (P, P), strides=(P, P), padding="VALID") with its default bias, then reshape to (T, D).
  2. Create pos_embed with self.param, nn.initializers.normal(stddev=0.02) and shape (T, D). Add it to the tokens.
  3. Call num_layers separate unnamed ViTBlock modules as specified above.
  4. Apply an unnamed LayerNorm, average over axis 0, and call an unnamed Dense(num_classes). Return the logits without softmax.

The top-level submodules are Conv_0, ViTBlock_0 through ViTBlock_{num_layers-1}, LayerNorm_0 and Dense_0. Initialize the entire module with jax.random.PRNGKey(int(seed)) and image, then apply it to the same image. Do not add wrappers that change the parameter paths.

The diagram illustrates attention masks. This model uses full bidirectional attention, without a causal or sliding-window mask.

Loading visualization…