PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Training of Language Models to Fill in the Middle

Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, Mark Chen

arXiv:2207.14255v1cs.CL

TL;DR

Autoregressive models need to generate text conditioned on both a prefix and suffix, a capability useful for applications such as coding assistance. The paper trains models with a simple middle-to-end data transformation and finds that FIM improves infilling without harming left-to-right performance, while identifying effective training practices and important limitations.

  • Problem

    Existing autoregressive models are limited at infilling because they condition only on prefixes, despite applications requiring context before and after the generated span.

  • Method

    The paper jointly trains causal decoder models on traditional left-to-right data and FIM-transformed data, then studies FIM rates, transformation structures, span selection, benchmarks, and finetuning.

  • Results

    FIM models match autoregressive models on left-to-right test loss at the same compute while achieving lower FIM loss, with benchmark results showing FIM-rate and implementation choices affect infilling.

  • Takeaways & Limitations

    Pretraining jointly with FIM is more effective for best FIM performance than finetuning a canonically pretrained left-to-right model, and recommended settings support robust infilling.

  • Takeaways & Limitations

    FIM sampling can generate spurious content or fail to produce <EOT> within the token budget because the desired infill length or style is unspecified.

Abstract

from arXiv · show

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.

1 Introduction

The paper adds fill-in-the-middle capability to causal autoregressive models through a simple data transformation, while preserving left-to-right generation. It frames FIM-for-free as an efficient way to expand model capability without changing architecture or pretraining compute.

  • Motivation: Causal decoder-based models are limited in infilling because they condition only on the prefix, despite practical contexts containing both prefix and suffix.Infilling is useful for coding tasks such as docstring generation, import completion, and partially written functions.
  • Approach: A simple training-data modification enables causal autoregressive models to learn infilling without changing the model architecture or compromising left-to-right generation.The transformation is applied to only a fraction of the training dataset.
  • Results: FIM models learn to condition on the suffix while predicting the middle, obtaining lower FIM test loss while matching baseline left-to-right loss.The comparison is described using the same model scales evaluated in the paper.
  • Approach: The method splits documents into prefix, middle, and suffix, then moves the middle span after the suffix and concatenates the pieces with sentinel tokens.This transformation is described as the key to the approach.
  • Results: FIM-for-free means jointly training on FIM-transformed and ordinary left-to-right data preserves left-to-right capability while adding infilling ability.The paper studies this property using scaling experiments, perplexity, and sampling-based evaluations.
  • Contributions: The paper contributes best practices, new infilling benchmarks, and evidence that sampling evaluations can reveal differences missed by FIM test losses.It also reports that finetuning existing autoregressive models for FIM is computationally inefficient relative to learning FIM during pretraining.

2 Evaluation

The evaluation combines autoregressive and infilling losses with standard capability benchmarks and code-generation tests. It uses carefully constructed HumanEval-based tasks to measure free-form infilling quality beyond perplexity.

  • Evaluation setup: AR evaluation measures whether FIM training affects left-to-right capability, while FIM evaluation measures infilling performance and scaling trends.AR loss is computed on ordinary left-to-right data, whereas FIM loss uses fully transformed data.
  • Evaluation setup: Test losses and standard benchmarks are evaluated across natural-language and code domains to assess both capability preservation and infilling behavior.The code evaluation includes HumanEval pass rates, while natural-language tasks include reasoning and comprehension benchmarks.
  • Loss evaluation: FIM test sets reuse AR examples after applying a 100% transformation, enabling direct comparison of AR and FIM losses.Masked evaluations measure loss on middle-span tokens under P(middle|prefix, suffix) for FIM models and P(middle|prefix) for AR models.
  • Sampling evaluation: The infilling benchmarks remove middle spans from HumanEval solutions and evaluate open-ended code completions with test suites.This setup supports correctness evaluation for long generated samples.
  • Sampling evaluation: The benchmark suite includes single-line, multi-line, and random-span tasks, with substantially more tasks than the original HumanEval dataset to reduce evaluation variance.The suite contains 1033, 5815, and 1640 tasks, respectively.
  • Evaluation setup: The paper reports SPM infilling results by default, except where PSM changes the conclusions.PSM and SPM are the two FIM formatting modes studied.

3 FIM training and inference

