What it actually does
Cut the image into a grid of 16-by-16 patches, flatten each one into a vector, and feed the resulting sequence to a plain text transformer with nothing changed. The model is told nothing about two-dimensional structure — not locality, not translation equivariance, not scale — beyond a learned position embedding per patch. That sounds like it should not work, and on a hundred thousand images it does not: below a certain data scale a convolutional network wins comfortably, because the biases a CNN hardcodes are true and free. Above that scale the transformer learns better biases than the ones we knew to build in, and overtakes.
- Enough pretraining data exists, directly or through a checkpoint, to learn the spatial structure the architecture does not assume
- The image can be meaningfully cut into fixed-size patches on a fixed grid — false for extreme aspect ratios and for resolution-critical fine detail
- Global relationships matter enough to justify attending over every patch pair from the first layer
- Inference resolution matches training resolution, or position embeddings are interpolated to match
- The preprocessing pipeline reproduces the checkpoint’s exactly, since the features are not robust to a normalization mismatch
Loss function
The loss is ordinary cross-entropy and carries no information about the architecture, which is exactly the point worth making: everything distinctive about ViT is in the forward pass and the training recipe, not in the objective. Two details in the attention expression do carry weight. The scale factor is not cosmetic — remove it and the dot products grow with head width until the softmax is effectively one-hot, at which point gradients stop flowing and training stalls silently. And the target is almost never one-hot: label smoothing, mixup and CutMix are not optional refinements here but part of what makes the model trainable at all at moderate data scale.
- the target, which is almost never one-hot in practice — label smoothing and mixup make it a distribution
- sequence length: patches per image. Attention cost is quadratic in this, which is the whole resolution problem
- the scale factor. Without it the dot products grow with head width and the softmax saturates into a near-one-hot that passes no gradient
- three linear projections of the same patch sequence — in practice one fused projection, then split
- softmax over classes from the class token, or from mean-pooled patch tokens
How it is minimized
AdamW with a long warmup, cosine decay, heavy augmentation and stochastic depth
- warmup steps — typically thousands. Skip them and the run diverges in the first few hundred updates
- decoupled weight decay, applied to the parameter and not through the gradient; this is the W in AdamW
- bias-corrected first and second gradient moments
- total steps. The schedule is defined against it, so changing the budget changes every learning rate
This is the entry where the optimizer is not a detail. A ViT trained with the recipe that works for a ResNet simply fails, and it took the field two years and a separate paper to establish that the gap was never architectural. Warmup is mandatory rather than helpful: the attention logits at initialization are near-uniform, early large updates saturate the softmax, and a saturated softmax passes no gradient — the failure is a loss curve that flattens in the first few hundred steps and never recovers. Decoupled weight decay matters because coupling decay into an adaptive optimizer scales it by the gradient magnitude, which is nonsense. Augmentation carries the inductive bias the architecture lacks: RandAugment, mixup and CutMix are how the model is told about translation and scale, and removing them costs several points of accuracy that no amount of extra capacity recovers. Stochastic depth is what allows depth past about twelve blocks without collapse.
patch size
Sets sequence length as the square of its inverse; halving it quadruples the tokens and roughly quadruples attention cost (8, 14 or 16 pixels)
embedding width
Model capacity; the projections are quadratic in it and dominate FLOPs at low resolution (384 to 1280)
depth
Blocks stacked. Past roughly 24 the returns need stochastic depth and careful initialization to appear at all (12 to 32)
heads
Attention subspaces. Head width is what matters, and below about 32 per head quality degrades noticeably (6 to 16)
warmup steps
Not optional. Too short and the run diverges immediately, with a characteristic flat loss curve (5k to 20k steps)
weight decay
Much larger than CNN practice; the main regularizer alongside augmentation (0.05 to 0.3)
stochastic depth rate
Probability of dropping a whole block during training; scales with depth and is what makes deep ViTs trainable (0.0 to 0.4)
Well-behaved once the recipe is right and brittle before then, which is an uncomfortable combination because the failures look like architecture problems. Three are characteristic. Skipping or shortening warmup gives a loss that flattens in the first few hundred steps as the attention softmax saturates, and it never recovers however long it runs. Attention collapse in deep models sends every head toward the same near-uniform distribution, so the model becomes an expensive MLP — visible as attention entropy climbing toward its maximum and constant across heads. And training from scratch on a mid-sized dataset simply underperforms a ResNet, which is not a bug but the data-scale threshold asserting itself: the honest response is to pretrain or to use a CNN, not to tune harder.
O(N²·d) for attention and O(N·d²) for the projections, with N the patch count. At 224 pixels and patch 16 there are 196 tokens and the projections dominate; at 1024 pixels there are over four thousand and the quadratic term takes over completely. That crossover is the entire reason hierarchical and windowed variants exist.
How, where, and why it is used
Time-Series Forecasting
AdaptedThe patching idea transfers directly and is the substance of PatchTST: cut the series into overlapping segments of a few dozen steps, treat each segment as a token, and run the same encoder. A patch is a far better token than a single timestep — it carries local shape, it cuts sequence length by the patch size and therefore attention cost by its square, and it gives the model something with internal structure to attend over.
- Long-horizon multivariate forecasting, where patching is what makes a thousand-step context affordable
- Channel-independent forecasting, where each series is encoded separately with shared weights — consistently stronger than mixing channels early
- Time-series foundation models, which are patch transformers pretrained across many domains
- Representation learning on sensor data, where masked-patch pretraining then feeds a small supervised head
The transfer is real and the reasoning is clean: a single timestep carries almost no information, so attention over timesteps spends its capacity relating noise, while a patch has local shape worth attending to. The quadratic saving is the other half — patch length P divides the sequence and therefore divides attention cost by P squared. Against it: on a single ordinary business series with a few hundred observations it loses decisively to exponential smoothing, and the honest trigger is many long series rather than model sophistication. The channel-independence result is also worth knowing because it is counterintuitive — encoding each series separately with shared weights beats letting the model mix channels early, which suggests cross-series attention mostly overfits.
- Use overlapping patches with a stride below the patch length, so a pattern straddling a boundary is not destroyed
- Normalize per series and per window (reversible instance normalization), which matters more here than the architecture does
- Keep channels independent with shared weights unless cross-series structure is demonstrably present
- Emit all horizons at once from the encoded sequence rather than decoding recursively, which stops errors compounding
Rolling-origin backtesting with MASE against seasonal-naive and against exponential smoothing, reported per horizon. The naive baselines matter more here than anywhere else in this reference, because long-horizon benchmark results in this literature have repeatedly failed to beat a linear model once evaluated carefully.
- Not normalizing per window, which leaves the model modelling the level rather than the shape
- Non-overlapping patches severing patterns exactly at boundaries
- Comparing against no linear baseline, which in this specific literature is how several published wins evaporated
- Leaking scaler statistics computed over the full series into a backtest
Anomaly Detection
AdaptedUse a pretrained ViT as a frozen feature extractor and do the detection in feature space. Patch tokens from a self-supervised model carry enough semantics that a simple nearest-neighbour or Gaussian model over a memory bank of normal patch features localizes defects precisely — and because the score is per patch, the output is a defect map rather than a single image-level number.
- Industrial visual inspection, where patch-level feature memory banks are the current standard approach on defect benchmarks
- Medical imaging triage, flagging studies whose features are unlike the normal population
- Content moderation and document fraud, where the anomaly is a local manipulation rather than a global property
- Wafer and PCB inspection, where defects are small, local and heterogeneous
It works because the hard part of visual anomaly detection is a good representation, and self-supervised ViT features are a very good one that required no defect labels — which is decisive, since defects are rare and labelling them exhaustively is impossible. Patch-level granularity gives localization for free. The caveats are real: the features are frozen and optimize nothing about anomalies, so a defect that is semantically invisible to the pretraining objective is invisible to the detector; the memory bank grows with the normal set and needs subsampling; and it is a genuinely heavy way to solve a problem that a classical detector handles when the defect is a simple intensity or geometry deviation.
- Take features from an intermediate block, not the last — the final layers are specialized to the pretraining objective and discard the texture detail defects live in
- Build the memory bank from confirmed-normal images only, and subsample it with a coverage-preserving selection or it grows without bound
- Aggregate neighbouring patches before scoring, so a defect spanning a patch boundary is not halved
- Keep the resolution high enough that the smallest defect of interest covers more than one patch
Image-level AUROC and pixel-level localization scores reported separately, since a model can flag the right image for the wrong region. Always compare against a classical baseline on the same data — on defects that are simple intensity deviations, a cheap detector matches this at a fraction of the cost and that comparison is routinely skipped.
- Using final-layer features, which are too semantic and miss texture-level defects entirely
- A memory bank built from images that quietly contain defects, which teaches the detector they are normal
- Domain shift from a lighting or camera change invalidating the whole bank at once
- Reporting only image-level AUROC, which hides that the localization is wrong
Optimization
AdaptedTwo separate stories. The training-recipe one is the more important: ViT is the clearest case in modern vision of the optimization procedure determining the result, where identical architectures differ by many points of accuracy purely through schedule, decay and augmentation. The systems one is attention cost — the quadratic memory of materializing an N-by-N attention matrix is the binding constraint on resolution, and the fix is a classic IO-aware restructuring that tiles the computation and recomputes the softmax online so the matrix is never written to memory at all.
- Warmup and cosine schedules as the difference between divergence and convergence, not as tuning
- Decoupled weight decay, which is a correction to how decay interacts with an adaptive optimizer
- Online softmax with tiling — the FlashAttention idea — as an IO-bound rather than compute-bound optimization
- Gradient checkpointing trading recomputation for activation memory, which is what makes large batches fit
Worth studying because it is the strongest available counterexample to architecture-first thinking: the field spent two years attributing to the architecture a gap that was entirely in the recipe. The attention-memory story is equally instructive for a different reason — the naive implementation is not slow because of arithmetic but because it writes and re-reads an N-by-N matrix, so the win comes from never materializing it rather than from doing less work. That distinction, compute-bound versus memory-bound, is the single most useful diagnostic habit in performance work.
- Never skip warmup; the attention softmax saturates in the first few hundred steps without it and no later schedule recovers
- Decouple weight decay from the adaptive update, or decay is silently scaled by gradient magnitude
- Tile attention and keep only running softmax statistics, so memory is linear in sequence length rather than quadratic
- Profile memory traffic before arithmetic on attention — the naive version is bandwidth-bound and optimizing FLOPs does nothing
Ablate the recipe one element at a time against a fixed architecture — warmup, decay, each augmentation — which is the only way to see how much of the result is the model. For the systems side, measure achieved memory bandwidth rather than FLOP utilization, since the naive attention kernel can look compute-idle while being entirely bandwidth-saturated.
- Attributing a recipe difference to the architecture, which is precisely the error the field made collectively
- Optimizing attention FLOPs when the kernel is memory-bound and the arithmetic was never the constraint
- Changing total steps without noticing that a cosine schedule redefines every learning rate along the way
Computer Vision
Primary fitThe default backbone for image understanding at scale. Patches become tokens, a stack of self-attention blocks mixes them globally from the first layer, and a class token or mean pooling produces the image representation. Detection and segmentation heads attach to the patch tokens directly, and the same encoder serves as the image tower of a vision-language model unchanged.
- Large-scale image classification and as the pretrained backbone nearly everything else fine-tunes from
- Self-supervised representation learning by masked-image modelling or self-distillation, which is where the strongest features come from
- The image encoder in vision-language models, where a shared token interface with the text side is the actual advantage
- Dense prediction — detection and segmentation — via hierarchical or windowed variants that restore some locality
Global receptive field from the first layer, and a token interface that composes with text, audio and video without adaptation — which is why multimodal models are built on it rather than on convolutions. The honest limit is data: below roughly ten million images a well-trained CNN matches or beats it, because convolution’s hardcoded locality and translation equivariance are true and cost nothing to assume. ViT wins when there is enough data to learn better biases than those, and when there is not, choosing it is a mistake that heavy augmentation only partly disguises.
- Normalize with the statistics the pretrained checkpoint used; a mismatch degrades features silently and substantially
- Interpolate position embeddings when changing resolution rather than reinitializing them, which is the standard and frequently forgotten step
- Use strong augmentation — RandAugment, mixup, CutMix — since it supplies the inductive bias the architecture does not have
- Prefer a hierarchical or windowed variant for dense prediction; a flat ViT at high resolution spends everything on quadratic attention
Top-1 accuracy for classification, with a matched-recipe CNN baseline at the same data scale — matched recipe is the whole point, since an unmatched comparison measures the schedule rather than the model. For transfer, linear probing and fine-tuning separately: they rank pretraining methods differently and reporting only one is how misleading claims get made.
- Training from scratch on a mid-sized dataset and concluding the architecture is weak, when it is the data scale
- Forgetting to interpolate position embeddings at a new resolution, which quietly destroys accuracy
- Quadratic attention making high resolution unaffordable, which a flat ViT has no answer to
- Comparing against a CNN trained with an old recipe, which measures the recipe gap rather than the architecture gap
Control & Operations
AdaptedUsed as the perception front end of a visuomotor policy: a frozen or lightly tuned ViT encodes camera frames into tokens, and a policy head consumes them. The appeal is that the representation was learned from data the robot never had to collect, and patch tokens give a spatially addressable interface a policy can attend over.
- Robot manipulation policies built on frozen self-supervised visual features
- Autonomous inspection and navigation, where a pretrained encoder replaces a hand-designed vision stack
- Quality-control automation, where the same encoder serves both inspection and the actuation decision
Pretrained features dramatically reduce the robot data needed, which is the dominant cost in this domain — and self-supervised ViT features transfer to manipulation better than supervised classification features, apparently because classification discards the spatial detail a policy needs. Against it: latency is a hard constraint in a control loop and a large ViT is not cheap; the features are frozen and optimize nothing about the control objective; and the distribution shift between internet images and a robot’s camera is larger than it looks.
- Use patch tokens rather than the pooled class token; a policy needs spatial addressing and pooling destroys it
- Match the preprocessing to the pretraining exactly, since a robot camera is already far from the pretraining distribution
- Distil into a smaller encoder if the control loop has a latency budget, rather than accepting a slower loop
- Freeze the encoder initially and fine-tune only if the data supports it — a small robot dataset will overfit a full backbone immediately
Task success rate on held-out object and scene configurations, not feature quality — the representation metric and the control metric diverge more than expected. Measure end-to-end loop latency, which is a hard constraint rather than a number to report.
- Encoder latency breaking the control loop, which no amount of accuracy compensates for
- Pooled features discarding the spatial information the policy depends on
- Overfitting a full backbone to a few hundred robot demonstrations
What it costs to run
Pretraining is a large-scale undertaking — hundreds to thousands of accelerator-days at the sizes that matter — which is why almost nobody pretrains and almost everybody fine-tunes. Fine-tuning is hours on a handful of GPUs, and the practical decision is which checkpoint to start from rather than how to train.
Tens of milliseconds per image on a GPU at 224 pixels for a base model; considerably more on CPU, which is usually prohibitive without distillation or quantization. Cost scales as the square of the token count, so resolution is the dominant lever and a doubling of image side roughly quadruples the attention term.
The backbone is effectively static — replaced when a better checkpoint appears, not retrained on a schedule. Heads are retrained as the label set or the domain shifts, which is a much cheaper and much more frequent operation.
- Track input statistics against the preprocessing the checkpoint expects; normalization drift degrades features without any error
- Monitor attention entropy per layer if training or fine-tuning deeply — uniformly high entropy across heads is attention collapse
- Watch per-class accuracy rather than the aggregate, since domain shift shows up in a few classes long before the average moves
- Alert on input resolution changes, which invalidate position embeddings unless they are interpolated
- Position embeddings are tied to a token grid. Change resolution without interpolating them and accuracy degrades badly with nothing logged
- Normalization constants are part of the model. Using ImageNet statistics with a checkpoint trained on different ones is a silent, substantial regression
- Attention memory is quadratic in tokens, so a resolution increase that looks modest can be an out-of-memory failure
- The class token and mean pooling are different interfaces and are not interchangeable between checkpoints, even when both exist
- Fine-tuning with the pretraining learning rate destroys the features; fine-tuning needs a rate one to two orders of magnitude smaller
Pros and cons, in context
Global receptive field from the first layer
A patch attends to every other patch immediately, whereas a CNN needs depth to relate distant regions. Decisive when the relationship between distant parts of an image is the signal; irrelevant when the task is local texture, where a CNN is both sufficient and cheaper.
One token interface across modalities
The reason vision-language models are built on ViT and not on convolutions: images, text and audio become the same kind of sequence and a single stack consumes all of them. Worth nothing in a single-modality pipeline.
Scales with data better than convolutional networks
Past roughly ten million images it learns spatial biases better than the ones we knew to hardcode, and keeps improving where CNNs flatten. Below that threshold this advantage is negative, and saying so plainly matters.
Exceptional self-supervised features
Masked-image modelling and self-distillation on ViTs produce representations that transfer better than supervised ones, including to dense tasks. This is where most current practical value comes from — the classification result is almost incidental.
Data-hungry by construction
Removing inductive bias means the bias has to come from data. On a mid-sized dataset from scratch a ResNet wins comfortably, and the honest response is to use a CNN or a pretrained checkpoint rather than to tune harder.
Quadratic attention caps resolution
Doubling the image side roughly quadruples attention cost, which makes high-resolution dense prediction unaffordable for a flat ViT. Hierarchical and windowed variants exist precisely because of this and give back some of the global reach.
Brittle to the training recipe
Warmup, decoupled decay, augmentation and stochastic depth are load-bearing, and getting them wrong looks like an architecture failure. This costs real practitioner time and is not visible from the model definition.
Resolution and preprocessing are baked in
Position embeddings are tied to a token grid and normalization constants to the checkpoint. Both cause silent, substantial degradation when mismatched — among the most common production errors with this model.
Attention maps are routing, not explanation
They are visually compelling and routinely presented as evidence of what the model looked at, which the interpretability literature has repeatedly cautioned against. In a regulated setting that gap between appearance and evidence is a genuine liability.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""One ViT encoder block and the patch embedding, written out literally.
The pipeline is short and there is nothing vision-specific in it after the
first step:
image -> patches -> linear projection -> + position -> [attention, MLP] xL
The only thing telling the model about two-dimensional structure is the
learned position embedding, one vector per grid cell. Everything else is a
plain text transformer.
Plain loops, no libraries. Attention is written as the definition reads, which
means the full N x N matrix is materialized - the exact thing the fast stage
refuses to do.
"""
import math
def patchify(image, patch_size):
"""(height, width, channels) -> (num_patches, patch_size^2 * channels).
Row-major over the grid, which fixes the token order and therefore what
the position embeddings mean. Changing it silently invalidates them.
"""
height, width, channels = len(image), len(image[0]), len(image[0][0])
if height % patch_size != 0 or width % patch_size != 0:
raise ValueError("image dimensions must divide the patch size")
patches = []
for row in range(0, height, patch_size):
for column in range(0, width, patch_size):
flat = []
for dy in range(patch_size):
for dx in range(patch_size):
for c in range(channels):
flat.append(image[row + dy][column + dx][c])
patches.append(flat)
return patches
def linear(vectors, weight, bias):
"""out[i][j] = sum_k vectors[i][k] * weight[k][j] + bias[j]."""
out_dim = len(weight[0])
output = []
for vector in vectors:
row = [bias[j] for j in range(out_dim)]
for k, value in enumerate(vector):
for j in range(out_dim):
row[j] += value * weight[k][j]
output.append(row)
return output
def softmax(scores):
"""Subtract the max before exponentiating, or a large logit overflows."""
peak = max(scores)
exponentials = [math.exp(score - peak) for score in scores]
total = sum(exponentials)
return [value / total for value in exponentials]
def attention_head(queries, keys, values, head_dim):
"""softmax(Q K^T / sqrt(d)) V for one head.
The scale factor is not cosmetic. Without it the dot products grow with
head width until the softmax is effectively one-hot, at which point the
gradient through it is zero and training stalls with no other symptom.
"""
num_tokens = len(queries)
scale = 1.0 / math.sqrt(head_dim)
# The full N x N score matrix, materialized because that is what the
# definition says. It is also the entire memory problem.
output = []
for i in range(num_tokens):
scores = []
for j in range(num_tokens):
dot = 0.0
for d in range(head_dim):
dot += queries[i][d] * keys[j][d]
scores.append(dot * scale)
weights = softmax(scores)
mixed = [0.0] * head_dim
for j in range(num_tokens):
for d in range(head_dim):
mixed[d] += weights[j] * values[j][d]
output.append(mixed)
return output
def multi_head_attention(tokens, wq, wk, wv, wo, bq, bk, bv, bo, num_heads):
"""Project, split into heads, attend per head, concatenate, project."""
queries = linear(tokens, wq, bq)
keys = linear(tokens, wk, bk)
values = linear(tokens, wv, bv)
model_dim = len(queries[0])
head_dim = model_dim // num_heads
concatenated = [[0.0] * model_dim for _ in tokens]
for head in range(num_heads):
lo, hi = head * head_dim, (head + 1) * head_dim
head_out = attention_head(
[row[lo:hi] for row in queries],
[row[lo:hi] for row in keys],
[row[lo:hi] for row in values],
head_dim,
)
for token, mixed in enumerate(head_out):
concatenated[token][lo:hi] = mixed
return linear(concatenated, wo, bo)
def layer_norm(vectors, gain, shift, epsilon=1e-6):
"""Per-token normalization: mean and variance over the feature axis."""
output = []
for vector in vectors:
mean = sum(vector) / len(vector)
variance = sum((value - mean) ** 2 for value in vector) / len(vector)
denominator = math.sqrt(variance + epsilon)
output.append([
(value - mean) / denominator * gain[i] + shift[i]
for i, value in enumerate(vector)
])
return output
def gelu(value):
"""The tanh approximation, which is what the reference models use."""
inner = math.sqrt(2.0 / math.pi) * (value + 0.044715 * value ** 3)
return 0.5 * value * (1.0 + math.tanh(inner))
def encoder_block(tokens, params, num_heads):
"""Pre-norm: normalize BEFORE each sublayer, add the residual after.
Pre-norm rather than post-norm is why deep ViTs train at all - it keeps a
clean identity path from input to output, so the gradient reaches the
early blocks without passing through every normalization on the way.
"""
normed = layer_norm(tokens, params['ln1_gain'], params['ln1_shift'])
attended = multi_head_attention(
normed, params['wq'], params['wk'], params['wv'], params['wo'],
params['bq'], params['bk'], params['bv'], params['bo'], num_heads,
)
tokens = [
[a + b for a, b in zip(token, delta)]
for token, delta in zip(tokens, attended)
]
normed = layer_norm(tokens, params['ln2_gain'], params['ln2_shift'])
hidden = linear(normed, params['w1'], params['b1'])
hidden = [[gelu(value) for value in row] for row in hidden]
projected = linear(hidden, params['w2'], params['b2'])
return [
[a + b for a, b in zip(token, delta)]
for token, delta in zip(tokens, projected)
]
def vit_forward(image, params, patch_size, num_heads, depth):
"""Patchify, embed, prepend the class token, add positions, then blocks."""
patches = patchify(image, patch_size)
tokens = linear(patches, params['patch_weight'], params['patch_bias'])
# The class token is a learned vector that belongs to no patch. It has
# nowhere to look but the patches, so its final state is the image
# summary - a trick borrowed wholesale from BERT.
tokens = [list(params['class_token'])] + tokens
# Position embeddings are the ONLY thing carrying 2-D structure, and they
# are tied to this exact token count. A different grid needs them
# interpolated, which is the most-forgotten step in deploying a ViT.
if len(params['position']) != len(tokens):
raise ValueError("position embeddings do not match the token count")
tokens = [
[value + offset for value, offset in zip(token, position)]
for token, position in zip(tokens, params['position'])
]
for layer in range(depth):
tokens = encoder_block(tokens, params['blocks'][layer], num_heads)
tokens = layer_norm(tokens, params['final_gain'], params['final_shift'])
return tokens[0] # the class token, now the image representation
O(N^2 * d) for attention and O(N * d^2) for the projections, with N the token count. Illustrative, not a measured benchmark: at 196 tokens and width 768 this is tens of millions of interpreter iterations per block, and the N x N score matrix is materialized per head — which is the memory problem, not the arithmetic one.
How this language actually wants it written.
"""The same block, typed, with the shape contract and grid geometry explicit.
What changes is the boundary rather than the arithmetic. Patch geometry
becomes a frozen dataclass that derives the token count instead of leaving it
implicit, position-embedding interpolation is provided rather than assumed,
and the preconditions that silently degrade a ViT in production - a mismatched
grid, a head width that does not divide, a normalization mismatch - are
checked rather than discovered as a quiet accuracy loss.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import NamedTuple, Sequence
Vector = list[float]
Matrix = list[Vector]
class ShapeMismatch(ValueError):
"""Raised on a shape violation instead of broadcasting past it."""
class GridMismatch(ValueError):
"""Raised when position embeddings do not match the current token grid.
Its own type because the correct response is specific: interpolate, not
reinitialize. This is the single most common silent ViT failure.
"""
@dataclass(frozen=True)
class PatchGeometry:
"""Frozen so the token count cannot drift from the grid it describes."""
image_height: int
image_width: int
patch_size: int
channels: int = 3
def __post_init__(self) -> None:
if self.patch_size <= 0:
raise ShapeMismatch("patch size must be positive")
if self.image_height % self.patch_size or self.image_width % self.patch_size:
raise ShapeMismatch(
f"a {self.image_height}x{self.image_width} image does not divide into "
f"{self.patch_size}-pixel patches"
)
@property
def grid(self) -> tuple[int, int]:
return self.image_height // self.patch_size, self.image_width // self.patch_size
@property
def num_patches(self) -> int:
rows, columns = self.grid
return rows * columns
@property
def sequence_length(self) -> int:
"""Patches plus the class token - what attention is quadratic in."""
return self.num_patches + 1
@property
def patch_dim(self) -> int:
return self.patch_size * self.patch_size * self.channels
class AttentionShape(NamedTuple):
"""Derived once, so no call site divides model_dim by num_heads again."""
model_dim: int
num_heads: int
@property
def head_dim(self) -> int:
return self.model_dim // self.num_heads
@property
def scale(self) -> float:
"""1/sqrt(d_k). Omitting it saturates the softmax and stalls training."""
return 1.0 / math.sqrt(self.head_dim)
def validate(self) -> None:
if self.num_heads <= 0:
raise ShapeMismatch("at least one head is required")
if self.model_dim % self.num_heads:
raise ShapeMismatch(
f"width {self.model_dim} does not divide into {self.num_heads} heads"
)
if self.head_dim < 16:
# Not fatal, but below roughly 32 quality degrades measurably and
# it is worth refusing to do it silently.
raise ShapeMismatch(f"head width {self.head_dim} is too narrow to be useful")
def patchify(image: Sequence[Sequence[Sequence[float]]], geometry: PatchGeometry) -> Matrix:
"""Row-major over the grid, which fixes what the position embeddings mean."""
if len(image) != geometry.image_height or len(image[0]) != geometry.image_width:
raise ShapeMismatch("image does not match the declared geometry")
size = geometry.patch_size
return [
[
image[row + dy][column + dx][channel]
for dy in range(size)
for dx in range(size)
for channel in range(geometry.channels)
]
for row in range(0, geometry.image_height, size)
for column in range(0, geometry.image_width, size)
]
def interpolate_positions(
positions: Matrix,
source_grid: tuple[int, int],
target_grid: tuple[int, int],
has_class_token: bool = True,
) -> Matrix:
"""Bilinearly resample patch position embeddings onto a new grid.
Provided rather than left to the caller because it is the step everyone
forgets: change resolution without it and accuracy collapses with nothing
logged anywhere. The class-token position is carried through untouched,
since it belongs to no grid cell.
"""
source_rows, source_columns = source_grid
target_rows, target_columns = target_grid
expected = source_rows * source_columns + (1 if has_class_token else 0)
if len(positions) != expected:
raise GridMismatch(f"expected {expected} position embeddings, got {len(positions)}")
if source_grid == target_grid:
return [row[:] for row in positions]
prefix = positions[:1] if has_class_token else []
grid = positions[1:] if has_class_token else positions
width = len(grid[0])
resampled: Matrix = []
for target_row in range(target_rows):
for target_column in range(target_columns):
source_y = (target_row + 0.5) * source_rows / target_rows - 0.5
source_x = (target_column + 0.5) * source_columns / target_columns - 0.5
y0 = min(max(int(math.floor(source_y)), 0), source_rows - 1)
x0 = min(max(int(math.floor(source_x)), 0), source_columns - 1)
y1 = min(y0 + 1, source_rows - 1)
x1 = min(x0 + 1, source_columns - 1)
wy, wx = source_y - y0, source_x - x0
corners = (
(grid[y0 * source_columns + x0], (1 - wy) * (1 - wx)),
(grid[y0 * source_columns + x1], (1 - wy) * wx),
(grid[y1 * source_columns + x0], wy * (1 - wx)),
(grid[y1 * source_columns + x1], wy * wx),
)
resampled.append([
sum(corner[index] * weight for corner, weight in corners)
for index in range(width)
])
return prefix + resampled
def linear(vectors: Sequence[Sequence[float]], weight: Matrix, bias: Vector) -> Matrix:
if len(weight) != len(vectors[0]):
raise ShapeMismatch(f"weight expects {len(weight)} inputs, got {len(vectors[0])}")
out_dim = len(weight[0])
return [
[
bias[j] + sum(value * weight[k][j] for k, value in enumerate(vector))
for j in range(out_dim)
]
for vector in vectors
]
def stable_softmax(scores: Sequence[float]) -> Vector:
"""Subtract the maximum first; a large logit overflows exp otherwise."""
peak = max(scores)
exponentials = [math.exp(score - peak) for score in scores]
total = math.fsum(exponentials)
return [value / total for value in exponentials]
def attention_head(
queries: Matrix, keys: Matrix, values: Matrix, shape: AttentionShape
) -> Matrix:
"""softmax(Q K^T * scale) V for one head."""
output: Matrix = []
for query in queries:
scores = [
sum(a * b for a, b in zip(query, key)) * shape.scale for key in keys
]
weights = stable_softmax(scores)
output.append([
sum(weight * value[d] for weight, value in zip(weights, values))
for d in range(shape.head_dim)
])
return output
def multi_head_attention(
tokens: Matrix, params: dict[str, Matrix | Vector], shape: AttentionShape
) -> Matrix:
shape.validate()
queries = linear(tokens, params['wq'], params['bq'])
keys = linear(tokens, params['wk'], params['bk'])
values = linear(tokens, params['wv'], params['bv'])
concatenated = [[0.0] * shape.model_dim for _ in tokens]
for head in range(shape.num_heads):
lo, hi = head * shape.head_dim, (head + 1) * shape.head_dim
head_out = attention_head(
[row[lo:hi] for row in queries],
[row[lo:hi] for row in keys],
[row[lo:hi] for row in values],
shape,
)
for token, mixed in enumerate(head_out):
concatenated[token][lo:hi] = mixed
return linear(concatenated, params['wo'], params['bo'])
def layer_norm(vectors: Matrix, gain: Vector, shift: Vector, epsilon: float = 1e-6) -> Matrix:
output: Matrix = []
for vector in vectors:
mean = math.fsum(vector) / len(vector)
variance = math.fsum((value - mean) ** 2 for value in vector) / len(vector)
denominator = math.sqrt(variance + epsilon)
output.append([
(value - mean) / denominator * gain[i] + shift[i]
for i, value in enumerate(vector)
])
return output
def gelu(value: float) -> float:
inner = math.sqrt(2.0 / math.pi) * (value + 0.044715 * value ** 3)
return 0.5 * value * (1.0 + math.tanh(inner))
def encoder_block(tokens: Matrix, params: dict, shape: AttentionShape) -> Matrix:
"""Pre-norm: normalize before each sublayer, add the residual after.
Pre-norm is why deep ViTs train at all - it leaves a clean identity path
from input to output, so gradients reach early blocks without passing
through every normalization on the way.
"""
normed = layer_norm(tokens, params['ln1_gain'], params['ln1_shift'])
attended = multi_head_attention(normed, params, shape)
tokens = [[a + b for a, b in zip(t, d)] for t, d in zip(tokens, attended)]
normed = layer_norm(tokens, params['ln2_gain'], params['ln2_shift'])
hidden = [[gelu(v) for v in row] for row in linear(normed, params['w1'], params['b1'])]
projected = linear(hidden, params['w2'], params['b2'])
return [[a + b for a, b in zip(t, d)] for t, d in zip(tokens, projected)]
What changed: The arithmetic is unchanged; the boundary is what moves. Patch geometry becomes a frozen dataclass that derives grid, token count and patch width instead of leaving them implicit at every call site, and the attention shape derives head width and the scale factor once so no caller divides by head count again — including a refusal to build heads narrower than is useful, which is not fatal but degrades quality and should not happen silently. Position-embedding interpolation is supplied rather than assumed, with its own exception type, because changing resolution without it is the single most common way a deployed ViT loses accuracy with nothing logged anywhere: the correct response is specific, so the error is specific. The softmax subtracts its maximum, which the literal version does not survive on a large logit, and layer normalization accumulates with fsum so a wide feature vector does not drift. Pre-norm ordering is stated in the block rather than left as an implementation detail, since it is the reason depth is trainable at all.
Identical asymptotics; the change is in what fails loudly rather than in speed. Illustrative, not a measured benchmark: comprehensions and fsum cost slightly more than the raw loops and buy numerical stability plus a shape contract that holds.
Faster, and the justification for each change.
"""Batched, fused, and tiled. The N x N attention matrix is never written.
Three changes, in order of how much they matter:
1. Q, K and V come from ONE projection with 3 * d columns, then a reshape.
Three separate GEMMs read the same activations three times; one reads
them once, and the split afterwards is free.
2. Attention is computed in tiles with an ONLINE softmax. Instead of
materializing scores and normalizing, each tile updates a running
maximum and running sum and rescales the accumulator. Memory drops from
quadratic in sequence length to linear, which is the FlashAttention
idea - and the win is bandwidth, not arithmetic.
3. Everything batches. Heads and batch fold into one leading dimension so
attention is a single batched matmul rather than a loop over heads.
"""
from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
FLOAT = np.float32
def patchify(images: NDArray[np.float32], patch_size: int) -> NDArray[np.float32]:
"""(B, H, W, C) -> (B, N, P*P*C) as a pure reshape and transpose.
No gather, no copy until the final ascontiguousarray: the patch grid is
already a strided view of the image, so the whole operation is metadata.
"""
batch, height, width, channels = images.shape
if height % patch_size or width % patch_size:
raise ValueError("image dimensions must divide the patch size")
rows, columns = height // patch_size, width // patch_size
blocked = images.reshape(batch, rows, patch_size, columns, patch_size, channels)
blocked = blocked.transpose(0, 1, 3, 2, 4, 5)
return np.ascontiguousarray(
blocked.reshape(batch, rows * columns, patch_size * patch_size * channels)
)
def fused_qkv(
tokens: NDArray[np.float32],
weight: NDArray[np.float32],
bias: NDArray[np.float32],
num_heads: int,
) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]:
"""One GEMM for all three projections, then a reshape into heads.
weight is (d, 3 * d). Three separate products would read `tokens` three
times; this reads it once, and the split is a view rather than a copy.
"""
batch, tokens_per_image, model_dim = tokens.shape
head_dim = model_dim // num_heads
projected = tokens.reshape(-1, model_dim) @ weight
projected += bias
projected = projected.reshape(batch, tokens_per_image, 3, num_heads, head_dim)
# (3, B, heads, N, head_dim): heads and batch lead, so the attention
# below is one batched matmul with no loop over heads.
projected = projected.transpose(2, 0, 3, 1, 4)
return projected[0], projected[1], projected[2]
def tiled_attention(
queries: NDArray[np.float32],
keys: NDArray[np.float32],
values: NDArray[np.float32],
block: int = 128,
) -> NDArray[np.float32]:
"""Online-softmax attention over key tiles. Never materializes (N, N).
The invariant: after processing tiles up to j, `accumulator` holds the
unnormalized weighted sum of values and (running_max, running_sum) hold
the softmax statistics for everything seen so far. A new tile with a
larger maximum rescales the accumulator by exp(old_max - new_max), which
is exactly the correction that keeps the running sum exact.
"""
*lead, num_tokens, head_dim = queries.shape
scale = FLOAT(1.0 / np.sqrt(head_dim))
accumulator = np.zeros(queries.shape, dtype=FLOAT)
running_max = np.full((*lead, num_tokens, 1), -np.inf, dtype=FLOAT)
running_sum = np.zeros((*lead, num_tokens, 1), dtype=FLOAT)
for start in range(0, num_tokens, block):
stop = min(start + block, num_tokens)
# One tile of scores: (..., N, tile). Quadratic in the TILE, not in
# the sequence, which is the entire point.
scores = queries @ keys[..., start:stop, :].swapaxes(-1, -2)
scores *= scale
tile_max = scores.max(axis=-1, keepdims=True)
new_max = np.maximum(running_max, tile_max)
# Rescale what we already have to the new maximum, then add the tile.
correction = np.exp(running_max - new_max, dtype=FLOAT)
np.subtract(scores, new_max, out=scores)
np.exp(scores, out=scores)
accumulator *= correction
accumulator += scores @ values[..., start:stop, :]
running_sum *= correction
running_sum += scores.sum(axis=-1, keepdims=True)
running_max = new_max
return accumulator / running_sum
def multi_head_attention(
tokens: NDArray[np.float32],
qkv_weight: NDArray[np.float32],
qkv_bias: NDArray[np.float32],
out_weight: NDArray[np.float32],
out_bias: NDArray[np.float32],
num_heads: int,
) -> NDArray[np.float32]:
batch, tokens_per_image, model_dim = tokens.shape
queries, keys, values = fused_qkv(tokens, qkv_weight, qkv_bias, num_heads)
mixed = tiled_attention(queries, keys, values)
# (B, heads, N, hd) -> (B, N, d): one transpose, one reshape, one GEMM.
mixed = mixed.transpose(0, 2, 1, 3).reshape(batch, tokens_per_image, model_dim)
return mixed @ out_weight + out_bias
def layer_norm(
tokens: NDArray[np.float32],
gain: NDArray[np.float32],
shift: NDArray[np.float32],
epsilon: float = 1e-6,
) -> NDArray[np.float32]:
"""Mean and variance over the feature axis, computed in one pass.
The variance uses the two-pass definition deliberately: the one-pass
E[x^2] - E[x]^2 form cancels catastrophically in float32 when the mean is
large relative to the spread, which is exactly the regime residual streams
drift into deep in a network.
"""
mean = tokens.mean(axis=-1, keepdims=True, dtype=np.float32)
centred = tokens - mean
variance = np.mean(centred * centred, axis=-1, keepdims=True, dtype=np.float32)
centred /= np.sqrt(variance + epsilon)
centred *= gain
centred += shift
return centred
def gelu_(values: NDArray[np.float32]) -> NDArray[np.float32]:
"""In-place tanh-approximation GELU: the MLP hidden state is 4x the model
width, so avoiding one temporary there is the largest single saving in
the block."""
cubed = values * values * values
cubed *= FLOAT(0.044715)
cubed += values
cubed *= FLOAT(np.sqrt(2.0 / np.pi))
np.tanh(cubed, out=cubed)
cubed += 1.0
values *= cubed
values *= FLOAT(0.5)
return values
def encoder_block(
tokens: NDArray[np.float32],
params: dict[str, NDArray[np.float32]],
num_heads: int,
) -> NDArray[np.float32]:
"""Pre-norm block, with both residual adds done in place."""
normed = layer_norm(tokens, params['ln1_gain'], params['ln1_shift'])
tokens = tokens + multi_head_attention(
normed, params['qkv_weight'], params['qkv_bias'],
params['out_weight'], params['out_bias'], num_heads,
)
normed = layer_norm(tokens, params['ln2_gain'], params['ln2_shift'])
hidden = normed @ params['w1']
hidden += params['b1']
gelu_(hidden)
tokens += hidden @ params['w2']
tokens += params['b2']
return tokens
What changed: The three changes are ordered by how much they matter, and the biggest one is not arithmetic. Fusing Q, K and V into a single projection means the token activations are read once instead of three times, and the split afterwards is a view; at 768 channels that is a real bandwidth saving for no extra FLOPs. Tiled attention with an online softmax is the substantive change: instead of materializing an N-by-N score matrix and normalizing it, each key tile updates a running maximum and running sum and rescales the accumulator by exp of the maximum difference, which keeps the result exact while making memory linear in sequence length rather than quadratic. That is the FlashAttention idea, and the point worth internalizing is that the naive version is slow because it writes and re-reads a large matrix, not because it does more arithmetic — the win is bandwidth. Batch and heads fold into leading dimensions so attention is one batched matmul rather than a loop, patchify becomes a pure reshape-and-transpose over the image’s existing strides, and the GELU runs in place because the MLP hidden state is four times the model width and is the largest temporary in the block. Layer normalization keeps the two-pass variance deliberately, since the one-pass form cancels catastrophically in float32 exactly where deep residual streams live.
Fuse adjacent operations so intermediates are never materialized
The N-by-N score matrix never exists — tiles are consumed as they are produced — and the GELU, both residual adds and the layer-norm scaling all run in place.
Trade-off: Attention weights are unavailable after the fact, so the attention maps people expect to visualize require a second, unfused pass; and in-place activations destroy the pre-GELU values a gradient check would need.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Every projection and both attention products are GEMMs over batch and heads folded into leading dimensions, so there is no Python loop over heads or over tokens.
Trade-off: The head-major transpose before attention is a genuine copy of the whole activation tensor, which on a short sequence can cost more than the loop it replaces.
Batch work to amortize interpreter overhead
Images, heads and tokens all become array dimensions, so a whole batch flows through a block with a constant number of interpreter-level operations.
Trade-off: Peak memory is the batch times sequence times four times the model width for the MLP hidden state alone, which is what sets the maximum batch size long before compute does.
Ensure contiguous memory layout and a single dtype
float32 throughout with an explicit ascontiguousarray after patchify keeps every GEMM on the BLAS fast path instead of silently copying a strided operand.
Trade-off: float32 accumulation in the online softmax drifts as sequence length grows, and the running-sum rescaling amplifies it — production kernels keep the accumulator in float32 but the statistics in higher precision for exactly this reason.
Attention memory drops from quadratic to linear in sequence length; the projections become four GEMMs per block for the whole batch. Illustrative, not a measured benchmark: past roughly a thousand tokens the tiled version is faster despite doing slightly more arithmetic, because the naive one is bound by writing and re-reading the score matrix.