Lookahead Decoding

The sequential wall: one token, one whole pass

Autoregressive generation is defined by a strict dependency: token x_i is chosen from p_θ(x_i | x_0, …, x_{i-1}), so you cannot compute x_i until x_{i-1} exists. Producing m tokens therefore means m forward passes, run strictly back-to-back. Nothing about the arithmetic forces this to be slow — the latency comes from the memory system.

During decode with batch size one, each step reads every weight in the model from memory to process a single token. A forward pass costs about 2P FLOPs per token for P parameters, but it must also move all P weights (plus the KV cache) across the memory bus. The ratio of compute to bytes moved — the arithmetic intensity — is tiny, so the processor stalls waiting on HBM or system RAM. This is the memory-bound regime: the compute units sit mostly idle on every decode step. Lookahead decoding’s entire premise is that this idle compute is free capacity waiting to be spent.

Decoding as a fixed point: the Jacobi view

Here is the reframing that unlocks everything. Instead of one token at a time, treat a block of m future tokens as the unknowns of a system of equations. Greedy decoding satisfies, for every position i = 1…m:

x_i = argmax p_θ( · | x_0, x_1, ..., x_{i-1} )   for i = 1..m

This is a triangular nonlinear system: the true solution is exactly what sequential decoding produces. But triangular systems can be attacked by Jacobi iteration — start from an arbitrary guess for all m tokens at once, then update every position simultaneously using the current guesses for the positions before it:

x_i^(t+1) = argmax p_θ( · | x_0, x_1^(t), ..., x_{i-1}^(t) )

The key fact: all m of these updates are computed in a single forward pass. A causal transformer already scores every position in parallel; feeding it the current guess sequence yields the next-token prediction at each of the m slots at once. We turned m sequential passes into repeated parallel passes.

Convergence: why iterating in parallel is safe

Does Jacobi iteration reach the right answer? Yes, and provably fast for the greedy case. Observe that position 1 depends only on the fixed prefix x_0, so after the first iteration x_1 is correct and frozen forever. On the next iteration, position 2 now sees the correct x_1, so x_2 locks in. By induction, each iteration guarantees at least one more token converges from the left, so the process reaches the exact autoregressive output in at most m steps.

Worst case, that is no better than sequential decoding. The win is that convergence is often much faster than one token per step: later positions guess correctly early because many continuations are locally predictable (whitespace, closing brackets, common phrases), so a single parallel step can confirm a run of tokens. Plain Jacobi decoding captures some of this but inconsistently. Lookahead decoding’s contribution is to stop discarding the partial guesses each iteration produces, and instead accumulate them into a reusable structure.

The n-gram pool: harvesting the Jacobi trajectory

Every Jacobi iteration produces, at each of its m positions, a short sequence of tokens — a trajectory through candidate continuations. Across iterations these fragments form n-grams: length-N token sequences the model has tentatively proposed. Lookahead decoding keeps a pool (a cache) of these n-grams, indexed by their first token.

The intuition is that the model’s own half-formed guesses are a free source of plausible drafts. When the next confirmed token is some w, the pool is queried for every stored n-gram beginning with w — each a ready-made continuation the model generated moments earlier. For structured text this pays off dramatically: code, markup, and repetitive prose contain the same n-grams over and over ( return, </div>, self.), so the pool fills with n-grams that will be needed again. It evolves: the lookahead branch refreshes it every step and committed tokens prune it, so it always reflects the current context.

Two branches, one forward pass

A lookahead decoding step fuses two jobs into one batched forward pass, kept apart by a custom attention mask so the branches do not attend across each other:

The lookahead branch runs the Jacobi update. It maintains a 2D window — W parallel trajectories (the window size) each looking N tokens ahead (the n-gram size, or levels). In one pass it refines all W×(N−1) positions and emits fresh n-grams into the pool.

The verification branch pulls up to G candidate n-grams from the pool whose first token matches the last confirmed token and lets the model score them. Wherever the model’s greedy prediction agrees with a candidate token, that token is accepted; the longest matching prefix is committed. Because both branches share the step’s single weight load, they are almost free relative to a normal decode — the weights were going to be read anyway. The step advances by 1 + (accepted length) tokens instead of by one.

A worked step: several tokens at once

Suppose the model is emitting Python and has just committed def. The pool, filled by earlier Jacobi iterations, happens to hold an n-gram whose first token is that anchor: def load ( path ) :. The verification branch scores the continuation in the same pass that advances the lookahead window.

anchor:     ... def          (already committed)
candidate:      load  (  path  )  :
model says:     load  (  path  )  :   <- greedy prediction at each slot
match:           Y    Y   Y    Y  Y
accepted 5 tokens in ONE step  →  commit "load(path):"

One memory-bound forward pass advanced the sequence by five tokens instead of one — the anchor def was already committed, and the five tokens after it are the gain. Had the candidate diverged at path — say the model wanted data — the branch would accept the matching prefix load (, commit those, and discard the rest. Nothing incorrect is ever emitted: acceptance requires the token to equal the model’s own greedy choice, so the output is bit-identical to ordinary greedy decoding. Lookahead changes the schedule, never the result.

The FLOPs-for-latency trade

The whole method is one economic bet: spend idle parallel FLOPs to buy back sequential steps. A baseline step processes one token; a lookahead step processes roughly W×(N−1) + G×(N−1) extra tokens. That is many times more FLOPs per step — but the step still loads the weights exactly once, and it was memory-bound with compute to spare. As long as the extra tokens do not push the step past the roofline into the compute-bound regime, they add almost no wall-clock time, yet each accepted token removes an entire future step.

The governing quantity is the average tokens accepted per step. If a run generates T tokens and accepts τ per step, it needs about T/τ steps instead of T. Speedup is bounded by τ from above and by the point where growing per-step FLOPs saturate the hardware: push the window too far and the step turns compute-bound, so the extra FLOPs now cost real time. The sweet spot sits where per-step compute time rises to meet the fixed weight-load time.