26  Chapter 25. Vision Transformers

Patch Tokens, Self-Attention, Positional Encoding, Inductive Bias, and Dense Vision

26.1 25.1 Why transformers for vision?

Chapters 6, 7, and 21 emphasized a central principle of classical and neural vision: local operators are powerful. Convolutional networks build representations by repeatedly applying local filters, nonlinearities, pooling, and normalization. This architecture has strong inductive bias: nearby pixels interact early, weight sharing enforces translation equivariance, and receptive fields grow gradually with depth.

Vision transformers begin from a different point of view. Instead of asking first for local convolutional filters, they ask:

Can an image be treated as a sequence of visual tokens, and can the interactions among those tokens be learned by attention?

Let an image be

\[ X\in \mathbb R^{H\times W\times C}. \]

A vision transformer partitions \(X\) into patches, embeds each patch into a vector, adds positional information, and repeatedly applies self-attention and feed-forward transformations. The basic representation has the form

\[ X \longmapsto (Z_0,Z_1,\dots,Z_L), \qquad Z_\ell\in \mathbb R^{N\times d}, \]

where \(N\) is the number of tokens and \(d\) is the embedding dimension. Each row of \(Z_\ell\) is a vector representation of one token after layer \(\ell\).

The key mathematical shift is from local filtering to learned pairwise interaction. In a convolutional layer, the value at one location depends on a fixed local neighborhood. In a self-attention layer, the representation of one token can depend on all other tokens, with weights determined by the data itself.

NoteCentral idea

A vision transformer represents an image as a sequence of patch tokens and learns a data-dependent weighted averaging operation among tokens. The weights come from inner products between learned query and key vectors.

The original transformer architecture was introduced for sequence modeling in natural language processing by Vaswani et al. (2017). The Vision Transformer, or ViT, adapted this idea to images by treating fixed-size image patches as tokens Dosovitskiy et al. (2020). Since then, transformers have become central in recognition, detection, segmentation, self-supervised learning, image-text modeling, and vision foundation models.

26.1.1 Learning goals

After this chapter, you should be able to:

1. convert an image into a sequence of patch tokens;
2. express patch embedding as a linear map or convolution;
3. derive scaled dot-product attention from query, key, and value matrices;
4. explain the role of the \(1/\sqrt d\) scaling factor;
5. compare attention with convolution from the viewpoints of locality, equivariance, complexity, and inductive bias;
6. describe positional encodings and why they are needed;
7. analyze multi-head attention as several learned interaction kernels;
8. explain the architecture of a ViT block;
9. understand class-token pooling, mean pooling, and dense-token prediction;
10. write Python experiments for patchifying images, computing attention maps, and building a small transformer block;
11. use AI tools to verify tensor shapes, attention formulas, and implementation choices.

26.2 25.2 From images to patch tokens

Let \(P\) be a patch size. Assume for simplicity that \(P\) divides both \(H\) and \(W\). The image is divided into

\[ N=\frac{H}{P}\frac{W}{P} \]

non-overlapping patches. Each patch has shape \(P\times P\times C\) and can be flattened into a vector in

\[ \mathbb R^{P^2C}. \]

Write the vectorized patches as

\[ x_1,x_2,\dots,x_N\in\mathbb R^{P^2C}. \]

A learned patch embedding maps each patch to a \(d\)-dimensional token:

\[ z_i=x_iE+b, \qquad E\in\mathbb R^{P^2C\times d}, \qquad b\in\mathbb R^d. \]

Stacking the patch vectors into a matrix

\[ X_{\mathrm{patch}} = \begin{bmatrix} x_1^T\\ x_2^T\\ \vdots\\ x_N^T \end{bmatrix} \in\mathbb R^{N\times P^2C}, \]

the embedded token matrix is

\[ Z=X_{\mathrm{patch}}E+\mathbf 1 b^T \in\mathbb R^{N\times d}. \]

This is a linear transformation applied independently to each patch. In software, patch embedding is often implemented as a convolution with kernel size \(P\) and stride \(P\). The convolution version and the matrix version are mathematically equivalent when patches do not overlap.

TipShape check

For an image of size \(224\times 224\times 3\) and patch size \(16\), the number of patches is

\[ N=14\cdot 14=196. \]

If \(d=768\), then the patch embedding has shape

\[ Z\in\mathbb R^{196\times 768}. \]

