Vision Transformers and Modern Vision Models

Convolutions use local neighborhoods and share weights across locations. These are useful preferences for many image tasks, although striding and boundary handling limit exact translation equivariance. A Vision Transformer (ViT) uses attention to mix information across image patches. It still groups nearby pixels into patches and shares a projection across them; its geometry is organized differently from a CNN’s, not discarded.

We will follow an image through patch vectors, positional information, attention blocks, and a classifier. The attention explanation below is enough to trace that route; the linked transformer article develops the mechanism in more detail. The later sections compare how Swin, DETR, and ConvNeXt change the design or the comparison.

Patch embedding

A 224-by-224 image contains 50,176 pixels. Dense attention over that many tokens has 2,517,630,976 pairwise scores per head per image; explicitly storing them in float32 takes about 9.38 GiB before other activations. That is expensive, not mathematically impossible. ViT groups pixels into non-overlapping \(P\times P\) patches. With \(P=16\), there are \(14\times14=196\) patch tokens. Each RGB patch contains \(16\times16\times3=768\) values, projected to a vector of model width \(D\).

A convolution with kernel and stride both \(P\), no padding, and shared weights implements this patch projection. The projection weights distinguish relative pixel positions within a patch; the same projection is applied at every patch location. Run the Python blocks in order. The following implementation deliberately supports one fixed square image size and rejects other sizes, including inputs that would otherwise lose border pixels during convolution.

import torch
import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch=16, in_ch=3, dim=768):
        super().__init__()
        if img_size <= 0 or patch <= 0 or img_size % patch:
            raise ValueError("img_size must be a positive multiple of patch")
        self.img_size, self.in_ch = img_size, in_ch
        self.n_patches = (img_size // patch) ** 2
        self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.empty(1, self.n_patches + 1, dim))
        nn.init.trunc_normal_(self.pos, std=0.02)

    def forward(self, x):
        if (x.ndim != 4 or x.shape[1] != self.in_ch
                or x.shape[-2:] != (self.img_size, self.img_size)):
            raise ValueError("input must match the configured channels and square image size")
        x = self.proj(x)
        x = x.flatten(2).transpose(1, 2)
        cls = self.cls.expand(x.size(0), -1, -1)
        x = torch.cat([cls, x], dim=1)
        return x + self.pos

pe = PatchEmbed()
print(pe.n_patches, pe(torch.randn(2, 3, 224, 224)).shape)
# 196 torch.Size([2, 197, 768])

The printed shape (2, 197, 768) means two images, 196 patch tokens plus CLS, and 768 features per token. The convolution produces a grid, then flatten(2).transpose(1, 2) turns it into a sequence ordered by spatial position. A learned [CLS] vector is prepended as token 0; attention will update it using the image tokens, and its final representation feeds the classifier. The learned positional table adds a different vector at each sequence position. With dropout disabled and no positional signals or position-dependent masks, self-attention is permutation-equivariant: permuting the tokens permutes their outputs. A classifier reading only the unchanged CLS slot is then invariant to permutations of patch tokens. Patch content can suggest location, but no explicit patch-grid coordinate has been supplied.

When fine-tuning at a new resolution, the patch grid and positional table must agree. A common approach reshapes the patch-position vectors into their old 2-D grid, interpolates to the new grid, and keeps the CLS position separate. Simply accepting a new token count—or a different grid with the same count—does not handle positional geometry correctly. This fixed-size example leaves interpolation to a model implementation that supports it.

Self-attention, masking, and the full block are derived in Transformers from Scratch: Self-Attention to Encoder–Decoder.

From patch tokens to a class prediction

For one attention head, let \(Z\) contain the token vectors. Learned projections produce queries \(Q=ZW_Q\), keys \(K=ZW_K\), and values \(V=ZW_V\). The head computes \(A=\operatorname{softmax}(QK^\top/\sqrt{d_k})\) row by row and returns \(AV\), where \(d_k\) is the query/key width. Each row of \(A\) weights the value vectors for one output token. With weights \([0.25,0.75]\) and scalar values \([2,6]\), the weighted output is 5. Multiple heads learn different mixtures, then a projection combines their outputs.

