module

ViT

17 members

lucid.models.vision.vit

Vision Transformer (ViT) model family — Dosovitskiy et al., 2020.

This sub-package exposes the canonical Vision Transformer architecture introduced in "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" (ICLR 2021, arXiv:2010.11929).

The architecture splits an image into non-overlapping square patches, linearly projects each patch into a $D$ -dimensional embedding, prepends a learnable [CLS] token, adds a learnable positional embedding, and processes the resulting sequence with a stack of pre-norm transformer encoder blocks. The final CLS-token representation is used as the global image feature.

Notes

Configuration:

ViTConfig — immutable dataclass describing any ViT variant.

Modules:

ViT — feature-extractor backbone (returns (B, dim) CLS).
ViTForImageClassification — backbone + linear head producing classification logits.

Backbone factories:

vit_base_16, vit_base_32
vit_large_16, vit_large_32
vit_huge_14

Classifier factories (CLS-token linear head):

vit_base_16_cls, vit_base_32_cls
vit_large_16_cls, vit_large_32_cls
vit_huge_14_cls

Image Classification

Paper

Dosovitskiy, Alexey, et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." International Conference on Learning Representations, 2021.

Overview

The Vision Transformer (ViT) shows that convolutions are not required for state-of-the-art image classification given enough data. An input image $x \in \mathbb{R}^{H \times W \times C}$ is reshaped into a sequence of non-overlapping square patches $x_p \in \mathbb{R}^{N \times (P^2 C)}$ with $N = HW / P^2$ , each patch flattened and projected to a $D$ -dimensional token by a single linear map. A learnable class token $x_{\text{cls}}$ is prepended, and learnable positional embeddings $E_{\text{pos}}$ are added:

z_0 = [\,x_{\text{cls}};\; x_p^1 E;\; \dots;\; x_p^N E\,] + E_{\text{pos}}.

The sequence then flows through $L$ stacked transformer encoder blocks — multi-head self-attention plus an MLP, each wrapped by LayerNorm and residual connections. Classification reads the final state of the class token through a single linear head.

The paper's central empirical claim is that the inductive biases convolutions provide (locality, translation equivariance) become unnecessary once pre-training data is large enough — ViT pre-trained on ImageNet-21k or JFT-300M matches or beats the best ResNet/EfficientNet baselines on downstream benchmarks while using fewer FLOPs. Three canonical variants from Table 1 (Base / Large / Huge) trade compute for accuracy by scaling $(D, L, H)$ — the hidden width, depth, and number of attention heads.