26.2.1 Patch size as a modeling choice

Patch size controls a trade-off. A smaller patch size gives more tokens and preserves fine spatial information. A larger patch size gives fewer tokens and lower computational cost, but may lose boundary and texture information.

The attention cost is roughly quadratic in the number of tokens:

\[ \mathrm{cost}\approx O(N^2d). \]

Since

\[ N=\frac{HW}{P^2}, \]

halving \(P\) multiplies \(N\) by \(4\) and multiplies the \(N^2\) part of attention cost by \(16\). This explains why transformer-based dense vision often requires design modifications such as windowed attention, hierarchical tokens, shifted windows, or hybrid convolution-transformer backbones.

26.3 25.3 A first interactive view: patch tokenization

The next interactive demonstration shows how patch size changes the number of visual tokens.

Patch tokenization

Change the patch size and observe how the number of tokens changes. The quadratic attention cost grows quickly when patches become small.

Patch size P: 16

26.4 25.4 Positional information

A transformer attention layer is naturally permutation equivariant. If we permute the input tokens, the output tokens are permuted in the same way. This is mathematically elegant, but images have spatial structure. A patch in the upper-left corner is not interchangeable with a patch in the lower-right corner.

To encode location, we add a positional vector \(p_i\in\mathbb R^d\) to each patch embedding:

\[ z_i^{(0)} = x_iE+b+p_i. \]

Equivalently, if

\[ P_{\mathrm{pos}} = \begin{bmatrix} p_1^T\\ \vdots\\ p_N^T \end{bmatrix} \in \mathbb R^{N\times d}, \]

then

\[ Z_0=X_{\mathrm{patch}}E+\mathbf 1b^T+P_{\mathrm{pos}}. \]

There are several common choices:

1. Learned absolute positional embeddings. Each location has a learned vector \(p_i\).
2. Sinusoidal positional embeddings. Location is encoded by deterministic sine and cosine functions.
3. Relative positional bias. Attention logits depend on relative displacement between two patches.
4. Two-dimensional factorized encodings. Row and column positions are represented separately.
5. Rotary or other geometric encodings. Positional information is injected through transformations of queries and keys.

For a two-dimensional patch grid, it is natural to index positions by \((r,c)\). A simple separable positional encoding has the form

\[ p_{r,c}=a_r+b_c, \]

where \(a_r\in\mathbb R^d\) is a row embedding and \(b_c\in\mathbb R^d\) is a column embedding.

26.4.1 Why attention needs position

Suppose there are two token matrices \(Z\) and \(\Pi Z\), where \(\Pi\) is a permutation matrix that shuffles patch order. A self-attention layer without positional information satisfies

\[ \operatorname{SelfAttn}(\Pi Z)=\Pi\operatorname{SelfAttn}(Z). \]

Thus, the model does not know which patch came from which spatial location except through the visual content of the patch. Positional encodings break full permutation symmetry and introduce spatial structure.

This is different from convolution. A convolution is translation equivariant by construction because the same local operator is applied at every spatial location. A basic transformer must be given positional information explicitly.

26.5 25.5 Scaled dot-product attention

Let

\[ Z\in\mathbb R^{N\times d} \]

be a token matrix. An attention layer constructs three matrices:

\[ Q=ZW_Q, \qquad K=ZW_K, \qquad V=ZW_V, \]

where

\[ W_Q,W_K\in\mathbb R^{d\times d_k}, \qquad W_V\in\mathbb R^{d\times d_v}. \]

The rows of \(Q\) are query vectors, the rows of \(K\) are key vectors, and the rows of \(V\) are value vectors. The attention scores are inner products:

\[ S=\frac{QK^T}{\sqrt{d_k}} \in\mathbb R^{N\times N}. \]

The attention weights are obtained row-by-row using softmax:

\[ A_{ij} = \frac{\exp(S_{ij})}{\sum_{m=1}^N\exp(S_{im})}. \]

Finally,

\[ \operatorname{Attention}(Q,K,V)=AV. \]

For each token \(i\), the output is

\[ y_i=\sum_{j=1}^N A_{ij}v_j. \]

Thus self-attention is a learned, data-dependent weighted average of value vectors.

NoteAttention is nonlinear

