Vision Transformer Deep Dive: ViT to Modern Descendants

An image transformer is not one trick called “attention.” It is a path: decode an image, turn its grid into tokens, repeatedly mix information, and read an output. Every descendant in this article changes one part of that path because vanilla ViT is data-hungry, expensive at high resolution, awkward for dense outputs, or too generous with tokens.

We will build the baseline first, keep one tiny PyTorch program running beside it, and then follow each branch using the same five questions: what limitation appeared, what changed, how tensors move, which code changed, and which costs moved elsewhere.

Vanilla ViT-B/16: pixels to logits

This chapter follows the original ViT paper and its official code. The local paper deep dive gives the experimental context; the account here stays with the execution path an engineer must trace from an input file to logits and, during training, back through the optimizer.

The complete machine has two boundaries worth keeping separate. The learned model begins at patch embedding and ends at class logits. The application begins earlier, with compressed image bytes, and ends later, after the logits have been converted into the output the caller consumes. Keeping both boundaries visible prevents a fast encoder from being mistaken for a fast image pipeline.

Decode and normalize: pixels are the input boundary

A JPEG or PNG arrives as compressed bytes, not as the floating-point tensor ViT consumes. The decoded image is an H × W × C array of uint8 RGB samples in [0, 255], so it is rank-3. Resize and crop choose the spatial field of view before any implementation-specific layout conversion.

Layout is implementation-specific. The educational PyTorch toy uses an unbatched C × H × W floating-point tensor and batched B × C × H × W; the official Google/JAX path uses channel-last H × W × C and B × H × W × C. In either layout, the unbatched image is rank-3 and adding a batch axis makes it rank-4. For the toy's running crop, that batched shape is B × 3 × 224 × 224.

Normalization is checkpoint-specific: the official Google pipeline maps byte-valued pixels to [-1, 1], while other checkpoints may use different scaling or per-channel mean and standard deviation. That preprocessing choice changes every downstream token, so the checkpoint's own recipe is part of the input contract.

Hold two quantities for the rest of the post. Everything else is a walk from the first to the second, then through twelve identical blocks.

1. 150,528 — 224 × 224 × 3. Raw RGB. No tokens, no embeddings, no CLS. Flatten the spatial grid and you have 50,176 sites, each still carrying 3 channels.
2. 197 — 14 × 14 patches plus one CLS. That is the sequence length ViT-B/16 actually attends over.

Attention is O(n²) in that length. 197² ≈ 39,000 pairs. 50,176² ≈ 2.5×10⁹. That ratio is why patches exist.

A CNN hides the tensor immediately: the first 3 → 64 kernel mixes the channels and you start talking about feature maps. ViT does not. The first honest sentence about this architecture is: the input is not an object. The input is a tensor.

The PureNumbers instrument samples the decoded uint8 boundary: click a pixel and you get three integers, with no "ear" or "cat" attached. The later normalized floating-point image tensor feeds the model's patch embedding in the layout its implementation expects; the classifier receives the final CLS representation, never the image tensor.

Primary sources for this figure: the ViT paper and official ViT code.

The next decision is forced by arithmetic, not aesthetics. If each pixel were a token, self-attention on 50,176 sites would dominate training. So we refuse pixels as tokens.

Patch size sets the token bill

A 16×16 patch is 16 × 16 × 3 = 768 raw numbers. On a 224 image that is a 14×14 grid — 196 patches.

In general the patch grid is (H/P) × (W/P), so the patch-token count is HW/P² when both spatial dimensions divide evenly by P. Halving patch size doubles each grid axis, creates four times as many tokens, multiplies token-linear block work by four, and can multiply the pair interactions by sixteen.

Flip the grain below. Pixel tokens explode. 8×8 is a middle tax. 16×16 is the bill ViT-B actually pays.

Primary sources for this figure: the ViT paper and official ViT code.

This is an explicit image-specific inductive bias: local rectangles exist. The shared non-overlapping patch projection keeps that patch-grid locality, but ViT does not build the growing receptive field of a staged hierarchy of overlapping 3×3 convolution blocks.

Two patches that share an edge are neighbors in the image and strangers in the sequence until attention says otherwise. Neighborhoods that matter across patch boundaries are learned, not free.

Patch size is a knob on that trade:

• Smaller patches → more tokens, more spatial fidelity, more token-linear encoder work, and a quadratic pair bill you may not want
• Larger patches → fewer, coarser tokens. They enlarge E's (P²C) × D parameter matrix and its per-token projection input, but the patch-token count is N_patch = HW/P², so for batch B the matrix-only total is BN_patch(P²C)D = BHWCD, independent of P at fixed H, W, C, and D. This N_patch excludes CLS and is distinct from the block sequence length N = 197. Fewer tokens reduce token-linear encoder and MLP work as well as attention-pair work.

16 on 224 is a historical default, not a law. Model names encode it: ViT-B/16 means Base, patch 16.

One learned projection maps each patch into model width

Text embeddings are a lookup. Token id 17 means "row 17 of the table". Images have no such id.

A patch is a point in a continuous 768-dimensional input space (16 × 16 × 3 flattened). ViT chooses a standard inexpensive learned linear map from that point onto the model's width. Call its matrix E.

With row vectors, a flattened patch x has shape 1 × (P²C), the learned matrix E has shape (P²C) × D, and xE has shape 1 × D. The expression xE is matrix-only notation and leaves out the learned implementation bias for clarity. For ViT-B/16 both sides happen to be 768, but that square shape is a coincidence.

The same E and bias are shared by every patch, image, and position; there are no independently learned per-location kernels.

For a batch, reshape the image into B × 196 × 768 patch rows and multiply the last dimension by E. A Conv2d with kernel and stride both equal to P is an implementation-equivalent way to perform the same shared projection; it is not a convolutional stem with overlapping local stages.

Primary sources for this figure: the ViT paper and official ViT code.

The same instrument's basis view shows why E is a learned linear map, not a lookup table and not a list of named semantic detectors. Output coordinate j is the dot product of the flattened patch with column j of E.

After this affine projection, every patch is just another 768-vector. The encoder no longer knows it came from pixels. Training can organize useful relationships among patch codes, but neither the linear map nor input similarity guarantees task-level nearness in the learned space. Treat E as a frozen codebook and you will go looking for a vocabulary that was never there.

The training section below follows how those projection weights move after the classifier makes an error.

Class and position tokens define the sequence

Two more additions happen before the first block. They are easy to mash together. They are not the same object.

1. CLS — a learned vector with no pixels behind it. Prepended at index 0. Attends to every patch (and they attend to it). After the last block you usually read only this vector and hand it to a linear classifier. It is not a patch, and it is not "the average of the image".
2. Positions — a second learned table P with one row per sequence index, P[0] through P[196]. Each row is added to the matching token (CLS included). Not concatenated (that would change the width). Not multiplied. Added.

Without position embeddings, the encoder's token outputs are permutation-equivariant: permuting the 196 patch tokens permutes their output rows. With the CLS slot held fixed, its classification readout is invariant to a permutation of the patch tokens. Learned absolute position rows tie content to named seats and break that symmetry, so moving a face from left to right can change the result.

This chapter and cost model use the released fine-tuned ViT-B/16 configuration with a linear classifier head (representation_size=None) shown in the official model and official B/16 config. The original paper used a one-hidden-layer tanh MLP head during pretraining, then a linear head for fine-tuning; those are two configurations, not one blended head.

Primary sources for this figure: the ViT paper and official ViT code.

The resulting input to block 1 has shape B × 197 × 768: one batch dimension, one sequence axis containing CLS plus 196 patches, and one model-width axis. Position rows and token rows are added element by element, so neither the token count nor width changes at this step.

The original ViT block is pre-norm

One encoder layer, in the order the tensor actually walks:

1. Layer-norm over the 768 dimensions of each token (not over the sequence)
2. Multi-head self-attention — the only step that mixes across tokens
3. Residual: x ← x + attn(LN(x))
4. Layer-norm again
5. MLP, applied independently to each token. Inner width 4 × 768 = 3072, GELU in the middle, project back to 768
6. Residual: x ← x + mlp(LN(x))

Shape in equals shape out: 197 × 768. Nothing in this block knows it is looking at an image. Stack it twelve times and you have the entire ViT encoder. No decoder. No cross-attention. No causal mask.

Writing the order as equations makes the residual ownership unambiguous: u_l = z_{l-1} + MSA(LN(z_{l-1})), then z_l = u_l + MLP(LN(u_l)). After block 12, ViT applies the final normalization to the CLS row before the classifier. LayerNorm and the MLP do not mix token positions; only attention does.

Primary sources for this figure: the ViT paper and official ViT code.

Attention talks across tokens. The MLP remixes inside one token. The skips are why you can stack twelve of these without the identity dying.

Now let's look at the only step that mixes across tokens.

Attention computes affinities, then mixes values

From each token — including itself — the layer builds three vectors with three learned matrices shared by the whole sequence:

1. q = x W_q — what am I looking for?
2. k = x W_k — how can I be found?
3. v = x W_v — what do I hand over if someone attends to me?

An ear-patch might seek eyes while offering fur. Those are different jobs. One raw x cannot do both. That is why there are three projections, not a comparison of the embeddings themselves.

After splitting heads, Q, K, and V each have logical shape B × H × N × d_h, where d_h = D/H. QKᵀ contracts the head-width axis and produces B × H × N × N; row-wise softmax preserves that shape; multiplying the result by V contracts the key-token axis and returns B × H × N × d_h.

Score matrix

Relevance of token j to query i is the dot product q_i · k_j. Doing that for every pair writes a 197 × 197 score matrix QKᵀ (5×5 in the toy).

The real ViT-B/16 head is 64-dimensional, so its sum has 64 terms and uses √64 = 8. The visible instrument is a six-dimensional toy head: its displayed dot products have six terms and use √6. In both cases, matching the divisor to the actual head width keeps large raw scores from shoving softmax toward one-hot saturation. Then softmax.

Mix of values

The new token is not a mix of keys. Keys only decided the invitations. The new token is Σ_j a_{ij} v_j — a mix of values.

Primary sources for this figure: the ViT paper and official ViT code.

The five beats are one head. Multiple heads repeat that path in narrower subspaces, concatenate the results, and apply W_o.

That is the whole formula:

Attention(Q, K, V) = softmax(QKᵀ / √d_h) V

One layer, global receptive field. A CNN spends many strides to reach across the image. Here an ear can read an eye on the first block. Deeper write-up: self-attention in ViT.

Two consequences that are easy to miss:

1. No causal mask. Token 196 may look at token 0. A single image is not a left-to-right story.
2. No retained decoder state. During inference the whole image sequence is known, so the encoder computes all tokens in parallel and can release its intermediate K and V tensors after the pass. During training, autograd may retain those tensors or an implementation may recompute them for backward. The later vision-language branch introduces a separate autoregressive decoder boundary.

Heads split width, not the token sequence

One head is one set of W_q, W_k, W_v and one softmax. If that were the whole layer, the network would have to pick a single notion of relevance for every token.

