Lazy loaded image11.1 Vision Transformer

Paper: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
notion image
Vision Transformer (ViT) is a model that applies the Transformer architecture to image classification tasks. The core idea of ViT is to divide an image into multiple small blocks (patches) and treat these small blocks as a sequence input to the Transformer Encoder. The ViT model is mainly composed of the following three modules:
  1. Linear Projection of Flattened Patches (Embedding Layer)
  1. Transformer Encoder
  1. MLP Head (Classification Head)

Input

The input requirement for a Transformer is a 2D matrix [num_token, token_dim], without considering the batch for now, while the image format is a 3D matrix [H, W, C]. To convert the image into the format required by Transformer, ViT adopts the following steps:
  1. Image Patchification: Divide the input image (e.g., 224x224) into multiple small patches of fixed size (e.g., 16x16). For example, in ViT-B/16, an image is divided into patches.
  1. Linear Projection: Each patch (shape [16, 16, 3]) is flattened into a vector of length 768 (token) and then passed through a linear projection. That is, [16, 16, 3] -> [768].
      2. Note that this flattening and projection operations can be simplified as a single convolution operation in practical implementation. The convolution kernel size is 16x16, and the stride is 16. Input [224, 224, 3] passes through the convolutional layer to get [14, 14, 768], and then is flattened to [196, 768].
  1. [class] token: Similar to BERT, ViT inserts a learnable [class] token at the beginning of the token sequence for classification tasks. This token is concatenated with other tokens, resulting in [197, 768].
  1. Position Embedding: To retain the position information of the image, ViT introduces a learnable 1D Position Embedding, which is directly added to the tokens. Its shape is the same as the tokens, i.e., [197, 768].
Code Example

Transformer Encoder

The Transformer Encoder is stacked with multiple Encoder Blocks, and each Encoder Block contains the following components:
  1. Layer Norm: Normalizes each token.
  1. Multi-Head Attention: Multi-head self-attention mechanism, used to capture relationships between tokens.
  1. Dropout/DropPath: Used to prevent overfitting. The original paper used Dropout, while some implementations used DropPath (stochastic depth) such as in the Swin Transformer.
  1. MLP Block: Composed of a fully connected layer, GELU activation function, and Dropout. The first fully connected layer expands the input node count by 4 times ([197, 768] -> [197, 3072]), and the second fully connected layer restores it ([197, 3072] -> [197, 768]).