For fixed \(A\), the map \(V\mapsto AV\) is linear. But in self-attention, the weight matrix \(A\) is itself computed from \(Z\) through \(QK^T\) and softmax. Therefore the map \(Z\mapsto \operatorname{Attention}(ZW_Q,ZW_K,ZW_V)\) is nonlinear.

26.5.1 The \(1/\sqrt{d_k}\) scaling

Assume for intuition that the components of \(q_i\) and \(k_j\) are independent with mean \(0\) and variance \(1\). Then

\[ q_i^T k_j=\sum_{m=1}^{d_k}q_{im}k_{jm} \]

has variance approximately \(d_k\). Without scaling, the scores grow in magnitude as \(d_k\) grows. Large scores make the softmax nearly one-hot, which can produce small gradients and unstable optimization. Dividing by \(\sqrt{d_k}\) stabilizes the variance:

\[ \operatorname{Var}\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right)\approx 1. \]

This is a statistical normalization argument. It is similar in spirit to variance-preserving initialization for neural networks.

26.6 25.6 Interactive attention weights

The next demo shows how attention weights depend on temperature. A lower temperature makes the distribution sharper. The standard transformer scaling behaves like a temperature chosen to control score variance.

Softmax attention and temperature

The query token attends to a row of keys. Lower temperature produces sharper attention.

Temperature: 1.00

26.7 25.7 Attention as a kernel smoother

Attention can be interpreted as a learned kernel smoother on a set of tokens. For token \(i\), define

\[ \kappa_\theta(i,j) = \exp\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right). \]

Then

\[ A_{ij} = \frac{\kappa_\theta(i,j)}{\sum_{m=1}^N\kappa_\theta(i,m)}. \]

The output is

\[ y_i=\sum_{j=1}^N A_{ij}v_j. \]

This resembles kernel regression:

\[ \widehat f(x_i) = \sum_{j=1}^n \frac{K(x_i,x_j)}{\sum_m K(x_i,x_m)}y_j. \]

The difference is that the kernel is learned and contextual. The query and key representations depend on the current token features, which themselves change across layers.

This view is useful for applied mathematics and statistics students because it connects transformers to:

• kernel methods;
• nonparametric regression;
• graph smoothing;
• Markov transition matrices;
• random-walk averaging;
• adaptive discretizations of integral operators.

Indeed, the attention matrix \(A\) is row-stochastic:

\[ A_{ij}\ge 0, \qquad \sum_{j=1}^N A_{ij}=1. \]

For each feature dimension of \(V\), multiplication by \(A\) averages values over a directed weighted graph on tokens.

26.8 25.8 Multi-head attention

A single attention head computes one type of token interaction. Multi-head attention computes several attention maps in parallel.

For head \(h=1,\dots,H_a\), define

\[ Q_h=ZW_Q^{(h)}, \qquad K_h=ZW_K^{(h)}, \qquad V_h=ZW_V^{(h)}. \]

The head output is

\[ Y_h=\operatorname{softmax}\left(\frac{Q_hK_h^T}{\sqrt{d_h}}\right)V_h. \]

The outputs are concatenated and linearly mixed:

\[ \operatorname{MHA}(Z) = \operatorname{Concat}(Y_1,\dots,Y_{H_a})W_O. \]

Different heads may learn different interaction patterns: local neighborhood relations, long-range symmetry, object-part relations, background suppression, boundary interactions, or semantic grouping. This interpretation is suggestive but should be used carefully. Individual attention heads are not always reliable explanations.

26.8.1 Parameter count

Assume \(d_h=d/H_a\). Then the combined query projection across all heads has about \(d^2\) parameters, and the same is true for keys and values. The output projection also has about \(d^2\) parameters. Thus a multi-head attention layer has roughly

\[ 4d^2 \]

projection parameters, ignoring biases. The main computational burden, however, is the attention matrix:

\[ QK^T\in\mathbb R^{N\times N}. \]

This is why token count matters so much.

26.9 25.9 The transformer block

A standard transformer encoder block contains:

1. layer normalization;
2. multi-head self-attention;
3. residual connection;
4. layer normalization;
5. feed-forward multilayer perceptron;
6. residual connection.

One common pre-normalization form is

\[ \widetilde Z_{\ell} = Z_{\ell-1}+\operatorname{MHA}(\operatorname{LN}(Z_{\ell-1})), \]

\[ Z_\ell = \widetilde Z_\ell + \operatorname{MLP}(\operatorname{LN}(\widetilde Z_\ell)). \]

