Vision Transformers and Patch Tokens

Patchify, embed, add CLS and position, run pre-LN blocks. Token counts, parameter math and any-resolution packing.

Concepts and terms
Tasks and commands
Annotated code
Diagrams
Defaults and numbers
Troubleshooting
Questions and tips

Source

ViT pipeline#

From Lesson 04.14

Shape and cost math#

Grid: $inline$ patches. 224 at patch 16 gives 14x14 = 196 tokens.
Patch vector: $inline$ values. At P = 16 that is 768.
Projection: nn.Conv2d(3, D, kernel_size=P, stride=P) equals unfold then linear.
Sequence: patches + 1 CLS + registers. ViT-B/16 at 224 is 197.
Attention cost: $inline$ . 196 tokens give 38,416 scores. 150,528 pixels give 22.6 billion.
Head dim: head_dim = hidden / num_heads. 768 / 12 = 64.
Block params: $inline$ + 3k. About 7.1M at D = 768. 12 blocks are 85M.
Patch embed params: 3 * P * P * hidden + hidden. 590,592 for ViT-Base/16.

From Lesson 12.01, Lesson 19.58, Lesson 19.59, Lesson 04.14

Tiny ViT forward#

vit = ViT(image_size=64, patch_size=16, num_classes=10, dim=192, depth=6, num_heads=3)

x = self.patch(x)(N, 16, 192)

cls = self.cls_token.expand(x.size(0), -1, -1)

x = torch.cat([cls, x], dim=1)(N, 17, 192)

x = x + self.pos_embedlearned, trunc_normal std=0.02

for blk in self.blocks:

x = blk(x)pre-LN attention, then MLP

return self.head(self.ln(x[:, 0]))classify from [CLS] only

From 04.14 code/main.py, Lesson 04.14

Pre-LN block at ViT-B scale#

cfg = ViTConfig()224, 16, hidden 768, depth 12, heads 12

qkv = self.qkv(x).reshape(b, n, 3, h, hd).permute(2, 0, 3, 1, 4)

q, k, v = qkv[0], qkv[1], qkv[2]

scores = (q @ k.transpose(-2, -1)) * self.scalescale = 1 / sqrt(head_dim)

attn = F.softmax(scores, dim=-1)no mask. Encoder is bidirectional

x = x + self.attn(self.ln1(x), store_attn=store_attn)

x = x + self.ffn(self.ln2(x))hidden to 4 * hidden to hidden

tokens, cls = encoder(img)(1, 197, 768) and (1, 768)

From 19.59 code/main.py, Lesson 19.59

Tasks in code#

Task	Code
Pretrained ViT-B/16	`timm.create_model("vit_base_patch16_224")`
Flat patches, stdlib	`patchify(image, patch_size)`
Fixed 2D position table	`sinusoidal_2d(grid_h, grid_w, dim)`
Check the Conv2d math	`unfold_then_linear(x, weight, bias, patch_size)`
Sequence length	`seq_length(cfg)`
FLOPs per forward	`flops_per_forward(cfg)`
Pack mixed sizes	`pack_batch(images, patch)`
Block-diagonal mask	`build_dense_mask(pack)`

From Lesson 04.14, 19.58 code/main.py, 12.01 code/main.py, 12.06 code/main.py, 07.09 code/main.py

Commands#

Command	What it does
`python3 code/main.py` (19.58)	Front end gives (1, 197, 768)
`python3 code/main.py` (19.59)	Full encoder. About 86M params
`python3 -m unittest code/test_main.py`	Runs the shape and equivalence tests
`code/main.py` (12.01)	Patch tokenizer and geometry calculator

From Lesson 19.58, Lesson 19.59, Lesson 12.01

Pooling and position choices#

[CLS] token: its final hidden state is the image. Original ViT and CLIP use it.
Mean pool: average of patch outputs. SigLIP and DINOv2 use it.
Register tokens: 4 to 16 extra tokens absorb high-norm artifacts. Drop them before the LLM.
Learned 1D position: ties the model to the training resolution.
2D sinusoidal: fixed, 0 parameters. Half the dim is row, half is column.
2D-RoPE, M-RoPE: rotate q and k by (row, col). No position table.
VLM handoff: skip pooling. Every patch becomes an LLM input token.

From Lesson 12.01, Lesson 07.09, Lesson 19.58, Lesson 12.06

Model sizes#

Setting	Default	Note
ViT-B/16	86M	dim=768, depth=12, heads=12
ViT-Large	~307M	hidden=1024, depth=24, heads=16
ViT-Huge	~632M	Feature dim 1280
SigLIP So400M	400M	hidden=1152, depth=27, heads=16
Tiny ViT (lesson 4.14)	2.8M	dim=192, depth=6, CPU
MLP expansion	4x	2x underfits, 8x overfits
Register tokens	4	DINOv2 and SigLIP 2
MAE mask	75%	Encoder sees the visible 25%

From Lesson 04.14, Lesson 07.09, Lesson 19.59, Lesson 12.01

Patch size and token count#

Setting	Default	Note
224, patch 16	196 tokens	Standard
224, patch 32	49 tokens	Faster, lower resolution
224, patch 8	784 tokens	Finer, quadratic cost
SigLIP 2 SO400m/14 at 384	729 patch tokens	2026 default VLM tower
NaFlex budgets	256, 729, 1024	One checkpoint, pick at inference
AnyRes 672x672, 2x2 + thumbnail	2880	4 * 576 + 576
OCR and documents	1024-4096 tokens	Natural photos need 256-576

From Lesson 04.14, Lesson 12.01, Lesson 12.06

Patch-n'-pack a mixed batch#

From Lesson 12.06, 12.06 code/main.py

Troubleshooting#

Symptom	Cause	Fix
ViT loses to ResNet on ImageNet-1k	No CNN priors, too little data	DeiT recipe or MAE pretraining
Deep stack fails to train	Post-LN past 6-8 layers	Use pre-LN
New resolution breaks positions	Learned 1D table	Interpolate, or use 2D-RoPE
Squished text, lost chart labels	Square resize	Encode at native resolution
50%+ of tokens are padding	Pad to the longest side	Patch-n'-pack with a mask
High-norm artifact patches	No sink token	Add register tokens
ValueError on the grid	`patch_size` does not divide `image_size`	Pick a dividing patch size
`sinusoidal_2d` raises	dim not divisible by 4	Use a dim that is a multiple of 4

From Lesson 04.14, Lesson 12.01, Lesson 12.06, 19.58 code/main.py, 19.59 code/main.py

Questions#

Why patches, not pixels?: A 224x224 image is 150,528 pixel tokens. Attention is quadratic.
Why pre-LN?: It trains deep stacks stably without warmup.
Patch 14 or 16?: Patch 14 gives more tokens. Better for OCR and dense tasks, slower.
Which backbone for dense features?: DINOv2. Use CLIP or SigLIP to match text.
Learned or sinusoidal position?: Learned wins at fixed resolution. Sinusoidal wins when resolution changes.

From Lesson 04.14, Lesson 07.09, Lesson 12.01, Lesson 19.58