ViT-B splits the 768-width into 12 heads of 64. Each head has its own projections. Concatenate the 12 outputs (back to 768) and a single W_o writes the mixture into model space.

Every head still processes all 197 sequence positions. Heads partition the feature width, not the token sequence: twelve heads do not create twelve groups of patches. With total width fixed, changing the number of heads changes d_h and the organization of the score tensors, not the leading dense-projection MAC formula.

Primary sources for this figure: the ViT paper and official ViT code.

This is not a committee vote on the same question. Separate matrices mean separate questions. The 64-wide subspace is also why √d_h = 8 and not √768. If you implemented "multi-head" as 12 copies of full-width 768-dimensional attention, you would build a different, higher-capacity layer: full-width heads change capacity, parameter count, and compute rather than merely reorganizing a fixed 768-wide projection.

After the head opens, one loss sends gradients through every branch

Before optimization, none of the learned state represents an "ear," but untrained does not mean every array is random noise. The official implementation zero-initializes the class token and classifier head, while the patch projection, positions, and encoder use their specified initializers.

That zero head kernel matters on the literal first backward pass: it cannot send a nonzero classifier gradient into the encoder, so the head updates before the lower stack can receive that route's signal. On a typical post-initial update, after the head kernel is nonzero, the tensor walks the full graph: CLS produces class logits, softmax gives p, cross-entropy is −log p(true class), and p − y seeds backpropagation through the encoder. The stage changes each beat below.

Primary sources for this figure: the ViT paper and official ViT code.

The ETrainFilm instrument is a sequence of local-gradient demonstrations, not one numerically connected forward-and-backward trace. Its value, softmax, activation, and projection stages teach distinct derivative routes; their displayed numbers should not be read as one shared computation graph.

Nobody typed "detect edges." Training can develop edge-sensitive features when they help the objective, but the label does not assign such a feature directly.

On a typical post-initial update, the complete training boundary is logits → cross-entropy loss → backpropagation through the classifier head → encoder blocks → patch projection → optimizer step. Backpropagation can also reach the learned position table, class token, attention and MLP weights, normalization parameters, and classifier. Inference stops after the application consumes the logits; training continues until the optimizer uses the accumulated gradients to update parameters and, for a stateful optimizer, optimizer state.

The backward pass is more concrete than "gradients flow". On the value branch, V reuses the forward attention weights to form its backward weighted sum: ∂L/∂V = Aᵀ(∂L/∂O), so those coefficients split the output gradient among value vectors. Q and K receive gradients through the attention probabilities and the softmax Jacobian instead: changing either changes QKᵀ, which changes A, which changes the value mixture. The simple “same weights backward” picture describes V only.

A column of E is an output feature or basis direction. Semantic responses can emerge from combinations of those coordinates and later layers; no named detector is assigned to one column.

Most of the stack is not linear. Between those mixes sit GELU (in the MLP) and softmax (in attention). Their forward roles differ: GELU smoothly gates inputs by magnitude, while softmax turns logits into a distribution that sums to 1. Backward, the wish arriving on the output of either nonlinearity is not the wish that should land on its input. The conversion factor is the local slope — the derivative.

That is why we differentiate activations. The chain rule is not decoration. If a = f(z) and a wish δ_a arrives on a, the only honest wish on z is δ_z = δ_a · f'(z). Skip the multiply and you pretend the layer was the identity. On the left of a ReLU that lie is total: f' = 0, the unit is dead, and this example never reaches E. Softmax has the same structure: p − y is already the derivative of cross-entropy-plus-softmax, which is why that residual is the seed of the whole backward pass.

Residuals exist for the opposite reason. The skip path is x + f(x). Its derivative through the + x branch is 1. A wish can ride that highway all the way down to E without being multiplied by twelve GELU slopes. We do differentiate the activations on the residual branch. We refuse to differentiate away the identity.

At the bottom, E never sees the label directly. For a row-vector patch, ∂L/∂E = xᵀ(∂L/∂(xE)): an outer product of the input coordinates and a gradient that has already crossed twelve blocks. A zero input coordinate contributes zero to that row for this example, while shared use of E makes the gradients from all patches accumulate.

One useful intuition is that repeated gradient signals can make task-relevant distinctions accumulate while idiosyncratic signals compete or cancel. That picture helps explain why representation geometry can change during training, but it is an interpretation of optimization rather than a guarantee about any two embeddings.

The optimizer step uses the descent convention E ← E - η∂L/∂E; this notation is plain gradient descent. For pretraining, the paper's training recipe used Adam with warmup, while fine-tuning used SGD with momentum. Adam is designed to minimize the training objective, but an individual Adam step is not guaranteed to lower the current batch loss.

E is not trained first and frozen. On typical updates after the zero-initialized head has moved, E, P, CLS, every W_q, the MLP, the LayerNorm scales, and the head can receive gradients in the same backward sweep. The randomized projection and encoder initializers help break symmetry; a column of E remains one patch-projection output feature rather than a named semantic detector.

On typical later updates, P and CLS can receive gradients directly as learned parameters. The paper observes that learned embeddings for closer image positions tend to be more similar. Interpreting that geometry as the result of correlated gradient histories is a useful hypothesis, not a causal mechanism established by the paper's experiment.

The complete execution and latency ledger

The literal vanilla inference path is decode → resize/crop → normalize → patch embedding → encoder → classifier head → postprocess. Decode changes compressed image bytes into RGB pixels. Patch embedding through the classifier is learned model work; the classifier reads the final normalized CLS vector and produces logits; postprocess turns those logits into the application output, for example probabilities or a ranked label set.

A model-only measurement begins at an already prepared tensor and normally ends at logits. An end-to-end measurement includes file decode, resize/crop, normalization, transfers, model execution, and postprocess. Later descendants may replace or extend the classifier head with a task projector, but the vanilla classification path here ends in the classifier head and contains no such extra projection stage.

For a row-vector patch x, the learned patch map has shape E ∈ ℝ^{(P²C) × D}, so xE produces one width-D token. For batch B, tokens N, width D, and MLP ratio r, one block's matrix MACs are

In the article's compact notation that is (4 + 2r)BND² + 2BN²D. QKV contributes 3BND², the output projection contributes BND², the two attention products contribute 2BN²D, and the two MLP projections contribute 2rBND². For r=4, this is 12BND² + 2BN²D.

At ViT-B/16's B=1, N=197, D=768, and r=4, QKV costs 348,585,984 MACs, the output projection 116,195,328 MACs, the two attention products 59,610,624 MACs together, and the MLP 929,562,624 MACs. The block total is 1,453,954,560 MACs; this substitution assumes B = 1. About 95.90% is in the linear-in-token projection and MLP terms at this resolution, while about 4.10% is in the quadratic attention products.

At this resolution, the projection and MLP terms are larger than the two quadratic attention products. Raising resolution changes that balance because the pair products grow with N²; “attention is quadratic” is not the same claim as “the attention matrix always dominates latency.” Normalization, activation, softmax, indexing, reshaping, allocation, and data movement are not included in this matrix-only equation and must be reported separately.

The storage view reports analytical tensor sizes, not a promise about process memory. A naive implementation may expose Q/K/V, a per-head N × N score tensor, probabilities, and the rD MLP activation, but fusion, recomputation, autograd, scheduling, and allocator behavior determine which objects coexist. The ledger therefore describes tensor accounting rather than a measured peak resident set or a universal latency.

Primary sources for this figure: the ViT paper and official ViT code.

The path view keeps non-model stages visible; the MAC view changes resolution and patch size through the tested cost model; and the storage view changes bytes per element without pretending that element count predicts a device runtime. Any real timing claim would still need the exact configuration, batch, precision, warm-up, synchronization, runtime, device, statistic, and measurement boundary.

Run the tiny model

The complete educational executable uses synthetic input, asserts every major shape, performs one classifier loss and backward pass, and never downloads a dataset or checkpoint:

python3 examples/vit_lineage_toy.py --demo vit --seed 7

PyTorch is optional for the website and must be installed separately for this command. The default path is CPU-only. The core data path is the same path used by the figures:

patches = self.patch_tokens(images)
tokens = self.prepend_tokens(patches)
tokens = self.add_positions(tokens)
tokens = self.encode_tokens(tokens)
cls = tokens[:, 0]
logits = self.head(cls)

Open the complete runnable toy.

The toy is deliberately small; it is a mechanism check, not a checkpoint reproduction or a latency benchmark.

Questions that close the baseline

Why exactly 16×16? A trade. Patch 8 → 784 tokens → finer grid, about 16× the attention pairs. Patch 32 → 49 tokens → cheap and coarse. Sixteen sits at the historical operating point for 224-pixel ViT-B/16; it is a configuration, not a law.

Why trust CLS with the answer? The classification loss is read from the CLS output, so CLS learns queries that collect class-relevant information. Mean-pooling the patch tokens is a different valid readout, but it changes which representation the head is trained to consume.

Why pre-norm, not post-norm? Pre-norm leaves a direct identity path through the residual stream. The original ViT paper specifies pre-norm even though the 2017 Transformer used post-norm.

Why GELU, not ReLU? GELU smoothly attenuates the negative side instead of clipping every negative value to zero. Its local derivative therefore usually passes a graded training signal rather than a binary gate.

Do softmax or the dot product have parameters? No. In matrix-only notation, attention's learned projection matrices are W_q, W_k, W_v, and W_o; the official dense layers also carry learned bias parameters. Scoring, scaling, and softmax are fixed operations over their outputs.

What if the image is not 224×224? Resize, or keep the patch size and accept a new token count. The learned absolute position table is then commonly interpolated over its 2-D patch grid before fine-tuning.

Where do color and texture live if attention only relates patches? Inside each token. The projection and per-token MLP encode within-patch appearance; attention relates those descriptions across positions.

That is the complete machine: numbers → patches → learned tokens with a class slot and positions → twelve rounds of cross-token attention and per-token MLP work → final normalization → read CLS → classifier logits → application output. The training route walks the same graph backward and ends only after the optimizer updates the shared parameters.

What ViT refuses to assume

Write down what is not in the machine. It will save you from borrowing the wrong story from CNNs or from GPT.

• No overlapping or staged convolutional hierarchy. Patch projection may be implemented as a shared Conv2d whose kernel and stride both equal P, but after that non-overlapping cut the encoder does not build a multistage convolutional feature pyramid.
• Limited pixel-level translation equivariance. E is shared across patch locations, but the non-overlapping stride and patch granularity plus learned absolute positions limit pixel-level translation equivariance: a sub-patch shift can change both patch contents and their named seats.
• No decoder, no causal mask. This is an encoder. Everyone may look at everyone.
• No retained decoder state on a single image. During inference the encoder sees the whole sequence, performs one parallel pass, and can release its attention intermediates. Training may retain or recompute them for backward. Decoder state belongs to the later autoregressive boundary, not to vanilla image encoding.
• No discretized codebook. Vanilla ViT never snaps a patch to a nearest neighbor in a VQ table. Continuous in, continuous through E. People add codebooks later, for generation, not for this classifier.

