ViT-22B: Scaling Vision Transformers to 22 Billion Parameters

Chapter 1: ViT Recap & Where 22B Lives

To see what changed, we need the standard Vision Transformer (Dosovitskiy et al., 2021) cleanly in mind. A ViT does not look at pixels directly. It chops the image into a grid of non-overlapping square patches, flattens each patch, and treats the resulting sequence exactly like a sequence of word tokens.

Concretely, take a 224×224 RGB image and a patch size of 14×14. That yields a 16×16 grid = 256 patches. Each patch is 14·14·3 = 588 raw numbers; a learned linear patch-embedding projects each into a d-dimensional vector. Prepend a learnable [CLS] token, add positional embeddings, and you have an input tensor of shape [T, d] with T = 257.

Then comes a stack of L identical Transformer blocks. The standard ViT block is a sequential composition — attention first, MLP second, each wrapped with a residual connection and a pre-LayerNorm:

Where Attn is multi-head self-attention and MLP is a two-layer feed-forward network with a GELU nonlinearity that expands d to 4d then back to d. Finally the [CLS] token's output feeds a classification head. That is all a ViT is.

Where does 22B sit on the ladder?

ViT-22B keeps the patch-token idea identical. It scales the three knobs that set a Transformer's size: the width d (hidden dimension), the depth L (number of blocks), and the MLP expansion. The configuration is roughly d = 6144, L = 48 blocks, 48 attention heads of dimension 128, and an MLP hidden size around 24576. Multiply those out across all blocks plus the patch embedding and head, and you land near 22 billion parameters.

A worked size estimate

Let us sanity-check the count from one block. A single block has two big sources of parameters. The attention projections W_Q, W_K, W_V, W_O are each [d, d] = 6144×6144 ≈ 37.7M parameters; four of them is about 151M. The MLP has two matrices [d, 4d] and [4d, d], each ≈ 6144×24576 ≈ 151M, so about 302M. One block is therefore roughly 151M + 302M ≈ 453M parameters. Across L = 48 blocks: 48 × 453M ≈ 21.7B. The patch embedding and head are a rounding error on top. The MLP, note, is two-thirds of every block — which is exactly why the parallel-layer reorganization in Chapter 4 pays off.

This sharding is not a footnote — it dictates the architecture. Every time a sharded matmul finishes, devices must exchange partial results (an "all-gather" or "reduce-scatter"). That cross-device communication is slow relative to on-chip matmul. The parallel-layer and asynchronous-linear tricks of Chapters 4–5 exist precisely to make those communication steps cheaper and to hide them behind compute. The architecture of ViT-22B is, in a real sense, co-designed with the TPU pod it runs on.

Chapter 2: The Divergence — Attention Entropy Collapse

Before we can appreciate the cure, we have to understand the disease precisely. Recall scaled dot-product attention. For a query q and key k (each a d_h-dimensional vector, d_h the per-head dimension), the pre-softmax logit is:

The attention weight is a_ij = softmax_j(ℓ_ij). Now here is the failure mode. Nothing in the standard block bounds the magnitude of q and k. They are produced by learned linear projections W_Q, W_K with no constraint on their norm. During training, if those weights grow, the dot product q·k grows, the logits grow, and the softmax sharpens.

The softmax has a runaway property. Consider one query attending over keys. The gradient that updates W_Q, W_K tends to push the largest logit larger (it is the one that "wins" and gets credit). Larger logits mean a sharper softmax, which concentrates the gradient even more on the winning key, which pushes that logit larger still. This is a positive feedback loop. At small scale the loop is gentle and self-limiting; at large d and large L, with thousands of these loops compounding across layers, it runs away.

The entropy view

A clean way to measure how sharp a softmax is, is its Shannon entropy. For an attention row a = (a₁, ..., a_T):

A uniform distribution over T keys has maximum entropy log T. A one-hot distribution has entropy 0. The paper observes that just before divergence, the average attention entropy collapses toward zero — the attention becomes nearly one-hot everywhere. This is the diagnostic signature.

A worked numerical example

Take a single query attending over three keys with raw logits before any scaling. Suppose the logit gap grows over training. Let the logits be (1, 0, 0), then (4, 0, 0), then (12, 0, 0).

Logits (1, 0, 0): exp gives (e¹, 1, 1) = (2.718, 1, 1), sum = 4.718, so a = (0.576, 0.212, 0.212). Entropy H = −(0.576 ln 0.576 + 2·0.212 ln 0.212) = −(0.576·(−0.552) + 0.424·(−1.551)) = 0.318 + 0.658 = 0.976 nats. Healthy — close to the maximum ln 3 = 1.099.