The MLP is applied independently to each token:

\[ \operatorname{MLP}(z) = W_2\sigma(W_1z+b_1)+b_2. \]

Typically the hidden dimension is larger than \(d\), for example \(4d\). The MLP mixes feature channels within each token, while attention mixes information across tokens.

ImportantToken mixing and channel mixing

In a transformer block:

• self-attention mixes tokens;
• the MLP mixes channels;
• residual connections stabilize optimization;
• layer normalization controls scale.

This structure is different from CNNs, where spatial mixing and channel mixing happen together inside convolutions. In a transformer, the two operations are separated.

26.10 25.10 Classification with a class token

For image classification, ViT usually adds a learned class token \(z_{\mathrm{cls}}\in\mathbb R^d\) to the sequence:

\[ Z_0= \begin{bmatrix} z_{\mathrm{cls}}^T\\ z_1^T\\ \vdots\\ z_N^T \end{bmatrix} +P_{\mathrm{pos}} \in\mathbb R^{(N+1)\times d}. \]

After \(L\) transformer blocks, the class token output is passed to a classification head:

\[ \widehat p(y=k\mid X) = \operatorname{softmax}(W z_{\mathrm{cls},L}+b)_k. \]

Alternatively, one can average the patch tokens:

\[ \bar z_L=\frac{1}{N}\sum_{i=1}^N z_{i,L}, \]

and classify from \(\bar z_L\). Class-token pooling and mean pooling are both used in practice.

The statistical loss is the usual cross-entropy:

\[ \mathcal L(\theta) = -\frac{1}{n}\sum_{r=1}^n \log p_\theta(y_r\mid X_r). \]

The architecture changes the representation map, but the classification objective remains the empirical negative log-likelihood.

26.11 25.11 Interactive token mixing

This demo illustrates how a row-stochastic attention matrix mixes token values. Each output token is a weighted average of value tokens.

Attention as token mixing

Move the locality slider. Small bandwidth gives local attention; large bandwidth gives global attention.

Bandwidth: 1.50

26.12 25.12 Comparing attention and convolution

Convolution and attention are both ways of mixing spatial information, but their mathematical structures are different.

26.12.1 Convolution

A convolution layer computes

\[ y_i=\sum_{r\in\mathcal N} w_r x_{i+r}, \]

where \(\mathcal N\) is a fixed local neighborhood. The weights \(w_r\) do not depend on the image content. They are learned parameters but fixed at inference time.

26.12.2 Attention

A self-attention layer computes

\[ y_i=\sum_{j=1}^N A_{ij}(X)v_j, \]

where the weights \(A_{ij}(X)\) depend on the input.

Thus convolution has fixed local weights, while attention has adaptive global weights.

Property	Convolution	Self-attention
Spatial mixing	local	global unless restricted
Weights	fixed after training	data-dependent
Translation structure	built in	must be learned or encoded
Cost in token count	roughly linear	roughly quadratic
Fine-scale bias	strong	weaker
Long-range interactions	require depth	immediate
Data efficiency	often strong on small data	often needs large data or pretraining

This comparison is not a claim that one architecture is always better. It is a mathematical description of different inductive biases. In many modern systems, convolutional and transformer ideas are combined.

26.13 25.13 Complexity and memory

Let \(N\) be the number of tokens and \(d\) the embedding dimension. A single attention head requires forming

\[ QK^T\in\mathbb R^{N\times N}. \]

The approximate computation is

\[ O(N^2d) \]

and the attention matrix requires

\[ O(N^2) \]

memory. This is manageable for image classification with \(196\) tokens, but difficult for dense prediction with high-resolution images.

For example, a \(1024\times 1024\) image with \(16\times 16\) patches has

\[ N=64^2=4096 \]

tokens. The attention matrix has

\[ N^2=16,777,216 \]

entries per head per layer. This motivates efficient variants.

26.13.1 Windowed attention

Windowed attention restricts each token to attend only to tokens in a local window. If each window contains \(M\) tokens, the cost becomes approximately

\[ O(NMd), \]

which is linear in the number of tokens when \(M\) is fixed.

The Swin Transformer uses shifted windows to exchange information across windows while preserving local efficiency Liu et al. (2021). Hierarchical designs also reduce resolution gradually, similar to CNN feature pyramids.

