Patch Embedding

The input image x ∈ R^(H×W×C) is divided into a grid of N = (H/P) × (W/P) non-overlapping 2D patches, each of size P×P×C. Each patch is then flattened into a 1D vector of length P²·C. For a standard ViT-B/16 on a 224×224 image, this produces 196 patches, each of dimension 768 (16×16×3). The patch size P is a critical hyperparameter that controls the sequence length and granularity of the visual tokens.

Each flattened patch vector is multiplied by a learnable embedding matrix E ∈ R^((P²·C)×D) to project it into a D-dimensional latent space. This operation is mathematically equivalent to a 2D convolution with kernel size P and stride P, using D output filters. The resulting sequence of N patch embeddings forms the input tokens to the Transformer, alongside a prepended learnable [class] token for classification tasks.

Since self-attention is permutation-invariant, spatial information must be explicitly added. Standard learnable 1D positional embeddings of shape (N+1)×D are summed with the patch tokens. Without this step, the model cannot distinguish between a patch in the top-left corner and one in the bottom-right. Alternatives include 2D-aware positional encodings, relative position biases, and Rotary Position Embeddings (RoPE) for improved spatial reasoning.

While the original ViT uses a single large-kernel projection, modern architectures employ overlapping convolutional patch embeddings for smoother feature extraction:

• ConvStem: A stack of small 3×3 convolutions replaces the single large stride, improving optimization stability.
• Overlapping patches: Using stride smaller than kernel size creates redundant tokens that capture finer local structure.
• Multi-scale embeddings: Hierarchical models like Swin Transformer extract patch tokens at multiple resolutions using patch merging layers.

The patch size P directly governs the sequence length N and thus the O(N²) cost of self-attention:

• Smaller P (e.g., 8): More tokens (e.g., 784 for 224×224), finer granularity, higher compute cost.
• Larger P (e.g., 32): Fewer tokens (e.g., 49), coarser granularity, faster inference. This trade-off is critical in medical imaging, where high-resolution scans (e.g., 1024×1024 pathology tiles) require careful patch sizing or hierarchical architectures to remain computationally tractable.

Patch embedding creates the initial token set, but many tokens—especially in background regions—are redundant. Token Merging (ToMe) is a training-free technique that operates downstream of patch embedding to progressively combine similar tokens, reducing sequence length by 40-60% with negligible accuracy loss. This is particularly valuable in medical image analysis, where large uniform tissue regions generate many near-identical patch embeddings that can be safely merged.

The core operation that processes the sequence of patch embeddings by computing a weighted sum of all token representations. For each patch, query (Q), key (K), and value (V) vectors are derived via learned linear projections. The attention weight between any two patches is the scaled dot-product of their Q and K vectors, enabling every patch to globally attend to every other patch.

• Computational complexity: O(N²) where N is the number of patches, making it expensive for high-resolution images
• Multi-head attention: Runs multiple attention operations in parallel, allowing the model to focus on different relationship types simultaneously
• Global receptive field: Unlike CNNs, every patch can directly interact with every other patch from the first layer

An architecture variant where the input to the Transformer is not raw image patches but feature maps extracted by a convolutional neural network. Instead of flattening pixel patches, a CNN backbone processes the image first, and the resulting feature map is treated as a grid of patch embeddings. This blends the local inductive bias of convolutions with the global modeling capacity of self-attention.

• Early convolution stages: Capture low-level features like edges and textures efficiently
• Reduced sequence length: CNN downsampling produces fewer tokens than raw patching, lowering attention cost
• Example: The original ViT paper explored a hybrid model using ResNet feature maps, showing competitive performance on smaller datasets