Logits (4, 0, 0): exp gives (54.6, 1, 1), sum = 56.6, a = (0.965, 0.0177, 0.0177). Entropy H = −(0.965·(−0.0357) + 2·0.0177·(−4.035)) = 0.0344 + 0.143 = 0.177 nats. The softmax has sharpened hard.

Logits (12, 0, 0): exp gives (162754, 1, 1), a = (0.99999, ~6e-6, ~6e-6). Entropy ≈ 0.0002 nats — effectively one-hot. The gradient through this softmax is essentially zero (the Jacobian of softmax is a·(diag − a a^T), which vanishes as a goes one-hot). The layer can no longer learn. This is entropy collapse.

Why does this only appear at scale? Two compounding reasons. First, wider models have larger d_h and more capacity in W_Q, W_K, so the achievable logit magnitudes are larger. Second, depth L matters: a small per-layer tendency to sharpen, repeated over 48 layers and millions of steps, compounds. The small-scale model never accumulates enough to tip over; the large one does. It is a classic case where a "works fine" recipe has a latent instability that only the scale exposes.

Chapter 3: QK-LayerNorm — Capping the Logits

The cure is almost embarrassingly small. Before computing the dot product, apply a LayerNorm to the queries and to the keys, per head. The standard attention logit was ℓ = (q·k)/√d_h. ViT-22B replaces q and k with their layer-normalized versions:

Recall what LayerNorm does to a vector v of dimension d_h: it subtracts the mean, divides by the standard deviation, then applies a learned per-dimension scale γ and shift β:

The decisive consequence: after LayerNorm, each of q and k has a controlled norm. Ignoring the learned γ, β for a moment, the normalized vector (v−μ)/σ has, by construction, second moment 1 per dimension, so its L2 norm is about √d_h. The dot product of two such vectors is bounded by the Cauchy–Schwarz inequality: |LN(q)·LN(k)| ≤ ||LN(q)|| · ||LN(k)|| ≈ d_h. Divide by √d_h and the logit is bounded on the order of √d_h — a constant set by head dimension, not something the weights can inflate without limit.

Why this stops the runaway loop

In the unnormalized block, the gradient could grow q·k by simply scaling up W_Q, W_K. With QK-LayerNorm in the path, scaling those weights up changes the mean and variance that LayerNorm immediately divides out. The norm of LN(q) is essentially invariant to the scale of q. So the positive-feedback channel — "make the weights bigger to make the winning logit bigger" — is cut. The logits cannot escape to infinity, the softmax cannot collapse, and the entropy stays healthy. Training is stable.

A worked example: LayerNorm caps a blown-up vector

Suppose mid-training the raw query for one head is q = (10, 2, −4, 0) (d_h = 4) and the weights have inflated it to 10q = (100, 20, −40, 0). Without QK-LN, the key k = (1, 1, 1, 1) gives a logit (100+20−40+0) = 80 before scaling — large, sharpening softmax. With QK-LN we normalize first.

For q: mean μ = (10+2−4+0)/4 = 2. Centered: (8, 0, −6, −2). Variance σ² = (64+0+36+4)/4 = 26, σ = 5.10. Normalized (taking γ=1, β=0): (1.569, 0, −1.177, −0.392). Now scale q by 10: mean 20, centered (80, 0, −60, −20), variance 2600, σ = 50.99, normalized = (1.569, 0, −1.177, −0.392) — identical. LayerNorm is invariant to the overall scale of its input. The weights inflating by 10× changed the logit by a factor of one. That is the entire mechanism.

Cost and placement

QK-LayerNorm adds two tiny LayerNorms per head per layer, operating on d_h-dimensional vectors. Relative to the [d, d] projections and the [d, 4d] MLP, this is a negligible amount of compute and parameters — fractions of a percent. It is applied only to q and k, not to v: v carries the content to be aggregated and we do not want to normalize away its magnitude information; only the similarity computation (q·k) needs taming.

python
import torch, torch.nn as nn

