A single-file-per-concept, heavily annotated reimplementation of the Vision Transformer (ViT) encoder [1] and of Contrastive Language–Image Pre-training (CLIP) [2]. Every module is written to be read: the comments derive the math, justify the design choices, and flag the failure modes that the equations alone do not reveal.
The goal is not another timm — it is a reference implementation whose
correctness you can verify line by line, with the training recipe (which
matters as much as the architecture) written out explicitly rather than hidden
behind a config system.
image ─▶ patchify ─▶ +[CLS] ─▶ +pos ─▶ [pre-norm block]×L ─▶ LN ─▶ pool ─▶ head
│
text ─▶ tokenize ─▶ +pos ─▶ [causal block]×L ─▶ LN ─▶ [EOT] ──┴──▶ InfoNCE (CLIP)
- Installation
- Quickstart
- Repository layout
- Part I — The Vision Transformer
- Part II — CLIP
- Part III — The training recipe
- Model presets
- Testing
- Extending this code
- References
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txttorch>=2.0 is the only hard dependency (needed for the fused
scaled_dot_product_attention kernels). torchvision and pillow are required
only for the real-dataset paths; every training loop runs end to end on
synthetic data with --dummy if they are absent.
# 1. Verify the whole stack on synthetic data (no downloads, runs on CPU)
python train.py --dummy --epochs 2 --model vit_tiny --img-size 32 --patch-size 4
# 2. Train a ViT on CIFAR-10
python train.py --dataset cifar10 --model vit_tiny --img-size 32 --patch-size 4 \
--epochs 200 --batch-size 128 --lr 1e-3 --drop-path 0.1
# 3. Train on an ImageNet-style folder tree (<root>/train/<class>/*.jpg)
python train.py --dataset folder --data-dir /path/to/data --model vit_base \
--img-size 224 --patch-size 16 --batch-size 256 --epochs 300 --ema
# 4. Contrastive image–text pre-training
python train_clip.py --dummy --epochs 6 --model clip_tiny --img-size 32 --patch-size 4
python train_clip.py --data pairs.json --image-root ./images --model clip_base \
--img-size 224 --batch-size 256 --epochs 30Using the library directly:
import torch
from vision_encoder import create_vision_encoder, create_clip, SimpleTokenizer
model = create_vision_encoder("vit_base", img_size=224, patch_size=16, num_classes=1000)
logits = model(torch.randn(2, 3, 224, 224)) # (2, 1000)
tokens = model.forward_features(torch.randn(2, 3, 224, 224)) # (2, 197, 768)
# Dense features for segmentation/detection decoders, as a (B, D, gh, gw) map
maps = model.get_intermediate_layers(torch.randn(1, 3, 224, 224), n=4, reshape=True)
# Zero-shot classification with CLIP
tok = SimpleTokenizer.build(["a photo of a cat", "a photo of a dog"])
clip = create_clip("clip_base", vocab_size=tok.vocab_size)
W = clip.build_zeroshot_classifier([["a photo of a cat"], ["a photo of a dog"]], tok)
preds = clip.zero_shot_predict(torch.randn(4, 3, 224, 224), W).argmax(-1)| Path | Contents |
|---|---|
vision_encoder/patch_embed.py |
Image → patch tokens; the Conv2d tiling identity |
vision_encoder/pos_embed.py |
Learnable and 2D sin–cos positions; resolution interpolation |
vision_encoder/attention.py |
Multi-head self-attention (fused SDPA + reference path) |
vision_encoder/layers.py |
MLP, DropPath (stochastic depth), LayerScale |
vision_encoder/block.py |
Pre-norm Transformer block |
vision_encoder/encoder.py |
VisionTransformer, config dataclass, model presets |
vision_encoder/text_encoder.py |
Causal Transformer text tower |
vision_encoder/tokenizer.py |
Dependency-free tokenizer with byte-level backoff |
vision_encoder/clip.py |
Dual-encoder CLIP + symmetric InfoNCE loss |
vision_encoder/utils.py |
Param groups, LR schedule, EMA, mixup/cutmix, metrics |
train.py |
Supervised classification training |
train_clip.py |
Contrastive image–text training |
tests/test_models.py |
27 shape, invariance and gradient tests |
Notation.
An image is a grid; a Transformer eats a sequence. ViT bridges the two in the
crudest way that works: tile the image into non-overlapping
The Conv2d identity. Extracting patches, flattening them and applying a
shared kernel_size = stride = P:
when the stride equals the kernel size the windows tile the plane without
overlap, and each output position is one patch's dot product with the kernel.
One fused op replaces three. This equivalence is asserted numerically in
test_patch_embed_equals_unfold_plus_linear.
Why
Self-attention is permutation-equivariant: for any permutation matrix
The model therefore cannot distinguish a patch's location from its content —
a property this repository tests explicitly
(test_attention_is_permutation_equivariant). Position must be added to the
tokens. Two schemes are implemented.
Learnable (pos_embed_type="learnable", the original ViT): a free
parameter
Fixed 2D sin–cos ("sincos", used by MAE/DINO): for a patch at grid
position
The reason sinusoids help is that a shift
so relative offsets are linear maps — precisely the structure the attention dot product can exploit.
Changing resolution. A model pretrained at 224px has a interpolate_pos_embed
reshapes the table to 2D, resizes it bicubically, and flattens it back. This
happens automatically in forward, so a 32px-trained model accepts 64px input
without any code change (test_vit_handles_new_resolution).
Each token emits a query, a key and a value; the output is a similarity-weighted average of values.
with
Why
Un-normalized scores therefore grow like
Why several heads. A single softmax yields one averaging pattern per token
— one relation per layer. Splitting into
QK-normalization (--qk-norm) applies LayerNorm to
Implementation. The forward pass calls F.scaled_dot_product_attention,
which dispatches to FlashAttention [8]. Flash never materializes the
test_fused_and_manual_attention_agree).
Attention is, once its weights are fixed, a linear map on the values — the
only nonlinearity lives inside the softmax that produces the weights. The MLP
supplies the network's actual nonlinear capacity, applied per token and
therefore costing
The activation is GELU [11],
Note where the parameters live: per block, attention holds
Pre-norm vs post-norm. The original Transformer [3] normalized after the
residual add,
so gradient reaches early layers through the identity term even when
Stochastic depth [5] drops an entire residual branch per sample with
probability
LayerScale [6] multiplies each branch by a learned per-channel gain
Two options, both supported:
The [CLS] token is a learned vector, identical for every image, that owns no
pixels and is therefore free to act as an accumulator across --global-pool avg. Prefix tokens are excluded
from the mean, since they carry no spatial meaning.
Register tokens [7] (--num-registers) add a few extra learnable tokens
with no position. Trained ViTs otherwise hijack low-information background
patches as scratch space for global computation, producing high-norm artifacts
that corrupt attention maps and hurt dense downstream tasks. Registers give the
model dedicated scratch space instead.
Per block, with
| Component | FLOPs (MACs) | Parameters | Scaling |
|---|---|---|---|
| QKV projection | |||
| Attention scores |
— | ||
| Attention |
— | ||
| Output projection | |||
| MLP | |||
| Total |
The crossover where attention overtakes the projections is at
| Model | Resolution | Tokens | GMACs | Attention share |
|---|---|---|---|---|
| ViT-B/16 | 224 | 196 | 17.6 | 4.1% |
| ViT-B/16 | 384 | 576 | 55.5 | 11.1% |
| ViT-B/16 | 512 | 1024 | 107.0 | 18.1% |
Memory is the other half of the story: the score matrix is
--grad-checkpointing, ~30% extra compute for near-constant
activation memory) attacks from the other side.
Two encoders — the ViT above and a causal text Transformer — are trained to map an image and its caption to the same point on a unit hypersphere. Supervision comes from pairing alone, so the training set is the web rather than a labeled corpus.
For a batch of
The correct match is always the diagonal, so the loss is symmetric cross-entropy against the identity permutation:
The first term is image→text retrieval, the second text→image; averaging keeps
neither modality privileged. Each is InfoNCE [9] with one positive and
Two sanity values follow immediately and are asserted in the test suite: at
initialization
so that --dummy training shows
The InfoNCE objective lower-bounds the mutual information between the two views,
which is the formal reason CLIP needs enormous batches: the bound itself is
capped at
Practical consequence. Gradient accumulation does not substitute for batch size here: it enlarges the gradient but each micro-batch still sees only its own negatives. The real remedies are multi-device training with an all-gather over embeddings before the loss, or a MoCo-style memory queue. Treat single-GPU runs as mechanism demonstrations, not accuracy reproductions.
Because the text tower embeds arbitrary strings, a trained CLIP classifies without a classifier. Embed a prompt per class and take the nearest:
The text encoder has synthesized the weights of a linear classifier from
language alone. Averaging several phrasings per class before renormalizing —
prompt ensembling — reduces phrasing variance and is worth a few points of
accuracy at no training cost. See CLIP.build_zeroshot_classifier.
Text tower details. It is architecturally the same Transformer, with two
differences: 1D positional embeddings, and a causal mask [EOT] token is the only one whose receptive field spans the whole caption.
A ViT trained with a CNN recipe underperforms badly. Most of the reported gain between ViT (2020) and DeiT (2021) [13] came from the recipe, not the architecture — so it is written out here rather than buried.
Why ViTs need this. A CNN hard-codes locality and translation equivariance. A ViT hard-codes almost nothing: any patch may attend to any other from layer 1, so it must learn that neighboring patches are related. That flexibility is exactly why ViTs scale better on large corpora and overfit worse on small ones. Every item below injects, through data and regularization, a prior the architecture lacks.
AdamW [14], not Adam. Adam's L2 penalty enters through the gradient and is
then divided by
Parameter groups matter. Decay is applied only to tensors with
[CLS] token, a LayerScale param_groups_weight_decay, tested in test_no_weight_decay_covers_1d_params).
Warmup is not optional for Transformers. At initialization, predictions are
arbitrary and gradients are large and poorly correlated with any useful
direction; meanwhile Adam's second moment starts at
| Technique | Default | Why |
|---|---|---|
| RandomResizedCrop | scale |
Teaches scale invariance the architecture lacks |
| RandAugment [15] |
|
Photometric/geometric diversity |
| Mixup [16] | Linear interpolation of inputs and labels | |
| CutMix [17] | Regional pasting; label mixed by area | |
| RandomErasing [18] | Forces distributed rather than local evidence | |
| Label smoothing | Caps target confidence, calibrates logits | |
| Stochastic depth [5] |
|
Regularizes depth itself |
| Weight decay |
|
On matrices only |
| Grad clipping | Bounds rare gradient spikes | |
| EMA | Averages the SGD trajectory |
Mixup and CutMix both draw
For CutMix, SoftTargetCrossEntropy.
EMA keeps
bf16 is preferred over fp16: it has fp32's 8-bit exponent, so it needs no
loss scaling. With fp16 (5 exponent bits), small gradients underflow to zero,
and torch.amp.GradScaler multiplies the loss by a large factor before backward
and divides it out before the step — note that gradients must be unscaled
before clipping, or you would clip the scaled values. The script selects
automatically and falls back if the GPU lacks bf16 support.
In train_clip.py the similarity matrix is cast back to fp32 before the loss:
a softmax over cosines scaled by up to
The zero-initialized classifier head gives an exact, checkable starting point:
If your first loss is not that number, something is wrong before training even
begins. The --dummy path is the second check: random labels are unlearnable
in general, so training loss falling well below
Encoder-only (no head), at 224px with
| Preset | Params | GMACs | |||
|---|---|---|---|---|---|
vit_tiny |
192 | 12 | 3 | 5.5M | 1.25 |
vit_small |
384 | 12 | 6 | 21.7M | 4.60 |
vit_base |
768 | 12 | 12 | 85.8M | 17.56 |
vit_large |
1024 | 24 | 16 | 303.3M | 61.55 |
vit_huge |
1280 | 32 | 16 | 630.9M | 127.31 |
CLIP presets (vocabulary 30k, 224px,
| Preset | Vision | Text | Joint dim | Params |
|---|---|---|---|---|
clip_tiny |
|
|
192 | 10.5M |
clip_small |
|
|
384 | 44.2M |
clip_base |
|
|
512 | 139.7M |
For small images, override the patch size — CIFAR at 32px with
python train.py --dataset cifar10 --model vit_tiny --img-size 32 --patch-size 4pytest -q # 27 tests, ~1.5s on CPUThe suite targets the bugs that do not raise. Transformer code fails silently
— a transpose that mismatches heads and channels, positional embeddings
misaligned with the patch raster, a causal mask that leaks the future — and
each one shows up only as a slightly worse final number. Pinned properties
include: the Conv2d patch embedding equals explicit unfold + Linear; attention
rows are probability distributions; the fused and reference attention paths
agree; attention is permutation-equivariant without positional embeddings;
sin–cos positions are mutually distinguishable; interpolation preserves prefix
tokens; the text encoder cannot see the future; the InfoNCE loss equals
-
Masked autoencoding (MAE) [19]: drop 75% of patch tokens after
_pos_embed, encode the visible ones, and attach a shallow decoder that reconstructs pixels. The encoder here already returns full token sequences, which is all that is required. -
Self-distillation (DINO) [20]: run two augmented views, add a projection
head, and train a student against a momentum-EMA teacher —
ModelEmais already implemented. -
Dense prediction: use
get_intermediate_layers(..., reshape=True)to get$(B, D, g_h, g_w)$ feature maps from several depths and feed a decoder. -
Multi-GPU CLIP: wrap in
DistributedDataParalleland all-gather embeddings before the loss, so negatives scale with the world size rather than the per-device batch. -
Better tokenization: replace
SimpleTokenizerwith a real BPE (open_clip.get_tokenizer, or HuggingFaceCLIPTokenizer).encode_textonly needs a(B, context_length)integer tensor, so nothing else changes.
- Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021. arXiv:2010.11929
- Radford et al. Learning Transferable Visual Models From Natural Language Supervision. ICML 2021. arXiv:2103.00020
- Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
- Xiong et al. On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745
- Huang et al. Deep Networks with Stochastic Depth. ECCV 2016. arXiv:1603.09382
- Touvron et al. Going Deeper with Image Transformers (CaiT). ICCV 2021. arXiv:2103.17239
- Darcet et al. Vision Transformers Need Registers. ICLR 2024. arXiv:2309.16588
- Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
- van den Oord et al. Representation Learning with Contrastive Predictive Coding. 2018. arXiv:1807.03748
- Dehghani et al. Scaling Vision Transformers to 22 Billion Parameters. ICML 2023. arXiv:2302.05442
- Hendrycks & Gimpel. Gaussian Error Linear Units (GELUs). 2016. arXiv:1606.08415
- Geva et al. Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021. arXiv:2012.14913
- Touvron et al. Training data-efficient image transformers & distillation through attention (DeiT). ICML 2021. arXiv:2012.12877
- Loshchilov & Hutter. Decoupled Weight Decay Regularization. ICLR 2019. arXiv:1711.05101
- Cubuk et al. RandAugment: Practical automated data augmentation. NeurIPS 2020. arXiv:1909.13719
- Zhang et al. mixup: Beyond Empirical Risk Minimization. ICLR 2018. arXiv:1710.09412
- Yun et al. CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features. ICCV 2019. arXiv:1905.04899
- Zhong et al. Random Erasing Data Augmentation. AAAI 2020. arXiv:1708.04896
- He et al. Masked Autoencoders Are Scalable Vision Learners. CVPR 2022. arXiv:2111.06377
- Caron et al. Emerging Properties in Self-Supervised Vision Transformers (DINO). ICCV 2021. arXiv:2104.14294
Released under the MIT License.