notion image
The output of the Transformer Encoder maintains the same shape as the input. Taking ViT-B/16 as an example, if the input is [197, 768], the output is still [197, 768]. There is also a Layer Norm layer at the end of the Encoder.
Code Example
  • MLP Block
    • Encoder

      MLP Head

      From the output of the Transformer Encoder, only the result corresponding to the [class] token needs to be extracted ([1, 768]). The MLP Head is used to map the [class] token to the final classification result. The structure of the MLP Head is as follows:
      1. When training on ImageNet21K: Linear + tanh activation function + Linear.
      1. When migrating to ImageNet1K or other datasets: Only one Linear layer is used.

      Model Parameters

      ViT (Vision Transformer) mainly includes three models: Base, Large, and Huge. Additionally, the source code also provides configurations for Patch Size of 16x16 and 32x32.
      • Layers: Determines the depth of the model. The more layers, the stronger the expressive power of the model, but the higher the computational complexity.
      • Hidden Size: Determines the dimension of the vector for each token. The larger the dimension, the stronger the expressive power of the model.
      • MLP Size: Determines the number of nodes in the first fully connected layer of the MLP Block, which is usually four times the Hidden Size.
      • Heads: Determines the number of heads in Multi-Head Attention. The more heads, the more feature information the model can capture.
      • Patch Size: Determines how many patches an image is divided into. The smaller the patch, the more detailed information the model can capture.
      notion image
      Prev
      11.0 Introduction
      Next
      11.2 CLIP
      Loading...
      Article List
      LLM Learning Roadmap
      ✨ Awesome-Anything
      Lazy loaded image Awesome-AI-Tutorials
      Lazy loaded image Awesome-AI-Researchs
      Lazy loaded image Awesome-GitHub-Repos
      Lazy loaded image Awesome-AI-Examples
      🖼️ Digital Image Processing
      Lazy loaded image DIP I Image Demosaicing and Histogram Manipulation
      Lazy loaded image DIP II Image Denoising
      Lazy loaded image DIP III Edge Detection
      Lazy loaded image DIP IV Digital Half-toning
      Lazy loaded image DIP V Geometric Image Modification
      🍃 LLM Components
      Lazy loaded image 1.1 Tokenization
      Lazy loaded image 1.2 Embedding
      Lazy loaded image 1.3 Attention
      Lazy loaded image 1.4 FFN, Residual Addition, and LN
      Lazy loaded image 1.5 Positional Encoding
      Lazy loaded image 1.6 Architecture & Decoding Policy
      🌱 LLM Pre-training
      Lazy loaded image 2.1 Pre-training Definition
      Lazy loaded image 2.2 Pre-training Data
      Lazy loaded image 2.3 Pre-training Workflow
      Lazy loaded image 2.4 Continued Pre-training (CPT)
      Lazy loaded image 2.5 Pre-training Evaluation
      ☘️ LLM Post-Training
      Lazy loaded image 3.1 Supervised Fine-Tuning (SFT)
      Lazy loaded image 3.2 Reinforcement Learning
      Lazy loaded image 3.3 RLHF: Reinforcement Learning from Human Feedback
      Lazy loaded image 3.4 Optimizing PPO-Based Algorithms
      Lazy loaded image 3.5 Direct Preference Optimization (DPO)
      🍀 LLM Popular Models
      Lazy loaded image 4.1 BERT & Variants
      Lazy loaded image 4.2 PaLM Series
      Lazy loaded image 4.3 GPT Series
      Lazy loaded image 4.4 LLaMA Series
      Lazy loaded image 4.5 GLM Series
      Lazy loaded image 4.6 Qwen Series
      Lazy loaded image 4.7 DeepSeek Series
      Lazy loaded image 4.8 MOE Series
      Lazy loaded image 4.9 Other
      🪴 LLM Applications
      Lazy loaded image 5.1 Prompt Engineering
      Lazy loaded image 5.2 LLM-based Agents
      Lazy loaded image 5.3 Retrieval-Augmented Generation (RAG)
      Lazy loaded image 5.4 Enterprise-level RAG Practice (I)
      Lazy loaded image 5.5 Enterprise-level RAG Practice (II)
      🌿 LLM Optimization
      Lazy loaded image 6.1 GPU Usage Analysis
      Lazy loaded image 6.2 Latency Analysis for Inference
      Lazy loaded image 6.3 FlashAttention
      Lazy loaded image 6.4 PagedAttention
      Lazy loaded image 6.5 Packing for Efficient Training
      Lazy loaded image 6.6 PEFT
      🌾 LLM Compression
      Lazy loaded image 7.1 Quantization
      Lazy loaded image 7.2 Pruning
      Lazy loaded image 7.3 Knowledge Distillation
      🌵 LLM Hands-on Practice
      Lazy loaded image 8.1 nanoGPT: Training a GPT from Scratch
      Lazy loaded image 8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
      Lazy loaded image 8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
      Lazy loaded image 8.4 OpenRLHF: Reward Model Training and Code Analysis
      Lazy loaded image 8.5 OpenRLHF: PPO Tutorial and Code Analysis
      🌴 LLM Must-read Papers
      Lazy loaded image 9.1 Technical Reports Explained
      Lazy loaded image 9.2 MedLLM
      Lazy loaded image 9.3 o1 Techinical Paper
      Lazy loaded image 9.4 Prompting
      Lazy loaded image 9.5 LLM-as-Judge
      🌳 LLM Q&A
      Lazy loaded image 10.1 Why GRPO training loss starts at 0
      Lazy loaded image 10.2 VLM-R1 Reproduction Results
      Lazy loaded image 10.3 KL(P||Q) or KL(Q||P)?
      Lazy loaded image 10.4 Left Padding or Right Padding?
      Lazy loaded image 10.5 What is QK-Norm?
      Lazy loaded image 10.6 Why LLM training rarely adopts Dropout?
      Lazy loaded image 10.7 Dropout and DropPath
      🐝 VLM Image Encoders
      Lazy loaded image 11.0 Introduction
      Lazy loaded image 11.1 Vision Transformer
      Lazy loaded image 11.2 CLIP
      Lazy loaded image 11.3 BLIP
      Lazy loaded image 11.4 EVA
      Lazy loaded image 11.6 BLIP-2
      Lazy loaded image 11.7 SigLIP
      📝 MISC.
      🤔 Large Language Models Pass the Turing Test.
      Lazy loaded image OpenAI ChatGPT Progression
      ⏳ Claude Sonnet 4 Test with Sandtris
      🔍 What DeepSeek-R1's Surge Teaches Us About SEO
      💰 CivitAI’s Payment Issue
      🎓 arXivSub, an Academia Assistant