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

Infilling requires conditioning on text both before and after a gap, a capability limited across existing model classes. This paper jointly trains causal decoders on left-to-right and FIM-transformed data, finding that at equal compute they match AR left-to-right test loss while achieving lower FIM loss.

  • Problem

    Existing model classes are limited at infilling, which requires conditioning on both a prefix and suffix and arises naturally in applications such as code completion.

    Source: §1 · p. 3

  • Method

    The paper trains causal decoder-based language models jointly on traditional left-to-right and FIM-transformed data.

    Source: §8 · p. 17

  • Results

    At equal compute, FIM models achieve the same left-to-right test loss as AR models while achieving lower FIM loss; left-to-right capabilities remain unaffected by FIM choices below a 100% rate.

    Source: §8 · p. 17

  • Takeaways & Limitations

    For best FIM performance, pretraining jointly from scratch with recommended hyperparameters is more effective than finetuning.

    Source: §8 · p. 17

    FIM sampling can fail to emit <EOT> within the allotted token budget, producing a bad sample that cannot easily be repaired by post-processing.

    Source: §6 · p. 15

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 addresses the difficulty of infilling—generating text between a prefix and suffix—with causal language models, which ordinarily condition only on the prefix. It proposes a data transformation that teaches infilling while preserving left-to-right capability, and studies the training choices and evaluations needed to make this efficient.

1 Introduction

  • The approach adds FIM capability to causal decoder models through a training-data change, without changing the architecture or compromising normal left-to-right generation.Training mixes FIM-transformed data with ordinary left-to-right data.Source: §1 · pp. 3–4

1.1 Our contributions

  • FIM-trained models achieve lower FIM test loss than baseline autoregressive models, indicating that they learn to condition on the suffix when predicting the middle.The paper identifies suffix conditioning as the capability measured by the FIM test set.Source: Figure 2 · p. 4, §1.1 · p. 5

1 Introduction

  • FIM training transforms some documents by moving a randomly selected middle piece after the prefix and suffix, then concatenating the pieces with sentinel tokens.The paper calls the property of learning FIM while preserving ordinary left-to-right capability “FIM-for-free.”Source: §1 · p. 4

1.1 Our contributions

  • The study ablates FIM rate, transformation variants, and middle-span selection, and reports that learning FIM during finetuning requires significant additional compute to approach pretraining performance.The authors present pretraining FIM as computationally free for left-to-right capability, in contrast to the cost of learning it during finetuning.Source: §1.1 · p. 4
  • The paper introduces random span infilling and random span infilling light benchmarks and finds that sampling evaluations can reveal hyperparameter differences missed by FIM test losses.The benchmarks complement single-line and multi-line code infilling evaluations; the authors argue that sampling evaluation better reflects use cases and should augment test-loss analysis.Source: §1.1 · p. 5
  • Across code and natural language, FIM models from an 8-model scaling study achieved similar left-to-right test loss to non-FIM models, with matching capability also observed throughout training.The study trained models from 50M to 6.9B parameters; matching learning curves were reported for left-to-right loss and HumanEval evaluations.Source: §1.1 · p. 5

2 Evaluation

The evaluation uses autoregressive (AR) and fill-in-the-middle (FIM) benchmarks to assess whether FIM training preserves left-to-right capabilities and how training choices affect infilling. It evaluates both test losses and benchmark performance, with code infilling assessed through correctness tests.

2 Evaluation

  • AR and FIM evaluations assess left-to-right capability retention and infilling behavior, respectively, while test losses and sampling benchmarks track model performance.AR and left-to-right are used interchangeably; losses are cross-entropy in nats per token, and sampling uses nucleus sampling with parameter 0.95 [Holtzman et al., 2020].Source: §2 · p. 5

2.1 Autoregressive evaluation

  • AR evaluation measures canonical-order test losses and standard natural-language benchmarks, plus HumanEval pass rates for code [Chen et al., 2021].Natural-language tasks cover reasoning, reading comprehension, and completion; all except DROP and QuAC use few-shot prompting.Source: §2.1 · pp. 5–6

2.2 Infilling evaluation

  • FIM test sets transform the same examples as AR sets, enabling loss comparisons; masked tests score only middle-span tokens to compare conditioning on prefix-and-suffix versus prefix alone.This measures how much information FIM models gain from conditioning on the suffix.Source: §2.2 · p. 6
  • Generative infilling focuses on code because test suites can check correctness for long, open-ended samples, unlike cloze-style tasks focused on single or few tokens.The suite includes partial function completions from HumanEval canonical solutions, including random middle spans selected uniformly at random.Source: §2.2 · p. 6
  • The infilling suite uses HumanEval solution spans for single-line, multi-line, and random-span tasks, with 1,033, 5,815, and 1,640 tasks, respectively.At least 100–200 samples per task are used for final snapshots; a 164-task random-span-light set tracks training trends. SPM results are reported except when PSM changes the conclusions.Source: §2.2 · p. 6

