Zhai et al. (2023) — ARXIV-2302.05442
Key takeaways
- Three changes carry ViT from 4B to 22B: parallel attention/MLP blocks
(
y = x + MLP(LN(x)) + Attention(LN(x)), which lets the QKV projections fuse with the MLP’s first linear and the out-projection with its second — PaLM reports 15% faster training from the same change), QK normalization, and omitted biases - The instability is diagnosed, not just fixed. Training diverged after a
few thousand steps at ~8B, caused by “extremely large values in attention
logits, which lead to (almost one-hot) attention weights with near-zero
entropy”. The fix follows from the mechanism: LayerNorm the queries and keys
before the dot product,
softmax[(1/√d)·LN(XW_Q)(LN(XW_K))ᵀ] - Dropping biases from QKV projections and LayerNorms improved accelerator utilisation by 3% with no quality loss — but biases are kept on the MLP dense layers, explicitly departing from PaLM, for observed quality with no speed cost. Bias removal is not uniform
- LayerNorms applied without bias and centering, citing Zhang & Sennrich — that is RMSNorm, keeping the gain
- 54.9% model FLOPs utilisation; 89.5% ImageNet; linear probing of ViT-22B can approach or exceed full fine-tuning of smaller models at high resolution
Standing in the anthology
Read — NOTE-058. Retagged to model-stability: the paper’s own framing is
efficiency and stability at scale, and the stability diagnosis is what survives
the modality gap.
This identified a hole, and the hole is now filled. This note said that
nothing in the record read a source establishing QK-norm, and called this
paper one of the earliest at scale — a hedge that was doing real work,
because it declined to claim the origin while the practice built on it
claimed exactly that. SOTA-192 was filed with source: LIT-088 and
introduced_by: LIT-088, so the practice asserted what the note had
carefully refused to.
The origin is Henry et al. (2020), LIT-640, which named the technique,
and SOTA-192’s introduced_by now points there. What this paper supplies
is unchanged and is the reason it stays in source:: a 22B model diverging
without the normalization and converging with it, the mechanism named, an
isolated before/after. One difference is worth carrying — Henry et al.
replace 1/√d with a learnable scalar, while the formula above keeps
it. Both bound the logits; only theirs makes the softmax temperature a
trained parameter.
The two arrivals do not cite each other, and LIT-449 is a third from
diffusion transformers that cites neither — see SOTA-192.
It also tightens SOTA-191 rather than promoting it. A 22B model dropping
LayerNorm bias and centering while keeping the gain is the RMSNorm half — the
opposite of the half SOTA-191 calls expendable — so it corroborates that
practice’s contested_by and does not meet its promote_when, which asks for
a decoder-only language model with the gain removed.