A Vision Transformer (ViT) applies a transformer to images by turning each image into an ordered sequence of patch tokens. Self-attention can then model relationships between distant parts of the image. The same visual encoder can support image classification or supply features to a text decoder for image captioning. section 3.1 section 4.0
For image classification, trace the image through four stages:
- Split the image into fixed-size patches, ordered from left to right and top to bottom. A 3 × 32 × 32 CIFAR-10 image split into 4 × 4 patches yields an 8 × 8 grid, or 64 patches. section 3.1.1 section 3.3.1
- Project each patch into a 48-value embedding and prepend a class (CLS) token. The resulting sequence has shape (65, 48): 64 patch tokens plus one token for the whole-image prediction. section 3.1.1
- Add learned positional embeddings. Self-attention alone does not encode where a patch came from; positional information lets the model use the image’s spatial arrangement. section 3.1.1
- Feed the embeddings to an encoder-only transformer. Train its class prediction against the image label using cross-entropy loss; use the trained model’s logits to select a class. section 3.1.2 section 3.4.3
One PyTorch implementation uses a Conv2d layer with kernel size and stride both set to 4. It projects nonoverlapping patches from (3, 32, 32) into (48, 8, 8); flattening and transposing gives the transformer a (64, 48) patch sequence. Its configuration uses four attention heads and four encoder layers. section 3.3.1
Listing 3.2 Converting image patches into a sequence
from torch import nn
class PatchEmbeddings(nn.Module):
def __init__(self, config):
super().__init__()
self.projection = nn.Conv2d( #1
config.num_channels,
config.hidden_size,
kernel_size=config.patch_size, #2
stride=config.patch_size) #3
def forward(self, x):
x = self.projection(x)
x = x.flatten(2).transpose(1, 2) #4
return x
The patch-embedding implementation shows how convolution, flattening, and transposition turn a batch of images into sequences of vectors. section 3.3.1
Inspecting a CIFAR-10 prediction
For a batch of 32 test images, the classifier returns class logits and an attention map per encoder layer. With four heads and 65 tokens, each layer’s maps have shape (32, 4, 65, 65). To inspect what receives attention, extract the CLS-to-patch values, average across heads and layers, reshape the 64 values to an 8 × 8 grid, and enlarge it to the image size. Compare the overlay with the predicted and actual labels. The reported test accuracy for this trained classifier is 0.5621. section 3.4.3
For image captioning, the ViT instead serves as an image encoder: a text decoder generates caption tokens from its visual features. One encoder implementation uses Unfold to extract 256 patches from a 128 × 128 image, then projects the patches and processes them with unmasked transformer blocks. Thus patch sequencing is reusable, while the output task determines what consumes the encoded image. section 4.0 section 4.3.1