26.13.2 Low-rank and sparse approximations

Other efficient attention ideas include:

• low-rank approximations to the attention matrix;
• sparse attention patterns;
• linear attention kernels;
• pooling or merging tokens;
• cross-attention from a smaller set of latent tokens;
• multi-scale token hierarchies.

The central numerical question is: how can we approximate global token interaction without explicitly storing all \(N^2\) pairwise weights?

26.14 25.14 Positional encoding demo

The next demo compares a one-dimensional sinusoidal positional encoding with a learned-looking random positional encoding. The sinusoidal encoding has smooth geometry: nearby positions tend to have similar vectors.

Positional encoding geometry

The heatmap shows pairwise dot products between position vectors. Smooth structure helps the model represent relative location.

Frequency scale: 1.00

26.15 25.15 Vision transformers for dense prediction

Image classification produces one label per image. Detection and segmentation require spatial outputs. A transformer can be adapted to dense vision in several ways.

26.15.1 Patch-token classification

For semantic segmentation, each patch token can be mapped to a class distribution:

\[ \widehat p_i=\operatorname{softmax}(Wz_{i,L}+b). \]

The token predictions are then upsampled to pixel resolution. This approach is simple but may lose boundary detail when patches are large.

26.15.2 Encoder-decoder segmentation

A transformer encoder can produce multi-scale features, which are then decoded by convolutional or transformer-based decoders. Skip connections or feature pyramids can recover spatial detail.

26.15.3 Cross-attention decoding

A set of learned query tokens can attend to image tokens. This idea appears in detection, instance segmentation, and promptable models. Cross-attention has the form

\[ Q=Z_{\mathrm{query}}W_Q, \qquad K=Z_{\mathrm{image}}W_K, \qquad V=Z_{\mathrm{image}}W_V, \]

\[ Y=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V. \]

Unlike self-attention, the queries and keys come from different token sets.

26.15.4 Detection as set prediction

The DETR architecture formulated object detection as a set prediction problem with transformer encoders and decoders Carion et al. (2020). Learned object queries attend to image features, and a bipartite matching loss assigns predictions to ground-truth boxes. This connects directly to Chapter 23.

26.15.5 Segmentation foundation models

Promptable segmentation models use image tokens, prompt tokens, and mask decoding. The Segment Anything model is one example of a vision foundation model with a transformer-based image encoder and prompt-conditioned mask decoder Kirillov et al. (2023).

26.16 25.16 Training vision transformers

ViT-style models often have less built-in locality than CNNs, so training them from scratch may require large datasets or strong regularization. Several strategies are common.

26.16.1 Data augmentation and regularization

Training may use random crop, color jitter, RandAugment-style transformations, mixup, CutMix, label smoothing, stochastic depth, dropout, and weight decay. The statistical purpose is to reduce variance and encode invariance.

26.16.2 Distillation

Data-efficient image transformers can be trained with teacher-student distillation Touvron et al. (2021). Suppose a teacher model produces soft labels \(q(y\mid X)\) and a student produces \(p_\theta(y\mid X)\). A distillation loss can use

\[ \mathcal L_{\mathrm{distill}}(\theta) = -\sum_{k=1}^K q(k\mid X)\log p_\theta(k\mid X). \]

Soft labels provide information about class similarity and uncertainty.

26.16.3 Self-supervised pretraining

Transformers are especially important in self-supervised vision. Two central examples are:

• DINO-style self-distillation with vision transformers Caron et al. (2021);
• masked autoencoders, which reconstruct missing image patches from visible patches He et al. (2022).

In masked image modeling, a subset \(M\) of patches is hidden and the model learns to reconstruct them:

\[ \mathcal L(\theta) = \sum_{i\in M}\|x_i-\widehat x_i\|^2. \]

This connects ViTs to representation learning, denoising, and generative modeling.

26.17 25.17 Statistical viewpoint: bias, variance, and inductive structure

For MA Applied Mathematics and MS Statistics students, the important question is not merely whether transformers are fashionable. The important question is what assumptions the architecture makes.

A model class \(\mathcal F\) controls approximation, estimation, and optimization:

\[ R(\widehat f)-R(f^*) = \underbrace{R(\widehat f)-\inf_{f\in\mathcal F}R(f)}_{\text{estimation and optimization}} + \underbrace{\inf_{f\in\mathcal F}R(f)-R(f^*)}_{\text{approximation}}. \]

