[Paper Review] Efficient Training of Language Models to Fill in the Middle

Mohammad Bavarian, Heewoo Jun+5 more|arXiv (Cornell University)|Jul 28, 2022

Topic ModelingComputer Science44 citations

TL;DR

Autoregressive language models can learn fill-in-the-middle (FIM) by a simple data transformation that moves a middle span to the end, with FIM training being effectively costless for left-to-right capabilities; finetuning FIM is compute-intensive, and the paper provides best practices and benchmarks.

ABSTRACT

We show that autoregressive language models can learn to infill text after we apply a straightforward transformation to the dataset, which simply moves a span of text from the middle of a document to its end. While this data augmentation has garnered much interest in recent years, we provide extensive evidence that training models with a large fraction of data transformed in this way does not harm the original left-to-right generative capability, as measured by perplexity and sampling evaluations across a wide range of scales. Given the usefulness, simplicity, and efficiency of training models to fill-in-the-middle (FIM), we suggest that future autoregressive language models be trained with FIM by default. To this end, we run a series of ablations on key hyperparameters, such as the data transformation frequency, the structure of the transformation, and the method of selecting the infill span. We use these ablations to prescribe strong default settings and best practices to train FIM models. We have released our best infilling model trained with best practices in our API, and release our infilling benchmarks to aid future research.

Motivation & Objective

• Motivate and enable infilling in causal language models by introducing a simple data transformation that moves a middle span to the end (FIM).
• Demonstrate that training with FIM does not harm left-to-right generation (the FIM-for-free property) across language and code domains.
• Characterize hyperparameters and practical guidelines (FIM rate, mode, context vs document-level, middle-span choice) to maximize FIM capabilities.
• Provide infilling benchmarks and release models to facilitate future FIM research and applications.

Proposed method

• Apply a dataset transformation that splits documents into prefix, middle, and suffix and moves the middle to the end, concatenating with sentinel tokens to form (prefix, suffix, middle).
• Train causal decoder models on a mixture of original left-to-right data and FIM-transformed data (FIM rate p, typically 50%), preserving autoregressive loss while enabling infilling.
• Explore two FIM implementations: PSM (prefix-suffix-middle order) and SPM (suffix-prefix-middle) and their joint training, with sentinel placement considerations.
• Investigate context-level vs. document-level FIM to mitigate data fragmentation during chunking, using 100B-token pretraining horizons and ablations across model scales.
• Evaluate left-to-right performance via standard autoregressive benchmarks and measure infilling via specialized benchmarks in language and code, including new random-span infilling tasks.

Experimental results

Research questions

• RQ1 Does FIM training affect the model's left-to-right autoregressive capabilities?
• RQ2 How does varying the FIM rate influence infilling performance and AR performance across scales?
• RQ3 What are the best practices for implementing FIM (PSM vs SPM, context vs document-level, middle-span selection) to maximize infilling without harming autoregression?
• RQ4 Can finetuning pretrained models achieve FIM capabilities as effectively as pretraining with FIM, and at what compute cost?

Key findings

• FIM training yields the FIM-for-free property: left-to-right (autoregressive) loss and performance remain unchanged when training with up to 50% FIM data across language and code domains.
• Higher FIM rates (up to 90%) do not degrade autoregressive perplexity but 100% FIM rate starts to degrade AR loss; infilling performance improves with higher FIM rate.
• Context-level FIM consistently outperforms document-level FIM in infilling benchmarks, with negligible AR loss difference.
• SPM mode is generally slightly stronger than PSM, and joint PSM+SPM training yields robust, flexible models with efficient inference modes.
• Finetuning AR models to acquire FIM capabilities requires substantial additional compute and does not generally reach the same performance as FIM-trained-from-scratch models.