Vision Transformers: patch equivalence, position and attention cost

A Vision Transformer turns an image into a sequence of patch features,then applies attention and tokenwise feed-forward blocks. Understanding the interfaces—patch order,projection,position,readout and resolution—is more useful than saying attention replaced convolution.

What the ViT result establishes

Dosovitskiy et al.'s An Image is Worth16×16 Words studies transformer image classifiers across pretraining scales. In the paper's experiments,large-scale pretraining lets a comparatively weak-spatial-prior model perform strongly. It is not a theorem that CNNs hit a universal ceiling,that ViT always loses below10M images,or that every vision team deployed a ViT within two years.

DeiT subsequently trains competitive image transformers with ImageNet-only data using a different recipe and teacher-distillation design. ConvNeXt modernizes both CNN architecture and training. These are counterexamples to an architecture ranking based only on dataset size. Compare exact model,pretraining data,training recipe,compute and evaluation before declaring a winner.

Outcomes and prerequisites. Starting from DL025's windowed dot products and DL029's augmentation choices,derive patchification two ways,build a complete small ViT,understand CLS versus average pooling,prove permutation equivariance,test resolution interpolation and count quadratic attention work. TransformerEncoderLayer supplies the attention block today; later lessons open its internals. CPU tests use random tiny tensors,not downloaded checkpoints or image benchmarks.

The full patch-to-logits contract

For B RGB images224×224 and square patch sideP16,there are14×14=196 patches. Each flattened patch contains3×16×16=768 numbers. A learned matrix maps those numbers to embedding widthD; D768 being equal to the input width does not make the matrix identity.

(B,3,224,224)
  -> non-overlapping patches (B,196,768)
  -> learned affine projection (B,196,D)
  -> prepend one CLS token (B,197,D)
  -> add position table (1,197,D), then dropout
  -> pre-norm transformer blocks (B,197,D)
  -> final normalization, select CLS
  -> linear classifier (B,classes), raw logits

The projection uses768 weights per output channel and D×768 weights overall,not3072 per channel. One patch can mix all its pixels/channels before attention. No spatial attention is needed to produce the initial embeddings.

Patch embedding as unfold plus linear or convolution

Both implementations below use row-major spatial patch order and channel-then-row-then-column order within each flattened patch. The explicit permutation after unfold is essential; a bare reshape of (B,C,gridH,gridW,P,P) interleaves the wrong axes. Matching weights and biases makes the functions equal; independent random initialization does not.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
 