3 FIM training and inference

The paper trains FIM models by rearranging document spans, with document-level and context-level implementations and an SPM variant. Context-level FIM boosts performance relative to document-level FIM and is used in the main runs.

3 FIM training and inference

  • With FIM rate p = 0.5 in the main suite, document-level FIM randomly splits documents into prefix, middle, and suffix, then orders them prefix, suffix, middle with sentinel tokens.The split occurs before tokenization, with each span one-third of the document in expectation; the resulting format is <PRE> ◦Enc(prefix) ◦<SUF> ◦Enc(suffix) ◦<MID> ◦Enc(middle).Source: §3 · pp. 6–7
  • Training retains loss on prefix, middle, and suffix to preserve the autoregressive learning signal, and <EOT> tokens must be trained to signal a successful join.The authors say sentinel masking does not change the FIM-for-free property; preliminary, unreported experiments suggest retaining loss on all three sections is crucial.Source: §3 · p. 7
  • At inference, the model receives the encoded prefix and suffix in PSM order and generates the infill until <EOT> signals that it has connected them.Failure to emit <EOT> within the allotted token budget often indicates difficulty connecting the spans and poorer samples, motivating EOT-aware best-of-n sampling.Source: §3 · p. 7

3.1 SPM mode

  • SPM swaps prefix and suffix order to improve key-value caching when prefixes are extended, though suffix changes can invalidate the prefix cache.The authors report a slight SPM edge over PSM on infilling benchmarks, while noting that caching superiority depends on the application; main runs train on both modes at 50% probability each.Source: §3.1 · p. 7

3.2 Context-level FIM

  • Context-level FIM can boost performance relative to document-level FIM, and the authors use it in all main FIM runs.It applies the transformation after chunking: documents in a context slice are split at <EOT>, selected for FIM using the FIM rate, rejoined, and trimmed to context length.Source: §3.2 · pp. 7–8

4 Pretraining results

FIM augmentation preserves left-to-right capability, while infilling performance depends on the training rate, transformation format, context scope, and middle-span selection.

4.1 Evaluation of left-to-right capabilities in downstream benchmarks

  • With 50% FIM augmentation, left-to-right test loss stayed unaffected and natural-language and code downstream benchmark performance matched within error.The benchmark evaluation strengthens the test-loss finding across models trained from 50M to 6.9B parameters.Source: Figure 3 · p. 8, §4.1 · p. 9

4.2 FIM rate

  • FIM rates up to 90% preserved left-to-right capability, but 100% degraded ordinary AR test loss; higher rates improved infilling pass rates without commensurate FIM-perplexity gains.HumanEval performance was similar across rates; the rate findings motivated a 6.9B code model trained at 90% FIM for stronger infilling, while 50% remained near optimal.Source: §4.2 · pp. 9–10, Figure 4 · p. 9, Figure 5 · p. 10

4.3 SPM vs PSM vs joint SPM+PSM training

  • SPM evaluation consistently slightly outperformed PSM, while joint 50% FIM training transferred capability across modes and matched pure 90% SPM training in SPM mode.Joint pretraining was more efficient and produced a flexible model supporting both inference modes; SPM’s contiguous prefix and middle may make continuation more natural.Source: §4.3 · pp. 10–11, Table 1 · p. 11

4.4 Context-level vs document-level FIM

  • Context-level FIM consistently and significantly outperformed document-level FIM across scales, despite only a 0.001 nats/token FIM-perplexity improvement.Document-level FIM can create fragmented examples, yet these did not hurt left-to-right evaluation, so its simpler implementation may still be preferable in some settings.Source: Figure 7 · p. 11, Figure 8 · p. 11, §4.4 · p. 12

4.5 Middle span selection

  • Character-level random spans best handled random-span infilling while remaining strong on single- and multi-line tasks; line-based spans aided in-distribution tasks but nearly failed on random spans.Token-level spans did slightly better than line-based spans on random infilling but missed cases where a BPE token crosses a boundary; character-level selection naturally trained on such cases.Source: §4.5 · p. 12, Table 2 · p. 12

5 Finetuning results