Convolutional networks encode locality and translation equivariance, which can reduce sample complexity when those assumptions are correct. Transformers encode flexible global interactions, which can reduce approximation error for tasks requiring long-range relationships. But flexibility can increase the amount of data and regularization needed.

26.17.1 Attention maps are not automatically explanations

It is tempting to interpret large attention weights as explanations. Sometimes attention maps are useful diagnostic tools, but attention weights alone do not always identify causal evidence. A model can attend to a region and still make its decision for another reason, because information is mixed across layers and heads.

A more careful interpretability study may combine attention visualization with:

• feature ablation;
• gradient-based saliency;
• counterfactual image edits;
• patch masking experiments;
• calibration and uncertainty checks;
• external validation on shifted data.

26.18 25.18 Interactive attention complexity

The next demo compares full attention and windowed attention as image resolution changes.

Full attention versus windowed attention

For high-resolution images, full attention grows quadratically in the number of tokens. Windowed attention is much cheaper when the window size is fixed.

Patch size: 16

26.19 25.19 Python example: patchifying an image

The following Python example shows patch extraction as a tensor reshaping operation.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(25)
H, W, C = 32, 32, 3
P = 8
X = rng.uniform(size=(H, W, C))

def patchify_numpy(X, P):
    H, W, C = X.shape
    assert H % P == 0 and W % P == 0
    h, w = H // P, W // P
    patches = X.reshape(h, P, w, P, C)
    patches = patches.transpose(0, 2, 1, 3, 4)
    patches = patches.reshape(h * w, P * P * C)
    return patches

patches = patchify_numpy(X, P)
print("image shape:", X.shape)
print("patch matrix shape:", patches.shape)

The output patch matrix has shape

\[ N\times P^2C. \]

This matrix can be multiplied by a learned embedding matrix \(E\in\mathbb R^{P^2C\times d}\).

d = 64
E = rng.normal(scale=0.02, size=(P * P * C, d))
b = np.zeros(d)
Z = patches @ E + b
print("token matrix shape:", Z.shape)

26.20 25.20 Python example: attention from scratch

This example computes one scaled dot-product attention layer using NumPy.

import numpy as np

rng = np.random.default_rng(7243)
N = 10
D = 32
dk = 16
dv = 24

Z = rng.normal(size=(N, D))
Wq = rng.normal(scale=1/np.sqrt(D), size=(D, dk))
Wk = rng.normal(scale=1/np.sqrt(D), size=(D, dk))
Wv = rng.normal(scale=1/np.sqrt(D), size=(D, dv))

Q = Z @ Wq
K = Z @ Wk
V = Z @ Wv

scores = Q @ K.T / np.sqrt(dk)
scores = scores - scores.max(axis=1, keepdims=True)
A = np.exp(scores)
A = A / A.sum(axis=1, keepdims=True)
Y = A @ V

print("Q shape:", Q.shape)
print("attention matrix shape:", A.shape)
print("row sums:", A.sum(axis=1))
print("output shape:", Y.shape)

A useful debugging check is

\[ \sum_{j=1}^N A_{ij}=1 \]

for every query token \(i\).

26.21 25.21 TensorFlow example: a small ViT block

The following TensorFlow example is intentionally small. It is meant to clarify shapes, not to compete with production vision libraries.

import tensorflow as tf

class PatchEmbedding(tf.keras.layers.Layer):
    def __init__(self, patch_size, embed_dim):
        super().__init__()
        self.patch_size = patch_size
        self.proj = tf.keras.layers.Conv2D(
            filters=embed_dim,
            kernel_size=patch_size,
            strides=patch_size,
            padding="valid"
        )

    def call(self, x):
        z = self.proj(x)                         # (B, H/P, W/P, d)
        b = tf.shape(z)[0]
        z = tf.reshape(z, [b, -1, tf.shape(z)[-1]])
        return z                                  # (B, N, d)

class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, mlp_ratio=4, dropout=0.0):
        super().__init__()
        self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.attn = tf.keras.layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=embed_dim // num_heads,
            dropout=dropout
        )
        self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.mlp = tf.keras.Sequential([
            tf.keras.layers.Dense(mlp_ratio * embed_dim, activation="gelu"),
            tf.keras.layers.Dropout(dropout),
            tf.keras.layers.Dense(embed_dim),
            tf.keras.layers.Dropout(dropout)
        ])

    def call(self, z, training=False):
        y = self.norm1(z)
        z = z + self.attn(y, y, training=training)
        z = z + self.mlp(self.norm2(z), training=training)
        return z