Within the paper's scaling experiments, BiT led with ImageNet pretraining and ViT overtook it with larger datasets. On JFT-300M, the tested ViT variants showed a better performance-versus-pretraining-compute trade-off and no saturation within the explored range. Those observations are bounded to the reported models, datasets, and compute budgets rather than establishing a universal causal ceiling.

Follow-up work makes targeted changes to the baseline path. DeiT can add a distillation token and head while changing the training recipe; Swin changes attention neighborhoods and staging through shifted local windows and a hierarchical representation. Each preserves transformer ingredients while changing part of the machine traced above.

The instruments do not depend on the scaling table. They only show the machine. The table tells you when the missing bias stops being a tax.

A practical trace checklist

1. Draw the tensor first — H × W × 3 — before you say "tokens". Count 150,528 before you count 197.
2. Pick patch size from the attention bill, not from aesthetics. 16×16 on 224 is 196 plus CLS. Write n² down.
3. Treat E as a linear map you will train, not as a frozen codebook. Check its shape: (P²C) × D.
4. Add positions. Measure what happens when you shuffle patches and keep P glued to the old seats. If nothing changes, P is not doing its job.
5. Read CLS after the last block for classification. Keep the patch tokens if the head is dense (segmentation, depth, a VLM connector).
6. When something is "black box", name the tensor that moves (E, P, CLS, W_q) and the scalar that moves it (the loss). If you cannot name both, you do not yet have a mechanism.

Data hunger: DeiT

Limitation. Vanilla ViT's weak image-specific inductive bias made the original recipe unusually dependent on large supervised datasets. A model that can relate every patch globally still has to learn useful visual regularities, and the first ViT results relied on substantially more labeled pretraining data than ImageNet alone.

Architecture delta. DeiT keeps the patch projection and transformer encoder, but changes the optimization recipe with strong augmentation, repeated augmentation, regularization, and careful distillation from a CNN teacher. The distilled variant adds a learned distillation token and a second classifier head; those additions are only part of the data-efficiency result, not a substitute for the training recipe.

Mechanism. ViT-B/16 supplies 196 patch tokens; the distilled student prepends CLS and DIST, so the encoder carries 196 patch tokens + CLS + DIST = 198 tokens. The CLS representation feeds a class head trained from the ground-truth label, while DIST feeds a distillation head trained from a hard or soft teacher target. Those two student heads share the encoder but own different supervision. For joint classification, the DeiT paper adds the two heads' softmax probability outputs. The released implementation in the official DeiT repository instead averages their logits, and this toy follows that released-code convention; it is not a universal fusion rule for every DeiT variant.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo deit --seed 7. DistilledViT adds DIST, a second head, detached teacher logits, and explicit hard and soft distillation losses; the synthetic run checks both training losses and the fused evaluation output.

def run_deit_demo(config: ToyConfig) -> ShapeReport:
    model = DistilledViT(config)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
# ...
    training_output = model(images)
# ...
    assert torch.allclose(fused_logits, expected_fused)
    return {
        "images": tuple(images.shape),
        "tokens": (
            config.batch_size,
            config.patch_tokens + 2,
            config.width,
        ),
        "class_logits": tuple(class_logits.shape),
        "distill_logits": tuple(distill_logits.shape),
    }

Latency ledger. The teacher forward, strong augmentation and regularization recipe, and distillation loss are training-only fan-out. They can reduce the labeled-data requirement without entering the deployed student. Under the released-code inference convention followed by the toy, the distilled student still carries 198 tokens through every block, evaluates two heads, and averages their logits; the paper's probability-output fusion also evaluates both heads. That token and head work is inference-visible, while a non-distilled DeiT uses the ordinary class-token path. Data efficiency is not an inference speedup, and whether the extra token or head matters in wall-clock time depends on the complete implementation and measurement boundary.

High-resolution cost: Swin, PVT, and MViT

At fixed patch size, doubling each image side roughly quadruples N and creates roughly sixteen times as many attention pairs. Dense prediction also needs spatial features at more than one scale instead of one final classification vector. Swin, PVT, and MViT therefore change the spatial execution path rather than merely shrinking one otherwise identical global block.

Swin: shifted locality

Limitation. Global attention connects every token immediately, but its pair matrix becomes expensive on a high-resolution feature map. Dense prediction needs fine early features as well as increasingly semantic later features, while a flat ViT keeps one token scale through the encoder.

Architecture delta. Swin partitions the feature map into non-overlapping local windows and alternates ordinary and shifted window blocks. A patch merge between stages halves each spatial axis and grows channel width, yielding a hierarchy rather than one fixed-resolution sequence.

Mechanism. A regular partition feeds W-MSA inside each window. The next block applies a cyclic shift, makes a shifted partition, and uses a pairwise mask so wrapped image edges do not become false neighbors during SW-MSA. A reverse shift restores the original coordinates; because the shifted windows crossed the previous boundaries, information now has cross-window flow. Patch merge then combines each 2×2 neighborhood before the next, wider stage. At 56×56 with window size 7, there are 64 windows of 49 tokens: the regular and shifted blocks both keep that count under the efficient cyclic-shift path. The small mask shown in the figure is an illustrative selected submatrix, not the literal full stage mask. The Swin paper specifies this alternating path, and the official Swin implementation exposes the corresponding stages and masks.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo swin --seed 7. The toy checks that window partition and reverse are exact inverses, applies ordinary attention, constructs a shifted-window mask with both allowed and blocked pairs, reverses the shift, and projects a merged 2×2 neighborhood into twice the width.

def run_swin_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
    patches = patch_embed(images)
# ...
    windows = window_partition(features, 4)
# ...
    assert torch.equal(partitioned_then_reversed, features)
# ...
    shifted = shifted_window_block(ordinary, attention, 4, 2)
    merged = patch_merge(shifted, merge_projection)
# ...
    return {
        "features": tuple(features.shape),
        "windows": tuple(windows.shape),
        "shifted": tuple(shifted.shape),
        "merged": tuple(merged.shape),
    }

Latency ledger. Window attention replaces one global pair matrix with a sum of smaller regular matrices, but fewer pair interactions do not guarantee lower wall-clock latency. Window partition and reverse, cyclic shift, the pairwise mask, reshape and layout work, padding when a grid does not divide evenly, and patch merge between stages all sit outside that pair-MAC count; later stage widths also change projection and MLP work. Swin locality trades immediate global connectivity for receptive field growth across alternating blocks, so the quality and execution trade-off must be measured at the chosen dense task and full pipeline boundary.

PVT: spatial reduction

Limitation. A flat ViT preserves a single token resolution until its final representation, while dense prediction consumes features at multiple scales. Applying full global attention to the finest grid makes the early pair bill especially large, but discarding the fine query grid would also remove the dense output locations the task needs.

Architecture delta. PVT builds a feature pyramid whose representative stage shapes are 56²×64 → 28²×128 → 14²×320 → 7²×512. Inside spatial-reduction attention, it keeps full-resolution queries but forms spatially reduced K/V context. Each stage can therefore expose a dense output map while later stages trade spatial resolution for channel width.

Mechanism. For a stage with N input locations and spatial-reduction ratio R, PVT projects N queries but groups the K/V input down to N/R² locations. Rectangular attention scores those N queries × N/R² keys, mixes the reduced values, and returns one output for every original query; it is not window attention because each query can read the whole compressed context. At the figure's default first stage, 56×56 = 3,136 queries read 7×7 = 49 K/V locations when R=8. The PVT paper defines the pyramid and spatial-reduction attention, and the official PVT implementation supplies the staged backbone.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo pvt --seed 7. SpatialReductionAttention projects Q from every input location, reduces the two-dimensional feature map before K/V projection, scores the resulting rectangular matrix, and asserts both the full-grid output and reduced-context shapes.

def run_pvt_demo(config: ToyConfig) -> ShapeReport:
# ...
    attention = SpatialReductionAttention(config.width, config.heads, 4)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
# ...
    output, reduced = attention(features)
# ...
    assert attention.query.weight.grad is not None
    return {
        "features": tuple(features.shape),
        "queries": tuple(output.shape),
        "key_values": tuple(reduced.shape),
    }

Latency ledger. Relative to dense attention, the K/V projection terms are each R² smaller because K and V see N/R² locations instead of N. The QKᵀ score product and A V value-mixing product are each R² smaller too, from N²C to N²C/R². The Q projection (NC²), output projection (NC²), and per-token MLP still run on the full query grid, so the whole block does not shrink by R². The learned spatial-reduction projection is also real matrix work: reshaping to (N/R²) × (R²C) and multiplying by Wˢ ∈ R^(R²C×C) costs (N/R²)(R²C)C = NC², ignoring bias. For the default N=3,136, C=64 case, the corrected figure/helper includes 12,845,056 learned spatial-reduction projection MACs in the 58,605,568 accounted SRA subtotal, while normalization and reshape/layout remain excluded. Compressed context can lose fine-grained key/value detail; that fidelity trade-off buys a smaller rectangular interaction but does not guarantee lower wall-clock latency once the complete stage, dense head, and memory path are measured.

MViT: pooled attention

Limitation. A flat encoder keeps a constant token scale and constant width even though useful visual structure is multiscale. Fine early features need many spatial locations, whereas later semantics can use fewer locations and more channels; paying one fixed geometry through every block wastes that opportunity.

Architecture delta. In the project-first path shown here, MViT projects Q, K, and V from the full spatial input grid, then can pool Q with one query stride and pool K/V with a different, more aggressive stride. Later stages reduce spatial resolution, grow channel width, and expose a hierarchy of outputs rather than a constant-width, constant-resolution sequence. The figure deliberately forms a synthetic image-only two-operation splice, not a default reference block: it places pooling and a separate 96 → 192 channel projection in one teaching view. The toy instead isolates constant-width, projection-first pooling. The original MViT is video-first, then presents an image adaptation by removing the temporal dimension.

Mechanism. In the shown 56×56×96 spatial slice, query stride 2 turns the 3,136 input positions into 28×28 = 784 queries, while K/V stride 4 creates a 14×14 = 196 context. These are spatial-token counts; the classifier's CLS token stays separate and unpooled. Query pooling sets the output grid, while K/V pooling independently sets the context grid, and the residual must align with the 28×28 output by pooling or projection as needed. In the reference MViT-B schedule, width grows from 96 to 192 in the final MLP of the preceding block; then the following pooling block receives the 56×56 grid at width 192 and applies query stride 2 and K/V stride 4 at that width. The separate 96 → 192 arrow in this figure represents that adjacent channel-growth operation rather than part of the reference pooling block. The MViT paper defines pooling attention and stagewise scaling, and the official MViT implementation provides the corresponding models.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo mvit --seed 7. PoolingAttention records that all three projections receive the full spatial input grid, applies query stride 2 and K/V stride 4 independently, and asserts that the output geometry follows the pooled queries rather than the smaller context grid.

def run_mvit_demo(config: ToyConfig) -> ShapeReport:
# ...
    attention = PoolingAttention(config.width, config.heads, 2, 4)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