FIM training randomly transforms data by rearranging document spans and marking them with sentinel tokens, while inference prompts the model with prefix and suffix context. The paper compares document- and context-level processing and supports both PSM and SPM formats.

  • Training transformation: Document-level FIM applies the transformation with probability p, splitting each document into randomly sized prefix, middle, and suffix character spans.The main model suite uses p = 0.5, with each span one-third of the document in expectation.
  • Training transformation: The transformed sequence orders the spans as prefix, suffix, and middle, with <PRE>, <SUF>, and <MID> sentinel tokens marking each section.The PSM sequence is <PRE> Enc(prefix) <SUF> Enc(suffix) <MID> Enc(middle).
  • Training objective: FIM training retains loss on prefix, middle, and suffix, preserving the autoregressive learning signal rather than removing supervision from reordered sections.The paper identifies this choice as important for the FIM-for-free property.
  • Inference: At inference, the model receives encoded prefix and suffix sections followed by <MID>, then samples until generating <EOT>.<EOT> signals that the model has connected the prefix and suffix.
  • Inference: Failure to generate <EOT> within the token budget often indicates difficulty connecting the prefix and suffix and can produce lower-quality samples.This motivates EOT-aware best-of-n sampling.
  • Formatting modes: SPM swaps the prefix and suffix order to improve key-value caching because appending to the prefix does not invalidate suffix keys and values.The caching advantage is not universal and may depend on the application.
  • Implementation variants: Context-level FIM applies transformation after chunking and is adopted in the main runs because it can boost performance relative to document-level FIM.Document-level FIM can fragment long documents during chunking, reducing the effective FIM rate.

4 Pretraining results

Pretraining with FIM adds infilling capability without degrading left-to-right performance across model scales, while FIM quality depends strongly on training design choices. Higher FIM rates improve infilling, and context-level FIM and joint SPM+PSM training provide strong practical configurations.

  • FIM-for-free property: 50% FIM augmentation leaves left-to-right test loss and standard autoregressive benchmark performance unchanged across natural-language and code models.The performance matches within error on standard benchmarks.
  • FIM rate: FIM rates up to 90% preserve left-to-right capability, whereas a 100% rate causes degradation in ordinary autoregressive test loss.The rate ablation trains models with FIM rates from 0 to 1.0 for 50B tokens.
  • FIM rate: Higher FIM rates consistently improve infilling pass rates even when FIM perplexity gains are negligible, showing that loss-based evaluation can miss sampling improvements.HumanEval performance remains similar across rates, while random-span infilling improves with higher rates.
  • SPM vs PSM vs joint training: SPM has a slight benchmark advantage over PSM, while joint SPM+PSM training transfers capability across modes and supports flexible inference.Joint pretraining with a 50% FIM rate performs roughly like pure SPM trained at 90% FIM in SPM mode.
  • Context-level vs document-level FIM: Context-level FIM consistently outperforms document-level FIM, despite an almost negligible 0.001 nats/token perplexity difference.Document-level FIM can produce fragmented examples, but these invalid examples do not affect left-to-right evaluation and may be simpler to implement.
  • Middle-span selection: Character-level random spans outperform token-level and line-based spans on random-span infilling because they naturally create subtoken boundary cases.Line-based spans slightly help line-based benchmarks but fail almost completely on random-span infilling.

5 Finetuning results

The paper tests whether existing autoregressive models can acquire FIM through finetuning instead of pretraining. Finetuning generally requires substantial compute and aggressive settings to match FIM-pretrained performance, although smaller runs can yield non-trivial FIM ability.

  • Finetuning existing AR models: Finetuning an autoregressive model pretrained without FIM generally fails to reach the performance of a model pretrained with FIM, even with significant additional compute.The comparison uses models finetuned for 25B or 50B tokens under varied hyperparameters.
  • Finetuning existing AR models: The pretrained-versus-finetuned gap closes only after 50B tokens with a 90% FIM rate and a learning-rate multiplier of 1.0.This was the only setting among 16 finetuned models that closed the gap.
  • Finetuning hyperparameters: Higher learning rates, higher FIM rates, and longer finetuning all appear helpful for improving FIM performance.These settings are associated with better finetuning outcomes in the reported experiments.
  • Finetuning hyperparameters: Matching pretrained performance requires high learning rates and lengthy finetuning, whereas smaller high-rate finetuning runs can still produce non-trivial FIM performance.The paper characterizes the compute needed to match pretraining as large.

6 Discussion