Finetuning an autoregressive model pretrained without FIM can produce non-trivial FIM performance, but matching the FIM-pretrained baseline requires substantial additional compute and an aggressive finetuning configuration.

5 Finetuning results

  • Only the 50B-token finetune at a 90% FIM rate and 1.0× pretraining learning rate closed the gap to the 100B-token, 50%-FIM pretrained baseline.Figure 9 compares 16 settings spanning four learning-rate multipliers, two FIM rates, and 25B- or 50B-token horizons.Source: Figure 9 · p. 13, §5 · p. 13
  • Matching the baseline required a large amount of additional compute, contrary to the expectation that FIM could be learned quickly through finetuning.The finetuned models generally did not match the FIM-pretrained model even with significant additional finetuning compute.Source: §5 · p. 13
  • Higher learning rates, higher FIM rates, and longer finetuning each appeared helpful for improving FIM performance.The tested learning-rate multipliers were 0.1, 0.2, 0.5, and 1.0; FIM rates were 0.5 and 0.9; and horizons were 25B and 50B tokens.Source: §5 · pp. 13–14
  • A small amount of finetuning, especially with high FIM and learning rate, was sufficient to achieve non-trivial FIM performance.Reaching performance similar to pretraining nevertheless required high learning rates and lengthy finetuning.Source: §5 · p. 14

6 Discussion

The discussion explains why FIM can be learned for free during pretraining yet remains harder than left-to-right generation. It links that difficulty to finetuning attention changes, suffix-connection failures in sampling, and lower effective FIM rates under document-level packing.

Pretraining vs finetuning

  • Moving the middle span to the end creates three left-to-right prediction pieces without increasing autoregressive tokens, explaining why FIM is free in pretraining despite costly finetuning.The authors suggest that FIM’s changed document-wide attention pattern, and possible ossification of pretrained attention, may make finetuning require a longer token horizon and higher learning rate.Source: §6 · p. 14

FIM loss, AR loss, and the difficulty of FIM task

  • FIM sampling must plan a continuation that fits both prefix and suffix and produce <EOT> on time; missing the token can yield truncated samples that do not connect to the suffix.The needed generation budget is difficult to determine, and trimming that can improve left-to-right samples does not readily fix this FIM failure.Source: §6 · pp. 14–15
  • Across scale, 50%-FIM models have consistently higher FIM than AR perplexity on matched documents, showing that the same document is harder to model in FIM.The comparison uses identical underlying documents for the AR and FIM test sets.Source: Figure 11 · p. 15, §6 · p. 15

Context-level vs document-level FIM and FIM rate

  • Document-level FIM has a lower effective rate than context-level FIM because chunking and packing can separate a document’s prefix, middle, and suffix across contexts.The authors identify this lower effective rate as the likely main reason context-level FIM performs more strongly, even when the nominal rates match.Source: §6 · pp. 15–16

7 Related work

Prior infilling work spans masked-language and flexible-order autoregressive models, as well as left-to-right models trained with infills moved to the context end. This paper emphasizes efficient data formatting as a practical way to add infilling while retaining left-to-right modeling.

7 Related work

  • Masked language models including T5 [Raffel et al., 2019], SpanBERT, and BART [Lewis et al., 2020] use contiguous masked spans, typically shorter than sentences or code lines.BERT [Devlin et al., 2019] initially masked tokens randomly; DOBF trains BERT on code, while HTLM [Aghajanyan et al., 2021] trains BART on HTML.Source: §7 · p. 16
  • Other autoregressive approaches enable flexible generation order: XLNet [Yang et al., 2019] changes the attention mask, while Insertion Transformer [Stern et al., 2019], KERMIT, InDIGO [Gu et al., 2019], and Blank Language models [Shen et al., 2020] generate by selecting locations or blanks.Source: §7 · p. 16
  • Zhu et al. [2019], Donahue et al. [2020], GLM [Du et al., 2022], CM3 [Aghajanyan et al., 2022], and InCoder similarly move infills to the context end.This paper focuses on practical single-span infilling, computational efficiency at scale, and character-level span selection for robustness; InCoder uses sampling benchmarks created from HumanEval [Chen et al., 2021]. The paper describes early research behind code-davinci-002 and compares it with its 6.9B models and InCoder 6.7B in Appendix 4.Source: §7 · pp. 16–17
  • Alternative infilling methods use GANs [Fedus et al., 2018] with REINFORCE for discrete text, or gradient search [Liu et al., 2019] to optimize span tokens and collapse them to nearest neighbors.Source: §7 · p. 16
  • Infilling can be added through new architectures or data formatting; the paper argues that introducing new training distributions is generally easier and more universal than hardwiring skills.Data formatting alters language-model behavior through control codes, as motivated by CTRL [Keskar et al., 2019]; DistAug jointly trains on transformed data conditioned on transformation type.Source: §7 · p. 16

