Skip to content

← Flax step 46 of 100

Hard Primitives

Vision Transformer with [CLS] Token

Jump to the task

A dedicated output token

A vision transformer turns image patches into a sequence of vectors. To classify the image, it needs one vector to pass to a classification layer. One approach averages the final patch vectors. This exercise instead prepends a trainable vector called the class token, or [CLS], and reads that token after the transformer blocks.

The initial class token is shared across images. It does not contain information about a particular image yet. Self-attention lets its output depend on the patch vectors. Training can then teach the model to collect information useful for classification at this position. Here you only initialize the model and run a forward pass, so the outputs are untrained logits, not meaningful predictions.

Sequence construction

For an image of shape (H, W, C) and square patches of side P, let T = (H // P) * (W // P). Assume H and W are divisible by P. A convolution with kernel and stride (P, P) maps each patch to D features. Flatten the patch grid in row-major order to get (T, D).

Prepend a parameter of shape (1, D), then add a position parameter of shape (T + 1, D). Position 0 belongs to the class token. Positions 1 through T belong to the image patches.

import jax.numpy as jnp

patches = jnp.array([[1., 2.], [3., 4.]])
cls = jnp.array([[0.1, 0.2]])
positions = jnp.array([[0., 0.], [10., 20.], [30., 40.]])
sequence = jnp.concatenate([cls, patches], axis=0) + positions
# [[0.1, 0.2], [11., 22.], [33., 44.]]

These chosen values illustrate sequence construction, not random model initialization. Concatenate along the sequence axis, not the feature axis. Omitting the class position from the position table causes a shape mismatch.

What positions contribute

The class token is a separate learned parameter and a designated output slot even without position embeddings. Position embeddings tell the model where the patches occur in the sequence. Without them, permuting the patch tokens leaves the class output unchanged in this unmasked model. With a position table, moving patch content between fixed positions can change that output.

Attention is bidirectional here: every token can attend to every token, including itself and the class token. The class token need not summarize every detail of an image. Its representation is shaped by the training objective, and this exercise does not train it.

Inside each transformer block

Use two residual branches, each with LayerNorm before its transformation:

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

Both branches preserve (T + 1, D). The attention uses num_heads heads, with D / num_heads features per head. Each block has its own parameters. After the last block, apply another LayerNorm, select row 0 and project it to num_classes logits. Do not apply softmax or average the patch rows.

For a (4, 4, 3) image, P=2 and D=8, the convolution produces (2, 2, 8), the patch sequence is (4, 8) and adding the class token gives (5, 8). The final class vector has shape (8,) regardless of the number of patches.

Problem

Implement vit_cls_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 one floating-point image with shape (H, W, C), without a 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.

Match this module structure so seeded initialization agrees:

  1. In a compact ViTCLS module, call an unnamed nn.Conv(d_model, (P, P), strides=(P, P), padding="VALID") with its default bias. Reshape its output to (T, D).
  2. Create cls_token with self.param, nn.initializers.normal(stddev=0.02) and shape (1, D). Prepend it to the patch sequence.
  3. Create pos_embed with the same initializer and shape (T + 1, D). Add it to the sequence.
  4. Call num_layers separate unnamed ViTBlock modules. Each compact block calls LayerNorm, nn.MultiHeadDotProductAttention with num_heads=num_heads and qkv_features=D, then the second LayerNorm, Dense(d_ff), ReLU and Dense(D), in the order shown above. Pass the normalized sequence as the attention's sole input. Use no mask or dropout; retain other Linen defaults, including biases.
  5. Apply an unnamed LayerNorm at the top level, take row 0, then apply an unnamed Dense(num_classes).
  6. Initialize the whole module with jax.random.PRNGKey(int(seed)) and image, then apply it to the same image.

The top-level submodules are Conv_0, ViTBlock_0 through ViTBlock_{num_layers-1}, LayerNorm_0 and Dense_0. Within each block they are LayerNorm_0, MultiHeadDotProductAttention_0, LayerNorm_1, Dense_0 and Dense_1. Adding wrapper modules or changing these paths changes the random parameter initialization.

With zero blocks the class logits cannot depend on the image: there is no attention operation to mix patches into the class token yet.

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

Loading visualization…