# ...
    output, queries, key_values = attention(features)
# ...
    assert projection_input_counts == {
        "query": config.patch_tokens,
        "key": config.patch_tokens,
        "value": config.patch_tokens,
    }
# ...
    return {
        "features": tuple(features.shape),
        "queries": tuple(queries.shape),
        "key_values": tuple(key_values.shape),
        "output": tuple(output.shape),
    }

Latency ledger. The figure's current same-width attention ledger at 96 channels is pedagogical, not a reference-block total. It prices the Q/K/V projections on the full 3,136 spatial positions, the output projection on 784 pooled queries, and the two attention products on the 784 × 196 rectangle. Pooling, layout conversion, residual alignment, and the separate channel projection remain outside that ledger, as does the unpooled CLS token. The corresponding reference pooling block receives width 192, so its projection and interaction work must be accounted at 192, not copied from this teaching splice. Channel growth increases representation capacity and projection work, a trade-off against the smaller spatial grid. The reduced rectangular interaction does not guarantee lower wall-clock latency once those operators and the downstream stages enter the boundary.

Across all three spatial branches, fewer pair interactions do not guarantee lower wall-clock latency. Partitioning, padding, reshaping, small operations, and irregular shapes remain outside the matrix MAC count, as do stage transitions and task heads; the full execution path, not one analytical term, decides the observed result.

Learning without labels: MAE, BEiT, DINO, and DINOv2

The next descendants change the signal used to train a visual encoder. Their pretraining systems may contain masks, tokenizers, teachers, extra views, or auxiliary heads even when deployment keeps one familiar-looking backbone. That makes the training boundary part of the architecture story: work saved or added before deployment is not automatically work saved or added on each later image.

MAE: visible-token pretraining

Limitation. Class labels are expensive, and dense encoder work on masked patches is unnecessary when their pixels have deliberately been hidden. Sending every placeholder through a wide encoder would spend most of the pretraining pass on positions that contain no image evidence.

Architecture delta. MAE samples a random 75% patch mask and keeps only the visible patches for the encoder, removing masked positions before the expensive backbone. Afterward, a narrower lightweight full-grid decoder combines the encoded visible tokens with learned mask tokens, restores their original positions, and predicts pixels for the hidden patches. The asymmetry puts representational capacity in the encoder that will survive deployment rather than in the reconstruction path that will be discarded.

Mechanism. At the ViT-B/16 analytical scale, 75% of 196 patches means 147 are hidden and 49 visible patches remain; adding CLS gives 50 encoder tokens during pretraining. The encoded visible rows are projected to the decoder width, mask tokens fill the missing coordinates, and the lightweight decoder sees the restored 197-token sequence before a masked-pixel loss scores only the 147 hidden patch targets. Fine-tuning or using the learned encoder downstream removes the masking machinery and decoder: ordinary 197-token inference again carries 196 image patches plus CLS. The MAE paper, its official implementation, and the local MAE deep dive describe that asymmetric pretraining path.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo mae --seed 7. The CPU toy is explicitly scale-reduced: its 32×32 synthetic image and 4×4 patches make 64 positions, of which 16 visible patches plus CLS produce 17 encoder tokens at the same 75% ratio. ToyMAE restores all positions in a one-block decoder and emits a 64-position pixel prediction whose last axis holds the 48 values in one 4×4 RGB patch. It asserts the mask count, positional restoration, masked-only loss, and gradients without downloading data or weights.

def run_mae_demo(config: ToyConfig) -> ShapeReport:
    model = ToyMAE(config).cpu()
    images = synthetic_images(config)
    encoder_tokens, prediction, mask = model(images)
    targets = patchify(images, config.patch_size)
# ...
    assert encoder_tokens.shape == (
        config.batch_size,
        expected_visible + 1,
        config.width,
    )
    assert prediction.shape == targets.shape
    assert mask.shape == (config.batch_size, config.patch_tokens)
# ...
    loss = mae_masked_loss(prediction, targets, mask)
    assert torch.isfinite(loss)
    loss.backward()
# ...
    return {
        "encoder_tokens": tuple(encoder_tokens.shape),
        "prediction": tuple(prediction.shape),
        "mask": tuple(mask.shape),
    }

Latency ledger. Random masking, the lightweight decoder, mask tokens, restoration, and pixel loss are pretraining-only work. The visible-only encoder reduces the backbone's analytical pretraining token bill from 197 to 50 for this reference mask, while the decoder and restoration add a separate, narrower path that must still be counted. Downstream inference does not retain that saving: it restores the ordinary 197-token full sequence and discards the decoder. This analytical encoder reduction does not guarantee lower wall-clock latency; the complete masking, gathering, decoder, backward, and input pipeline must be measured, and the reconstruction objective trades pixel fidelity for the learned representation.

BEiT: discrete visual targets

Limitation. Pixel reconstruction is not the only possible self-supervised target. Asking a model to reproduce every color value can emphasize local appearance even when the desired representation should retain higher-level semantics, so BEiT asks which visual code belongs at a hidden position instead.

Architecture delta. A separate clean-image tokenizer converts the uncorrupted patches into discrete visual IDs, one target per patch. On the encoder side, learned mask embeddings replace selected inputs, but the full set of patch positions remains: unlike MAE, BEiT does not remove the masked rows before attention. A vocabulary head then classifies the hidden positions rather than regressing their RGB values.

Mechanism. Start with the clean image and obtain discrete targets, one target ID for every patch. Build the masked encoder input by replacing selected patch embeddings with the learned mask embedding, prepend CLS, and still send 197 encoder tokens through the backbone, with no masking token savings. The output at each selected position feeds a vocabulary loss against its clean-image target; unselected positions do not contribute to that masked prediction loss. Vocabulary size belongs to the tokenizer and checkpoint—the 32-ID synthetic vocabulary in this toy is only a compact executable stand-in. The BEiT paper, the official BEiT implementation, and the local BEiT deep dive keep the clean target path distinct from the corrupted encoder path.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo beit --seed 7. ToyBEiT keeps all 64 patch positions and CLS, so the scale-reduced run has 65 encoder tokens even when its 40% mask selects 26 patches. It constructs deterministic synthetic IDs, returns a B × 64 × 32 vocabulary logits tensor, and applies cross-entropy at masked-only positions. Assertions prove that input replacement does not shorten the sequence, targets stay inside the 32-entry vocabulary, and gradients reach the shared patch projection.

def run_beit_demo(config: ToyConfig) -> ShapeReport:
    model = ToyBEiT(config).cpu()
    images = synthetic_images(config)
    mask_ratio = 0.4
# ...
    tokens, logits, targets = model(images, mask)
    assert tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
# ...
    loss = beit_masked_loss(logits, targets, mask)
# ...
    return {
        "tokens": tuple(tokens.shape),
        "logits": tuple(logits.shape),
        "targets": tuple(targets.shape),
        "mask": tuple(mask.shape),
    }

Latency ledger. The clean-image tokenizer, mask replacement, vocabulary head, and discrete-target loss are training-only branches; they disappear when one pretrained encoder is selected for deployment. During pretraining, however, the encoder still processes all 197 positions, so BEiT masking produces no encoder token saving by itself. The tokenizer and target pipeline also sit outside the backbone matrix-MAC equation and must be included when that wider boundary is evaluated. Discrete targets can trade exact pixel detail for codebook-level structure, and changing the objective does not guarantee lower wall-clock latency; the chosen tokenizer, data path, full-sequence encoder, backward pass, and quality target must be measured together.

DINO: self-distillation

Limitation. Useful semantic features should be learnable without labels and without a reconstruction objective over pixels or discrete visual codes. The difficulty is constructing a target that does not collapse to the same constant output for every image while no external annotation says which distinctions matter.

Architecture delta. DINO builds two copies of the same network around multiple augmented views. A trainable student sees global and local crops; a stop-gradient online EMA teacher sees global crops only. Centered, temperature-sharpened teacher probabilities become cross-view targets for the student, and the teacher parameters follow the student by an exponential moving average instead of direct backpropagation.

Mechanism. In the released-code default analytical schedule, the 2 global and 8 local student crops enter the trainable network, while the 2 global teacher crops enter the target network. Each teacher output supervises every differently indexed student view, excluding its matching global crop, so 2 × (2 + 8) − 2 = 18 valid cross-view loss terms remain. The paper's appendix instead describes 2 global and 6 local student crops; the 2-plus-8 inventory is therefore a released-code default, not a universal paper configuration. Before the loss, centering subtracts the old class-wise running center and sharpening applies the teacher temperature; the result is detached, so the teacher is a stop-gradient target. Backpropagation changes the student, the student optimizer step happens before the teacher EMA update, and the new center is retained for the next batch. The semantic object-like attention maps reported by DINO are an emergent representation outcome, not a supervised segmentation target. The DINO paper, its official implementation, and the local DINO deep dive establish those routes and stabilization choices.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo dino --seed 7. The CPU toy reduces the crop inventory to 4 student views—2 global and 2 local—and 2 teacher global views, yielding 6 valid cross-view pairs after the two aligned global pairs are removed. DINOStudentTeacherToy forms centered and sharpened detached targets, checks that gradients reach only the student, takes an optimizer step, and then verifies the EMA equation against the post-step student parameters. Its synthetic logits demonstrate the mechanism rather than reproducing a trained attention map.

def run_dino_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config)
    student_views = _ssl_student_views(images, config)
    teacher_views = student_views[:2]
# ...
    student_logits = [model.student_logits(view) for view in student_views]
    teacher_logits = [model.teacher_logits(view) for view in teacher_views]
# ...
    loss = dino_cross_view_loss(
        student_logits,
        teacher_logits,
        model.center,
        student_temperature=model.student_temperature,
        teacher_temperature=model.teacher_temperature,
        center_momentum=model.center_momentum,
    )
    assert torch.isfinite(loss)
# ...
    return {
        "student_views": (len(student_views),),
        "teacher_views": (len(teacher_views),),
        "student_logits": tuple(student_logits[0].shape),
    }

Latency ledger. The released-code default pretraining forward bill contains 10 student crop forwards and 2 teacher crop forwards, followed by student backward and optimizer work; all of that fan-out is training-only. Crop generation, centering, sharpening, cross-view loss construction, and the EMA update sit outside the encoder matrix-MAC subtotal even though they remain part of the training path. One student-shaped backbone and one view remain at inference, with no teacher, crop fan-out, or self-distillation head. That removal does not guarantee lower wall-clock latency relative to another pretrained encoder; the deployed model configuration and full input-to-output boundary still require measurement, while pretraining pays extra work in exchange for label-free representation quality.

DINOv2: curation and objectives

Limitation. Scaling self-supervision is not solved by enlarging the encoder alone. It also requires reliable data and stable training objectives: duplicated or low-value samples can consume the training budget, while one global image-level target can leave patch representations less useful for dense transfer.

