Transformer

A transformer is a neural network built from a stack of identical blocks, each of which lets every token exchange information with every other token through self-attention and then transforms each token separately with a small multilayer perceptron (MLP). Both steps sit inside a residual connection with layer normalization. Vaswani et al. (2017) introduced it as an encoder–decoder for machine translation. In vision the tokens are image patches or feature-map positions: the Vision Transformer (ViT) classifies images with the encoder alone, DETR detects objects with the full encoder–decoder, and Swin Transformer serves as a hierarchical backbone.

Core Idea

A convolutional network mixes neighboring pixels with fixed learned filters. A transformer mixes all tokens with weights computed from their content. Each block has two halves: token mixing by multi-head self-attention, the only place where tokens interact, and channel mixing by a position-wise MLP that applies the same weights to every token. The network stacks such blocks at a constant width $d$. Nothing in a block refers to where a token sits, so position must be supplied separately.

Building Blocks

Multi-head self-attention. Each of $h$ heads computes scaled dot-product attention on its own learned projections of width $d/h$; the head outputs are concatenated and projected back to width $d$. The self-attention article covers the operation, masks and cost.

Position-wise MLP. Two linear layers with a nonlinearity between them, applied to each token independently. Vaswani et al. (2017) used a ReLU and an inner width of 2,048 for $d = 512$; ViT uses a GELU (Dosovitskiy et al., 2021). Both expand the width by a factor of 4 (ViT-Base: $d = 768$, MLP size 3,072).

Residual connections and layer normalization. Every sub-layer is wrapped in a residual connection, so all produce outputs of width $d$. Layer normalization (Ba, Kiros & Hinton, 2016) normalizes each token’s $d$ features to zero mean and unit variance, then applies a learned scale and shift; unlike batch normalization, it uses no batch statistics.

Positional information. Position enters in one of three ways:

• Sinusoidal encodings. Vaswani et al. (2017) added fixed sines and cosines of geometrically spaced wavelengths to the input embeddings, $PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d})$ and $PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d})$. Learned embeddings gave nearly identical results.
• Learned absolute embeddings. ViT adds a learned vector per token position, one-dimensional in patch order, since 2D-aware variants brought no significant gain in its experiments (Dosovitskiy et al., 2021).
• Relative positions. Shaw, Uszkoreit and Vaswani (2018) let attention use learned representations of the offset $j - i$ between tokens, clipped to a maximum distance, and reported gains of 1.3 and 0.3 BLEU over absolute encodings for their big models on WMT 2014 English–German and English–French.

How Data Flows

Encoder. A stack of $N$ blocks ($N = 6$ in the original model) maps input embeddings to contextual representations of the same length; every token attends to every other.

Decoder. Also $N$ blocks, each with three sub-layers. The first is masked self-attention: a causal mask keeps position $i$ from attending to later positions, so each prediction depends only on earlier outputs. The second is cross-attention (Vaswani et al. call it encoder–decoder attention): queries come from the decoder, keys and values from the encoder’s output, so each output position can draw on the whole input. The third is the MLP.

Three configurations. Later models keep what they need:

• Encoder–decoder, for mapping one sequence to another: the original translation model, and DETR, which maps image features to a set of object detections.
• Encoder-only, for inputs available all at once, with no causal mask: BERT (Devlin et al., 2019) in language, ViT in vision.
• Decoder-only, for generating sequences: masked self-attention and the MLP, with no encoder and no cross-attention. iGPT (Chen et al., 2020) uses the GPT-2 form of such a decoder to predict the pixels of low-resolution images one at a time in raster order, as unsupervised pre-training.

Mathematical Formulation

Write the tokens entering block $\ell$ as the rows of $Z_{\ell-1} \in \mathbb{R}^{n \times d}$. The original transformer is post-norm: it normalizes after each residual sum,

