Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Code LLM Training › Fill-in-the-Middle Training

Michael BrenndoerferJanuary 29, 20269 min read

Language AI HandbookAll chapters

Part 47 of 67Code Generation

Chapter 357Code LLM Training › Fill-in-the-Middle Training

Reading section 4 of 13 · Fill-in-the-Middle Training · 9 minSection 4 of 13 · 9 min

Section 4 of 13 · 9 min read

Code LLM Training › Fill-in-the-Middle TrainingOne section · Sections in this chapter4 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The most significant training innovation specific to code models is fill-in-the-middle (FIM) training. It fundamentally changes what completion models can do and is worth understanding in depth, both in terms of the mechanics of how it works and the intuition for why it is so effective.

The Problem with Left-to-Right CompletionLink Copied

Standard autoregressive language models predict the next token given all previous tokens. This makes them excellent at completing code from a prefix: given the beginning of a function, they continue it forward. But software developers spend a large fraction of their time not completing code forward from a cursor, but rather inserting code in the middle of existing code.

Consider these scenarios that require infilling:

Writing the body of a function when the signature and docstring are already present
Filling in a missing function between two complete ones
Implementing a try block when the except clause is already written
Completing a line in the middle of a larger codebase where surrounding context is already established
Filling in argument lists where the function call structure is already partially written

A pure left-to-right model cannot use the suffix (code after the cursor) as context. It is forced to treat completion as if nothing comes after, which produces code that may not connect correctly with what follows. A model generating a function body without knowing what comes after it may choose variable names or return types that are incompatible with the calling code that already exists below the cursor.

This is not a minor usability issue. For an IDE completion assistant, infilling is arguably the dominant use case. Developers rarely write code from scratch in a single forward pass; they more commonly add a function somewhere in the middle of a file, fill in an argument to a function call, or implement a method in a class that already has several methods defined.

Fill-in-the-middle training teaches the model to use both prefix and suffix context when generating the infill region. The key insight is that this capability can be trained into a standard autoregressive model without modifying its architecture, simply by rearranging the order in which tokens are presented during training.

The FIM TransformationLink Copied

FIM works by transforming the document order of training examples. Given a complete source document, three special tokens are introduced:

<PRE> (prefix sentinel): marks the start of the prefix context
<SUF> (suffix sentinel): marks the start of the suffix context
<MID> (middle/infill sentinel): marks where the model should begin generating the middle content

A random split point is chosen within the document, dividing it into prefix, middle, and suffix portions. The document order is then rearranged to present the task as: given prefix and suffix, predict middle. The transformed document looks like:

<PRE> prefix text <SUF> suffix text <MID> middle text

The model is trained autoregressively on this rearranged sequence. By the time the model reaches <MID>, it has attended to both the prefix (before the split point) and the suffix (after the split point) during the forward pass. The middle tokens are then predicted causally given this combined context.

At inference time, the model receives <PRE> {prefix} <SUF> {suffix} <MID> and must generate the middle portion. This is the same format as training, so the model knows exactly what is expected of it.

There are two main variants of the FIM transformation, originally described in the Starcoder paper and related work:

PSM (Prefix-Suffix-Middle): The arrangement shown above. The model sees the prefix first, then the suffix, then generates the middle. PSM has the advantage that the prefix is processed first and may set up important context before the suffix is read.

SPM (Suffix-Prefix-Middle): The suffix is presented first, followed by the prefix, and then the model generates the middle. In SPM format, the document looks like:

<PRE> <SUF> suffix text <MID> prefix text [middle text]

SPM was found to be more effective in some evaluations because it allows the model to immediately attend to the suffix without having first processed a potentially long and distracting prefix. When the suffix is short and highly constraining (for example, return result), presenting it first allows the model to condition everything that follows on this constraint.

In practice, most code models train on a mixture of SPM and PSM examples, with FIM applying to approximately 50 percent of training examples. The remaining 50 percent are trained in standard left-to-right order to preserve the model's completion capability from prefix-only contexts. This mixed training ensures the model does not lose generative quality in the standard completion scenario, which is still common.

Why FIM Works: The Mathematical IntuitionLink Copied