Architecture delta. DINOv2 adds a data pipeline that first filters and post-processes raw inputs, deduplicates the resulting uncurated pool, and only then retrieves and matches samples against curated sources. Its official training system combines a DINO CLS objective for image-level agreement, an iBOT masked-patch objective, and KoLeo regularization over representations, together with the online teacher stabilization inherited from self-distillation. After large-model pretraining, a separate offline model-distillation stage can train smaller backbones against a frozen larger teacher.

Mechanism. Multiple student views and clean global teacher views first enter the online self-distillation graph. DINO matches image-level CLS distributions across views; iBOT makes masked patch outputs match clean-teacher targets at corresponding patch positions; and KoLeo spreads L2-normalized raw student CLS representations. The paper appendix describes KoLeo on the first global crop. The official code applies one KoLeo term to each of the two global crops and sums both per-crop terms into the optimization loss; it divides that sum by two only for the logged metric. The online stop-gradient EMA teacher evolves during pretraining, whereas offline model distillation later uses a frozen larger teacher to supervise a selected smaller model. Curation happens before these tensor routes, so it belongs to the full training system even though it is not a layer in the encoder. The original DINOv2 architecture has no register tokens; those arrive in the separate follow-up below. The DINOv2 paper, its official implementation, and the local DINOv2 deep dive provide the primary training-system boundary.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo dinov2 --seed 7. The CPU toy implements a scale-reduced centered-softmax path, not full Sinkhorn-Knopp. It routes 4 student views—2 masked global and 2 clean local—and 2 teacher views, both clean globals, through 6 DINO cross-view pairs and 2 same-crop iBOT patch pairs. For KoLeo, the CPU toy uses both global crops but averages its two per-crop KoLeo terms in the loss itself instead of reproducing the official optimization sum. That shared two-crop selection is not the same reduction convention. It verifies separate CLS and patch centers, detached teacher targets, gradients for every student objective, the optimizer-before-EMA order, and exact masks. Dataset curation and offline distillation are narrative full-system stages and are not executed by this synthetic demo.

def run_dinov2_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config)
    student_views = _ssl_student_views(images, config)
    global_views = student_views[:2]
    local_views = student_views[2:]
# ...
    step = dinov2_objective_step(model, student_views, masks=masks)
    assert step.student_routes == (
        ("global-0", "masked student"),
        ("global-1", "masked student"),
        ("local-0", "clean student"),
        ("local-1", "clean student"),
    )
# ...
    return {
        "student_cls": tuple(step.student_global_cls[0].shape),
        "teacher_cls": tuple(step.teacher_cls_values[0].shape),
        "patch_logits": tuple(step.student_patch_logits[0].shape),
        "mask": tuple(step.masks[0].shape),
    }

Latency ledger. Data curation—filtering and post-processing raw inputs, deduplicating the uncurated pool, then retrieving and matching against curated sources—crop construction, online teacher forwards, DINO, iBOT, and KoLeo objectives, backward, EMA, and later offline distillation are training-only work; none recurs for every deployed image. Inference keeps one selected backbone. The figure's illustrative cost scenario uses 256 patch positions plus CLS—257 tokens—and no registers, but that is a declared analytical configuration rather than a rule for every DINOv2 checkpoint. Removing the training system does not guarantee lower wall-clock latency for that backbone; input resolution, patch size, width, depth, task head, and the end-to-end boundary still require measurement. The trade is substantial pretraining-system complexity for transferable representation quality, not an automatic inference optimization.

Register tokens

Limitation. Large visual transformers can produce high-norm, low-information background patch artifacts. Those outlier patch rows appear to act as scratch space for global computation even though their coordinates tell a downstream dense head that they should describe real image regions, creating a conflict between an internal workspace and a spatial output.

Architecture delta. The later follow-up Vision Transformers Need Registers adds four learned non-spatial register tokens between CLS and the patch sequence in its common analytical example. Attention may write global scratch work into these dedicated slots while the number and coordinates of patch outputs stay unchanged. A downstream head can discard the registers rather than interpreting a background patch that the encoder repurposed.

Mechanism. Registers enter with CLS and the patches, participate in every encoder block during both training and inference, and can exchange information through ordinary self-attention. Only the register output vectors are excluded from downstream heads; removing them at the output does not erase the work already performed inside the stack, and CLS or spatial patch outputs continue according to the task. For ViT-B/16, the analytical sequence changes from 197 to 201 tokens; those four registers preserve 196 spatial patch outputs. This is a separate follow-up: registers were not part of the original DINOv2 architecture. The Vision Transformers Need Registers paper and its official implementation in DINOv2 document the artifact and the added-token path.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo registers --seed 7. The scale-reduced CPU toy subclasses the same tiny backbone as RegisteredViT: 64 patches, CLS, and 4 registers produce 69 tokens through the inherited encoder, after which slicing exposes 4 register rows and 64 spatial outputs. It keeps positional embeddings on CLS and the patches while leaving registers non-spatial, then verifies logits and nonzero gradients for the patch projection, register parameters, and attention. The run uses synthetic input and does not claim a deployed measurement.

def run_registers_demo(config: ToyConfig) -> ShapeReport:
    model = RegisteredViT(config).cpu()
    images = synthetic_images(config)
    tokens, cls = model.forward_features(images)
    registers = tokens[:, 1 : 1 + model.register_count]
    patches = tokens[:, 1 + model.register_count :]
    logits = model.head(cls)
# ...
    assert tokens.shape == (
        config.batch_size,
        1 + model.register_count + config.patch_tokens,
        config.width,
    )
# ...
    return {
        "tokens": tuple(tokens.shape),
        "registers": tuple(registers.shape),
        "patches": tuple(patches.shape),
        "logits": tuple(logits.shape),
    }

Latency ledger. Register work is inference-visible when registers are included in the deployed backbone: the extra tokens enter projections, MLPs, and attention interactions in every block. In the tested analytical ViT-B/16 scenario, the 197 → 201 sequence adds 369,082,368 encoder MACs over twelve blocks and 19,104 score elements per block across the twelve heads; these are derived matrix and tensor counts, not observed timings. Discarding the four output vectors before a dense head does not recover that encoder work, though the spatial output count remains 196. Cleaner spatial features are the intended fidelity benefit, and adding registers does not guarantee lower wall-clock latency; the model, task head, sequence layout, and full pipeline still require measurement.

MAE, BEiT, DINO, and DINOv2 therefore alter the learning signal in four different ways: visible-pixel reconstruction, discrete masked targets, online self-distillation, and a curated multi-objective training system. Registers solve a later representation artifact inside the deployed encoder. Keeping those boundaries separate is what prevents a pretraining saving from being reported as an inference saving—or a training-only teacher from being mistaken for a second deployed network.

Beyond classification: CLIP, SigLIP, ViTDet, and SAM

A class logit is only one possible consumer of a visual representation. Retrieval needs an image to meet language in a shared space; detection needs every spatial row rather than only CLS; interactive segmentation needs many answers from one image. The next four branches keep a ViT in the system but move its output boundary, its training objective, or the work reused by the application.

CLIP: dual encoder

Limitation. A fixed classifier head binds the model to a closed label vocabulary and emits scores only for the classes installed in that head. Open-vocabulary retrieval instead needs an image to be comparable with previously unseen text, and it needs either side to be encoded without running the other side at the same moment.

Architecture delta. CLIP replaces the fixed class head with distinct image and text towers. Each tower feeds its own learned modality-specific projection into a common comparison width, then L2-normalizes the projected vector. A learned logit scale multiplies the cosine similarities. The image path is therefore image → ViT → image projection → L2 normalization; the text path is text → text encoder → text projection → L2 normalization. The projection heads are separate learned parameters, not one shared adapter and not normalization by itself.

Mechanism. For a batch of B matched image-text examples, the normalized projected vectors form a B × B similarity-logit matrix. The diagonal entries name matches; off-diagonal entries supply alternatives. CLIP computes a mean row cross-entropy for image-to-text retrieval and a mean column cross-entropy for text-to-image retrieval, then takes the average of those two directional losses. That matrix construction and symmetric loss are training-only. At inference, an application may encode and cache a candidate collection with either tower, encode only the changing side, and perform similarity search; it does not replay the training loss. The CLIP paper defines the learned projections, normalization, scale, and symmetric objective, while the official CLIP implementation exposes the two independent encoding interfaces.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo clip --seed 7. The synthetic CPU check builds a tiny ViT, a text embedding tower, and two distinct bias-free projection heads. It L2-normalizes each projected modality exactly once, constructs a 2 × 2 logit matrix, backpropagates through both projection heads, and checks nonzero gradients. The toy uses a fixed temperature at 0.07; reference CLIP instead learns its logit scale. Its exact shape report is:

def run_clip_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config)
    tokens, image_features = vision.forward_features(images)
# ...
    text_features = text_tower(text_ids)
    image_projection = nn.Linear(config.width, config.width, bias=False).cpu()
    text_projection = nn.Linear(config.width, config.width, bias=False).cpu()
    assert image_projection is not text_projection
# ...
    image_projected = image_projection(image_features)
    text_projected = text_projection(text_features)
    image_normalized = F.normalize(image_projected, dim=-1)
    text_normalized = F.normalize(text_projected, dim=-1)
# ...
    loss, logits = _clip_symmetric_loss_from_normalized(
        image_normalized,
        text_normalized,
    )
# ...
    return {
        "image_features": tuple(image_features.shape),
        "logits": tuple(logits.shape),
        "tokens": tuple(tokens.shape),
    }

{"demo": "clip", "shapes": {"image_features": [2, 64], "logits": [2, 2], "tokens": [2, 65, 64]}}

Latency ledger. The figure's illustrative 4 × 4 batch creates 16 dot products: 4 diagonal matches and 12 off-diagonal alternatives. Those counts omit both encoder towers, both learned projections, L2 normalization, the learned scaling operation, loss reduction, and backward work. They are a training matrix bill, not a serving measurement. Inference has a different boundary: image-tower work, text-tower work, embedding storage, index construction, candidate transfer, and similarity search can be scheduled separately. Reusing one side can avoid repeated encoding of that side, but it does not guarantee lower wall-clock latency for the full retrieval system; corpus size, indexing, input processing, and the changing tower still require measurement.

SigLIP: independent pairs

Limitation. CLIP's batch-softmax objective couples every selected logit to the other logits in its image row and text column through shared denominators. That probability normalization makes the loss of one matched pair depend on the other examples that happen to be compared with it, even though the dual-encoder architecture itself can encode examples independently.

Architecture delta. SigLIP keeps normalized image and text embeddings but changes the training objective to binary sigmoid classification over image-text pairs. A target z_ij ∈ {+1, −1} marks a matched or unmatched pair. The paper applies a learned positive scale and a learned bias to each similarity logit before the sigmoid loss; the learned scale and bias are objective parameters, not replacements for the two encoders.

