NLP Foundations
Fill-in-the-Middle
A way to teach a plain causal decoder to infill text by rearranging training documents, not the model, so a single architecture serves both left-to-right generation and code-editor-style completion.
intermediate · 6 min read
A code editor's autocomplete needs to predict a token given both the code above the cursor and the code already written below it, for example the closing brace of a function whose body you are still typing. A causal decoder cannot do that natively; it only ever conditions on the left (see causal-lm-pretraining-task). Bavarian et al., 2022 showed you do not need a different architecture to fix this, only a different way of presenting the same training documents.
The rearrangement
Fill-in-the-middle (FIM) takes a document, splits it into three pieces at random positions, prefix, middle, suffix, and then reorders them at training time so a causal model always reads them in an order that lets the true "middle" appear last, where causal generation naturally happens:
?wzxhzdk:0?
The model reads a special <PRE> token, the prefix, a <SUF> token, the suffix, then a <MID> token, and only then generates the middle, one token at a time, exactly as it would generate the next token of any other document. Nothing about the transformer, the attention mask, or the loss changes: it is still ordinary next-token cross-entropy over a causally masked sequence. The only difference is that the sentinel tokens and reordering give the model, at generation time, access to "future" context (the suffix) by placing it earlier in the linear token sequence than the span being generated.
Bavarian et al. also test SPM order (suffix before prefix) and find it performs comparably; the important structural fact is that all context the model needs is moved before the generation point, not the specific ordering of prefix and suffix.
Why this is cheaper than it sounds
Span corruption (Raffel et al., 2019) achieves a related capability but requires an encoder-decoder architecture and bidirectional attention in the encoder (see span-corruption). FIM gets infilling behaviour out of a stock decoder-only model with zero architectural changes, by transforming a fraction of pretraining documents into the PSM or SPM layout and mixing them in with ordinary left-to-right documents. This is what makes FIM attractive for large generative models that are decoder-only for other reasons (inference simplicity, KV-cache friendliness, see kv-cache): you get an editor-friendly capability without paying for a second stack.
The "FIM-for-free" result
The central empirical claim in the paper is that applying the FIM transformation to a substantial fraction of training documents (they explore rates up to around 90%) does not measurably hurt the model's ordinary left-to-right generation quality, as long as it is trained jointly with un-transformed documents rather than replacing them. In other words, the model learns to infill essentially "for free," without trading away the capability it would have had from standard causal pretraining alone. This is why production code models (Codex-style completions, StarCoder, CodeLlama) fold FIM into pretraining by default rather than treating it as a separate fine-tuning stage.
When it falls down
- Span selection strategy matters. Splitting at random character or token offsets can produce prefixes and suffixes that break in the middle of a syntactic unit (a variable name, a string literal); character-level splits and token-level splits give measurably different results.
- The model still cannot condition on suffix content it has not been shown. FIM makes context ordering flexible during training, but at inference the caller must still explicitly provide the suffix; there is no mechanism for the model to discover relevant "future" context on its own.
- Sentinel token overhead. Every FIM example spends extra tokens on
<PRE>,<SUF>,<MID>markers, a small but real tax on effective context length compared to a plain document. - Not a substitute for bidirectional understanding tasks. FIM teaches generation with hindsight of a fixed suffix; it does not give the model the same continuously bidirectional representations that masked-language-modelling produces for classification or retrieval.
Further reading
- Bavarian et al., 2022, Efficient Training of Language Models to Fill in the Middle - the PSM/SPM transforms, FIM rate ablations, and the "FIM-for-free" result.
- Raffel et al., 2019, Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer - the encoder-decoder alternative for the same underlying problem.