The discussion explains why FIM is harder than ordinary left-to-right generation despite preserving autoregressive token processing, and identifies effective FIM rate and sampling budget as practical constraints. It also connects context-level FIM’s advantage to document packing and highlights the distinct attention pattern introduced by FIM.

  • Pretraining vs finetuning: FIM can be learned for free in pretraining because transformed documents still require predicting the same total number of tokens autoregressively.The transformation effectively creates three smaller documents while preserving the total autoregressive token count.
  • Pretraining vs finetuning: FIM introduces a different document-wide attention pattern, which may help explain why finetuning requires a relatively long token horizon and high learning rate.The passage presents this as a possible explanation involving ossification in learned document-wide attention.
  • FIM loss, AR loss, and the difficulty of FIM task: FIM is harder than ordinary left-to-right generation because it must connect a plausible middle span to both prefix and suffix while ending at the right time.The model must plan a narrative bridge, match the prefix, and generate <EOT> when the middle ends.
  • FIM loss, AR loss, and the difficulty of FIM task: FIM sampling often fails when the model does not produce <EOT> within the allotted budget, yielding truncated samples that connect poorly to the suffix.A larger budget may help, but the appropriate maximum middle length is unclear and a practical limit is often necessary.
  • FIM loss, AR loss, and the difficulty of FIM task: FIM perplexity is consistently higher than AR perplexity across model scales, even when both test sets derive from the same underlying documents.The comparison uses FIM models trained with a 50% FIM rate.
  • FIM loss, AR loss, and the difficulty of FIM task: The FIM middle-span loss is lower when conditioning on both prefix and suffix than when conditioning only on the prefix.The figure attributes this advantage to FIM’s ability to attend to the suffix.
  • Context-level vs document-level FIM and FIM rate: Document-level FIM can lose effective FIM exposure after chunking and packing, especially when documents exceed the context size.The resulting reduction in effective FIM rate is identified as a likely reason context-level FIM performs better.

7 Related work

Related work approaches infilling through masked or flexible-order modeling, new architectures, or autoregressive data formatting. The paper positions data formatting as a simpler and more universal way to add infilling and relates its work to prior sentinel-based autoregressive systems.

  • Masked and flexible-order modeling: Masked language models such as BERT, T5, SpanBERT, and BART treat contiguous masked runs as spans to infill, but primarily target representation learning.The passage contrasts these models with autoregressive infilling approaches.
  • Masked and flexible-order modeling: Flexible-order approaches modify attention or predict insertion locations so text can be generated in orders other than standard left-to-right decoding.Examples include XLNet, Insertion Transformer, KERMIT, and InDIGO.
  • Autoregressive data formatting: Prior autoregressive systems including Donahue et al., GLM, CM3, and InCoder move infill regions to the end of context and separate regions with sentinel tokens.Donahue et al. also study spans at word, sentence, and paragraph granularities.
  • Alternative methods: Infilling has also been attempted with GANs and gradient search, using REINFORCE or optimization of tokens within the infilled span.These methods provide alternatives to both architectural and data-formatting approaches.
  • Architectures versus data formatting: The paper distinguishes architectural infilling methods from data formatting, arguing that introducing new training distributions is generally easier and more universal than hardwiring additional skills.Data formatting is framed as altering model behavior through control codes or transformed data.
  • Positioning and benchmarks: The work reports early research behind code-davinci-002 and compares that system with the paper’s 6.9B models and InCoder 6.7B on infilling benchmarks.The comparison is presented in an appendix.

8 Conclusion

The paper shows that causal decoder models gain FIM capability through joint training while preserving left-to-right performance, and identifies practical training recommendations and open limitations.

  • FIM models achieve lower FIM loss while matching AR models on left-to-right test loss at the same compute.The result is further supported by non-loss-based evaluations.
  • Finetuning pretrained left-to-right models does not fully acquire FIM capability, making joint pretraining from scratch more effective for best FIM performance.
  • FIM capabilities depend on FIM rate and implementation, whereas left-to-right capabilities remain unaffected below a 100% FIM rate.
  • Recommended FIM hyperparameters: Recommended practice is character-level FIM with some random character spans, joint PSM and SPM training, context-level implementation, and a 50%-90% FIM rate.The paper notes that document-level FIM is simpler to implement and that PSM can improve mid-token robustness.
  • Future directions: The authors identify smarter span selection, steerable generation, and broader evaluation of the FIM-for-free property as future directions.They cannot completely rule out underperformance on benchmarks not considered.

A Architecture and datasets

The study uses eight causal transformer decoder models with Codex- and GPT-3-like configurations, trained from scratch on a large Python dataset under a fixed token budget.

  • Architecture: The experiments use 8 causal transformer decoder models with architectures, optimization hyperparameters, and encodings similar to Codex and GPT-3.Relative attention replaces learned positional embeddings, and all models use a context size of 2048.
  • Architecture: The models use relative attention rather than learned positional embeddings, with negligible parameter-count increase and improved performance.
  • Datasets: The code models are trained on Codex's 159 GB Python dataset scraped in May 2020 and initialized randomly from scratch.
  • Training setup: All main-scan models train for 100B tokens regardless of size, leaving the largest models expected to be undertrained under the fixed budget.
  • Architecture: The six largest models follow architectures similar to GPT-3 models Small through 6.7B, with parameter counts reported in Table 3.