8 Conclusion

The conclusion presents joint FIM training as adding infilling capability while preserving left-to-right performance in the reported evaluations, recommends training settings, and outlines open research directions.

8 Conclusion

  • Jointly trained causal decoder models can infill while one FIM model handles importing modules, writing docstrings, and completing functions, exceeding AR models within the reported evaluations.The single model subsumes specialized models finetuned for individual tasks [Chen et al., 2021].Source: §8 · p. 17
  • At equal compute, FIM models match AR test loss on left-to-right evaluation while achieving lower FIM loss, with non-loss evaluations further supporting the FIM-for-free finding.The authors note that the property is supported but not proven across benchmarks they did not evaluate.Source: §8 · p. 17, §8.2 · p. 18
  • Careful finetuning does not give a pretrained left-to-right model the full FIM capability of its model size, making joint pretraining from scratch more effective for best FIM performance.This held despite significant finetuning compute relative to pretraining.Source: §8 · p. 17
  • In the reported studies, perplexity does not reflect true infilling performance; FIM rate and implementation affect FIM ability, while left-to-right ability remains unaffected below a 100% FIM rate.Character-level FIM also improves robustness to subtoken boundaries and supports deployment such as coding assistance.Source: §8 · p. 17

8.1 Recommended FIM hyperparameters

  • Apply FIM at the character level, include character-level random spans, and jointly pretrain with PSM and SPM; context-level FIM is preferred, and 50%–90% FIM rates are reasonable.Inference in PSM mode may be more robust to mid-token boundaries than the explored SPM mode, while document-level FIM offers a simpler implementation.Source: §8.1 · p. 17

8.2 Future directions

  • Future work includes smarter span selection, steerable generation, further testing of FIM-for-free, multi-slot infilling, better natural-language infilling, and systematic study of attention and architectures.The authors suggest meaningful or shorter spans may help, and note that free-form generation is harder to evaluate than code functional correctness.Source: §8.2 · pp. 17–18
  • For research on jointly learned capabilities, set a budget for acceptable loss of original capability, then maximize the new capability within that budget.The authors argue this trade-off framework can support practical decisions about integrating capabilities.Source: §8.2 · p. 18

A Architecture and datasets

The study uses eight causal decoder models with a relative-attention modification and trains code and natural-language models on established datasets. Its architecture follows GPT-3 designs, with training choices varying by model size.

A Architecture and datasets

  • Eight causal transformer decoders use relative attention instead of learned positional embeddings, with 2048-token contexts and doubled learning rates for the three largest models.The authors report that relative attention negligibly increases parameter count and improves performance; the models otherwise have similar architecture and optimization hyperparameters to Codex and GPT-3.Source: A Architecture and datasets · p. 23
  • Code models train from scratch on Codex’s 159 GB Python dataset, while natural-language models use the GPT-3 dataset; main-scan models train for 100B tokens.The code data was scraped in May 2020, preceding HumanEval’s public release; the fixed token budget may leave the largest models undertrained and able to benefit from longer training.Source: A Architecture and datasets · p. 23
  • The architecture table covers model parameter counts and non-embedding counts; scaling plots use the latter, and embedding and unembedding weights are untied.The six largest models follow architectures similar to GPT-3’s Small-to-6.7B models.Source: Table 3 · p. 23

B Scaling trends for FIM rate ablations · C Details of FIM implementation · D Details of SPM encoding

The appendix reports how FIM rates affect model evaluations and describes practical choices for implementing FIM and encoding SPM. It recommends an SPM encoding that supports transfer between PSM and SPM training, while noting an alternative’s advantage for prefix subtokens.

B Scaling trends for FIM rate ablations

C Details of FIM implementation

  • For document-level FIM before packing, choose two random split positions and format the prefix, suffix, and middle with their corresponding sentinels.Character-level and token-level transformations differ only in whether splitting happens before or after encoding.Source: C Details of FIM implementation · pp. 24–25
  • For context-level FIM, token-level augmentation is straightforward after splitting contexts at <EOT>; character-level augmentation requires decoding and checking for broken UTF-8 fragments, then trimming or padding to context length.BPE chunking can leave UTF-8 characters as token fragments that fail to decode.Source: C Details of FIM implementation · p. 25