Mechanism. For each pair, SigLIP evaluates −log σ(z_ij(s x_iᵀy_j + b)). There are B² binary terms in a B × B batch matrix. The published reduction sums those terms and divides that sum by B, not by B², so it is not the ordinary mean over matrix cells. This removes the batch-softmax probability normalization only, not all batch concerns: training still chooses and scores negative pairs, aggregates gradients, communicates model updates when the training system is distributed, and depends on the sampled data. The SigLIP paper specifies the pair labels and reduction; its official implementation in Big Vision provides the training path.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo siglip --seed 7. The scale-reduced synthetic function L2-normalizes two feature matrices, uses a fixed unit logit scale and an additive bias, constructs +1 diagonal and −1 off-diagonal labels, and divides the sum of all losses by batch size. The report is 2 × 2; it checks objective mechanics rather than reproducing the paper's learned scale and bias:

def run_siglip_demo(config: ToyConfig) -> ShapeReport:
# ...
    vision = TinyViT(config).cpu()
    images = synthetic_images(config)
    tokens, image_features = vision.forward_features(images)
    text_tower = nn.Embedding(config.batch_size, config.width).cpu()
    text_ids = torch.arange(config.batch_size)
    text_features = text_tower(text_ids)
    loss, logits, labels = siglip_pair_loss(image_features, text_features)
# ...
    assert logits.shape == (config.batch_size, config.batch_size)
    assert labels.shape == logits.shape
# ...
    return {
        "labels": tuple(labels.shape),
        "logits": tuple(logits.shape),
    }

{"demo": "siglip", "shapes": {"labels": [2, 2], "logits": [2, 2]}}

Latency ledger. At the figure's B=4, SigLIP computes 16 binary pair losses and divides their sum by 4; all 16 terms and the loss are training-only. Switching from shared softmax denominators does not remove the B² pair scores, the negative examples, either encoder, or model-update communication. At inference the same dual-encoder separation permits precomputed image or text embeddings, but that property comes from the architecture rather than the sigmoid reduction. The objective changes normalization and training-system flexibility; by itself it does not guarantee lower wall-clock latency for encoding or retrieval, which must be measured across the complete application path.

ViTDet: dense backbone

Limitation. Reading only CLS discards the dense spatial rows that a detector needs, while applying global attention to every retained high-resolution position makes the interaction term grow quadratically. A detection backbone must preserve the patch grid, exchange information beyond a small neighborhood, and present multiple output scales to a dense head.

Architecture delta. ViTDet retains the dense patch tokens instead of collapsing the backbone to CLS. Most blocks use regular non-overlapping, non-shifted windows; a sparse set of global-attention blocks bridges those local groups. The final sequence is restored to a 2-D feature map, and simple upsampling and downsampling adapters form a feature pyramid for detector heads. Unlike a hierarchical transformer, the plain ViT backbone itself can keep one resolution and width while adapters create the output scales.

Mechanism. In a window block, each query reads tokens only inside its regular window. Most blocks follow that path, while a global bridge block lets every retained position exchange information across windows. This is not the Swin shift-and-mask mechanism: ViTDet's local windows are not shifted, and cross-window communication comes from the sparse global blocks. After the last backbone block, the N rows reshape to their patch grid as a backbone feature map; feature-pyramid adapters then produce higher, base, and lower spatial outputs. The ViTDet paper describes the plain-backbone adaptation and sparse global blocks, and the official Detectron2 implementation supplies the dense backbone and pyramid components.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo vitdet --seed 7. The synthetic demo isolates the retained-token handoff, 2-D reshape, and three-level pyramid: 64 patch rows become an 8 × 8 base map, then nearest-neighbor upsampling and average pooling produce 16 × 16 and 4 × 4 maps. It deliberately does not reproduce the window/global block schedule; its tiny backbone uses the shared global-attention toy, so the demo validates output geometry rather than claiming a ViTDet checkpoint implementation.

def run_vitdet_demo(config: ToyConfig) -> ShapeReport:
    model = TinyViT(config).cpu()
    images = synthetic_images(config)
    tokens, _ = model.forward_features(images)
    pyramid = vitdet_feature_pyramid(tokens[:, 1:], config.grid_size)
# ...
    assert pyramid["base"].shape == (
        config.batch_size,
        config.width,
        config.grid_size,
        config.grid_size,
    )
# ...
    return {name: tuple(value.shape) for name, value in pyramid.items()}

{"demo": "vitdet", "shapes": {"base": [2, 64, 8, 8], "high": [2, 64, 16, 16], "low": [2, 64, 4, 4]}}

Latency ledger. The figure's declared twelve-block schedule has 8 window blocks and 4 global bridges over an 8 × 8, width-64 teaching grid. Its attention interactions total 2,359,296 MACs; adding the modeled dense projections and MLPs gives 40,108,032 MACs. Windowing changes the attention interaction but does not remove linear projection work, per-token MLP work, activation and normalization work, softmax, layout conversion, pyramid-adapter work, or detector-head work; all remain in the dense task boundary. The retained high-resolution maps can also increase activation traffic. The smaller interaction subtotal does not guarantee lower wall-clock latency, because the complete backbone, pyramid, head, and input/output path must be measured together.

SAM: reuse the image embedding

Limitation. Interactive segmentation asks many prompts about the same image. Rerunning the heavy image encoder for every point or box would repeat identical visual work even though only the prompt and requested mask changed.

Architecture delta. SAM separates a heavy image encoder, an application-owned image embedding retained once, a prompt encoder, and a lightweight mask decoder. The image encoder runs once per image. The prompt encoder plus mask decoder run per query, combining each new point, box, or mask prompt with the same dense image representation.

Mechanism. Let E be image-encoder work, P prompt-encoder work, D mask-decoder work, and q the number of prompts for one image. The first result costs E + P + D; q results cost E + q(P + D); amortized work per result is E/q + P + D. The application owns the retained image embedding and chooses its lifetime. There is no retained decoder state here: each prompt creates its own small prompt path and mask result while the dense image representation is reused. The SAM paper defines the promptable segmentation system, and the official Segment Anything implementation exposes separate image-setting and prompt-prediction boundaries.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo sam --seed 7. ToySAM uses one synthetic image encode followed by three prompt queries. Each prompt projection changes the resulting mask vector, while an explicit counter proves that the image tower executed once. This smaller three-prompt executable is distinct from the four-prompt analytical slider default:

def run_sam_demo(config: ToyConfig) -> ShapeReport:
    model = ToySAM(config).cpu()
    images = synthetic_images(config)
    image_tokens = model.encode_image(images)
    points = [
        torch.tensor(
            [[float(index), float(index + 1)] for _ in range(config.batch_size)]
        )
        for index in range(3)
    ]
    masks = [model.decode_prompt(image_tokens, point) for point in points]
    assert model.image_encode_calls == 1
    assert all(mask.shape == (config.batch_size, config.patch_tokens) for mask in masks)
    assert not torch.allclose(masks[0], masks[1])
# ...
    return {
        "image_tokens": tuple(image_tokens.shape),
        "masks": (len(masks), *tuple(masks[0].shape)),
    }

{"demo": "sam", "shapes": {"image_tokens": [2, 64, 64], "masks": [3, 2, 64]}}

Latency ledger. The teaching ledger assigns 100 illustrative work units to E and 4 to each combined P + D. With q=4, the first result is 104, four results total 116, amortized work is 29 per result, and retaining the embedding avoids 300 units of repeated image encoding. These are illustrative work units, not a timing or a claim about any deployed configuration. Retention also consumes application memory and introduces lifetime management, while prompt preprocessing and output handling remain outside the arithmetic. Avoiding three repeated image passes does not guarantee lower wall-clock latency for every interaction; the actual image, prompt mix, decoder, and application boundary must be measured.

Connecting vision to language: BLIP-2 and LLaVA

Vision-language systems add a second sequence boundary. The visual encoder can emit hundreds of patch rows, but an autoregressive language model consumes vectors at its own width and pays attention work according to the sequence it receives. BLIP-2 compresses that interface with learned queries; LLaVA projects a retained visual sequence directly into the language stream.

BLIP-2: query bottleneck

Limitation. Handing a long patch sequence directly to a frozen language model lengthens every downstream language-layer sequence. Fully retraining both large pretrained towers would also discard the modularity and training economy that made them useful starting points.

Architecture delta. BLIP-2 keeps a frozen vision encoder and a frozen language model, placing a trainable Q-Former and projection between them. The reference design uses 32 learned queries. Those queries are a fixed-size bottleneck: they read the variable visual source and emit one downstream vector per query, regardless of the original patch count.

Mechanism. With N source tokens and M learned queries, each Q-Former cross-attention layer exposes M × N query-source interactions and emits exactly M visual outputs. In the reference architecture, cross-attention appears in every other transformer block. Fixed M bounds the downstream visual sequence but creates a detail trade-off: too few queries may not carry all evidence required by a fine-grained task. BLIP-2 trains the bridge in two stages. The first performs vision-language representation learning with the frozen image encoder, using objectives that teach queries to extract image-text information. The second performs vision-to-language generative learning by connecting Q-Former outputs to the frozen language model. The BLIP-2 paper specifies both stages and 32-query design; the official LAVIS implementation provides the Q-Former training and model interfaces.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo blip2 --seed 7. The toy reduces the bottleneck to 8 learned queries reading 65 source tokens—64 patches plus CLS. It freezes the vision module, checks an 8 × 65 normalized attention map per example, and backpropagates only through the learned queries and cross-attention. It does not execute a frozen language model or either full BLIP-2 pretraining stage; it is a local tensor-contract check.

def run_blip2_demo(config: ToyConfig) -> ShapeReport:
    query_tokens = BLIP2_TOY_QUERY_TOKENS
# ...
    vision = TinyViT(config).cpu()
# ...
    images = synthetic_images(config)
    with torch.no_grad():
        source_tokens, _ = vision.forward_features(images)
    resampler = LearnedQueryResampler(
        config.width,
        config.heads,
        query_tokens=query_tokens,
    ).cpu()
    output, weights = resampler(source_tokens)
# ...
    assert output.shape == (config.batch_size, query_tokens, config.width)
# ...
    return {
        "attention": tuple(weights.shape),
        "query_tokens": tuple(output.shape),
        "source_tokens": tuple(source_tokens.shape),
    }

{"demo": "blip2", "shapes": {"attention": [2, 8, 65], "query_tokens": [2, 8, 64], "source_tokens": [2, 65, 64]}}

Latency ledger. The reference teaching bill starts with 257 visual source tokens—256 spatial patches plus CLS—and emits 32 downstream query vectors. The shown rectangle has 8,224 query-source pairs per cross-attention layer because 32 × 257 = 8,224. That is not the full Q-Former total: cross-attention appears in every other transformer block, and its self-attention, feed-forward layers, projections, normalization, and two-stage training objectives are additional work. The language model receives 32 visual vectors instead of all 257 source rows, but the image encoder and Q-Former still execute before that downstream sequence. The fixed query count trades visual bandwidth for a shorter language interface; it does not guarantee lower wall-clock latency or preserved task quality, both of which require measurement at the complete generation boundary.

LLaVA: visual prefix