$$Z' = \operatorname{LN}\big(Z_{\ell-1} + \operatorname{MSA}(Z_{\ell-1})\big), \qquad Z_\ell = \operatorname{LN}\big(Z' + \operatorname{MLP}(Z')\big).$$

ViT, iGPT and Swin are pre-norm: they normalize the input of each sub-layer and leave the residual path untouched (Dosovitskiy et al., 2021, Eqs. 2–3),

$$Z' = Z_{\ell-1} + \operatorname{MSA}\big(\operatorname{LN}(Z_{\ell-1})\big), \qquad Z_\ell = Z' + \operatorname{MLP}\big(\operatorname{LN}(Z')\big),$$

with one final layer normalization after the last block. Xiong et al. (2020) showed with mean-field theory that at initialization the post-norm transformer has large expected gradients for the parameters near the output, so large learning rates make training unstable; they offer this as the reason a learning-rate warmup is needed. Pre-norm gradients are well behaved at initialization, and in their translation and BERT pre-training experiments pre-norm models trained without warmup reached comparable results in less time. In our reading, pre-norm gives the transformer the clean identity path of the pre-activation residual block.

Parameters of a block. With $d_k = d_v = d/h$, the query, key, value and output projections are four $d \times d$ matrices, whatever the number of heads: $4d^2$ weights and $4d$ biases. The MLP has $d \cdot 4d$ weights and $4d$ biases in its first layer and $4d \cdot d$ weights and $d$ biases in its second: $8d^2 + 5d$. Two layer normalizations add a scale and a shift each, $4d$. In total,

$$P_{\text{block}} = 12d^2 + 13d \approx 12d^2 .$$

The MLP holds two thirds of the weights. For ViT-Base ($d = 768$) a block has 7,087,872 parameters, and its 12 blocks 85.05 million. Adding the patch embedding of $16 \times 16 \times 3$ pixels (590,592), the class token, 197 position embeddings, the final normalization and a linear 1,000-class head gives about 86.6 million, consistent with the 86M Dosovitskiy et al. list for ViT-Base (our arithmetic). A decoder block adds cross-attention, for $16d^2$ weights in all. The script below (run with PyTorch 2.14.1) checks the count against nn.TransformerEncoderLayer, and checks that a hand-written pre-norm block computes the same function as that layer with norm_first=True:

import math
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

class Block(nn.Module):
    """Pre-norm transformer block: x + MSA(LN(x)), then x + MLP(LN(x)), MLP expansion 4."""
    def __init__(self, d, heads):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
    def forward(self, x):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))

def n_params(m):
    return sum(p.numel() for p in m.parameters())

print(f"torch {torch.__version__}")
for d, heads in [(32, 4), (768, 12)]:
    ref = nn.TransformerEncoderLayer(d, heads, dim_feedforward=4 * d, dropout=0.0,
                                     activation="gelu", norm_first=True, batch_first=True)
    print(f"d={d:3d}: 12d^2+13d = {12 * d * d + 13 * d:,}, Block = {n_params(Block(d, heads)):,}, "
          f"nn.TransformerEncoderLayer = {n_params(ref):,}")

# Same weights, same function: copy Block's weights into the PyTorch layer and compare outputs.
torch.manual_seed(0)
blk = Block(32, 4).eval()
ref = nn.TransformerEncoderLayer(32, 4, 128, dropout=0.0, activation="gelu",
                                 norm_first=True, batch_first=True).eval()
ref.load_state_dict({"self_attn.in_proj_weight": blk.attn.in_proj_weight, "self_attn.in_proj_bias": blk.attn.in_proj_bias,
                     "self_attn.out_proj.weight": blk.attn.out_proj.weight, "self_attn.out_proj.bias": blk.attn.out_proj.bias,
                     "linear1.weight": blk.mlp[0].weight, "linear1.bias": blk.mlp[0].bias,
                     "linear2.weight": blk.mlp[2].weight, "linear2.bias": blk.mlp[2].bias,
                     "norm1.weight": blk.ln1.weight, "norm1.bias": blk.ln1.bias,
                     "norm2.weight": blk.ln2.weight, "norm2.bias": blk.ln2.bias})
x = torch.randn(5, 17, 32)
with torch.no_grad():
    print(f"max |Block(x) - TransformerEncoderLayer(x)| = {(blk(x) - ref(x)).abs().max():.1e}")

torch 2.14.1
d= 32: 12d^2+13d = 12,704, Block = 12,704, nn.TransformerEncoderLayer = 12,704
d=768: 12d^2+13d = 7,087,872, Block = 7,087,872, nn.TransformerEncoderLayer = 7,087,872
max |Block(x) - TransformerEncoderLayer(x)| = 0.0e+00

Cost in the number of tokens. For $n$ tokens a block spends about $12nd^2$ multiply-adds in its linear layers and $2n^2d$ in the attention products $QK^\top$ and the weighted sum of values (our count). The quadratic term dominates once $n > 6d$. A $224 \times 224$ image cut into $16 \times 16$ patches gives $n = 197$ tokens with the class token, far below $6d = 4{,}608$ for ViT-Base; the same patches on a $1024 \times 1024$ image give 4,097 tokens, close to it. The self-attention article works through the memory this takes.

Why It Works

Self-attention gives every block a global receptive field and mixing weights that adapt to each image, but no built-in knowledge of images. Dosovitskiy et al. (2021) point out that a CNN builds locality, 2D neighborhood structure and translation equivariance into every layer, whereas in ViT only the MLP is local and spatial relations must be learned from data.

The price of fewer built-in assumptions is data. Fine-tuned on ImageNet at $384 \times 384$, ViT-L/16 reached 76.5% top-1 accuracy when pre-trained on ImageNet itself (1.3 million images), 85.2% when pre-trained on ImageNet-21k (14 million), and 87.1% on the in-house JFT-300M dataset (303 million); on ImageNet alone it trailed the smaller ViT-B/16 (77.9%). Pre-trained on random subsets of JFT-300M and evaluated by linear few-shot classification, ViT-B/32 did much worse than a slightly slower ResNet-50 with 9 million images and better with 90 million or more. The authors concluded that the convolutional inductive bias helps on smaller datasets, but on larger ones learning the patterns from data suffices.

Training Considerations

Transformers are usually trained with Adam or AdamW, a learning-rate warmup and strong regularization (training a neural network covers schedules and regularization in general):

• Original transformer. Adam with $\beta_2 = 0.98$; the learning rate rose linearly for 4,000 steps, then decayed with the inverse square root of the step number; dropout 0.1 (base model) on each sub-layer output and on the input embeddings; label smoothing 0.1 (Vaswani et al., 2017).
• ViT. Pre-trained with Adam, batch size 4,096, weight decay 0.1 and a 10,000-step linear warmup, then fine-tuned with SGD and momentum (Dosovitskiy et al., 2021).
• DeiT. Touvron et al. (2021) trained ViT-B on ImageNet-1k alone, on one 8-GPU machine in two to three days, with AdamW (weight decay 0.05, 5 warmup epochs, cosine decay, 300 epochs), heavy augmentation (RandAugment, Mixup, CutMix, random erasing, repeated augmentation), stochastic depth and label smoothing, without dropout. It reached 81.8% top-1 accuracy at $224 \times 224$ and 83.1% fine-tuned at $384 \times 384$, against 77.9% reported for ViT-B/16 trained on ImageNet alone (at $384 \times 384$). Pre-training with SGD instead of AdamW lowered the first figure to 74.5%. A distillation token trained to reproduce a convolutional teacher’s predictions raised the best result to 85.2% (at $384 \times 384$, after 1,000 epochs).

Variants

• Hierarchical transformers. Swin Transformer starts from $4 \times 4$ patches and merges each $2 \times 2$ group of neighboring tokens between stages, giving feature maps at 1/4, 1/8, 1/16 and 1/32 of the input resolution, like a CNN backbone. It restricts attention to local windows that shift between blocks, so the cost grows linearly with image size (Liu et al., 2021).
• Hybrid CNN–transformer models. The input tokens can be positions of a CNN feature map instead of raw patches. Dosovitskiy et al. (2021) fed ResNet features to ViT and found hybrids slightly better than pure ViTs at small compute budgets, with the difference vanishing for larger models. DETR uses a ResNet backbone in the same way.

Use in Computer Vision

ViT (Dosovitskiy et al., 2021) applies an unmodified encoder to images. It cuts an $H \times W$ image into $N = HW/P^2$ patches of $P \times P$ pixels and maps each flattened patch to width $d$ by a learned linear patch embedding. A learned class token is prepended, as in BERT; its final output, layer-normalized, feeds the classifier. The learned position embeddings are interpolated in 2D when the model is fine-tuned at a higher resolution.

DETR (Carion et al., 2020) treats object detection as set prediction with an encoder–decoder. The positions of a ResNet feature map at 1/32 resolution, reduced to width $d = 256$, are the encoder’s tokens, with fixed positional encodings added in every attention layer. The decoder takes $N = 100$ learned embeddings called object queries and decodes them in parallel rather than one at a time, using self-attention among the queries and cross-attention to the encoded image. Each output becomes a class and a box, or “no object”. Training matches predictions one-to-one to ground-truth objects with the Hungarian algorithm, so DETR needs no anchors or non-maximum suppression. On COCO validation, DETR with a ResNet-50 (41M parameters, 500 training epochs) reached 42.0 AP, equal to an enhanced Faster R-CNN with FPN (42M); it scored 61.1 against 53.4 AP on large objects and 20.5 against 26.6 on small ones. Query-based trackers such as TrackFormer and MOTR extend it to multi-object tracking.

Swin Transformer’s hierarchical feature maps plug into detection and segmentation heads designed for CNNs, and attention between two images drives feature matching methods such as SuperGlue and LoFTR.

A small ViT on digit images. Continuing the script above, the code builds a ViT-style classifier for scikit-learn’s $8 \times 8$ digits: $2 \times 2$ patches give 16 tokens of 4 pixels, embedded to $d = 32$, plus a class token and learned position embeddings, followed by two pre-norm blocks with 4 heads and a GELU MLP. It is trained for 690 steps with AdamW, a 50-step warmup and cosine decay, and compared with a three-layer CNN of nearly the same size, on all 1,437 training images and on 150 of them. A third model omits the position embeddings.

digits = load_digits()
X = torch.tensor(digits.images / 16.0, dtype=torch.float32).unsqueeze(1)   # (1797, 1, 8, 8)
y = torch.tensor(digits.target)
i_tr, i_te = train_test_split(np.arange(len(y)), test_size=0.2, random_state=0, stratify=digits.target)

class TinyViT(nn.Module):
    """ViT-style classifier: 2x2 patches -> 16 tokens + class token, learned positions, pre-norm blocks."""
    def __init__(self, patch=2, d=32, depth=2, heads=4, pos=True):
        super().__init__()
        self.p, n = patch, (8 // patch) ** 2
        self.embed = nn.Linear(patch * patch, d)                       # patch embedding
        self.cls = nn.Parameter(torch.zeros(1, 1, d))                  # class token
        self.pos = nn.Parameter(0.02 * torch.randn(1, n + 1, d)) if pos else None
        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(depth)])
        self.norm, self.head = nn.LayerNorm(d), nn.Linear(d, 10)
    def forward(self, img, perm=None):
        B, p = len(img), self.p                                         # cut into patches, row by row
        t = img.reshape(B, 8 // p, p, 8 // p, p).permute(0, 1, 3, 2, 4).reshape(B, -1, p * p)
        if perm is not None:                                            # shuffle the image's patches
            t = t[:, perm]
        t = torch.cat([self.cls.expand(B, -1, -1), self.embed(t)], dim=1)
        if self.pos is not None:
            t = t + self.pos
        return self.head(self.norm(self.blocks(t))[:, 0])               # classify from the class token

def small_cnn():
    return nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(),
                         nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 4x4
                         nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 2x2
                         nn.Flatten(), nn.Linear(256, 10))

