DL0055 Vision Transformer

Written by

Explain the Vision Transformer (ViT). How does it convert an image into a class prediction?

Answer

A Vision Transformer converts an image into a sequence of fixed-size patch tokens and processes them with a Transformer encoder. A learned class token is prepended, positional embeddings preserve patch order, and global self-attention lets every patch exchange information with every other patch. After the encoder stack, the class-token representation is passed to a prediction head. Compared with a CNN, a ViT has weaker built-in locality and translation bias, but it can model long-range interactions directly and scales effectively with data and compute.

Vision Transformer architecture from image patchification through token encoding and classification.

Figure 1: ViT architecture with tensor shapes, patch tokenization, the Transformer encoder stack, and the classification path.

(1) Tokenization: An image of shape $H\times W\times C$ is divided into $N=(H/P)(W/P)$ non-overlapping $P\times P$ patches; each flattened patch is projected to width $D$ .
(2) Global Context: Multi-head self-attention mixes information across all patch positions, while residual connections and MLP sublayers refine each token.
(3) Classification: A learned [CLS] token aggregates evidence across the encoder stack; its final state is normalized and mapped to class logits.

Vision Transformer inference flowchart showing the ordered transformation from pixels to class logits.

Figure 2: ViT inference flow from image validation and patch embedding to encoder processing and class prediction.

Mathematical Formulation:
$z_0=[x_{\mathrm{cls}};x_p^1E;x_p^2E;\ldots;x_p^NE]+E_{\mathrm{pos}}$
$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^T}{\sqrt{d_h}}\right)V$

Where:

$z_0$ is the initial token sequence supplied to the Transformer encoder.
$x_{\mathrm{cls}}$ is the learned class token, and $x_p^i$ is flattened image patch $i$ .
$E\in\mathbb{R}^{P^2C\times D}$ projects each $P\times P\times C$ patch to width $D$ , while $E_{\mathrm{pos}}$ supplies positional embeddings.
$N=(H/P)(W/P)$ is the patch count for an image of height $H$ , width $W$ , and channel count $C$ .
$Q$ , $K$ , and $V$ are query, key, and value matrices; $d_h$ is the per-head query/key width.