Vision Transformer › ViT Architecture

With patches converted to embeddings and positional information injected, the remainder of the architecture follows the standard transformer encoder we detailed in Encoder Architecture. The embeddings pass through $L$ identical transformer blocks, each consisting of several key components working in concert.

The four core components of each transformer block are:

1. Multi-Head Self-Attention (MSA): As covered in Multi-Head Attention, this allows each patch to attend to all other patches in the sequence, modeling global relationships from the very first layer. This is in stark contrast to CNNs, where receptive fields grow gradually with depth and only achieve global context in the final layers. In ViT, a patch in the upper-left corner of an image can directly attend to a patch in the lower-right corner in layer 1, allowing the model to capture long-range visual relationships that would require many convolutional layers to achieve.

2. Feed-Forward Network (MLP): A two-layer MLP with GELU activation and hidden dimension typically set to $4D$, following the transformer convention from Feed-Forward Networks. Each position in the sequence passes through this MLP independently. This provides non-linear transformation capacity beyond what the attention mechanism alone can achieve. The MLP can be understood as transforming each patch representation in place, while attention mixes information across patches.

3. Layer Normalization: Applied before each sublayer (Pre-Norm configuration), as we discussed in Pre-Norm vs Post-Norm. Pre-Norm stabilizes training by ensuring that the inputs to each sublayer have controlled scale, which is particularly important for deep networks where vanishing or exploding gradients can derail training.

4. Residual Connections: Applied around both the attention and MLP blocks, creating gradient highways essential for training deep networks, as explored in Residual Connections. These connections allow gradients to flow directly back through the network to early layers without passing through the non-linearities of intermediate layers.

The equations for layer $\ell$ process the input through self-attention and feed-forward sublayers with residual connections:

$$\begin{aligned} \mathbf{z}'_\ell &= \text{MSA}(\text{LN}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1} \\ \mathbf{z}_\ell &= \text{MLP}(\text{LN}(\mathbf{z}'_\ell)) + \mathbf{z}'_\ell \end{aligned}$$

where:

• $\mathbf{z}_{\ell-1} \in \mathbb{R}^{(N+1) \times D}$: input to layer $\ell$, the output from layer $\ell-1$ (or $\mathbf{z}_0$ for the first layer)
• $\mathbf{z}'_\ell \in \mathbb{R}^{(N+1) \times D}$: intermediate representation after the Multi-Head Self-Attention sublayer and its residual connection
• $\mathbf{z}_\ell \in \mathbb{R}^{(N+1) \times D}$: output of layer $\ell$ after the MLP sublayer and its residual connection
• $\ell$: layer index, ranging from $1$ to $L$ (total number of layers)
• $\text{MSA}$: Multi-Head Self-Attention function applied to all $N+1$ tokens in parallel
• $\text{MLP}$: Feed-Forward Network applied position-wise with GELU activation and expansion ratio 4
• $\text{LN}$: Layer Normalization applied before each sublayer (Pre-Norm configuration)

The symmetry between the NLP transformer and ViT is striking: after the patch embedding step, the equations are identical. The same PyTorch TransformerEncoderLayer module used for BERT or GPT can be used for ViT without modification. This architectural equivalence is what makes unified multimodal models possible.

After $L$ layers, the final hidden state of the CLS token $\mathbf{z}_L^{(0)} \in \mathbb{R}^D$ is extracted and passed to a classification head. This head is typically a single linear layer during fine-tuning, and a two-layer MLP during pre-training:

$$\hat{\mathbf{y}} = \mathbf{z}_L^{(0)} \mathbf{W}_\text{head} + \mathbf{b}_\text{head}$$

where $\mathbf{W}_\text{head} \in \mathbb{R}^{D \times K}$ maps from the embedding dimension to $K$ class logits. The CLS token is used here because it has no direct input content (it begins as a learned null token), which means its final representation is entirely determined by what it has aggregated through attention over the other tokens. It acts as a summary of the entire image, specialized through training to carry information relevant to the classification objective.