Class Token [CLS]

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Easy

Problem

Theory

The [CLS] token is a learnable embedding prepended to the sequence of patch embeddings in the Vision Transformer (ViT). After passing through all transformer layers, the state of this single token at the output serves as the image representation for classification. This idea was borrowed directly from BERT, where a similar token aggregates sentence-level information.

What It Is

In ViT, an input image is split into fixed-size patches (e.g., 16 \times 16 pixels), each linearly projected into a D-dimensional embedding. This produces a sequence of N patch tokens, each of shape (1, D). The [CLS] token is an additional learnable parameter of shape (1, 1, D) that is prepended to this sequence before the transformer encoder processes it.

The [CLS] token is not derived from any region of the image. It carries no spatial information at initialization. Instead, it starts as a randomly initialized vector and learns, through backpropagation, to aggregate information from all patches via the self-attention mechanism. By the final transformer layer, the [CLS] token's state encodes a global summary of the entire image.

The classification head (typically a small MLP) is attached exclusively to the [CLS] token's output, ignoring the patch token outputs entirely. This makes the [CLS] token the bottleneck through which all image-level information must flow.

Key Equations

Let the patch embeddings after linear projection be \mathbf{z}_1, \mathbf{z}_2, \ldots, \mathbf{z}_N, each \in \mathbb{R}^D. The learnable [CLS] token is a parameter \mathbf{x}_{\text{class}} \in \mathbb{R}^D.

Tiling across the batch. The [CLS] token is a single parameter shared across all images, so it must be replicated to match the batch size B:

\mathbf{x}_{\text{class}}^{\text{tiled}} = \text{tile}(\mathbf{x}_{\text{class}}, B) \in \mathbb{R}^{B \times 1 \times D}

Concatenation. The tiled [CLS] token is concatenated to the front of the patch sequence along the sequence dimension:

\mathbf{z}_0 = [\mathbf{x}_{\text{class}}^{\text{tiled}} \;;\; \mathbf{z}_1, \mathbf{z}_2, \ldots, \mathbf{z}_N] \in \mathbb{R}^{B \times (N+1) \times D}

The sequence length changes from N to N + 1. This augmented sequence is then passed through the transformer encoder (after adding positional embeddings).

Output extraction. After L transformer layers, the output is \mathbf{z}_L \in \mathbb{R}^{B \times (N+1) \times D}. The [CLS] token's final state is extracted from position 0:

\mathbf{y} = \mathbf{z}_L[:, 0, :] \in \mathbb{R}^{B \times D}

This vector \mathbf{y} is fed to the classification head to produce class logits.

Why a [CLS] Token

Transformers produce a sequence of output tokens (one per input token), but image classification requires a single vector per image. The [CLS] token solves this aggregation problem.

Borrowed from BERT. Devlin et al. (2019) introduced the [CLS] token for sentence-level tasks. Dosovitskiy et al. (2020) directly adapted this for vision: "Similar to BERT's [class] token, we prepend a learnable embedding whose state at the output serves as the image representation."

Fixed position for classification. The [CLS] token is always at position 0, so the classification head always knows exactly where to look. There is no ambiguity about which token to use -- extraction is simply indexing position 0 of the output sequence.

Attends to all patches through self-attention. In each transformer layer, the [CLS] token computes attention scores against every patch token. It selectively attends to patches most informative for classification. Over multiple layers, it progressively refines its representation by gathering information from all spatial locations.

Aggregates image-level information. Unlike any single patch token (which is biased toward a local region), the [CLS] token is position-agnostic. It has no spatial prior and must learn to integrate global context. By the final layer, it represents a learned aggregation of the entire image, weighted by what the attention layers deemed relevant.

Why Learnable

The [CLS] token is a trainable parameter, not a fixed constant. This is a deliberate design choice.

Initialized randomly. The [CLS] token is initialized from a normal distribution scaled by 0.02 (i.e., randn * 0.02). The small scale prevents the token from dominating attention scores at the start of training, allowing gradients to flow normally.

Learned during training. As the model trains on classification loss, gradients flow back through the classification head, through the transformer layers, and into the [CLS] token parameter. The token learns an initialization state that, when processed by self-attention, produces the most useful image-level representation. In effect, it learns what "question" to ask the patch tokens.

