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.
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.
The KV-cache complication
Plain causal decoding caches each token's key and value vectors once and reuses them forever (kv-cache), because a causal token's representation, once computed, never changes as later tokens are added. That guarantee breaks inside a prefix-LM's bidirectional block: a prefix token's key and value depend on every other prefix token, so appending a new prefix token can, in principle, change the representation of every earlier prefix token too. In practice this means the prefix's cache can only be trusted once the prefix is finalised; extending it requires recomputation in a way pure causal decoding does not.
When it falls down
- Needs a committed prefix boundary. Streaming or multi-turn settings, where "the input" keeps growing, do not map cleanly onto a mask that assumes a fixed prefix/continuation split decided in advance.
- Complicates caching. The bidirectional block inside the prefix loses the free incremental-caching property that makes plain causal decoding cheap to serve, exactly where the win over encoder-decoder was supposed to come from.
- It underperformed full encoder-decoder in T5's own controlled comparison. Raffel et al. found the full two-stack encoder-decoder architecture ahead of prefix-LM at matched compute, meaning prefix-LM is better understood as a cheaper compromise than a categorically superior design.
- The field mostly moved past it anyway. As pure decoder-only models scaled and their in-context conditioning improved, the specific gap prefix-LM was built to close, the input being processed causally instead of bidirectionally, mattered less in practice than the T5-era comparisons suggested it would, and few frontier models adopted the prefix-LM mask.
Further reading
- Raffel et al., 2019, Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (T5), arXiv:1910.10683 - the controlled comparison of LM, prefix-LM, and encoder-decoder objectives.
- Dong et al., 2019, Unified Language Model Pre-training for Natural Language Understanding and Generation (UniLM), arXiv:1905.03197 - one shared stack trained under a mixture of mask types.
4 flashcards for this concept
Click a card to reveal the answer.