Limitation. A vision tower and language decoder usually have different embedding widths, and the decoder cannot consume raw vision features as text tokens without a learned bridge. Keeping many visual tokens preserves spatial evidence, but those tokens lengthen language-model prefill and all state derived from that prefix.

Architecture delta. LLaVA connects a frozen pretrained vision tower to an autoregressive language model through a learned projector, turning retained image features into a visual prefix at the language width. Its original feature-alignment stage freezes both the vision tower and language model and trains the projector. During the original instruction-tuning stage, the vision tower remains frozen while the projector and language model are updated. Those freeze schedules are training choices; inference still executes the vision tower, projector, and autoregressive decoder.

Mechanism. KV cache is not part of encoder-only ViT inference: every encoder token attends bidirectionally during one parallel forward pass, and the encoder can release its attention intermediates afterward. It begins at the autoregressive decoder boundary. LLaVA runs the vision tower once, projects each retained image feature to language width, inserts the visual prefix into the prompt, performs one parallel causal prefill over visual and text tokens, and creates prefix K/V state for each decoder layer and K/V head. Sequential decode then appends one generated token at a time; each token reads the retained prefix and prior outputs, so stored state and the one-token attention span grow with generation. This article's teaching layout is image-first for clarity, but it is not universal placeholder placement: reference prompts insert image features at the image placeholder, which may occur before or after the question. The LLaVA paper describes the two training stages and visual projection, while the official LLaVA implementation provides placeholder insertion and generation paths.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo llava --seed 7. The synthetic path projects 64 visual tokens from width 64 to width 80, prepends them to 6 text tokens, and performs one 70-token prefill. A tiny two-layer decoder-state module then appends 2 generated tokens sequentially, growing the stored length to 72 while proving that prefix keys and values remain unchanged and that K and V are distinct tensors.

def run_llava_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config).requires_grad_()
    image_tokens, _ = vision.forward_features(images)
# ...
    prefix = llava_visual_prefix(image_tokens[:, 1:], text_embeddings, projector)
# ...
    prefill_key, prefill_value = cache.prefill(prefix)
# ...
    for embedding in generated_embeddings:
        post_decode_key, post_decode_value = cache.append(embedding)
# ...
    assert cache.prefill_calls == 1
    assert cache.append_calls == generated_tokens
    assert cache.cache_length_history == [70, 71, 72]
    assert cache.projection_token_counts == [70, 1, 1]
# ...
    return {
        "combined_prefix": tuple(prefix.shape),
        "image_tokens": tuple(image_tokens.shape),
        "post_decode_key": tuple(post_decode_key.shape),
        "post_decode_value": tuple(post_decode_value.shape),
        "prefill_key": tuple(prefill_key.shape),
        "prefill_value": tuple(prefill_value.shape),
        "visual_prefix": tuple(prefix[:, : config.patch_tokens].shape),
    }

{"demo": "llava", "shapes": {"combined_prefix": [2, 70, 80], "image_tokens": [2, 65, 64], "post_decode_key": [2, 2, 2, 72, 8], "post_decode_value": [2, 2, 2, 72, 8], "prefill_key": [2, 2, 2, 70, 8], "prefill_value": [2, 2, 2, 70, 8], "visual_prefix": [2, 64, 80]}}

Latency ledger. The default analytical controls use 256 visual plus 128 text tokens, so prefix length is 384. Parallel causal prefill exposes 73,920 logical analytical pairs per layer. Generating 32 tokens adds 12,784 sequential pairs per layer—first step 384, last step 415—through G·P + G(G−1)/2; this bill does not repeat the parallel prefill. At prefix creation, the declared four layers, four K/V heads, head width 16, and one-example batch hold 98,304 key elements, 196,608 K-plus-V elements, or 393,216 bytes at two bytes per element. After 32 generated tokens, 416 retained positions hold 106,496 key elements, 212,992 K-plus-V elements, or 425,984 bytes. These are analytical causal pairs and illustrative tensor bytes, not a resident-memory measurement; implementations may compute larger tiles or store state differently. Visual encoding and projection are separate bills, and fewer visual tokens trade evidence for a shorter prefix. These reductions do not guarantee lower wall-clock latency; the full encode, project, prefill, sequential decode, and output path must be measured.

Making ViTs cheaper: MobileViT, EfficientViT, DynamicViT, and ToMe

The previous branches changed what a ViT learns or where its tokens go. These ask which representation or execution path should carry less work. MobileViT mixes local and global operators, EfficientViT-CGA reorganizes a block, DynamicViT prunes selected tokens, and ToMe merges them; each moves a different bill and fidelity or systems constraint.

MobileViT: local-global hybrid

Limitation. Convolutional maps provide local spatial structure, while image tasks also need global context. A plain global transformer block can lose local operators and the two-dimensional representation expected downstream. The problem is to expose distant locations to one another, then return their information to the same coordinates for later convolutional stages.

Architecture delta. MobileViT keeps a local convolutional route and inserts a transformer route inside a convolution-shaped block. Its full dataflow is C → d → unfold → transformer → fold → d → C → fusion: local features at width C are projected to transformer width d, rearranged into sequences, globally mixed, restored to a map, projected back from d to C, and fused with the local representation. That named seven-stage hybrid path makes the entry feature map and final fusion explicit; unfold and fold are rearrangements between them, not learned reductions that silently throw spatial sites away.

Mechanism. The teaching case begins with an 8 × 8 feature map of width 64, hence 4,096 elements per example. A 2 × 2 patch contains four positions. Unfold groups equal within-patch positions across the grid into 4 sequences, each containing 16 tokens; all 64 spatial locations and all 64 channels still exist. The transformer mixes each position-index sequence globally. Unfold preserves the number of elements even though the two tensor views index them differently, and fold then restores every value to its original spatial coordinates before the post-fold projection and fusion. The MobileViT paper describes the local-global block, and the official code supplies its feature-map implementation and model family.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo mobilevit --seed 7. ToyMobileViTBlock executes the local convolution, channel projection, unfold, transformer, fold, post-fold projection, and fusion in that order. The conservation assertion binds the projected B × d × H × W tensor to its B·P × N × d unfolded view; those have the same number of elements. The default command has C = d = 64, which masks that distinction. A separate unequal-width check uses C = 6 and d = 8: its input and unfolded tensors have different element counts, while its projected and unfolded tensors agree. The input is synthetic and the exercise is not a checkpoint benchmark, so it cannot establish performance; it exists to make every reshape and learned boundary runnable without downloads.

def run_mobilevit_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config).requires_grad_()
    patches = vision.patch_tokens(images)
    feature_map = patches.transpose(1, 2).reshape(
        config.batch_size,
        config.width,
        config.grid_size,
        config.grid_size,
    )
    output, unfolded = block(feature_map)
