Vision Transformer › Patch Embeddings and the CLS Token

Once we have flattened the image patches into vectors, we need to project them into the transformer's embedding dimension $D$. This process is analogous to word embeddings in NLP, where we map discrete vocabulary indices to dense continuous vectors. In ViT, we apply a trainable linear projection $\mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}$ to each patch vector.

The embedding for the $i$-th patch is computed as:

$$\mathbf{z}_0^{(i)} = \mathbf{x}^{(i)} \mathbf{E}$$

where:

• $\mathbf{z}_0^{(i)} \in \mathbb{R}^{D}$: the embedding vector for the $i$-th patch after projection
• $\mathbf{x}^{(i)} \in \mathbb{R}^{P^2 \cdot C}$: the flattened $i$-th patch (input vector), where $P$ is the patch size and $C$ is the number of channels
• $\mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}$: the trainable linear projection matrix mapping from patch dimension to embedding dimension $D$
• $i$: the patch index, ranging from $1$ to $N$ (total number of patches)

This linear projection is a learned weight matrix, initialized randomly and trained end-to-end along with the rest of the model. Conceptually it serves the same purpose as an embedding table in NLP: it maps the raw input representation to a space where meaningful relationships can be captured and transformed by the downstream attention layers. Unlike a discrete embedding table that has one learned vector per vocabulary entry, the patch embedding matrix is a general linear transformation applied uniformly to every patch, regardless of its content. The model must therefore learn a transformation that works for all possible patches it might encounter, a more challenging requirement that again explains why more pre-training data is beneficial.

An important practical equivalence exists between this explicit flatten-then-project operation and a $P \times P$ convolutional layer applied with stride $P$ and no padding. Both operations extract non-overlapping windows of size $P \times P$, apply the same linear map to each window, and produce a spatial grid of feature vectors. ViT implementations therefore sometimes use nn.Conv2d(in_channels=3, out_channels=D, kernel_size=P, stride=P) as the patch embedding layer, which is computationally equivalent to the explicit formulation but benefits from highly optimized convolution routines on modern hardware.

Following the convention established by BERT (which we covered in BERT Architecture), ViT prepends a special learnable embedding $\mathbf{x}_\text{class} \in \mathbb{R}^{D}$ to the beginning of the sequence. This classification token (CLS token) serves a unique role in the architecture. It begins as a learnable parameter initialized randomly, with no inherent content. As information flows through the transformer layers via self-attention, the CLS token can attend to every patch in the sequence, aggregating global information from all spatial locations.

Why use a separate CLS token for classification rather than, say, averaging all patch representations? The averaging approach, called global average pooling, is also used in practice and works well. The CLS token approach is conceptually cleaner for sequence classification because the model can learn to selectively aggregate the information that is most relevant for the classification decision, rather than uniformly combining all spatial information. Some patches might be background clutter that contributes little to the class label; a learnable aggregation via the CLS token can in principle learn to downweight such patches. In practice, both approaches yield similar accuracy, and the choice is sometimes made for architectural consistency with BERT-style models.

After prepending the CLS token, we add learned position embeddings to inject spatial information. The complete input sequence to the transformer is:

$$\mathbf{z}_0 = [\mathbf{x}_{\text{class}};\; \mathbf{x}^{(1)}\mathbf{E};\; \mathbf{x}^{(2)}\mathbf{E};\; \cdots;\; \mathbf{x}^{(N)}\mathbf{E}] + \mathbf{E}_{\text{pos}}$$

where:

• $\mathbf{z}_0 \in \mathbb{R}^{(N+1) \times D}$: the complete input sequence to the transformer, with position embeddings added element-wise
• $\mathbf{x}_{\text{class}} \in \mathbb{R}^{D}$: the learnable CLS token embedding prepended to the sequence at position 0
• $\mathbf{x}^{(i)} \mathbf{E} \in \mathbb{R}^{D}$: the linearly projected $i$-th patch embedding
• $\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}$: the matrix of position embeddings, with one row per sequence position including the CLS token
• $[\cdot;\cdot]$: vertical concatenation (stacking vectors row-by-row to form a matrix)
• $N$: total number of patches ($14 \times 14 = 196$ in the running example)

The $+$ here denotes element-wise addition: each row of the concatenated patch matrix is added to the corresponding row of the position embedding matrix. This means every token in the sequence, including the CLS token, receives a position-specific bias vector that encodes its index in the sequence. After this addition, the model's self-attention layers can use these position signals to infer spatial relationships.