def train(model, idx, steps=690, bs=64, lr=3e-3, warmup=50):
    """AdamW with weight decay 0.05, linear warmup then cosine decay; returns test accuracy."""
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.05)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min((s + 1) / warmup,
                0.5 * (1 + math.cos(math.pi * (s - warmup) / (steps - warmup)))))
    step = 0
    while step < steps:
        for b in torch.from_numpy(idx)[torch.randperm(len(idx))].split(bs):
            if step == steps:
                break
            loss = nn.functional.cross_entropy(model(X[b]), y[b])
            opt.zero_grad(); loss.backward(); opt.step(); sched.step(); step += 1
    model.eval()
    with torch.no_grad():
        return (model(X[i_te]).argmax(1) == y[i_te]).float().mean().item()

SEED = 8
print(f"{len(i_tr)} training and {len(i_te)} test images")
i_small = np.random.default_rng(SEED).choice(i_tr, 150, replace=False)
models = {}
for name, make in [("ViT", TinyViT), ("ViT, no positions", lambda: TinyViT(pos=False)), ("CNN", small_cnn)]:
    torch.manual_seed(SEED)
    models[name] = make()
    acc = train(models[name], i_tr)
    line = f"{name:18s} {n_params(models[name]):6,d} parameters  test accuracy: all {len(i_tr)} images {acc:.3f}"
    if name != "ViT, no positions":
        torch.manual_seed(SEED)
        line += f", 150 images {train(make(), i_small):.3f}"
    print(line)

