ViT: Vision Transformer

An input image \(\mathbf{x} \in \mathbb{R}^{H \times W \times C}\) is divided into non-overlapping 2D patches \(\mathbf{x}_{p} \in \mathbb{R}^{N \times (P^{2}\cdot C)}\), where:

• \((P,P)\) is the patch resolution,
• \((H,W)\) is the image resolution,
• \(N = \frac{HW}{P^{2}}\) is the effective sequence length (the number of patches).

Notably, we directly flatten image patches into a \(P^{2}\cdot C\)-dimension vector.

We map each flattened image patch to the latent model dimention \(D\) via a trainable linear projection \(\mathbf{E} \in \mathbb{R}^{(P^{2}C) \times D}\).

A learnable token embedding \(\mathbf{x}_{\text{class}} \in \mathbb{R}^{1 \times D}\) is prepended to the patch sequence, serving as the aggregate visual representation for classification, whose state at the output of the Transformer blocks serves as the image representation \(\mathbf{y}\).

Patches are arranged sequentially, and standard 1D learnable positional embeddings¹ are applied \(\mathbf{E}_{\text{pos}}\in\mathbb{R}^{(N+1)\times D}\)

Since ViT is first proposed as an image classifier, an MLP acts as an classification head. The model architecture looks like

\[\begin{aligned} \mathbf{z}_{0} &= [\mathbf{x}_{\text{class}}; \mathbf{x}_{p}^{1}\mathbf{E}; \dots; \mathbf{x}_{p}^{N}\mathbf{E}] + \mathbf{E}_{\text{pos}} \\ \mathbf{z}'_{\ell} &= \texttt{MultiHeadSelfAttn}(\texttt{LayerNorm}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1} \\ \mathbf{z}_{\ell} &= \texttt{MLP}(\texttt{LayerNorm}(\mathbf{z}'_{ell})) + \mathbf{z}'_{\ell} \\ \mathbf{y} &= \texttt{LayerNorm}(\mathbf{z}_{L}^{0}) \end{aligned}\]

where \(\mathbf{x}_{p}^{i}\) denotes the non-overlapping flattened 2D patch, \(\ell=1\dots L\) stands for Transformer block layers.