ViT uses pre-norm blocks: normalize tokens before attention and add its output to the incoming tokens; normalize again before a GELU feed-forward network and add that result. The feed-forward network applies the same transformation to each token independently; attention mixes tokens. The original 2017 Transformer used post-norm, so the normalization order is a real difference. ViT-B/16 uses 12 layers, 12 heads, width 768, and feed-forward width 3072, with roughly 86 million parameters depending on the classifier. At 224-by-224 resolution, its encoder keeps 197 tokens and width 768 throughout. Reading CLS is one aggregation choice; other variants pool patch tokens.

class TinyViT(nn.Module):
    def __init__(self, img_size=32, patch=8, dim=32, heads=4, depth=2, classes=3):
        super().__init__()
        self.embed = PatchEmbed(img_size, patch, 3, dim)
        self.blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=dim, nhead=heads, dim_feedforward=4 * dim,
                dropout=0.0, activation="gelu", batch_first=True, norm_first=True)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, classes)

    def forward(self, images):
        tokens = self.embed(images)
        for block in self.blocks:
            tokens = block(tokens)
        return self.head(self.norm(tokens[:, 0]))

torch.manual_seed(0)
model = TinyViT()
images = torch.randn(2, 3, 32, 32)
labels = torch.tensor([0, 2])
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
model.train()
optimizer.zero_grad(set_to_none=True)
logits = model(images)
loss = nn.functional.cross_entropy(logits, labels)
before = model.embed.proj.weight.detach().clone()
loss.backward()
optimizer.step()
print(tuple(logits.shape))
print("patch weights updated:", not torch.equal(before, model.embed.proj.weight))
# (2, 3)
# patch weights updated: True

model.eval()
with torch.no_grad():
    predictions = model(images).argmax(dim=1)
print(tuple(predictions.shape))
# (2,)

This small model turns 16 patch vectors plus CLS into three class logits for each image. Cross-entropy receives logits directly and integer class indices; the example verifies one update from the classification loss back to the patch projection. Random inputs and arbitrary labels do not measure recognition quality. Dropout is disabled to keep the example simple, and no causal mask is used: image tokens can attend to every token in the sequence.

Data and training recipes

The original ViT experiments found that large-scale pretraining helped plain ViT compete with CNNs, while their smaller-data settings favored stronger convolutional biases. This was an experimental result under particular recipes, not a fixed data threshold. Patch construction already supplies local structure; the attention blocks leave more spatial relationships to be learned. The usefulness of that design depends on pretraining, augmentation, optimization, and the downstream task.

The ViT paper observed a mix of short- and long-range attention in earlier layers and more global attention later. That pattern is not a rule every head must follow. CNNs can also aggregate global context through depth, pooling, or additional modules; attention supplies direct content-dependent interactions between distant tokens within a layer.

DeiT demonstrated competitive ImageNet-only training with a stronger recipe, including models trained without distillation; a teacher and an additional distillation token improved results further. MAE pretrains by reconstructing masked image patches from visible ones, avoiding manual class labels for that pretraining objective. Supervised downstream fine-tuning can still use labels, and pretraining still requires data and compute.

Swin: reintroducing hierarchy

Dense prediction often benefits from high-resolution or multi-scale features. Plain ViT keeps one token grid, and dense attention’s pairwise work grows quadratically with token count. Swin limits attention to windows of \(M\times M\) tokens and alternates regular and shifted partitions. For \(T\) tokens with fixed window size, there are \(T/M^2\) windows and \(M^4\) scores per window, or \(TM^2\) scores per head in total, assuming complete windows. Its attention cost is linear in \(T\) when window size and width are fixed.

Shifting the partition allows tokens formerly separated by a window boundary to interact. The cyclic-shift implementation uses a mask to prevent artificial wraparound connections; it is not just rolling a tensor and applying unrestricted attention. Patch merging between stages halves each spatial dimension and increases feature width, providing a hierarchy for detection or segmentation heads. Framework integration still needs compatible feature scales and channels; Swin is one approach to dense prediction, not the sole route.

DETR: detection as set prediction

Original DETR combines a CNN backbone with a transformer encoder–decoder and learned object-query embeddings. Each query produces a box and class scores, including a no-object class. It uses a fixed query count, so that count limits how many objects can be returned. Its set-prediction objective supports inference without hand-designed anchors or an NMS stage; DETR is not simply a plain ViT backbone.

During training, Hungarian matching selects a minimum-cost one-to-one assignment from ground-truth objects to predictions, assuming enough query slots. Matched predictions receive class and box losses; unmatched ones receive the no-object class target. This discourages duplicates but does not guarantee their absence at inference. Matching is a discrete operation: gradients are taken through the losses for the selected pairs, not through the assignment itself. “End-to-end training” does not mean every operation is differentiable.