B, H, W, C = 4, 64, 64, 3
patch_size = 8
embed_dim = 64
num_heads = 4

x = tf.random.normal([B, H, W, C])
patch = PatchEmbedding(patch_size, embed_dim)
block = TransformerBlock(embed_dim, num_heads)

z = patch(x)
z = block(z)
print(z.shape)  # (B, 64, 64) because N=(64/8)^2=64

The model above uses no class token. For classification, one could add a learned class token or use mean pooling over the token dimension.

26.22 25.22 Python example: attention visualization

Attention matrices can be visualized as heatmaps. The following example creates synthetic attention weights and plots them.

import numpy as np
import matplotlib.pyplot as plt

N = 14 * 14
coords = np.array([(i, j) for i in range(14) for j in range(14)])
query_index = 7 * 14 + 7
q = coords[query_index]

dist2 = ((coords - q) ** 2).sum(axis=1)
weights = np.exp(-dist2 / (2 * 2.0 ** 2))
weights = weights / weights.sum()

plt.figure(figsize=(5, 4))
plt.imshow(weights.reshape(14, 14), cmap="viridis")
plt.colorbar(label="attention weight")
plt.title("synthetic local attention from center patch")
plt.axis("off")
plt.show()

When interpreting real attention maps, remember that attention weights are internal mixing coefficients, not guaranteed causal explanations.

26.23 25.23 Mini-project: ViT on a small dataset

A good student project is to compare a small CNN and a small ViT on the same image classification task.

Suggested workflow:

1. Use a small dataset such as Fashion-MNIST, CIFAR-10, or a course-provided medical-image subset.
2. Train a CNN with a comparable number of parameters.
3. Train a small ViT with patch sizes \(4\), \(8\), and \(16\).
4. Compare training accuracy, validation accuracy, calibration, runtime, and sensitivity to augmentation.
5. Plot attention maps for selected examples.
6. Write a statistical conclusion: which architecture has better bias-variance behavior for this dataset?

The point is not to declare a universal winner. The point is to understand how architecture, dataset size, inductive bias, and regularization interact.

26.24 25.24 AI-assisted study prompts

Use an AI assistant as a mathematical checking tool, not as a substitute for reasoning. Good prompts include:

1. Shape checking.
   “For a ViT with image size \(224\times224\times3\), patch size \(16\), embedding dimension \(768\), and \(12\) heads, list the shapes of \(X_{\mathrm{patch}}\), \(Q\), \(K\), \(V\), \(QK^T\), and the output of one attention block.”

2. Formula verification.
   “Explain why attention weights sum to one across key tokens. Identify the dimension over which softmax is applied.”

3. Conceptual comparison.
   “Compare convolution and self-attention as operators on image tokens. Discuss locality, equivariance, and computational complexity.”

4. Debugging.
   “This TensorFlow ViT block has a shape mismatch after patch embedding. Find the most likely cause and suggest a fix.”

5. Statistical critique.
   “A ViT trained from scratch on a small dataset overfits badly. Give a mathematical explanation in terms of inductive bias, variance, and regularization.”

AI-generated explanations should be checked by doing the tensor algebra yourself.

26.25 25.25 Exercises

26.25.1 Conceptual exercises

1. Explain why a basic transformer without positional information is permutation equivariant.
2. Compare the inductive biases of CNNs and ViTs for medical image classification.
3. Explain why patch size affects both computational cost and boundary precision.
4. Why might a ViT require more training data than a CNN when trained from scratch?
5. Give one reason why attention maps are useful and one reason why they can be misleading.

26.25.2 Mathematical exercises