D Details of SPM encoding

  • SPM variant 1 orders sentinels and content as <SUF> suffix <PRE> prefix <MID> middle <EOT>, whereas variant 2 uses <PRE> <SUF> suffix <MID> prefix middle <EOT>.The paper uses variant 2.Source: D Details of SPM encoding · p. 25
  • Variant 2 supports PSM–SPM transfer because empty-prefix PSM examples naturally match it; variant 1 may handle prefix-ending subtokens better and may suit pure SPM, but the paper chooses variant 2 for joint training.Pure PSM runs achieve strong SPM-mode performance with variant 2.Source: D Details of SPM encoding · p. 25

E Random span infilling benchmark

The section introduces a HumanEval infilling benchmark that selects spans between random positions in canonical solutions, extending line-based benchmarks to more varied code-editing cases. It uses 10 random tasks per problem, for 1,640 tasks, balancing evaluation variance against sampling time.

E Random span infilling benchmark

  • The benchmark selects middle spans between two random positions in HumanEval canonical solutions, addressing the limitation of line-based single- and multi-line benchmarks [2022].It uses 10 random tasks per problem, yielding 1,640 tasks as a balance between evaluation variance and sampling time.Source: E Random span infilling benchmark · pp. 25–26
  • Random spans capture varied practical corner cases, including multi-line infills ending mid-subtoken that can break BPE- and token-based FIM approaches.One example requires recognizing an undefined variable, handling a prefix ending in a subtoken without causing an indentation error, and stopping completion in-line.Source: E Random span infilling benchmark · p. 26

F Dynamics and learning curves of finetuning … H.1 Successful infilling examples

Finetuning gradually enables random-span infilling while ordinary HumanEval recovers from an early decline. Qualitative examples show FIM using both preceding and following context for code and natural-language completion.

F Dynamics and learning curves of finetuning

G Top models comparison

  • Table 4 compares 6.9B FIM models trained at 50% and 90% rates with similar-size InCoder and code-davinci-002 across three infilling benchmarks; InCoder results are self-reported.The authors note that implementation differences may cause slight discrepancies between evaluation frameworks.Source: Table 4 · p. 27, G Top models comparison · p. 27

H Qualitative evaluation

  • Qualitative evaluation finds infilling works better for code than language, while remaining harder than extending a prefix; the authors show challenges and possible mitigations.Source: H Qualitative evaluation · p. 28

H.1 Successful infilling examples

  • FIM uses both preceding and following context to infer code components in one model, including imports, function names, arguments, docstrings, and definitions, unlike Codex [Chen et al., 2021].The example uses a prefix and suffix containing subtokens, which traditional language models without methods such as stochastic BPE [Provilkov et al., 2019] are known to mishandle.Source: §H.1 · p. 28
  • FIM also adapts to an existing writing style and completes a passage naturally while taking its ending into account.The example completes a Wikipedia passage.Source: §H.1 · p. 28

H.2 Limitations · H.3 Mitigations

FIM can struggle to connect unrelated suffixes, stop at the right point, and avoid repetition; prompt cues and selecting among multiple EOT-ending samples can mitigate, but not eliminate, these failures.

H.2 Limitations

  • Unrelated suffixes can prompt very long, often unsuccessful middle spans, and even plausible connections vary in quality.The model attempts to construct a trajectory joining the ending, but limited context often prevents it from succeeding.Source: §H.2 · p. 28
  • Even straightforward infills may continue too long because competing completion lengths reduce EOT probability, and sampling can miss the terminal token.The model may generate valid but spurious content or write its own ending while ignoring the supplied suffix.Source: §H.2 · p. 29
  • Failure to emit EOT can make the model copy the suffix and repeat the prompt indefinitely, including in large models.The passage’s example ends with “the the heart” as generation continues inside the missing span.Source: §H.2 · p. 29

H.3 Mitigations

  • Hints that constrain the expected middle can improve EOT generation and help the model connect to the suffix within a reasonable token budget.The general difficulty of deciding when to stop also affects left-to-right completion, but is less problematic there because no suffix must be joined.Source: §H.3 · p. 29, §H.2 · p. 29
  • Numbered examples at both ends help track position and discourage spurious content, while leading examples without cues can worsen ambiguity about joining the suffix.Numbered few-shot prompting helps considerably but does not fully prevent the model from starting a new list.Source: §H.3 · pp. 29–30
  • EOT-aware best-of-n sampling recommends generating multiple samples and preferring those ending with EOT to increase the chance of a sample that joins the ending.When multiple samples end in EOT, they can be reranked by likelihood or other heuristics.Source: §H.3 · p. 30

Related papers

Loading 2207.14255v1…