← Concept library

NLP Foundations

Prefix Language Models

A single decoder-only stack that grants bidirectional attention to a prefix segment before switching to causal generation, and why the field mostly passed on this compromise anyway.

advanced · 8 min read · Premium

Between "fully causal" (decoder-only) and "fully bidirectional plus a separate causal decoder" (encoder-decoder) sits a design that gets bidirectional context on the input without paying for a second stack: the prefix language model. It is a single set of transformer weights, identical in shape to a decoder-only model, distinguished entirely by its attention mask.

The mask, precisely

Split every training sequence into a prefix (the "input," analogous to what an encoder would read) and a continuation (the "output," what gets generated). A prefix-LM's attention mask allows full, bidirectional attention within the prefix, every prefix token can attend to every other prefix token, both earlier and later, exactly as in BERT, while the continuation attends causally: each continuation token can see the entire prefix plus every earlier continuation token, but no later one, and, critically, prefix tokens never attend forward into the continuation. One stack, one set of weights, a block-structured mask instead of a strict lower triangle.

Implementing this requires no architectural change to a decoder-only codebase, only a different mask function, which is a large part of its appeal: you can experiment with the prefix-LM objective on the same code that trains a plain causal model.

Why bother: a cheaper middle ground

Raffel et al., 2019 directly compared language-model, prefix-LM, and full encoder-decoder architectures under matched parameter counts and compute in the T5 paper, precisely to answer whether the extra machinery of a separate encoder stack was earning its cost. Prefix-LM gets part of what full encoder-decoder gets, dense bidirectional context on the input, using half the parameters, because it shares one stack instead of running two. Dong et al., 2019 (UniLM) generalised the idea further, training a single shared transformer under a mixture of mask types, unidirectional, bidirectional, and sequence-to-sequence, so the same weights could be specialised toward any of the three regimes at fine-tuning time depending on which mask was applied.

The tension it resolves, and the one it creates

The tension prefix-LM resolves is real: plain decoder-only wastes the input's context by forcing causal masking even where nothing forces it (see decoder-only-architecture); full encoder-decoder pays double parameters and adds a whole extra cross-attention sublayer (see encoder-decoder-models-t5) to fix that. Prefix-LM sits in between at close to decoder-only's parameter cost.

But it creates a new tension of its own: the prefix/continuation boundary must be known, and in practice fixed, before the mask can be applied. That is a natural fit for tasks with a clean input/output split, translate this, summarise this, but an awkward fit for open-ended, multi-turn interaction where what counts as "the prompt so far" keeps growing turn by turn, and where treating the entire growing history as one ever-expanding bidirectional prefix would mean recomputing attention over it from scratch every turn.

Keep reading with Pro.

You're reading the preview. Unlock the full concept plus the library, study plans, the AI mentor, and daily emails.

Sign in to save and react.
Share Copied