B Scaling trends for FIM rate ablations

Increasing the FIM rate improves infilling benchmarks without materially harming left-to-right capability below 100%, while encoding choices create implementation and transfer trade-offs.

  • FIM rate effects: The only noticeable perplexity degradation occurs at a 100% FIM rate; lower rates preserve left-to-right scaling trends.
  • FIM rate effects: Higher FIM rates significantly improve infilling benchmark performance, with no noticeable improvement after 90% FIM.
  • Evaluation: Higher FIM rates do not produce a commensurate reduction in FIM losses, showing that perplexity does not always capture real-world infilling performance.
  • FIM structure: Document-level FIM randomly splits a character sequence into prefix, middle, and suffix before tokenization, while context-level FIM applies transformation after packing and chunking.
  • Implementation: Character-level context FIM requires decoding tokenized documents before augmentation and careful handling of characters that may fail to decode.
  • SPM encoding: The selected SPM variant places prefix and suffix sentinels so SPM data occurs naturally during PSM training, increasing transfer between formats.The alternative SPM variant can be stronger for subtokens at the end of the prefix.

E Random span infilling benchmark

The paper extends HumanEval-based infilling evaluation with random-span tasks that test varied, including mid-token, code-completion scenarios beyond line-based benchmarks.

  • Benchmark motivation: Existing single-line and multi-line HumanEval infilling benchmarks remove non-empty solution lines but do not capture general real-world spans.
  • Benchmark construction: Random span infilling selects the missing middle span between two uniformly random positions in each canonical HumanEval solution.
  • Challenge cases: The benchmark requires models to infer undefined variables, handle subtoken boundaries, and stop completion at the appropriate inline position.
  • Challenge cases: Random-span tasks include cases where the prefix ends in a subtoken and the missing span crosses multiple lines.Such cases can cause indentation or tokenization failures for token-based FIM systems.
  • Evaluation design: The evaluation uses 10 random tasks per HumanEval problem, totaling 1640 tasks to balance variance reduction and sampling time.

F Dynamics and learning curves of finetuning

Finetuning reveals a trade-off between ordinary completion and infilling: HumanEval initially degrades, while random span infilling improves from zero. FIM supports broader code and language completion capabilities, though infilling remains harder than prefix completion.

  • Finetuning dynamics: HumanEval degrades significantly at the beginning of finetuning, especially with higher learning rates, but catches up to similar levels by training's end.The dynamics are reported for FIM rate 0.9; similar plots for FIM rate 0.5 are omitted.
  • Finetuning dynamics: Random span infilling starts at zero as expected and slowly rises during finetuning.
  • Qualitative evaluation: Infilling works better on code than language, while remaining more difficult than ordinary prefix extension.The paper qualitatively evaluates samples to identify strengths, challenges, and possible mitigations.
  • Capabilities: FIM lets one model condition on context before and after generation, supporting code completions that infer imports, names, arguments, docstrings, and definitions.The paper contrasts this with Codex's separate docstring model.
  • Capabilities: FIM also adapts completions to existing writing style and takes the passage ending into consideration.

H.2 Limitations

FIM sampling can fail when the model cannot infer a suitable middle span or decide when to stop, especially when connecting unrelated prefixes and suffixes. Prompting and EOT-aware best-of-n sampling help, but do not eliminate these failures.

  • Difficult prompts: Unrelated suffixes can trigger excessively long middle sections that fail to connect within the limited context size.The model may generate a plausible trajectory joining the ending, but usually cannot complete the join.
  • Deciding when to stop: Even straightforward prompts can make predicting the terminal <EOT> token difficult because valid completions vary in length.Longer candidates reduce the probability of selecting <EOT>.
  • Repetition: Missing <EOT> can cause the model to copy the suffix and repeat it indefinitely, including repeated words.This failure mode remains possible even for large models.
  • Prompting: Numbered few-shot prompts can improve EOT generation and suffix connection by constraining the desired output, but they do not completely fix the problem.The model may still accidentally begin a new list of items.
  • Sampling: EOT-aware best-of-n sampling generates multiple candidates, prioritizes samples ending with <EOT>, and can rerank them by likelihood or other heuristics.This increases the chance of selecting a sample that joins the ending.
Loading 2207.14255v1…