Vision Transformers

Vision Transformer (ViT)

The Vision Transformer (ViT) is a deep learning model that applies the Transformer architecture—originally designed for language processing—to visual data. Unlike CNNs, which operate on local pixel neighborhoods, ViT divides an image into patches and models global relationships among them via self-attention.

1. Image to Patch Embeddings

The input image:

$$ \mathbf{x} \in \mathbb{R}^{H \times W \times C} $$

is divided into non-overlapping patches of size \( P \times P \), giving a total of

$$ N = \frac{H \times W}{P^2} $$

patches. Each patch \( \mathbf{x}^{(i)} \) is flattened and linearly projected into a \( D \)-dimensional embedding:

$$ \mathbf{e}^{(i)} = \mathbf{W}_{\text{embed}} \, \text{vec}(\mathbf{x}^{(i)}) \in \mathbb{R}^D, \quad i = 1, \dots, N $$

After stacking all patch embeddings, we form:

$$ \mathbf{E} = [\mathbf{e}^{(1)}, \dots, \mathbf{e}^{(N)}]^\top \in \mathbb{R}^{N \times D} $$

Notation:
\( H, W \): image height and width
\( C \): number of channels (e.g., 3 for RGB)
\( P \): patch size
\( N \): number of patches
\( D \): embedding dimension
\( \mathbf{W}_{\text{embed}} \): learnable patch projection matrix

2. Positional Embeddings

Transformers have no inherent notion of spatial order, so positional embeddings are added to retain the location of each patch:

$$ \mathbf{Z}_0 = \mathbf{E} + \mathbf{P} $$

\( \mathbf{P} \in \mathbb{R}^{N \times D} \): learnable positional embeddings

3. Multi-Head Self-Attention (MHSA)

The self-attention mechanism models global dependencies among all patches. For each head \( h = 1, \dots, H \):

$$ \mathbf{Q}_h = \mathbf{Z} \mathbf{W}_Q^{(h)}, \quad \mathbf{K}_h = \mathbf{Z} \mathbf{W}_K^{(h)}, \quad \mathbf{V}_h = \mathbf{Z} \mathbf{W}_V^{(h)} $$

$$ \text{Attention}_h = \text{Softmax}\!\left( \frac{\mathbf{Q}_h \mathbf{K}_h^\top}{\sqrt{d_k}} \right)\mathbf{V}_h $$

$$ \text{MHSA}(\mathbf{Z}) = [\text{Attention}_1; \dots; \text{Attention}_H] \mathbf{W}_O $$

Notation:
\( H \): number of attention heads
\( d_k = D / H \): dimension per head
\( \mathbf{W}_Q^{(h)}, \mathbf{W}_K^{(h)}, \mathbf{W}_V^{(h)} \): projection matrices for queries, keys, and values
\( \mathbf{W}_O \): output projection matrix

4. Feed-Forward Network (FFN)

Each Transformer block contains a position-wise FFN applied to every patch embedding independently:

$$ \text{FFN}(\mathbf{z}) = \text{GELU}(\mathbf{z}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2 $$

Each encoder layer applies Layer Normalization and residual connections:

$$ \mathbf{Z}' = \text{LayerNorm}(\mathbf{Z} + \text{MHSA}(\mathbf{Z})) $$ $$ \mathbf{Z}_{\text{next}} = \text{LayerNorm}(\mathbf{Z}' + \text{FFN}(\mathbf{Z}')) $$

5. Classification Head

A special learnable classification token \( \mathbf{z}_{\text{CLS}} \) is prepended to the patch embeddings. After the final Transformer layer, its output representation is used for classification:

$$ \text{logits} = \mathbf{z}_{\text{CLS}} \mathbf{W}_{\text{cls}} \in \mathbb{R}^C $$

\( C \): number of classes, \( \mathbf{W}_{\text{cls}} \): classification weight matrix.

6. Training Objective

$$ \mathcal{L} = - \sum_{i=1}^{C} y_i \log(\text{Softmax}(\text{logits})_i) $$

Cross-entropy loss is minimized to train the model for classification.

References

Dosovitskiy, A., Beyer, L., Kolesnikov, A., et al. (2020). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv preprint, arXiv:2010.11929.