1437 training and 360 test images
ViT                26,538 parameters  test accuracy: all 1437 images 0.975, 150 images 0.781
ViT, no positions  25,994 parameters  test accuracy: all 1437 images 0.725
CNN                25,866 parameters  test accuracy: all 1437 images 0.981, 150 images 0.928

With all training images, the ViT and the CNN are close (97.5% and 98.1% test accuracy). With 150 images, the CNN keeps 92.8% while the ViT falls to 78.1%, the data dependence that Dosovitskiy et al. describe at a much larger scale. Without position embeddings the ViT reaches only 72.5%: it can use which pixel patterns occur, but not where. Over ten seeds the ViT reached 96.1–98.3% (all images) and 70.8–81.7% (150 images), the CNN 97.5–98.3% and 89.7–94.7%, and the model without positions 68.9–75.3%; seed 8 is typical.

Permutation check. Continuing, the code shuffles the 16 patches of every test image before they are embedded:

perm = torch.randperm(16)
with torch.no_grad():
    for name in ("ViT, no positions", "ViT"):
        m = models[name]
        out, out_shuffled = m(X[i_te]), m(X[i_te], perm=perm)
        print(f"{name:18s} max |change in logits| after shuffling patches: {(out - out_shuffled).abs().max():.1e}, "
              f"predictions changed: {(out.argmax(1) != out_shuffled.argmax(1)).sum().item()} of {len(i_te)}")

