NNX Vision Transformer

Why this matters

A vision transformer treats an image as a sequence of patch vectors. Attention lets each patch combine information from other patches. For classification, this exercise averages the final patch vectors and passes that average to a linear head.

You will build a small, freshly initialized model and return its logits. There is no training or checkpoint loading in this task. The model uses ReLU feed-forward layers and zero-initialized position embeddings; these choices define this exercise rather than a complete reproduction of a particular pretrained ViT.

Turn patches into vectors

Input image has shape (height, width, channels), without a batch axis. Use nnx.Conv with kernel size (P, P), stride P, and padding="VALID" to produce D features per patch.

This applies the same learned affine projection to every full patch. The kernel covers a patch's spatial positions and channels, while its output axis selects the feature. The convolution does not flip the kernel. Its output has shape (height // P, width // P, D). Incomplete patches at the bottom and right edges are dropped.

Here is a complete example with known weights. Each patch becomes its sum and its negated sum:

import jax.numpy as jnp
import numpy as np
from flax import nnx

image = jnp.arange(1., 9.).reshape(2, 4, 1)
conv = nnx.Conv(1, 2, kernel_size=(2, 2), strides=2,
                padding="VALID", rngs=nnx.Rngs(0))
conv.kernel[...] = jnp.broadcast_to(
    jnp.array([1., -1.]), (2, 2, 1, 2)
)
conv.bias[...] = jnp.zeros(2)
patches = conv(image)
sequence = patches.reshape(-1, 2)
np.testing.assert_allclose(sequence, [[14., -14.], [22., -22.]])
np.testing.assert_allclose(sequence.mean(axis=0), [18., -18.])

The full model reshapes the patch grid in the same row-major order, giving (N, D) where N = (height // P) * (width // P). Keep the D feature values of each patch together.

Add positions, then mix patches

Create pos_embed = nnx.Param(jnp.zeros((N, D))) and add pos_embed[...] to the patch sequence. Each row is a trainable position vector. At initialization all rows are zero, so they provide no location information yet. Training can make the rows different.

Each encoder block has two pre-normalized residual branches:

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

Attention uses separate D-to-D query, key, value and output linear projections, all with biases. Split queries, keys and values into num_heads heads of width Dh = D // num_heads. Within each head, compute softmax(Q @ K.T / sqrt(Dh), axis=-1) @ V. Softmax runs over keys. Concatenate head outputs in their original order, then apply the output projection. Every patch can attend to every patch; there is no causal mask or dropout.

The feed-forward layers have widths D -> F -> D, with ReLU only between them. Both LayerNorms normalize the D features of each row, using epsilon=1e-6 and default trainable scale and bias.

Pool and classify

Apply a final LayerNorm to each patch row, then compute pooled = jnp.mean(x, axis=0). This averages patches and retains the D features. Finally, use nnx.Linear(D, K) to produce K logits. Return the logits directly, without softmax.

Normalize before pooling: LayerNorm is nonlinear, so normalizing the pooled vector generally gives a different result. Do not use the first patch as the classifier input; this model has no class token.

Without position information, self-attention reorders its outputs when its input rows are reordered. Averaging removes that ordering. Consequently, with the initial zero position embeddings, rearranging whole patches does not change this model's logits. The model still uses the arrangement of pixels within each patch through the convolution kernel. It is not free of architectural assumptions.