class AttnQKNorm(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h = n_heads
        self.dh = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        # the two extra LayerNorms — over the per-head dim only
        self.q_norm = nn.LayerNorm(self.dh)
        self.k_norm = nn.LayerNorm(self.dh)

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(B, T, self.h, self.dh)
        k = k.view(B, T, self.h, self.dh)
        v = v.view(B, T, self.h, self.dh)
        # QK-LayerNorm: cap the logits BEFORE the dot product
        q = self.q_norm(q)
        k = self.k_norm(k)            # note: v is NOT normalized
        scores = torch.einsum('bthd,bshd->bhts', q, k) / self.dh**0.5
        a = scores.softmax(-1)
        out = torch.einsum('bhts,bshd->bthd', a, v).reshape(B, T, D)
        return self.proj(out)

Chapter 4: Parallel Layers — Attention and MLP, Side by Side

This is the architectural reorganization at the heart of the labs. Recall the standard, sequential ViT block: attention runs, its output is added back, that result is normalized and fed into the MLP, whose output is added back again. The MLP's input depends on the attention's output.

ViT-22B (following GPT-J and PaLM) uses a parallel block. Attention and the MLP both read the same normalized input, run independently, and their outputs are summed into a single residual:

Read those two formulas carefully — the difference is small on the page but consequential. In the sequential version, the MLP sees LN₂(z'), which includes the attention update. In the parallel version, the MLP sees LN(z), the original input, with no knowledge of what attention computed this layer. The two sublayers are now independent functions of the same input, and only meet when their outputs are added.

Why would you give up the dependency?

Three concrete wins, all about systems efficiency, with negligible quality cost at scale:

(1) One LayerNorm instead of two. The sequential block normalizes twice per block (before attention, before MLP). The parallel block shares a single LayerNorm feeding both branches. Over 48 layers that halves the LayerNorm count.

(2) Fused QKV and MLP-input matmuls. Because attention and MLP both consume the same LN(z), their first linear projections can be combined into one larger matrix multiply. Instead of a [d, 3d] QKV matmul and a separate [d, 4d] MLP-up matmul, you do a single [d, 3d+4d] = [d, 7d] matmul. Likewise the output side fuses W_O [d, d] and the MLP-down [4d, d] into one combined linear. Bigger matmuls run at higher hardware utilization than many small ones — fewer kernel launches, better tiling, more arithmetic per byte moved.

(3) Fewer communication barriers. On a sharded model, each linear layer ends with a cross-device reduction. Sequential attention-then-MLP forces those reductions to happen one after another (the MLP cannot start until attention's reduction lands). Fusing them lets the two reductions be combined and overlapped, which is exactly the asynchronous-linear idea of the next chapter. The paper reports the parallel layout gives a substantial throughput improvement at 22B scale.

A worked example: same input, summed outputs

Let z be one token's representation, LN(z) = u = (1, −1) for a tiny d = 2. Say attention contributes a = Attn(u) = (0.4, 0.2) and the MLP contributes m = MLP(u) = (−0.1, 0.5), both computed from the same u. Sequential would instead feed the MLP the post-attention vector z + a then re-normalize; parallel just adds: out = z + a + m. If z = (2, 3), out = (2+0.4−0.1, 3+0.2+0.5) = (2.3, 3.7). The attention and MLP never saw each other's work this layer — they communicate only through the residual stream, one layer later.

The cost: a small loss of expressivity

Is anything lost? In principle, yes: in the parallel block the MLP can no longer react to this layer's attention output — it must wait for the next layer to see it through the residual stream. Empirically, at large width and depth this barely matters: the residual stream already mixes both contributions, and the next layer's LN sees their sum. The paper, like PaLM before it, found the quality difference negligible while the speedup is real. This is a recurring theme in large-scale training: trade a sliver of theoretical expressivity for a large, concrete systems win.

python
# Parallel block (ViT-22B / PaLM style)
class ParallelBlock(nn.Module):
    def __init__(self, d, n_heads, mlp_mult=4):
        super().__init__()
        self.norm = nn.LayerNorm(d)        # ONE shared norm
        self.attn = AttnQKNorm(d, n_heads)  # with QK-LayerNorm inside
        self.mlp  = nn.Sequential(
            nn.Linear(d, mlp_mult * d), nn.GELU(),
            nn.Linear(mlp_mult * d, d))

    def forward(self, z):
        u = self.norm(z)            # both branches read the SAME u
        # out = z + Attn(u) + MLP(u)  — summed into one residual
        return z + self.attn(u) + self.mlp(u)

Chapter 5: Asynchronous Parallel Linear Operations

The parallel block sets up the third trick. On a model sharded across a device mesh, a single big linear layer y = xW is not one operation — it is a local matmul on each chip followed by a collective communication to assemble the result. The two costs are different in kind: matmul is compute-bound and lives in the fast on-chip arithmetic units; the collective is bandwidth-bound and crosses the slow inter-chip network.

If you run them strictly in order — matmul, then wait for the all-reduce, then the next matmul — the expensive arithmetic units sit idle during every communication step, and the network sits idle during every matmul. Utilization is poor on both. The idea behind asynchronous parallel linear operations is to overlap them: start the communication for one tile of the output while still computing the matmul for the next tile, so compute and communication run at the same time.

If compute and communication take roughly equal time, overlapping them nearly halves the wall-clock cost of the layer. The parallel block makes this far easier: because attention's and the MLP's first matmuls are fused and read the same input, there is a single large matmul whose output tiles can be streamed into a single overlapped collective, rather than several small matmuls each with its own barrier.

A worked example: serial vs overlapped

Say a layer's matmul takes 10 ms and its cross-device all-reduce takes 8 ms. Done serially, the layer takes 10 + 8 = 18 ms. With perfect overlap (start reducing the early output tiles while computing the late ones), the layer takes max(10, 8) = 10 ms — a 1.8× speedup. Over L = 48 layers and millions of steps, that compounds into days of training time saved and is a large fraction of why 22B was feasible on the available hardware budget.

Why the architecture and the systems are co-designed

This chapter is where it becomes undeniable that ViT-22B's "architecture" cannot be cleanly separated from the hardware it runs on. The parallel block was chosen partly because it makes the matmuls fusible and the collectives overlappable. QK-LayerNorm was chosen partly because it adds almost no cost to those matmuls. The model is shaped by the shape of the TPU pod. The paper is as much a distributed-systems contribution as an architecture one — it shows a concrete recipe (in JAX, with GSPMD-style sharding annotations) for training a dense 22B vision model efficiently.

Chapter 6: Frozen-Feature Transfer

Now the payoff. The cleanest way to ask "how good is a representation?" is to freeze the backbone — never update a single weight — and train only a tiny linear classifier on top of its features. This is the linear probe. If the frozen features are linearly separable for a downstream task, a single matrix can read them out; if not, no amount of linear readout will help. Linear-probe accuracy is therefore a direct measure of representation quality, with no fine-tuning to muddy the comparison.

Why care about frozen transfer specifically? Because a 22B model is expensive to fine-tune — you would have to store optimizer state and gradients for all 22B parameters per downstream task. If the frozen features are already excellent, you skip all of that: extract features once, train a cheap linear head per task, and reuse the same frozen backbone for everything. ViT-22B reports exactly this — its frozen features are strong across a wide battery of classification tasks, often rivaling fully fine-tuned smaller models. Scale converts the backbone into a reusable, general feature service.

A worked example: why frozen probing is a clean test

Imagine a task with 1000 classes and a feature dimension d = 6144. A linear probe is a single [6144, 1000] matrix — about 6.1M parameters, trivially cheap to train. If ViT-22B's frozen 6144-dim CLS feature lets that linear layer reach, say, the high-80s top-1 on ImageNet, it tells you the 22B backbone has already arranged its feature space so that classes are linearly separable — the hard work is done by the representation, not the readout. Contrast a weaker backbone whose linear probe lags its fine-tuned number by many points: its features need nonlinear untangling, so fine-tuning is doing real work. The smaller the probe-vs-finetune gap, the better the frozen representation.

Robustness and out-of-distribution transfer

The paper also evaluates ViT-22B features on out-of-distribution and robustness benchmarks (ImageNet variants like ImageNet-R, ImageNet-A, and others), on dense tasks like depth estimation and segmentation via frozen features, and on fairness/calibration. The recurring finding is that the large frozen backbone transfers gracefully — it does not just memorize ImageNet, it learns features that generalize to shifted distributions and to tasks it was never trained on. This is the practical argument for one very large, frozen, shared backbone over many task-specific fine-tunes.

Chapter 7: Shape Bias & Human-Perception Alignment

Here is the most surprising result, and it is not about accuracy at all. When you classify a "cat-shaped object painted with elephant skin texture," what is it? Humans overwhelmingly say cat — we classify by shape. Standard ImageNet-trained CNNs say elephant — they classify by texture (Geirhos et al., 2019). The fraction of such cue-conflict images a model labels by shape is its shape bias; the rest is texture bias. Humans sit around 90–96% shape; a typical ResNet sits around 20–30% shape (i.e. mostly texture).

ViT-22B's headline perceptual finding: its shape bias rises dramatically compared to smaller models and CNNs — the paper reports it reaching levels far closer to the human range than prior vision models, a substantial shift along the shape–texture axis. The largest model is, in this specific and measurable sense, the most human-like vision model reported at the time.

How shape bias is measured

The cue-conflict test set (Geirhos et al.) takes an image with one object's shape and another object's texture, restricted to the cases where the model picks one of the two conflicting labels. Among those, shape_bias is the fraction where it chose the shape label. A pure texture classifier scores 0; a pure shape classifier scores 1. Plot a model as a single point on a shape-vs-texture axis and you can place it relative to the human band.

A worked example

Suppose on 16 cue-conflict categories a model, restricted to shape-or-texture answers, picks the shape label 13 times and the texture label 3 times. Its shape bias is 13/16 = 0.81 = 81%. A ResNet on the same set might pick shape only 4 times: 4/16 = 25%. Humans pick shape ~15 of 16: ~94%. The number is a single fraction, but it encodes which cue the model trusts when forced to choose — and ViT-22B's number moved far toward the human end of that line.

Beyond shape: error consistency and calibration

Shape bias is one slice of a broader claim: ViT-22B's behavior aligns more with human perception. The paper also looks at error consistency — do the model and humans make mistakes on the same images? — and finds ViT-22B's errors are more human-consistent than smaller models'. It examines calibration (are the model's confidences well-matched to its accuracy?) and reports improvements. The thread tying these together: scaling did not merely make a better ImageNet scorer; it nudged the entire perceptual profile of the network toward the way people see.

Why might scale do this? One hypothesis the community offers: shape is a more global, relational cue than texture, and self-attention over the whole image — given enough capacity and data — can integrate global structure that small or convolutional models, biased toward local texture statistics, miss. A 22B ViT has the capacity to represent and rely on global shape, so it does. This is a representational consequence of the architecture-plus-scale combination, not something explicitly trained for.

Chapter 8: The ViT-22B Block Explorer

This is the payoff simulation — the paper's two core architectural ideas, live and interactive. The explorer runs a single attention head over a 6-token patch sequence and lets you (a) toggle QK-LayerNorm on or off and inflate the query/key weights to watch entropy collapse appear and be cured, and (b) switch the block layout between sequential and parallel to see how the residual output is assembled. Everything is computed from real arithmetic in the browser — the bars are actual softmax weights, the entropy is the actual Shannon entropy.

Drive it like an experimentalist. First, turn QK-LN off and crank the weight scale: the attention bars sharpen to one-hot and the entropy readout crashes toward zero — that is the divergence of Chapter 2, frozen as a control. Then flip QK-LN on and crank the same slider: the bars stay spread, the entropy holds up — the logits are capped (Chapter 3). Finally, toggle the block layout and watch which input the MLP branch reads (the post-attention vector in sequential, the shared normalized input in parallel) and how the residual output is summed (Chapter 4).

What you should take away from playing with it: the two ideas are independent levers. QK-LayerNorm is about numerical stability — it governs whether the run survives. Parallel layers are about composition and systems efficiency — they govern how fast each surviving step is. ViT-22B needed both: stability to start, efficiency to finish in a reasonable budget. Remove either and 22B is out of reach — without QK-LN it diverges; without the parallel-layer / async-linear efficiency it is too slow to train on the available hardware.

Chapter 9: Limitations, Connections & Cheat Sheet

ViT-22B is a milestone, not a finish line. Honest limitations:

Dense and expensive. 22B dense parameters are activated for every patch of every image. Sparse mixture-of-experts vision models can reach larger nominal parameter counts at a fraction of the per-image FLOPs — ViT-22B deliberately chose dense to keep the architecture simple and the result a clean scaling data point, but density is a cost.

Data and compute access. The model is trained on a very large proprietary image dataset on a TPU pod. The recipe is reproducible in principle, but the resources are out of reach for most labs — a recurring equity concern with frontier-scale work.

Quadratic attention remains. QK-LayerNorm fixes the stability of attention but not its O(T²) cost in the number of tokens. High-resolution images mean many patches; the attention matrix still grows quadratically. Long-context / efficient-attention work is orthogonal and still needed.

Bias is shifted, not solved. Higher shape bias and better human-error-consistency are improvements in degree, not a guarantee of human-like robustness or fairness everywhere. The paper's own fairness analysis is a starting point, not a clean bill of health.

Where this sits in the literature

The lineage is direct: the original Transformer introduced self-attention and the sequential block; ViT applied it to image patches; ViT-22B is the scaling study that re-engineered the block to push an order of magnitude further. QK-LayerNorm and the parallel block both have roots in the language-model scaling literature (PaLM, GPT-J), which ViT-22B imports into vision.

Related lessons

• Attention Is All You Need — Veanor: scaled dot-product attention, multi-head, the sequential encoder/decoder block.
• Vision Transformer — Gleam: how images become patch tokens; the base ViT.
• The Transformer — Gleam: the block, residuals, pre-LN vs post-LN.
• Normalization — Gleam: BatchNorm / LayerNorm / RMSNorm, the math QK-LN builds on.
• Attention Variants — Gleam: efficient and stabilized attention, where QK-LN fits.

Cheat sheet