ViT, no positions  max |change in logits| after shuffling patches: 5.2e-06, predictions changed: 0 of 360
ViT                max |change in logits| after shuffling patches: 1.3e+01, predictions changed: 304 of 360

Without position embeddings the logits agree to float32 rounding and no prediction changes: the model sees an unordered set of patches. With them, each patch arrives at the wrong position and 304 of 360 predictions change.

Limitations

• Quadratic cost in the number of tokens, which limits patch size and resolution.
• Data hunger. Without image-specific assumptions, transformers need large pre-training sets, or heavy augmentation and distillation, to match CNNs.
• Single-scale features. ViT keeps one low-resolution token grid throughout; Liu et al. (2021) judged this, together with quadratic cost, unsuitable for a general-purpose backbone for dense tasks or high-resolution inputs.
• Training sensitivity. Post-norm transformers depend on learning-rate warmup (Xiong et al., 2020), and Touvron et al. (2021) found vision transformers sensitive to initialization and optimizer settings.

Important Models

Model	Year	Configuration	Key idea	Reported result
Transformer	2017	Encoder–decoder, post-norm, $d = 512$, 6 + 6 layers (base)	Attention and MLP blocks without recurrence or convolution	WMT 2014 English–German 28.4 BLEU (big model)
ViT	2021	Encoder-only, pre-norm, image patches	Image patches as tokens, class token, large-scale pre-training	ImageNet top-1 88.55% (ViT-H/14, JFT-300M pre-training)
DeiT	2021	ViT architecture	Data-efficient training recipe, distillation token	ImageNet top-1 83.1% (DeiT-B, ImageNet-1k only, 384²); 85.2% with distillation
iGPT	2020	Decoder-only, pre-norm, pixels as tokens	Autoregressive pixel prediction as pre-training	CIFAR-10 linear-probe accuracy 96.3% (iGPT-L)
DETR	2020	CNN backbone + encoder–decoder	Object queries, set prediction with bipartite matching	COCO val 42.0 AP (ResNet-50)
Swin Transformer	2021	Hierarchical encoder, windowed attention	Shifted windows, patch merging	ImageNet-1K top-1 81.3% (Swin-T, ImageNet-1K only); COCO test-dev 58.7 box AP (Swin-L with HTC++, ImageNet-22K pre-training, multi-scale testing)

All results are as reported by the authors, under different datasets, resolutions and protocols, so they are not directly comparable.