How does Adaptive Softmax speed up LLMs?

What you need to know

Why the output layer can be expensive

At each position, the final hidden vector (size d) is multiplied by a matrix of size d × V to get one score per vocabulary word. With a word-level vocabulary of about 267,000 words and d = 1,024:

267,000 x 1,024 = about 273 million multiply-adds per token

That can cost more than the rest of a small network.

How adaptive softmax works

1. Sort the vocabulary by frequency and split it: a head of, say, the 20,000 most frequent words, and tail clusters for the rest.
2. Head step — every token computes scores for the 20,000 head words plus one extra score per tail cluster ("the word is somewhere in tail 1").
3. Tail step — only if the tail cluster is likely (or during training, only if the target is in it) compute scores inside that cluster, using a smaller projection such as d/4.
4. Combine — probability of a tail word = p(tail cluster) × p(word inside the cluster).

With the numbers above (illustrative): the head costs about 20,002 × 1,024 ≈ 20 million multiply-adds. A 40,000-word tail cluster at d/4 = 256 adds about 10 million, but only for the small share of tokens that need it. Most steps cost under a tenth of the full softmax.

Related techniques

• Hierarchical softmax — a tree of binary decisions, cost about log V.
• Sampled softmax / negative sampling — during training, score the correct word plus a random sample of wrong ones.
• Adaptive input — the same frequency idea applied to the input embedding table.

Why modern LLMs rarely use it

• Subword tokenizers keep vocabularies at 30,000 to about 250,000.
• GPUs compute one large matrix multiply very efficiently; branching into clusters is awkward on them.
• The output matrix is still big — 128,000 × 4,096 is about 524 million parameters — so modern stacks use fused kernels that compute the loss in chunks without storing all logits, and speculative decoding to reduce the number of expensive steps.

A real-life example

A company builds an on-device next-word suggestion model for a Hindi–English keyboard app. It uses a word-level vocabulary of 200,000 entries (Hindi words, English words and Romanised spellings like "kya", "kyaa" and "kia"), and the phone has a tight compute budget.

A full softmax made each prediction too slow for typing speed. With adaptive softmax, the 15,000 most common words ("hai", "the", "kar", "ok") sit in the head and resolve almost every keystroke cheaply; rare words like brand names are in tail clusters that are computed only when the head points there. Latency drops enough for real-time suggestions, with a small accuracy loss on rare words. A server-side LLM for the same company uses a normal softmax, because its subword vocabulary and GPU make the plain version fast enough.

Follow-up questions to expect

• "Is adaptive softmax exact?" — It defines a valid probability over the full vocabulary, but it is a different factorisation, so the model learns slightly differently; rare words tend to be modelled less precisely.
• "How does it differ from hierarchical softmax?" — Hierarchical softmax uses a deep binary tree; adaptive softmax uses a shallow two-level split tuned to GPU efficiency and word frequency.
• "What limits the vocabulary size of modern LLMs?" — Output-layer and embedding memory versus tokenization efficiency; larger vocabularies make text shorter in tokens, especially for non-English languages.