class PatchEmbeddingUnfold(nn.Module):
    def __init__(self,img_size=224,patch_size=16,in_ch=3,embed_dim=768):
        super().__init__()
        if any(type(v) is not int or v<1 for v in (img_size,patch_size,in_ch,embed_dim)):
            raise ValueError('image, patch, channels and embedding sizes must be positive integers')
        if img_size%patch_size: raise ValueError('image must tile exactly')
        self.img_size=img_size;self.patch_size=patch_size;self.in_ch=in_ch
        self.num_patches=(img_size//patch_size)**2
        self.proj=nn.Linear(in_ch*patch_size**2,embed_dim)
    def forward(self,x):
        if x.ndim!=4 or x.shape[1:]!=(self.in_ch,self.img_size,self.img_size):
            raise ValueError('wrong fixed image shape')
        B=x.shape[0];p=self.patch_size
        windows=x.unfold(2,p,p).unfold(3,p,p)
        patches=windows.permute(0,2,3,1,4,5).reshape(B,self.num_patches,-1)
        return self.proj(patches)
 
class PatchEmbeddingConv(nn.Module):
    def __init__(self,img_size=224,patch_size=16,in_ch=3,embed_dim=768):
        super().__init__()
        if any(type(v) is not int or v<1 for v in (img_size,patch_size,in_ch,embed_dim)):
            raise ValueError('image, patch, channels and embedding sizes must be positive integers')
        if img_size%patch_size: raise ValueError('image must tile exactly')
        self.img_size=img_size;self.patch_size=patch_size;self.in_ch=in_ch
        self.num_patches=(img_size//patch_size)**2
        self.proj=nn.Conv2d(in_ch,embed_dim,patch_size,stride=patch_size)
    def forward(self,x):
        if x.ndim!=4 or x.shape[1:]!=(self.in_ch,self.img_size,self.img_size):
            raise ValueError('wrong fixed image shape')
        return self.proj(x).flatten(2).transpose(1,2)
 
torch.manual_seed(30)
u=PatchEmbeddingUnfold(8,2,3,6).double()
c=PatchEmbeddingConv(8,2,3,6).double()
with torch.no_grad():
    c.proj.weight.copy_(u.proj.weight.reshape(6,3,2,2));c.proj.bias.copy_(u.proj.bias)
x=torch.randn(2,3,8,8,dtype=torch.float64,requires_grad=True)
x2=x.detach().clone().requires_grad_()
a,b=u(x),c(x2)
torch.testing.assert_close(a,b,rtol=1e-12,atol=1e-12)
g=torch.randn_like(a);(a*g).sum().backward();(b*g).sum().backward()
torch.testing.assert_close(x.grad,x2.grad,rtol=1e-12,atol=1e-12)
torch.testing.assert_close(u.proj.weight.grad,c.proj.weight.grad.reshape(6,-1),rtol=1e-12,atol=1e-12)
torch.testing.assert_close(u.proj.bias.grad,c.proj.bias.grad,rtol=1e-12,atol=1e-12)
wrong=x.detach().unfold(2,2,2).unfold(3,2,2).reshape(2,16,-1)
assert not torch.allclose(u.proj(wrong),a.detach())
try:
    c(torch.randn(1,3,9,8,dtype=torch.float64));raise AssertionError('silent border drop')
except ValueError:
    pass
print('patch forward/backward equivalence and wrong-order control passed')

A strideP,kernelP convolution is a concise implementation of the same shared patch projection. Whether it is faster than unfold+linear depends on device,kernel,shape and memory layout; a3–5× promise is not part of the identity. “No convolution” in the original paper describes removal of a convolutional backbone; the patch projection can be written as a convolution. Overlapping patch embeddings or convolutional stems introduce additional local structure and different token counts.

CLS is a learned readout, not a required law

The class token is a shared learned vector with shape(1,1,D),expanded over B and prepended at position0. Through attention it combines information from patch tokens. Classification gradients reach both the token and the patches it attends to. It does not independently know the image before attention.

Mean pooling over the final patch features is a legitimate alternative. ViT's appendix finds a properly tuned global-average-pooling variant can work comparably; there is no universal.5%-worse rule. Selecting the first image patch as readout is also architecturally possible after global attention,but gives that position a special role. For segmentation,dense patch outputs may be retained instead of discarded.

Position, permutation equivariance and invariance

For unmasked self-attention without positional structure,permuting token rows permutes output rows. This is equivariance. A symmetric mean readout is invariant; a fixed CLS readout is invariant to permutation of the other patch tokens when no position signal distinguishes them. A CNN instead applies weights at relative grid offsets; it need not encode absolute corner identity either.

If Q,K,V are token projections and P is a permutation matrix,then Q′=PQ,K′=PK,V′=PV and row-softmax(PQKᵀPᵀ)=P softmax(QKᵀ)Pᵀ. Multiplying by PV yields P times the original attention result. Tokenwise LayerNorm/MLPs preserve this symmetry. Fixed positional signals or masks can break it; stochastic dropout also complicates exact samplewise comparisons,so test in eval mode.

torch.manual_seed(30)
attn=nn.MultiheadAttention(16,4,batch_first=True,dropout=0.).double().eval()
tokens=torch.randn(2,7,16,dtype=torch.float64); perm=torch.tensor([2,5,0,6,1,4,3])
a,_=attn(tokens,tokens,tokens,need_weights=False)
b,_=attn(tokens[:,perm],tokens[:,perm],tokens[:,perm],need_weights=False)
torch.testing.assert_close(a[:,perm],b,rtol=1e-10,atol=1e-10)
position=torch.randn(1,7,16,dtype=torch.float64)
a,_=attn(tokens+position,tokens+position,tokens+position,need_weights=False)
b,_=attn(tokens[:,perm]+position,tokens[:,perm]+position,tokens[:,perm]+position,need_weights=False)
assert not torch.allclose(a[:,perm],b)

A learned position table is added,not concatenated,so widthD is unchanged. Its CLS row is separate from the spatial grid. It is a learned parameter,not a new random position vector each forward. Dropout may act after the addition; do not forget the addition and interpret successful compilation as proof of spatial awareness.

Resolution changes and a tested interpolation policy

224/P16 uses14×14 spatial positions;384/P16 uses24×24. Keep the CLS position row unchanged,reshape only spatial rows to a2D grid,interpolate,then flatten in the same order. The original paper describes2D interpolation; implementations choose bilinear or bicubic,alignment and antialiasing policies. Below is an explicit bicubic,align_corners=False choice,not a claim that all loading libraries do this automatically.

def resize_position(table,old_grid,new_grid):
    if any(type(v) is not int or v<1 for v in (old_grid,new_grid)):
        raise ValueError('grid sides must be positive integers')
    if table.ndim!=3 or table.shape[0]!=1 or table.shape[1]!=old_grid**2+1:
        raise ValueError('invalid position table')
    if old_grid==new_grid: return table.clone()
    cls=table[:,:1];spatial=table[:,1:].reshape(1,old_grid,old_grid,-1).permute(0,3,1,2)
    resized=F.interpolate(spatial,size=(new_grid,new_grid),mode='bicubic',align_corners=False)
    return torch.cat([cls,resized.permute(0,2,3,1).reshape(1,new_grid**2,-1)],dim=1)
 
table=torch.randn(1,197,8)
resized=resize_position(table,14,24)
assert resized.shape==(1,577,8) and torch.equal(resized[:,:1],table[:,:1])
assert torch.equal(resize_position(table,14,14),table)
constant=torch.ones(1,197,8)
torch.testing.assert_close(resize_position(constant,14,24),torch.ones(1,577,8))

Changing patch size also changes projection weights; position interpolation alone does not solve it. A fixed-size model should reject a mismatched input instead of dropping border pixels or filling unknown positions with zeros. For a real checkpoint,adapt the new model's projection/position contract explicitly and validate its new resolution rather than merely making tensors fit.

A complete small ViT with independent block initialization

The pre-norm block is attention on normalized tokens plus a residual,then an MLP on normalized tokens plus a residual. QKᵀ scores are scaled by the square root of head width before softmax. LayerNorm normalizes features per token,not batch statistics. GELU is one activation choice; many modern LLMs instead use SiLU-based gated feed-forwards,so “every LLM uses GELU” is false.

The old implementation used TransformerEncoder clones without discussing its same-initial-parameters warning. This version constructs each block separately,so weights initialize independently. They still share the same architectural hyperparameters,not the same learned matrices; deterministic LayerNorm scale/bias initial values may intentionally match. It returns logits and accepts either CLS or mean readout for controlled comparison. The mean option averages the final patch outputs but retains CLS inside the encoder, where it can influence those patches through attention. It is a readout ablation, not a separately implemented CLS-free model. All image/patch/grid sizes in these helpers are positive integer square-side contracts; rectangular position grids need two explicit dimensions, not a guessed square root.

class ViT(nn.Module):
    def __init__(self,img_size=224,patch_size=16,in_ch=3,embed_dim=768,depth=12,
                 num_heads=12,mlp_ratio=4,num_classes=1000,dropout=.1,readout='cls'):
        super().__init__()
        if any(type(v) is not int or v<1 for v in (embed_dim,num_heads,depth,num_classes)):
            raise ValueError('width, heads, depth and class count must be positive integers')
        if not math.isfinite(mlp_ratio) or int(embed_dim*mlp_ratio)<1:
            raise ValueError('MLP hidden width must be positive and finite')
        if embed_dim%num_heads or readout not in ('cls','mean'):
            raise ValueError('invalid width, depth or readout')
        self.patch_embed=PatchEmbeddingConv(img_size,patch_size,in_ch,embed_dim)
        self.cls_token=nn.Parameter(torch.zeros(1,1,embed_dim))
        self.pos_embed=nn.Parameter(torch.zeros(1,self.patch_embed.num_patches+1,embed_dim))
        self.dropout=nn.Dropout(dropout);self.readout=readout
        self.encoder=nn.ModuleList([nn.TransformerEncoderLayer(
            d_model=embed_dim,nhead=num_heads,dim_feedforward=int(embed_dim*mlp_ratio),
            dropout=dropout,activation='gelu',batch_first=True,norm_first=True)
            for _ in range(depth)])
        self.norm=nn.LayerNorm(embed_dim);self.head=nn.Linear(embed_dim,num_classes)
        nn.init.trunc_normal_(self.cls_token,std=.02);nn.init.trunc_normal_(self.pos_embed,std=.02)
    def forward(self,x):
        x=self.patch_embed(x)
        x=torch.cat([self.cls_token.expand(len(x),-1,-1),x],dim=1)
        x=self.dropout(x+self.pos_embed)
        for block in self.encoder: x=block(x)
        x=self.norm(x)
        return self.head(x[:,0] if self.readout=='cls' else x[:,1:].mean(1))
 
torch.manual_seed(30)
model=ViT(32,4,3,24,2,4,num_classes=3,dropout=0.).double()
assert not torch.equal(model.encoder[0].linear1.weight,model.encoder[1].linear1.weight)
xb=torch.randn(3,3,32,32,dtype=torch.float64);labels=torch.tensor([0,1,2])
opt=torch.optim.AdamW(model.parameters(),lr=.001)
old=model.head.weight.detach().clone()
logits=model(xb);assert logits.shape==(3,3)
loss=F.cross_entropy(logits,labels);opt.zero_grad();loss.backward();opt.step()
assert not torch.equal(old,model.head.weight) and torch.isfinite(model.pos_embed.grad).all()
model.eval()
with torch.no_grad():
    assert torch.equal(model(xb),model(xb))
    model.readout='mean';assert model(xb).shape==(3,3)
# Parameter inventory without allocating an86M-value CPU tensor model.
with torch.device('meta'):
    base=ViT()
counts={name:sum(p.numel() for p in module.parameters()) for name,module in base.named_children()}
total=sum(p.numel() for p in base.parameters())
assert counts['encoder']==85054464 and total==86567656
print('ViT base counts',counts,'total including CLS/position',total)

This is a complete forward/backward/update example. It is not pretrained ViT accuracy. To use local image data,plug the model into DL027's local train/validation loop with its own hyperparameter sweep and no architecture-quality promise.

Parameters, token cost and data tradeoffs

For embedding widthD and MLP ratio4,one biased attention+MLP block with two LayerNorms has $12D^2+13D$ parameters. ViTBase's12 blocks atD768 hold85,054,464. Patch projection590,592;position151,296;CLS768;final norm1,536;1000-class head769,000;total 86,567,656. Named-child iteration excludes standalone CLS/position parameters,so add them before comparing totals. Halving depth removes six blocks;it does not prove “more data needed per parameter.”

For imageH×W and patchP,patch count is HW/P² when divisible. HalvingP from16 to8 at224 changes196→784 patch tokens:4× as many. The full attention score matrix grows 16×,not4×. Including CLS gives197²→785²,about15.88×. QKV/MLP work grows roughly4× at fixedD/depth;not every operation is quadratic. Memory-efficient attention can avoid materializing the full matrix,while exact full attention still performs quadratic token interactions.

A CNN's local receptive field grows across layers. Full ViT attention can connect all patch tokens in one block,but this says nothing about how strongly a trained model uses those connections. Patches have internal pixel order even without a position table,so ViT is not literally devoid of all spatial prior. Pooling/windowed/hierarchical variants change these tradeoffs.

For limited labels,compare a suitable pretrained representation, a modest CNN and—if budget permits—a carefully tuned small ViT. If no matching pretrained model exists,CNN locality may be valuable. If transfer is available,pretraining data is part of the comparison. There is no scientifically fixed100M-image crossover or rule that scratch ViTs are always irrational at tens of thousands of examples.

Descendants and vision-language interfaces

• DeiT: teacher-student attention/distillation token plus a data-efficient training recipe. Swin: hierarchical shifted windows permit cross-window interaction with local-window complexity. Hiera: removes many vision-specific additions from a hierarchical transformer and leverages masked pretraining.
• DINOv2: combines self-supervised methods and curated large-scale data to learn reusable image/pixel features;the released family is not a blanket10B+ model. MAE reconstructs missing pixels,whereas EVA02 reconstructs strong language-aligned teacher features. These are distinct targets.
• SAM/SAM2: promptable image/video segmentation systems,not ordinary image classifiers. Dense heads need spatial features and prompt/memory interfaces.
• ViT22B: investigates a22B-parameter vision transformer with architecture/training changes,not simply this notebook made wider. Scale results are conditioned on its data and compute.
• PerceiverIO: uses cross-attention to map arbitrary inputs into a latent array and queries to structured outputs. This separates input/output size from latent processing,not just “patchify every modality.”
• CLIP: image/text contrastive alignment supports prompt-based zero-shot classification and retrieval;it does not itself become a generative captioning LLM. SigLIP: pairwise sigmoid objective. SigLIP2: adds captioning/self-supervised/data-curation techniques and multilingual/native-aspect-ratio variants.
• Perception Encoder (2025): studies contrastive vision-language encoders and using intermediate representations for different tasks,then alignment stages. Its paper does not establish that it is a drop-in replacement inside Llama4's released vision path.

For a VLM,an encoder emits image features,a projector adapts their width/sequence contract,and a language model conditions on those tokens. The projector and training objectives matter. Do not infer the exact encoder from a family name: PaliGemma2's primary abstract explicitly uses SigLIP-So400m,not SigLIP2;Molmo and Qwen2-VL have their own documented designs. Nor does every model start with kernel/stride16—DINOv2's familiar /14 variant already contradicts that claim.

Choose a backbone with task evidence: classification quality and latency;dense feature resolution;image–text retrieval alignment;or mobile operator support. MobileNetV4/FastViT are candidates for device-constrained evaluation,not guaranteed faster than all attention models. There is no single permanent “2025 default” for every task.

Worked answers and retained exercises

1. Full forward: patchify→shared affine projection→CLS→positions→encoder blocks→normalized readout→class logits. Average pooling is an alternative readout,not a failure to implement ViT.
2. Equivalent patchification: unfold,permute grid before channels,flatten,linear;or Conv2d(kernel=stride=P). Copy reshaped weights and biases,then compare outputs and gradients. Independent initialization is not an equivalence test.
3. Why positions? Unmasked position-free attention is permutation-equivariant. A symmetric readout cannot distinguish patch arrangement;position signals break that symmetry. CNN relative-grid structure is a different spatial prior,not free absolute coordinates.
4. What happens to CLS? A shared learned initial token aggregates through attention and its final normalized vector is read by the head. Patch outputs may instead be used for dense tasks.
5. When does each architecture win? Only under a stated data,pretraining,recipe and compute comparison. The original ViT scaling result and DeiT's ImageNet-only result cannot be compressed into a universal threshold.
6. What does self-supervision say about labels? Rich learning signals can come from views,masking and text pairing;downstream labels and trustworthy evaluation remain important. It does not prove all supervised pretraining is obsolete.
7. **Patch16→8 cost:**196→784 spatial tokens,16× score elements without CLS;4× linear tokenwise work. The original4×-attention answer confused these terms.
8. Resolution stretch: run both exact patch implementations and the interpolation checks. If a local pretrained timm/torchvision model is already available,inspect its patch stem and loading policy without downloading anything. No such checkpoint is used here.

Tomorrow reconstruct the shape chain;next week derive the attention permutation equation and base-model parameter count without opening the page. The next lesson changes image patches to word IDs and studies what a learned embedding geometry does—and does not—mean.

Sources and execution boundary

Primary reading: ViT, DeiT, PyTorch2.9 TransformerEncoder initialization warning, Swin, Hiera, DINOv2, ViT22B, PerceiverIO, SigLIP2, Perception Encoder, PaliGemma2.

The original optional videos (H5vpBCLo74U, wjZofJX0v4M) and cs231n.github.io/transformers/ carried specific title/channel/content claims that are unverified here. They are not evidence for this lesson. Use the primary ViT/DeiT papers and the displayed tested implementations for those objectives; no video identity or timestamp is inferred from its URL.

Executed: exact tiny CPU ViT update,patch output/input/weight-gradient parity,wrong-axis control,permutation/position controls,interpolation invariants and meta-device parameter inventory. Not executed: ImageNet/CIFAR training,pretrained-model loading,GPU throughput or paper reproduction. Previous: augmentation · Next: Word2Vec.