Adapts to the task. When fine-tuning ViT on different downstream tasks, the [CLS] token adapts along with all other parameters. A fixed, non-learnable token (e.g., a zero vector) would provide no task-specific signal and would rely entirely on the transformer layers to compensate.

Why Prepend, Not Append

The [CLS] token is placed at position 0 (the front of the sequence), not at the end. While self-attention is permutation-equivariant (position is encoded via positional embeddings, not by ordering), the prepend convention matters for practical reasons.

Convention from BERT. BERT always places [CLS] at the beginning. ViT follows this directly. Since ViT's contribution was showing that a standard transformer works for vision with minimal modifications, keeping conventions identical to NLP transformers reinforced that message.

Position 0 means consistent extraction. Regardless of the number of patches N (which changes with image resolution or patch size), the [CLS] token is always at index 0. If the [CLS] token were appended, its position would be at index $$, which varies with input configuration. Position 0 is the simplest, most robust convention.

Positional embedding alignment. The positional embedding at position 0 is specifically learned for the [CLS] token. The model learns that position 0 is special (a classification aggregator, not a spatial patch), helping the attention layers treat it differently from the patch tokens.

Paper Context

The [CLS] token was introduced for vision in "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" (Dosovitskiy et al., 2020).

BERT inspiration. The authors explicitly state: "Similar to BERT's [class] token, we prepend a learnable embedding to the sequence of embedded patches, whose state at the output of the Transformer encoder serves as the image representation." The mechanism is identical in structure to BERT's approach.

Alternative: Global Average Pooling. The paper also tested applying global average pooling (GAP) over the patch token outputs instead of using a [CLS] token. GAP computes the mean of all N patch embeddings at the output, producing a single D-dimensional vector without an extra token.

Empirical comparison. The paper found that [CLS] and GAP perform comparably but require different learning rate schedules. When using [CLS], the learning rate schedule from BERT worked well. The paper defaults to [CLS] to stay consistent with the original transformer design.

Pre-training and fine-tuning. During pre-training on large datasets (JFT-300M, ImageNet-21k), the [CLS] token learns a general-purpose image representation. During fine-tuning, the classification head is replaced with a new zero-initialized linear layer, but the [CLS] token parameter is kept, providing a strong initialization.

Numerical Example

Consider a small ViT with batch size B = 2, number of patches N = 4, and embedding dimension D = 3.

Step 1: Define the [CLS] Token Parameter

The [CLS] token is a learnable parameter of shape (1, 1, D) = (1, 1, 3), initialized with small random values (e.g., randn * 0.02):

\mathbf{x}_{\text{class}} = [[\;0.01, \;-0.02, \;0.03\;]]

Step 2: Patch Embeddings

After splitting images into patches and projecting them, we have N = 4 patch embeddings per image. For a batch of B = 2 images, the patch embedding tensor has shape (2, 4, 3):

\mathbf{Z}_{\text{img1}} = \begin{bmatrix} 0.5 & -0.3 & 0.1 \\ 0.2 & 0.4 & -0.1 \\ -0.2 & 0.6 & 0.3 \\ 0.1 & -0.1 & 0.5 \end{bmatrix}, \quad \mathbf{Z}_{\text{img2}} = \begin{bmatrix} -0.4 & 0.2 & 0.3 \\ 0.3 & -0.5 & 0.1 \\ 0.1 & 0.3 & -0.2 \\ -0.1 & 0.4 & 0.2 \end{bmatrix}

Step 3: Tile the [CLS] Token Across the Batch

The single [CLS] token must be replicated B = 2 times so each image gets its own copy:

\mathbf{x}_{\text{class}}^{\text{tiled}} = \begin{bmatrix} [0.01, -0.02, 0.03] \\ [0.01, -0.02, 0.03] \end{bmatrix} \in \mathbb{R}^{2 \times 1 \times 3}

Both images receive the same [CLS] token values. Self-attention will produce different outputs per image because the patch embeddings differ. During backpropagation, gradients from all batch elements are summed to update the single shared parameter.

Step 4: Concatenate to the Front

Concatenate along dimension 1 (sequence dimension). For image 1, the resulting sequence of shape (5, 3) is:

• Position 0 ([CLS]): [0.01, -0.02, 0.03]
• Position 1 (Patch 1): [0.5, -0.3, 0.1]
• Position 2 (Patch 2): [0.2, 0.4, -0.1]
• Position 3 (Patch 3): [-0.2, 0.6, 0.3]
• Position 4 (Patch 4): [0.1, -0.1, 0.5]

