Skip to content

← Flax NNX step 46 of 100

Hard Primitives

NNX ViT with CLS Token

Jump to the task

Why this matters

An image classifier needs one prediction for the whole image. A class token gives a transformer one dedicated sequence position to read for that prediction. Prepend a trainable row to the patch embeddings, let it attend to the patches, then send its final representation to a classification head.

This exercise builds a small model with that pattern. It runs a freshly initialized model, without training or loading a checkpoint. Its logits are not useful class predictions yet. The ReLU feed-forward layers and zero-initialized tokens below define this exercise's architecture, not a complete reproduction of a pretrained ViT.

From image to sequence

The image has shape (height, width, channels), with no batch axis. Apply nnx.Conv with kernel size (P, P), stride P, and "VALID" padding. Each non-overlapping patch becomes a vector of width D. Incomplete patches at the bottom or right edge are dropped.

The output has shape (height // P, width // P, D). Reshape it to (N, D), keeping patches in row-major order, where N = (height // P) * (width // P).

Create two trainable arrays with nnx.Param:

  • cls_token: zeros of shape (1, D).
  • pos_embed: zeros of shape (N + 1, D).

Concatenate the class token before the patches, then add the position array. Position row 0 belongs to the class token; rows 1 through N belong to the patches. Concatenation requires two-dimensional inputs here. It does not broadcast a vector of shape (D,) into a row.

How the class token gathers information

Use unmasked self-attention: every row can attend to every row, including itself. The class token therefore receives a weighted sum of value vectors from both the patches and the class position.

Here is a complete example of that averaging operation with one head. It isolates attention, without the projections or residual blocks:

import jax
import jax.numpy as jnp
import numpy as np

# One class row followed by two patch rows.
values = jnp.array([[0., 0.], [2., 4.], [4., 2.]])
queries = jnp.zeros_like(values)
keys = jnp.zeros_like(values)
scores = queries @ keys.T / jnp.sqrt(2.)
weights = jax.nn.softmax(scores, axis=-1)
result = weights @ values
np.testing.assert_allclose(result[0], [2., 2.])

Equal scores give equal weights over all three rows, so the class row receives their mean. Learned projections can produce unequal weights. Training can make the final class representation useful for prediction; the existence of a class token alone does not guarantee that outcome.

Self-attention without position information is permutation-equivariant: reordering its input rows reorders its output rows. If the class row stays first and only patch rows move, its output stays the same. With the zero position array used at initialization here, the model cannot distinguish absolute patch locations. Training can change the position rows so they carry location information. Adding zeros to zeros does not itself break this symmetry.

Encoder and classifier

Each encoder block uses pre-normalization and two residual additions:

  1. x = x + attention(ln1(x)).
  2. x = x + ff2(jax.nn.relu(ff1(ln2(x)))).

For attention, project queries, keys, and values with separate nnx.Linear(D, D) layers. Split the last axis into num_heads heads of width Dh = D // num_heads. For each head, compute softmax(Q @ K.T / sqrt(Dh), axis=-1) @ V. Join heads in their original order and apply another nnx.Linear(D, D) output projection. Do not apply a causal mask or dropout.

The feed-forward layers have widths D -> F -> D. Both LayerNorms normalize each row over its D features, using epsilon=1e-6 and their default trainable scale and bias.

After all blocks, apply a final LayerNorm with the same settings. Read row 0 and apply nnx.Linear(D, K) to return K class logits. Do not average patch rows or apply softmax to the logits.

Problem

Write vit_cls_forward(seed, image, patch_size, d_model, num_heads, d_ff, num_layers, num_classes) using NNX modules.

Cast the seed and size arguments to Python integers. Let P, D, F, L, and K denote patch size, model width, feed-forward width, block count, and class count. Assume positive sizes, P no larger than either image dimension, L >= 0, and D divisible by the positive head count.

Use one nnx.Rngs(int(seed)) throughout construction. To reproduce the initialized model, create modules and parameters in this order:

  1. Patch convolution from image channels to D features.
  2. Zero class token, then zero position array, as specified above.
  3. L distinct encoder blocks stored in nnx.List. Within each block, create ln1, query projection, key projection, value projection, attention output projection, ln2, ff1, then ff2.
  4. Final LayerNorm, then the D-to-K classification head.

Use default NNX convolution and linear initializers, with biases enabled. Read parameter arrays with parameter[...]. Return a one-dimensional array of shape (K,).

Zero-block check: with L=0, the freshly initialized class row remains zero. The default final normalization and head biases are also zero, so every logit is zero, regardless of the image. With encoder blocks, attention provides a path from patches to that row.

The diagram below compares allowed attention connections. This model uses the bidirectional pattern; the diagram does not show learned attention weights or this model's exact sequence length.

Loading visualization…