Bridges: DSA to LLMs

Sorting and selection to top-k sampling

Picking the best few from a distribution

~10 minLesson 40 of 57Written by Claude Opus 4.8

Greedy decoding is argmax; top-k is "find the $k$ largest logits" - the exact job a heap and quickselect do without ever sorting the whole 50,000-token vocabulary.

A transformer's final layer hands you one raw score - a logit - for every token in its vocabulary. To speak, it must pick one. That pick is a selection problem, the same family that quickselect and heaps were built for. Greedy decoding is argmax: the single largest logit. Top-k decoding wants the $k$ largest. Neither needs the other ~50,000 tokens put in order - and refusing to order them is where the speed comes from.

Decoding is selection in disguise

You have 50,257 candidate tokens (GPT-2's vocabulary) and you want the best 40. Sorting all of them into a perfect ranking is an

O (n lo g n)

job that answers a question nobody asked - you throw away 49,960 of the answers immediately. Selection asks the cheaper question: "which few are the largest?", and never bothers to order the losers among themselves. Greedy is

k = 1

; top-k is any small

k

Selection is cheaper than sorting

Three tools sit on a cost ladder. argmax is a single left-to-right scan tracking the biggest so far - $O (n)$ , no ordering at all. Quickselect generalises it: it does one quicksort-style partition around a pivot, then recurses into only the one side that contains the $k$ -th element - so it never sorts the other side, giving $O (n)$ on average. A size-k min-heap takes a different route that also keeps the survivors ranked: sweep once through all $n$ logits, keeping only the $k$ best seen so far. Its smallest entry is a live cut-off; anything below it is rejected on sight. That is $O (n lo g k)$ - and crucially, $lo g k$ , not $lo g n$ .

Task	Method	Cost
Single best token (greedy)	argmax scan	$O (n)$
Top-k, in any order	quickselect (partition)	$O (n)$ avg
Top-k, kept ranked	size-k min-heap	$O (n lo g k)$
Full ranking of the vocab	comparison sort	$O (n lo g n)$

The bottom row is the one decoding never needs. Once you have the top- $k$ survivors, you run softmax over just those few to turn logits back into a probability distribution, and sample one token from it. Select, re-normalize, draw - that is the entire top-k pipeline.

Softmax over the survivors, and why log k wins

After selection keeps the set $T$ of the top- $k$ tokens, softmax re-normalizes over only that set:

p_{i} = \frac{e ^{z_{i}}}{j \in T \sum e ^{z_{j}}}, i \in T .

The denominator sums $k$ terms, not $n$ - the discarded tail contributes nothing. On the selection cost itself, with GPT-2's $n = 50, 257$ we have $lo g_{2} n \approx 15.6$ . For $k = 40$ , $lo g_{2} k \approx 5.3$ . The heap's per-token work is a third of a full sort's, and quickselect skips the logarithm entirely with $O (n)$ average work. The bigger the vocabulary and the smaller the $k$ , the more lopsided the win.

See it: the size-k min-heap sweep

Below are 12 vocabulary tokens with fixed logits, in vocabulary order - not sorted, exactly as they leave the model. Set $k$ , then press Run heap sweep to watch a size- $k$ min-heap crawl left to right: each token is pushed if the heap has room, or must beat the heap's minimum (the amber cut-off) to evict it, otherwise it greys out on sight. The survivors are softmax-normalized into a sampling wheel on the right - spin it to draw one token, weighted by probability.

Top-k by min-heap → softmax wheel. set k · run the sweep · spin

survivor (top-k)currently in heaprejected

empty

the size-k min-heap - its smallest entry ( amber) is the cut-off every new logit must beat

Ready - press Spin.

top-k = 4

The heap-min is the bouncer

Watch the amber cut-off during a sweep. A newcomer only gets in if its logit beats that minimum; otherwise it is turned away after a single comparison. Because the heap never holds more than

k

items, each of the

n

tokens costs at most one

O (lo g k)

push-or-evict. That moving threshold - never a full ranking - is the entire reason top-k is

O (n lo g k)

. Drag

k

to 1 and the sweep is argmax: one survivor, greedy decoding.

import heapq, numpy as np

def top_k_sample(logits, k=40):
    # 1) SELECT the k largest logits - no full sort.
    #    nlargest keeps a size-k min-heap: O(n log k).
    top = heapq.nlargest(k, range(len(logits)), key=lambda i: logits[i])
    #    numpy alt: top = np.argpartition(logits, -k)[-k:]  # quickselect, O(n)

    # 2) SOFTMAX over just the survivors, re-normalized.
    z = np.array([logits[i] for i in top])
    p = np.exp(z - z.max())
    p /= p.sum()                          # denominator sums k terms, not n

    # 3) DRAW one token, weighted by p - spin the wheel.
    return np.random.choice(top, p=p)

The whole top-k pipeline in one function

Check yourself

Why is a size-k heap cheaper than sorting the whole vocabulary?

During the sweep, what does the size-k heap's minimum represent?

Recall: why is top-k decoding O(n log k) with a heap, not O(n log n)?

Pin down where the speed comes from. Try to state it, then check.

Lock it in

Decoding is a selection problem: greedy is argmax ( $k = 1$ ), top-k wants the $k$ largest logits - neither needs the vocabulary sorted.
A size- $k$ min-heap sweeps once in $O (n lo g k)$ , its minimum acting as a live cut-off; quickselect does it in $O (n)$ average with no ordering.
Sorting the whole vocabulary is $O (n lo g n)$ and answers a question nobody asked - the discarded tail is thrown away immediately.
Softmax then re-normalizes over just the $k$ survivors, and you sample one token: select, re-normalize, draw.

Primary source

Robert Sedgewick and Kevin Wayne's algorithms booksite (algs4) is the canonical treatment of the machinery underneath this lesson - binary heaps and priority queues, and the quickselect / partition-based selection that finds the

k

-th largest in linear average time. For the LLM half - how the transformer produces that vector of logits and turns the survivors into a next token - Jay Alammar's The Illustrated GPT-2 walks the pipeline visually.

Ask your teacher

Here is the connection worth keeping: top-k decoding is the heap and quickselect at work - you never need to sort the whole 50,000-token vocabulary, just find the top few. Want me to visualise quickselect's partition finding the

k

-th logit directly, or show how top-p (nucleus) trades the fixed

k

for a prefix-sum plus binary search in the very next lesson? Ask and I'll build the follow-up.