At first glance, FIM might seem like a hack: you are presenting tokens out of their natural order and asking the model to treat it as if this were normal. But the mathematical reason FIM works is straightforward. A standard autoregressive model learns the distribution:

P(t_1, t_2, \ldots, t_n) = \prod_{i=1}^{n} P(t_i \mid t_1, \ldots, t_{i-1})

By reordering the document as <PRE> prefix <SUF> suffix <MID> middle, FIM trains the model to learn a conditional:

P(\text{middle} \mid \text{prefix}, \text{suffix})

This conditional distribution is exactly what you want for infilling. The model learns that when it sees <PRE>...<SUF>...<MID>, the correct action is to generate content that bridges the prefix and suffix. Because the reordered document is still a valid token sequence, the standard autoregressive training loop applies without modification. No architectural changes are needed.

The softness here is that the model does not have true bidirectional attention to the suffix. When generating middle token $t_m$ , the model has attended to the suffix in its prefix context, but it cannot attend to the middle tokens it has already generated in relation to the suffix. This contrasts with masked language models, which use bidirectional attention. In practice, this limitation is largely mitigated by the model learning to "plan" the infill based on the suffix it has already read.

Span-Level vs. Line-Level FIMLink Copied

The original FIM formulation splits at a randomly chosen character position, which creates infill spans of arbitrary size. In practice, the training distribution of span types matters for downstream performance:

Character-level splits: Split at a uniformly random character position. This creates a wide variety of span lengths, from single characters to most of the document. Character-level splits are maximally diverse but produce many awkward examples where the split falls inside a token (for example, splitting calculateMean at the c after cal).

Line-level splits: Split at line boundaries. This is better aligned with how code editors implement completion, where the cursor typically sits at the end of a line or between lines. Line-level splits produce more realistic training examples and are used in most production code models.

Identifier-level splits: Split at identifier or token boundaries so that the prefix and suffix are syntactically complete. This avoids training the model on examples where splits cut through tokens, which are unrealistic completion scenarios.

CodeLlama used a combination of line-level and character-level splitting, finding that including some character-level examples improved robustness to mid-line cursor positions common in IDEs. The intuition is that a model trained only on line-level splits will not have seen the task of completing a partially written identifier, a common situation in real usage.

Repository-Level Context in FIMLink Copied

A significant enhancement to basic FIM is providing repository context. When filling in a function body, the most useful context is often not the immediately surrounding lines but the signatures of related functions in the same file, the types defined in imported modules, or the docstrings of functions that the infill will call.

Repository-level FIM constructs the prefix and suffix from the current file and from a retrieval-augmented view of the repository. A retrieval system identifies the most relevant adjacent files based on import relationships, semantic similarity, or co-occurrence in the version history, and their contents are prepended to the context window.

This was a core innovation in models like Copilot Chat, where the IDE extension retrieves semantically relevant code from the open project and includes it in the context before calling the model. The retrieval is done client-side in the IDE, so the model itself does not need to learn retrieval, it just needs to use the retrieved context effectively.

Worked Example: Tracing a FIM TransformationLink Copied

To make FIM concrete, let's trace through a specific example. Suppose the training document is:

In[3]:

Code

def compute_area(radius):
    """Compute the area of a circle."""
    area = 3.14159 * radius**2
    return area


result = compute_area(5.0)
print(f"Area: {result:.2f}")

Out[3]:

Console

Area: 78.54

A line-level FIM split at line 3 (zero-indexed) produces:

Prefix: def compute_area(radius):\n """Compute the area of a circle."""
Middle: area = 3.14159 * radius ** 2
Suffix: return area\n\nresult = compute_area(5.0)\nprint(f"Area: {result:.2f}")

The PSM-transformed training example becomes:

<PRE>def compute_area(radius):
    """Compute the area of a circle."""<SUF>    return area

result = compute_area(5.0)
print(f"Area: {result:.2f}")<MID>    area = 3.14159 * radius ** 2<EOT>

The model is trained to predict the middle line area = 3.14159 * radius ** 2 given the function signature, docstring, and the return statement plus caller. At inference time, a developer who has already written the function signature, the docstring, and the return statement would receive exactly this infill as a suggestion, which is both natural and useful.