The sequence length increased from N = 4 to N + 1 = 5.

Step 5: Shape Summary

• cls_token parameter: (1, 1, 3)
• After tiling: (2, 1, 3)
• Patch embeddings: (2, 4, 3)
• After concatenation: (2, 5, 3) -- i.e., (B, N+1, D)

After the transformer encoder, the classification output is extracted from position 0: \mathbf{y} = \mathbf{z}_L[:, 0, :] \in \mathbb{R}^{2 \times 3}.

[CLS] Token vs Global Average Pooling

The [CLS] token is not the only way to aggregate patch information. Global Average Pooling (GAP) is the main alternative, and the ViT paper explicitly compared both.

How GAP works. Instead of prepending a [CLS] token and extracting position 0, GAP averages all N patch token outputs:

\mathbf{y}_{\text{GAP}} = \frac{1}{N} \sum_{i=1}^{N} \mathbf{z}_L^{(i)} \in \mathbb{R}^{B \times D}

This is parameter-free and does not increase the sequence length.

ViT paper findings. Both achieve similar accuracy when properly tuned, but require different hyperparameter settings, particularly learning rate schedules. The paper chose [CLS] as the default to maintain consistency with BERT.

Trade-offs:

• Sequence length. [CLS] adds one token, making self-attention O((N+1)^2) vs O(N^2). For typical ViT configurations (N = 196 for ViT-B/16), the overhead is negligible.
• Simplicity. GAP is simpler: no extra parameter, no tiling, no concatenation. Just a mean operation after the transformer.
• Learned vs uniform aggregation. The [CLS] token learns to attend selectively to informative patches through attention. GAP gives equal weight to all patches. However, the transformer layers before GAP can still modulate patch representations, so the distinction is nuanced.
• Convention. [CLS] is the default in ViT, DeiT, BEiT, and MAE. Some later architectures prefer GAP: Swin Transformer uses average pooling. The choice is often convention rather than a clear winner.

Pitfalls

Wrong Initialization Scale

The [CLS] token should be initialized with standard deviation 0.02. Using standard normal (randn without scaling) produces values too large, destabilizing early attention scores. Zero initialization is also problematic: a zero vector produces zero attention logits, making the [CLS] token invisible to softmax. The randn * 0.02 convention matches the initialization used for other ViT parameters.

Forgetting to Tile Across the Batch

The [CLS] token parameter has shape (1, 1, D), but the patch embeddings have shape (B, N, D). Concatenating directly without tiling fails when B > 1 because the batch dimensions do not match. The [CLS] token must be expanded or tiled to shape (B, 1, D) before concatenation. Some frameworks support broadcasting with expand(), but the operation must be explicit.

Appending Instead of Prepending

The [CLS] token must go at position 0, not position N. Appending breaks the extraction logic (which expects position 0) and misaligns with positional embeddings. The positional embedding at position 0 was learned to correspond to [CLS] during training; placing it elsewhere means it receives the wrong positional signal.

Wrong Sequence Length After Concatenation

After prepending, the sequence length is N + 1, not N. The positional embedding table must have N + 1 entries (one for [CLS] plus one per patch). Creating a table of size N causes a shape mismatch when adding positional embeddings to the augmented sequence.

Extracting the Wrong Position at Output

The classification output is the [CLS] token at position 0: output[:, 0, :]. Using position -1 (the last patch token) or position 1 (the first patch) gives the wrong representation. The [CLS] token is always at position 0, and this is the only correct extraction index.

Examples

Hints

1. Broadcast cls_token to the input batch size.
2. Concatenate along the sequence axis.

Requirements

• Use NumPy.
• Broadcast the supplied classification token across the batch.
• Place the classification token before every patch token.
• Preserve every patch value and its order.
• Return a float64 NumPy array with shape (B, N + 1, D).

Constraints

• patches has shape (B, N, D) and dtype float64.
• cls_token has shape (1, 1, D) and dtype float64.
• B, N, and D are positive.

Starter Code

import numpy as np

def prepend_class_token(patches: np.ndarray,
                        cls_token: np.ndarray) -> np.ndarray:
    """
    Returns the float64 sequence with the class token at position zero.
    """
    pass

Test Cases