The original DETR used a long training schedule and had weaker small-object results than some contemporary detectors. Deformable DETR introduced attention over a limited set of learned sampling locations, including multiple feature scales, and reported faster convergence and improved small-object performance. These results address the observed limitations without guaranteeing their elimination on every dataset.

ConvNeXt: the control experiment

ConvNeXt studied both training and architectural changes starting from a ResNet baseline. The experiments modernized optimization, training length, augmentation, and regularization, then changed the convolutional design with elements such as depthwise larger kernels, an inverted bottleneck, and LayerNorm. These are two different sources of improvement: the recipe and the architecture.

The resulting convolutional models were competitive with Swin in the paper’s comparisons. That supports testing a modern convolutional baseline; it does not isolate training as the cause of most transformer gains across papers. Parameter count or FLOPs alone also does not determine deployment latency. A comparison should state the checkpoint, training data, recipe, input resolution, and measured resource budget.

The development of earlier convolutional backbones is traced in CNN Architectures Compared: LeNet to EfficientNet.

Choosing

A checkpoint trained on relevant data can matter more than a small architectural difference, but fine-tuning cost and results still need measurement. Use a representative validation split, preserve each checkpoint’s preprocessing requirements, and give candidates comparable tuning opportunities. A subset can screen obvious failures; confirm the choice at the data scale and resolution you intend to use.

Masked reconstruction and other pretraining objectives, including text–image pairing, are covered in Self-Supervised Learning: SimCLR, BYOL, MAE, and CLIP.

Exercises

1. Why patches. Compute the number of tokens and the size of the attention matrix for a \(224\times224\) image treated as individual pixels, then for \(16\times16\) patches. Report both and the ratio, excluding CLS, batch size, and multiple heads from this calculation.

You should get: a pixel-level attention matrix with billions of entries and a patch-level one with tens of thousands.

px = 224*224
patch = (224//16)**2
for n, name in ((px, "pixels"), (patch, "16x16 patches")):
    print(f"{name:14s} tokens {n:7d}  attention entries {n*n:,}")
print("ratio", round((px*px)/(patch*patch)))
# pixels          tokens   50176  attention entries 2,517,630,976
# 16x16 patches   tokens     196  attention entries 38,416
# ratio 65536

2. Patch embedding is a convolution. Show that a strided convolution with kernel size equal to stride produces exactly the same result as flattening each patch and applying a linear layer. Verify numerically.

You should get: two tensors that match to floating-point precision.

import torch, torch.nn as nn
torch.manual_seed(0)
P, C, D = 4, 3, 8
conv = nn.Conv2d(C, D, kernel_size=P, stride=P)
x = torch.randn(1, C, 8, 8)                        # 2x2 = 4 patches

via_conv = conv(x).flatten(2).transpose(1, 2)      # (1, 4, D)

lin = nn.Linear(C*P*P, D)
with torch.no_grad():
    lin.weight.copy_(conv.weight.reshape(D, -1))
    lin.bias.copy_(conv.bias)
patches = x.unfold(2, P, P).unfold(3, P, P).reshape(1, C, 4, P, P)
patches = patches.permute(0, 2, 1, 3, 4).reshape(1, 4, C*P*P)
print(torch.allclose(via_conv, lin(patches), atol=1e-5))
# True

3. What ConvNeXt controlled for. A paper reports that architecture X beats a ResNet-50 by 4 points on ImageNet. The ResNet baseline was trained with the original 2015 recipe. List three confounds this comparison does not control, and state what a fair comparison requires.

You should get: three training-recipe differences, none of which is architectural.

References

• Carion et al. (2020). End-to-End Object Detection with Transformers. ECCV.
• Dosovitskiy et al. (2021). An Image Is Worth 16×16 Words: Transformers for Image Recognition at Scale. ICLR.
• Touvron et al. (2021). Training Data-Efficient Image Transformers and Distillation Through Attention.
• Liu et al. (2021). Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows. ICCV.
• Zhu et al. (2021). Deformable DETR: Deformable Transformers for End-to-End Object Detection. ICLR.
• He et al. (2022). Masked Autoencoders Are Scalable Vision Learners. CVPR.
• Liu et al. (2022). A ConvNet for the 2020s. CVPR.
• Radford et al. (2021). Learning Transferable Visual Models From Natural Language Supervision.