1. Let \(X\in\mathbb R^{224\times224\times3}\) and \(P=14\). Compute the number of tokens and the dimension of each flattened patch.
2. Let \(Q,K\in\mathbb R^{N\times d_k}\). Prove that \(QK^T\in\mathbb R^{N\times N}\).
3. Prove that the attention matrix \(A=\operatorname{softmax}(QK^T/\sqrt{d_k})\) is row-stochastic.
4. Assume the components of \(q\) and \(k\) are independent with mean \(0\) and variance \(1\). Compute the variance of \(q^Tk\).
5. Show that, without positional encodings, self-attention is permutation equivariant.
6. Derive the approximate attention cost as a function of image size \(H\times W\), patch size \(P\), and embedding dimension \(d\).
7. Let \(A\) be a row-stochastic attention matrix. Show that each output token is in the convex hull of the value tokens when values are scalar or vector-valued.
8. Compare a \(3\times3\) convolution on a \(14\times14\) feature map with full self-attention on \(196\) tokens. Count the number of spatial interactions.

26.25.3 Python exercises

1. Write a NumPy function patchify(X, P) for an image tensor \(X\).
2. Implement scaled dot-product attention from scratch and verify that each row of the attention matrix sums to one.
3. Visualize attention matrices for random tokens, locally correlated tokens, and repeated-pattern tokens.
4. Build a small TensorFlow ViT block and print the tensor shapes at each step.
5. Compare training a small CNN and a small ViT on the same dataset.
6. Create a plot showing how full-attention memory grows with image resolution.

26.25.4 AI/modeling exercises

1. Ask an AI assistant to explain why the \(1/\sqrt{d_k}\) factor appears in attention. Then write your own derivation using variance.
2. Ask an AI assistant for three ways to make ViTs efficient for high-resolution images. Classify each idea as sparse, hierarchical, low-rank, or token-reduction based.
3. Ask an AI assistant to generate pseudocode for a ViT classifier. Check every tensor shape manually.
4. Ask an AI assistant to compare class-token pooling and mean pooling. Identify claims that require empirical evidence.
5. Ask an AI assistant to critique an attention visualization. Then propose a stronger causal test using patch masking.

26.26 25.26 Notes and references

The transformer architecture was introduced in Vaswani et al. (2017). Vision Transformers adapted the transformer encoder to image patches and demonstrated strong results when trained at scale Dosovitskiy et al. (2020). Data-efficient image transformers and distillation are discussed in Touvron et al. (2021). Hierarchical windowed transformers such as Swin address the quadratic cost of full attention in dense vision Liu et al. (2021). DETR uses transformers for object detection as set prediction Carion et al. (2020). DINO and masked autoencoders show the importance of transformer architectures in self-supervised vision Caron et al. (2021); He et al. (2022). Vision-language models and promptable segmentation systems build on transformer-style token interaction and cross-attention mechanisms Radford et al. (2021); Kirillov et al. (2023).

Carion, Nicolas, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. 2020. “End-to-End Object Detection with Transformers.” European Conference on Computer Vision, 213–29. https://arxiv.org/abs/2005.12872.

Caron, Mathilde, Hugo Touvron, Ishan Misra, et al. 2021. “Emerging Properties in Self-Supervised Vision Transformers.” IEEE/CVF International Conference on Computer Vision, 9650–60.

Dosovitskiy, Alexey, Lucas Beyer, Alexander Kolesnikov, et al. 2020. “An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale.” arXiv Preprint arXiv:2010.11929. https://arxiv.org/abs/2010.11929.

He, Kaiming, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. 2022. “Masked Autoencoders Are Scalable Vision Learners.” IEEE/CVF Conference on Computer Vision and Pattern Recognition, 16000–16009.

Kirillov, Alexander, Eric Mintun, Nikhila Ravi, et al. 2023. “Segment Anything.” Proceedings of the IEEE/CVF International Conference on Computer Vision. https://arxiv.org/abs/2304.02643.

Liu, Ze, Yutong Lin, Yue Cao, et al. 2021. “Swin Transformer: Hierarchical Vision Transformer Using Shifted Windows.” IEEE/CVF International Conference on Computer Vision, 10012–22. https://arxiv.org/abs/2103.14030.

Radford, Alec, Jong Wook Kim, Chris Hallacy, et al. 2021. “Learning Transferable Visual Models from Natural Language Supervision.” International Conference on Machine Learning. https://arxiv.org/abs/2103.00020.

Touvron, Hugo, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. 2021. “Training Data-Efficient Image Transformers and Distillation Through Attention.” International Conference on Machine Learning, 10347–57. https://arxiv.org/abs/2012.12877.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, et al. 2017. “Attention Is All You Need.” Advances in Neural Information Processing Systems 30. https://arxiv.org/abs/1706.03762.