# ...
    assert unfolded.shape == (
        config.batch_size * block.patch_height * block.patch_width,
        (config.grid_size // block.patch_height)
        * (config.grid_size // block.patch_width),
        config.width,
    )
# ...
    return {
        "feature_map": tuple(output.shape),
        "unfolded": tuple(unfolded.shape),
    }

{"demo": "mobilevit", "shapes": {"feature_map": [2, 64, 8, 8], "unfolded": [8, 16, 64]}}

Latency ledger. The selected-operation ledger assigns local convolution to 2,359,296 MACs, C → d channel projection to 262,144 MACs, unfolded transformer to 2,228,224 MACs, d → C post-fold projection to 262,144 MACs, and local/global fusion to 4,718,592 MACs; the analytical total is 9,830,400 MACs. The reshape carries 4,096 elements through 2 layout boundaries. Unfold and fold preserve arithmetic elements, but possible copies can appear when the required layout is not already contiguous; layout conversion, normalization, activation, and framework scheduling are outside the selected MAC subtotal. Parameters alone do not predict latency, and the lower analytical subtotal does not guarantee lower wall-clock latency; the complete local, layout, transformer, fusion, and application path must be measured.

EfficientViT: cascaded group attention

Limitation. Arithmetic counts omit memory traffic, operator boundaries, layouts, and scheduling, so they cannot establish runtime. A smaller attention expression may still materialize costly tensors or fragment execution. An efficiency-oriented block needs an explicit operator graph alongside a count, with proxy arithmetic separated from whole-model claims.

Architecture delta. This section means Liu et al.'s CVPR 2023 model: an FFN → one CGA → FFN sandwich whose attention divides channels into channel groups and uses a cascade between them. Each later group incorporates the previous attention head's output before computing its own head. It is not the separate MIT EfficientViT family built around MSLA; sharing the name “EfficientViT” does not make their mechanisms or evidence interchangeable. Here, CGA always names cascaded group attention from the Liu/Cream branch.

Mechanism. CGA first splits the C channels into head-sized groups. The first head consumes its own group; the previous head output is added to the next channel group, so each later head receives a cascaded dependency rather than an independent slice alone. After all groups have attended, their outputs concatenate and pass through one output projection. That previous-head-to-next-head route is a sequential dependency, even though work inside an individual head may have parallel structure, and it must remain visible when interpreting the grouped projection count. The EfficientViT-CGA paper motivates the sandwich and redundancy reduction, while the official code defines the Liu/Cream implementation used by this lineage.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo efficientvit-cga --seed 7. The toy operates on dense B × N × C tensors to isolate the channel split, previous-head cascade, concatenation, and final projection. It is not a spatial EfficientViT block: it deliberately omits the surrounding spatial and depthwise convolutions, attention bias machinery, and paper schedule. Assertions bind each head to its own channel group plus the prior output, preserve the original dense input/output shape, and verify gradients through all groups. The synthetic run checks this narrow dependency graph rather than reproducing trained weights.

def run_efficientvit_cga_demo(config: ToyConfig) -> ShapeReport:
# ...
    tokens = torch.randn(
        config.batch_size,
        config.patch_tokens,
        config.width,
        requires_grad=True,
    )
    output = model(tokens)
    assert output.shape == tokens.shape
# ...
    baseline_groups = model.cascade(tokens.detach())
    perturbed = tokens.detach().clone()
    perturbed[..., : model.group_width] += 1
    changed_groups = model.cascade(perturbed)
    for index in range(1, model.groups):
        assert not torch.equal(baseline_groups[index], changed_groups[index])
# ...
    return {"input": tuple(tokens.shape), "output": tuple(output.shape)}

{"demo": "efficientvit-cga", "shapes": {"input": [2, 64, 64], "output": [2, 64, 64]}}

Latency ledger. For 64 tokens at width 64, the dense QKV projection comparator is 786,432 MACs, while the grouped QKV projection proxy is 196,608 MACs. The concatenated output projection is 262,144 MACs, and each efficient FFN side is 1,048,576 MACs × 2 sides. The grouped projection proxy is not a full block or a latency measurement. It excludes spatial and depthwise convolutions, attention interactions and attention bias, normalization, activation, layout conversion, cascade serialization, and scheduling. Consequently, the one-quarter QKV projection ratio does not guarantee lower wall-clock latency; the actual block, model, input shape, framework path, and full measurement boundary must be benchmarked together.

DynamicViT: learned pruning

Limitation. A fixed-depth encoder gives every image every token at every depth, even when later patch representations are unnecessary. Redundant regions still pay projection, attention, and MLP work. An adaptive route must choose from the image which patches leave while protecting CLS and accounting for decision cost and quality.

Architecture delta. DynamicViT inserts a learned importance predictor at several depths and performs progressive pruning instead of choosing one static token budget for the whole dataset. The fixed-shape differentiable gates used for training let gradients reach the predictors while attention masking preserves tensor dimensions. Physical inference instead uses top-k selection and gather to retain the chosen patch rows, with CLS carried forward. This distinction matters because masked training tensors demonstrate learnability but do not realize the later-block work reduction produced by a physically shorter inference sequence.

Mechanism. An early block first processes all 64 patch tokens plus CLS. The predictor then scores the 64 patches without treating CLS as a deletion candidate; at the teaching keep ratio, repeated decisions produce 64 → 32 → 16 patches, while physical sequences move through 65 → 33 → 17 once CLS is included. Each decision affects only blocks after that point, so no saving may be credited to the early full-sequence work or to the predictor itself. The selected token identities depend on image content, while fixed-ratio top-k keeps the configured count at each decision. The DynamicViT paper develops the differentiable masking and hierarchical policy, and the official code provides the training and inference routes.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo dynamicvit --seed 7. At the first 50% decision, the demo runs a straight-through fixed-shape training gate that preserves all positions and sends gradient into the scorer; its physical inference gather then retains 32 patches plus CLS for 33 tokens. It covers one decision, not the full schedule of progressive decisions or a deployed batching policy. The synthetic check keeps the differentiable training route distinct from the shorter inference tensor rather than pretending that hard top-k indices themselves are differentiable.

def run_dynamicvit_demo(config: ToyConfig) -> ShapeReport:
# ...
    images = synthetic_images(config).requires_grad_()
# ...
    training_tokens, gates = dynamic_prune_tokens(
        early, scorer, 0.5, training=True
    )
    assert training_tokens.shape == early.shape
    assert gates.shape == (config.batch_size, config.patch_tokens)
# ...
        pruned, indices = dynamic_prune_tokens(
            early, scorer, 0.5, training=False
        )
        later_inference = vision.blocks[1](pruned)
# ...
    assert indices.shape == (config.batch_size, kept_patches)
    assert pruned.shape == (
        config.batch_size,
        kept_patches + 1,
        config.width,
    )
# ...
    return {"indices": tuple(indices.shape), "tokens": tuple(pruned.shape)}

{"demo": "dynamicvit", "shapes": {"indices": [2, 32], "tokens": [2, 33, 64]}}

Latency ledger. Starting from 64 patches, the retained counts are 32 and then 16: the retained path is 64 → 32 → 16 patches, or physical sequences 65 → 33 → 17; these are later-block savings only. The full-sequence later-block comparator is 3,735,680 MACs; the modeled block after decision one is 1,761,408 MACs, and after decision two is 872,576 MACs. The analytical work saved across those two later blocks is 4,837,376 MACs. That subtraction excludes the importance predictor, top-k, gather, padding, batching, compilation, and accuracy evaluation; padding a mixed batch toward its longest retained sequence can restore work that a per-example count removed. Dynamic shapes may also change available execution paths. The MAC reduction does not guarantee lower wall-clock latency, so selection overhead, shape policy, complete model path, and task quality must be measured together.

ToMe: token merging

Limitation. Redundant tokens waste later work, but dropping them can lose needed information. Merging shortens the sequence by combining rather than simply discarding rows, so it needs a similarity rule and bookkeeping for how much original content each representation carries.

Architecture delta. ToMe partitions tokens into bipartite source and destination sets, chooses high-scoring edges using normalized similarity, and replaces selected pairs with weighted aggregation. Every token also carries a represented size, so a destination that already summarizes several inputs contributes proportionally in a later merge. Unlike pruning, a selected source is folded into a destination rather than discarded, although the resulting feature is still an approximation and no longer preserves every original token independently.

Mechanism. The teaching grid splits eight tokens into two sets of four, creating 4 × 4 = 16 comparisons. Selecting two source-to-destination edges makes 8 → 6 tokens while preserving represented mass of 8. The weighted rule is easiest to audit numerically: destination [2, 4] with size 1 plus source [8, 10] with size 3 produces [(2×1 + 8×3)/4, (4×1 + 10×3)/4] = [6.5, 8.5] with size 4. The values aggregate information rather than simply drop a matched source, and future weighted merges use the updated size. The ToMe paper defines bipartite soft matching, and the official code implements token-size-aware merging in existing transformer blocks.

Toy-code delta. Run python3 examples/vit_lineage_toy.py --demo tome --seed 7. The toy uses normalized synthetic representations for bipartite matching, selects two merges, and applies represented-size-weighted aggregation. Its 8 input rows become 6 output rows while total mass remains 8; assertions also reproduce the [6.5, 8.5] weighted example. The demonstration isolates matching and merging rather than approximating a pretrained model's feature geometry or quality.

def run_tome_demo(config: ToyConfig) -> ShapeReport:
    tokens = torch.randn(
        config.batch_size,
        8,
        config.width,
        requires_grad=True,
    )
    sizes = torch.ones(
        config.batch_size,
        8,
        requires_grad=True,
    )
    merged, merged_sizes = tome_bipartite_merge(tokens, sizes, merge_count=2)
    assert merged.shape == (config.batch_size, 6, config.width)
    assert merged_sizes.shape == (config.batch_size, 6)
    assert torch.equal(merged_sizes.sum(dim=1), sizes.sum(dim=1))
# ...
    return {"sizes": tuple(merged_sizes.shape), "tokens": tuple(merged.shape)}

{"demo": "tome", "shapes": {"sizes": [2, 6], "tokens": [2, 6, 64]}}

Latency ledger. The matching stage pays 4 × 4 = 16 comparisons; 8 input tokens become 6 output tokens, and represented mass is 8 before and 8 after. For one modeled later block, the later-block comparator is 401,408 MACs before and 299,520 MACs after, a 101,888 MACs analytical saving. Those counts exclude matching, weighted aggregation, and scatter overhead, as well as any layout work around the shortened tensor. Because merged features change, fidelity must be re-measured for the downstream task. The smaller later-block bill does not guarantee lower wall-clock latency; matching cost, merge placement, implementation, sequence length, and quality target must be measured at the same boundary.

Optimization synthesis

The right first move follows the observed bottleneck; it does not follow a universal ranking of tricks. The same token reduction can matter greatly in one boundary and barely register in another, so diagnosis comes before choosing a descendant or changing a model knob.

Synthesis sources for this figure: the FlashAttention paper, transformer data movement, kernel fusion, and quantization.

1. Excess token count → resolution, patch size, windows, pruning, merging, or resampling.
2. Projection/MLP work → depth, width, MLP ratio, or distillation.
3. Large intermediates → memory-efficient attention or recomputation choices.
4. Data movement/kernel boundaries → layout and fusion.
5. Representation bytes → mixed precision or quantization with quality checks.
6. Shape/control overhead → static shapes, compilation, batching, and pipeline overlap.
7. Re-measure decode → resize/crop → normalize → patch embedding → encoder → head/projector → postprocess before accepting capability loss.

Open the complete runnable toy.

Analytical MACs can identify where matrix work moves, while element counts and boundary counts expose possible movement; they cannot predict every runtime. A change earns its place only after the same workload, semantics, and measurement boundary are used again. Keep the evidence categories separate: parameter count describes stored weights, MACs describe selected arithmetic, tensor shapes expose intermediates, and an end-to-end measurement includes work that the model equation omits.

Resolution and patch size jointly determine the encoder token sequence: at fixed image dimensions, a smaller patch creates more spatial rows, while reducing resolution or increasing patch size discards spatial detail before the encoder. Windows change who interacts without necessarily changing how many tokens survive. Resolution reduction, patch-size increase, and token merging each require approximation. Dynamic pruning requires both approximation and dynamic shapes. Visual resampling requires approximation and a visual-token downstream. Model right-sizing and student distillation each require approximation. Static windows are the only default-safe token action. These eligibility gates expose candidate changes; they do not promise task quality.

When a visual-token downstream exists, query count makes resampling an eligible interface choice. A query-based resampler maps N encoder outputs to M downstream outputs, but its M × N rectangle counts query-source interactions per query-based cross-attention layer, not full resampler work: projections, value aggregation, feed-forward layers, and repeated layers remain. A shorter downstream visual sequence does not shorten the image encoder, which still produces its full token sequence; report encoder and downstream boundaries separately. Fewer visual outputs can lose fine detail. In a vision-language system, include optional visual-token projection, resampling, and language-model prefill instead of stopping the bill at the vision encoder.

Distillation changes training and deployment in different places. The teacher forward is training-only, and the smaller deployed student is the only source of deployment savings; copying a teacher objective onto an unchanged student does not shrink deployment work. Depth, width, and MLP ratio similarly require a compatible trained model rather than a serving-only switch, and every such change needs a capability comparison at the intended task boundary.

For large attention intermediates, FlashAttention is the primary example of an exact, memory-efficient attention algorithm whose tiling changes data movement without changing the mathematical attention result. Its relevance must still be established from the observed intermediate or traffic bottleneck. The local transformer data movement treatment explains why bytes crossing memory boundaries can matter when MACs stay fixed; kernel fusion covers eliminating compatible operator boundaries. For representation bytes, mixed precision or quantization must carry an explicit quality check and supported execution path.

Finally, the following must be measured and recorded together: model and configuration, input resolution, patch size, visual query count when present, batch, precision, warm-up, synchronization, framework, compiler, runtime, device, boundary, statistic, and quality delta. The complete measurement path is decode → resize/crop → normalize → patch embedding → encoder → head/projector → postprocess. For a vision-language consumer, expand the middle to encoder → visual-token projection → resampling → language-model prefill, then include generation and postprocess if they are in scope. Observed performance depends on the implementation and workload, so change one knob, inspect execution and quality, and re-measure before accepting capability loss.

Where the branch continues

References

• ViT — paper, official code
• DeiT — paper, official code
• Swin Transformer — paper, official code
• PVT — paper, official code
• MViT — paper, official code
• MAE — paper, official code
• BEiT — paper, official code
• DINO — paper, official code
• DINOv2 — paper, official code
• Vision Transformers Need Registers — paper, official code
• CLIP — paper, official code
• SigLIP — paper, official code
• ViTDet — paper, official code
• SAM — paper, official code
• BLIP-2 — paper, official code
• LLaVA — paper, official code
• MobileViT — paper, official code
• EfficientViT-CGA — paper, official code
• DynamicViT — paper, official code
• ToMe — paper, official code

Local deep dives

• An Image Is Worth 16×16 Words
• Swin Transformer
• Masked Autoencoders
• BEiT
• DINO
• DINOv2
• CLIP
• ViTDet
• Segment Anything
• Why SAM returns multiple masks
• BLIP-2
• LLaVA
• Transformer data movement
• Kernel fusion
• Quantization