Function-Aware Fill-in-the-Middle as Mid-Training for Coding Agent Foundation Models

Yubo Wang, Jiarong Liang, Yuxuan Zhang, Xuye Liu, Cong Wei, Yuyu Zhang, Ping Nie, Wenhu Chen

Function-aware fill-in-the-middle mid-training aligns base models with agentic reasoning structures.

How can we improve coding agents by training them to better handle the "action-observation-continuation" loop inherent in function calls?

Coding agents struggle to integrate external tool returns into their reasoning because standard left-to-right pretraining fails to expose the structural dependencies of an action-observation-continuation loop. The authors introduce function-aware fill-in-the-middle (FIM) mid-training, which masks function call sites selected via program dependency graph analysis and embeds chain-of-thought rationales within the masked span to force the model to learn bidirectional reasoning. This mid-training stage consistently improves SWE-Bench performance across multiple base models and post-training pipelines while simultaneously mitigating the capability erosion typically caused by agentic fine-tuning.

Paper Primer

The core mechanism treats a function call site as an internet-scale proxy for an agent's interaction loop: the pre-call context maps to history, the call to an action, the return value to an observation, and the downstream code to the continuation. By masking these function sites and requiring the model to reconstruct them using a chain-of-thought rationale, the model learns to reason about dependencies that are otherwise obscured in standard training.

Mid-training provides consistent, robust gains on agentic coding benchmarks.

On SWE-Bench-Verified, the method yields +2.8 to +3.0 point gains for Qwen2.5-Coder (7B/14B) and +3.2 points for Qwen3-8B. Consistent improvement across three distinct base models and two different post-training pipelines (R2E-Gym, SWE-Smith).

Mid-training mitigates the "capability erosion" typical of agentic post-training.

Post-training alone regresses non-agent benchmarks (e.g., LiveCodeBench, BFCL); adding mid-training restores performance, with LiveCodeBench recovering +11.1 points. The structural prior installed during mid-training survives post-training, enabling cross-domain transfer even though the mid-training corpus contains only Python code.

Why is this mid-training stage necessary if standard FIM is already used in pretraining?

Standard FIM uses arbitrary span boundaries that lack function-level dependency signal and lack reasoning supervision. This method targets specific function call sites and embeds chain-of-thought rationales to explicitly align the model with the think-then-act structure required by agents.

Does this method require a frontier model to generate the training data?

The default recipe uses Gemini-3-Flash to generate rationales, but ablations show that self-generated rationales (self-CoT) recover most of the performance gains, indicating the method is not merely a distillation pipeline.

The Agentic Tool-Use Gap

Standard left‑to‑right code pretraining misses the tool‑use loop that coding agents need.

Coding agents need to consume a tool’s output and then continue reasoning, but the dominant pretraining paradigm—predicting the next token from left to right—only ever sees the forward direction of code.

When an agent calls a tool, it first decides what to ask (action), receives the tool’s answer (observation), and then writes code that uses that answer (continuation); standard code pretraining never sees the observation step.

The core problem is the mismatch between standard code pretraining and the conditioning structure required for agentic tool‑use.

Limitations of Random-Span FIM

We outline why Random‑Span FIM misaligns with agentic reasoning and present a function‑aware FIM mid‑training pipeline that consistently improves coding agents.

Random‑Span FIM is widely used in code‑LLM pretraining, yet it fails to provide the structured reasoning needed for agentic tool use.

It masks a randomly chosen contiguous token span and asks the model to generate the missing code directly, without any intermediate reasoning step.

We address these three shortcomings by (i) selecting mask targets at function granularity using a program‑dependency‑graph analysis, (ii) embedding chain‑of‑thought rationales inside the masked span, and (iii) applying the objective at a dedicated mid‑training stage right before agentic post‑training.

Evaluation spans three robustness axes: (1) two model sizes (7 B and 14 B) on the SWE‑Bench‑Verified metric, (2) two post‑training pipelines (R2E‑Gym and SWE‑Smith), and (3) transfer to an alternative base model (Qwen3‑8B with SWE‑Lego).

Beyond the verification metric, adding mid‑training before standard post‑training restores large drops in other capabilities: LiveCodeBench improves by +11.1, BFCL by +2.4, and $\tau$‑bench by +3.9.

Function-Aware FIM Mechanism

Method introduces Function‑Aware FIM to align coding‑agent steps with function calls, selecting mask targets via structural scores.

Random‑span FIM only occasionally masks code that resembles an agent step, leaving a systematic gap: the model never learns to predict a function body after observing its call‑site return. By treating function calls as the four‑stage loop of a coding agent, we can deliberately target those points.

Think of a programmer debugging: they read surrounding code (context), invoke a helper function (action), inspect its return value (observation), then edit the surrounding code based on that result (continuation).

How does this differ from a naïve sequential execution of code?

In a naïve sequence the model only predicts the next token; it never sees the intermediate return as a separate observation that can influence subsequent predictions. The Agentic Loop explicitly separates the call (action) from its result (observation) before the next reasoning step.

Imagine a teacher picking a quiz question that is both challenging (complex) and answerable from the lecture notes (inferable). The teacher scores each function on complexity and inferability, then selects those that score high on both.

Why is Function‑Aware FIM better than the random‑span variant?

Random‑span FIM masks spans without regard to code structure, so many masked spans are either trivial or impossible to infer. Function‑Aware FIM deliberately selects functions that are both non‑trivial (high $\hat{H}$) and recoverable from surrounding code (high $\hat{I}$), giving the model a meaningful reasoning target.

Complexity $\hat{H}(\text{add}) = w_{\ell}\phi(2,c_{\ell}) + w_{c}\phi(1,c_{c}) + w_{d}\phi(1,c_{d}) = 0.4$ (example weights).

Complexity $\hat{H}(\text{total}) = w_{\ell}\phi(4,c_{\ell}) + w_{c}\phi(2,c_{c}) + w_{d}\phi(2,c_{d}) = 0.78$.

Inferability $\hat{I}(\text{add}) = 0.45$ (three cues weighted).

Inferability $\hat{I}(\text{total}) = 0.92$ (all five cues weighted).

FIM $\text{add}= \frac{0.4\times0.45}{0.4+0.45+\epsilon}\approx0.22$; FIM $\text{total}= \frac{0.78\times0.92}{0.78+0.92+\epsilon}\approx0.41$.

Threshold $\tau=0.08$ passes both, but

The harmonic‑mean penalizes a function that is easy to infer but trivial to write (low $\hat{H}$), ensuring the model practices non‑trivial reasoning.

Parse each source file into an AST and build the program‑dependency graph (call and sibling edges).

For every function node $v$, compute $\hat{H}(v)$ using lines‑of‑code, cyclomatic complexity, and nesting depth.

Compute $\hat{I}(v)$ by aggregating the five context signals defined in Eq. 2.

Combine $\hat{H}(v)$ and $\hat{I}(v)$ with the harmonic‑mean formula, apply the difficulty penalty $\rho$, and compare against threshold $\tau$.

Select the highest‑scoring single functions or small groups (size 2–3) as FIM mask targets.

Generate a CoT rationale and candidate body for each target, filter by a judge model, and format as an FIM middle span.

**Figure 1.** **Left:** A function call site and a single step of a coding agent are structurally similar, decomposing into the same four stages: context, call/action, return/observation, continuation. **Middle:** We exploit this analogy via function-aware FIM mid-training. A function $B$ is selected from the program dependency graph using complexity ($\hat{H}$) and inferability ($\hat{I}$) scores; the model is then mid-trained to fill in $B$'s body together with a CoT rationale, given the surrounding file as an FIM-formatted prompt. **Right:** Mid-training yields consistent gains across both Qwen2.5-Coder-Instruct (7B, 14B) and Qwen3 (8B) on SWE-Bench-Verified (solid bars) and SWE-Bench-Lite (hatched bars).

**Figure 2.** Function-aware FIM target selection on a small calculator example. (a) Program dependency graph parsed from the AST: solid arrows are call edges $\mathcal{E}_{call}$, dashed lines are sibling edges $\mathcal{E}_{sib}$ between same-class methods. (b) Stacked bars decompose the complexity score $\hat{H}=0.40$ (Eq. 1; LoC, CC, depth) and the inferability score $\hat{I}=0.48$ (Eq. 2; five context signals) for Calculator.total, yielding FIM $\approx 0.22 \ge \tau = 0.08$ (a hyperparameter; see Appendix B for the full list).

Agentic Performance Gains

FIM Mid-Training delivers consistent multi‑percent gains on SWE benchmarks across models and pipelines.

FIM Mid‑Training improves SWE‑Bench‑Verified by +3.00 % on the 14B‑Instruct model and by +2.80 % on the 7B‑Instruct model.

Table 1 shows the +2.80 % and +3.00 % lifts when adding FIM‑Midtrain to the R2E‑Gym pipeline.

R2E‑Gym fine‑tunes a pretrained coding model on a curated set of real‑world repository tasks, aligning its behavior with tool‑use scenarios.

How does R2E‑Gym differ from ordinary fine‑tuning on code?

Ordinary fine‑tuning optimizes for next‑token prediction on static code snippets, while R2E‑Gym explicitly presents the model with a tool‑return and asks it to continue reasoning, forcing a learnable action‑observation‑continuation loop.

**Table 2.**

**Table 2.** Capability preservation and cross-domain transfer at 14B with R2E-Gym. Bold marks the better trained variant per column (post-only vs. ours); the Instruct row is shown as a reference ceiling. All cells are percentages. “Terminal” denotes Terminal-Bench 2.0.

General Capability Preservation

Mid‑training with Function‑Aware FIM restores lost coding capabilities while preserving SWE‑Bench gains.

Mid‑training with Function‑Aware FIM closes most of the capability gap introduced by agentic post‑training.

It raises the six‑benchmark average from 16.04 to 19.56, a +3.52‑point gain, offsetting the 4.81‑point loss of post‑training alone.

Mid‑training restores the most affected benchmarks: LiveCodeBench improves by +11.10 points, OJBench by +1.94 (within 0.46 of the Instruct ceiling), and FullStackBench‑EN by +0.53.

**Table 3.** Ablations on the 7B model with R2E-Gym post-training. (A) rationale source. (B) function-selection algorithm. (C) mask granularity (single vs. multi-function groups). All blocks share the baseline (w/o mid-train) and a controlled 200K-target budget; CoT is fixed to Gemini-3-Flash in (B) and (C). Bold marks the best configuration per block; the 80%/15%/5% mixture in (C) is the recipe used in our main results (Table 1), which is trained on the full corpus rather than the 200K budget here. Absolute numbers across all blocks therefore lie below the main-table recipe; the relative orderings within each block, rather than the absolute values, are the object of comparison.

Even without any tool‑use trajectories in the mid‑training corpus, $\tau$‑bench gains +3.90 and BFCL gains +2.40, and Terminal‑Bench 2.0 recovers +1.25, demonstrating that the structural prior learned at mid‑training transfers to unrelated tool‑use domains.

Ablation Analysis

Ablations isolate each component’s impact and analysis shows where gains arise.

We run three controlled ablations on Qwen2.5‑Coder‑7B‑Instruct with R2E‑Gym post‑training, keeping a shared baseline that skips any mid‑training modifications.

Removing the chain‑of‑thought rationale drops the average score by 1.18 points, confirming that the function‑aware FIM signal works even without reasoning supervision.

Baseline with full pipeline vs. no‑CoT variant differ by +1.18 points.

The full function‑selection pipeline improves the average score by 1.65 points over random masking.

ComparisonGrid shows 13.95 → 15.60.

Using the 80 %/15 %/5 % mask mix raises the average by 0.5 points relative to pair‑only masking.

Mask granularity grid: 15.60 → 16.10.

Mid‑training raises the recovery rate from 24.8 % to 28.8 %, a +4.0‑pp gain.

Table 4 headline metrics.

On multi‑function patches the pass rate climbs from 13.6 % to 22.7 %, a +9.1‑pp improvement.

Figure 5 grouped bar chart.

The full recipe solves 15 more tasks than the baseline.

Outcome distribution in Figure 6 shows a +15 solved increase.

**Table 4.** Headline trajectory metrics on SWE-Bench-Verified. Full metrics in Appendix D.

**Figure.** Pass rate by gold-patch shape (SWE-Bench-Verified)

**Figure 6.** Outcome distribution per evaluation run on SWE-Bench-Verified (14B, R2E-Gym), averaged over three runs.

Related Work

Survey of prior work on mid‑training, code‑masking objectives, and coding‑agent pipelines.

Mid‑training has emerged as a distinct stage placed after the bulk of pretraining but before final task‑specific fine‑tuning, allowing models to absorb specialized inductive biases that generic web text does not provide.

Recent large‑scale models such as MiniCPM, OLMo, DeepSeek‑V3, and Code‑Llama schedule this stage near the end of pretraining, often to improve context‑length generalization or to inject domain‑specific data.

Fill‑in‑the‑Middle (FIM) is now a standard pretraining objective for code LLMs; extensions such as AST‑T5 and AST‑FIM mask abstract‑syntax‑tree subtrees, Horizon‑Length Prediction adds a planning horizon, and Instruction‑aware FIM appends a developer‑comment slot.

Repository‑level retrieval methods like GraphCoder and DRACO exploit program dependencies at inference time, but they do not alter the pretraining objective.

Coding‑agent foundations have progressed along three dimensions: benchmarks have expanded from SWE‑Bench to Multi‑SWE‑Bench and SWE‑Bench‑Pro; scaffolds such as SWE‑agent, OpenHands, and Agentless expose diverse action interfaces; and trajectory‑centric pipelines (SWE‑Gym, R2E‑Gym, SWE‑Smith, SWE‑Lego, Skywork‑SWE) curate full agent trajectories for post‑training.

Our work builds on this lineage by inserting a function‑aware FIM stage that selects whole functions via program‑dependency‑graph analysis, embeds an explicit chain‑of‑thought inside the masked region, and runs as a dedicated mid‑training phase.

Distillation from frontier models further enriches the signal: Gemini‑3‑Flash supplies chain‑of‑thought rationales that are distilled alongside the code, and ablations show the approach retains most of its gain even without a proprietary teacher.

Algorithmic Implementation

Full algorithmic pipeline and scoring formulas for single‑function FIM target selection.

This appendix expands the pipeline sketched in Section 2.3, giving the exact Python‑style algorithm and all default hyper‑parameters used for single‑function FIM target selection.

Program Dependency Graph (PDG) captures call relationships between functions as directed edges $uarrow v$ in the set $E_{\text{call}}$.

Complexity score $\hat{H}(v)$ combines three normalized code‑complexity measures—lines of code, cyclomatic complexity, and nesting depth—each capped and weighted.

Inferability score $\hat{I}(v)$ aggregates five signals that measure how well the surrounding code can reveal the target function.

Difficulty penalty $\rho(\Delta)$ down‑weights functions whose complexity exceeds inferability, using a Gaussian decay beyond a threshold $\tau_{d}$.

Hard filters discard functions outside length windows $[L_{\min},L_{\max}]$, $[\ell_{\min},\ell_{\max}]$, dunder methods, low‑complexity ($\hat{H}<\tau_{\hat{H}}$) or low‑FIM ($FIM<\tau_{\text{FIM}}$) candidates.

Algorithm 1 – single‑function FIM target selection pipeline (simplified)

Donut chart showing the distribution of software licenses.

**Table 6.** Numerical constants inside the $\hat{I}$ components.

Multi-Function Scoring Details

Technical details for scoring and selecting multi‑function groups.

The multi‑function extension of Function‑Aware FIM scores a set of functions $G$ so that the training pipeline can mask them jointly while preserving useful context.

Group inferability $\hat I(G)$ is recomputed assuming every member of $G$ is masked simultaneously: all intra‑group call, callee, and sibling contributions are removed, docstrings are zeroed, and the ``__init__`` bonus is granted only if ``__init__`` lies outside $G$. Signatures remain because they are never masked.

The topology taxonomy enumerates eight connected‑subgraph patterns that a candidate group may follow: for $k\!=\!2$ we have caller‑callee, co‑callee, sibling‑coupled, and mutual‑call; for $k\!=\!3$ we have call‑chain, hub, fan‑in, and class‑triad.

From all enumerated groups we keep only those that are large enough, sufficiently coupled, and pass size‑specific score floors, then greedily pick the highest‑scoring non‑overlapping groups per file.

Algorithm 2 – Multi‑function group FIM target selection for one file.

Coupling: one call edge ⇒ $|E_{\text{call}}^{G}|/(k(k-1)) = 1/2$. With $(w_c,w_s,w_{st})=(0.50,0.20,0.30)$ we get $\text{Coup}(G)=0.25$.

Complexity average: $\hat H(\text{normalize})\approx0.36$, $\hat H(\text{word\_freq})\approx0.30$ ⇒ $\hat H(G)\approx0.33$.

Joint‑mask inferability: after masking both bodies, only signatures remain, yielding $\hat I(G)\approx0.18$.

Harmonic blend: $(0.33\times0.18)/(0.33+0.18)\approx0.116$.

Final score: $0.25 \times 0.116 \approx 0.029$, below the pair floor $0.04$ → group rejected.

Joint masking dramatically reduces inferability for tiny, tightly coupled pairs; only larger, class‑method groups retain enough $\hat I$ to pass the threshold.

In the full corpus, selected groups are typically class‑method pairs or small method clusters where shared‑state edges and Jaccard similarity boost $\text{Coup}(G)$, and signatures plus class‑level signals keep $\hat I(G)$ high enough to survive joint masking.

Chain‑of‑Thought generation prompt (Appendix B.9).

Completion‑quality filtering prompt (Appendix B.10).

Mid‑training sample formatting template (Appendix B.11).

Negative‑Observation patterns (Appendix B.12) are simple string matches used to flag tool outputs that indicate failure: Python stack‑trace prefixes, common exception names, shell error markers, and a few harness‑specific failure strings. Matching is case‑insensitive.

Extended Behavioral Analysis

Function‑aware FIM mid‑training markedly improves pass rates and reduces empty‑patch failures on SWE‑Bench‑Verified.

Function‑aware FIM mid‑training boosts pass rate on multi‑function single‑file tasks by 9.1 percentage points.

D.1 reports a +9.1 pp gain (n = 88) compared to the baseline.

Mid‑training shifts the agent’s action mix: search actions fall from 15.3 % to 11.0 % while `execute_bash` actions rise from 19.5 % to 24.6 %. Consequently, average steps per solved task increase from 15.1 to 22.4 and edits per solved task from 3.3 to 5.4, with errors encountered rising from 3.5 to 5.8.

The recovery rate on solved trajectories climbs to 24.8 % (unsolved 28.8 %) under FIM‑midtrain, while overall localization correctness improves modestly from 70.6 % to 73.4 %. Prompt context length at success more than doubles, reaching 26.2 tokens (solved) versus 11,826 tokens in the baseline (unsolved 29.2 vs 16,397).

Both checkpoints solve roughly 11.3 % of multi‑file tasks, indicating no differential benefit from function‑aware FIM at the cross‑file level. Mid‑training also eliminates most no‑patch failures, cutting them from about 11 trajectories per run to roughly one.

The underlying mechanism is the FIM signal: during training the model is always conditioned to emit a non‑empty token span between the surrounding context, which persists through the post‑training pipeline and forces the agent to produce at least one edit instead of emitting alone.

**Table 9.** Trajectory-level behavioral metrics on SWE-Bench-Verified (14B, R2E-Gym), run-means over three evaluation runs of each checkpoint. Steps, edits, and errors-encountered are means per task. “Empty patch” is the fraction of failed trajectories whose final `output_patch` is empty; “Loc. correct” is the fraction of all trajectories that locate at least one file overlapping the gold patch; “Ctx @ success” is the mean prompt context length on solved trajectories.

Failure Mode Analysis

Limits of mid‑training: how it reshapes agent stopping and error patterns.

Recall that standard left‑to‑right code pretraining cannot incorporate tool returns, so the paper injects Function‑Aware FIM during mid‑training to close the action‑observation‑continuation loop.

D.7 A Concrete Contrast shows the baseline aborting after a single step on scikit‑learn‑26323, emitting with an empty edit and a wrong target file, while our mid‑trained agent executes ~41 steps, applies ~20 `str_replace` operations, observes negative tool returns, and ultimately produces a passing patch.

**Table.** Failure-mode counts on Lite move in the same direction as on Verified: the no-patch mode is eliminated (~8 → 0 trajectories per run), localization errors fall by about 11, and patch errors rise by about 7, reflecting the same iterate-and-verify shift.

E Behavioral Analysis on SWE‑Bench‑Lite replicates the Verified trends: mid‑training lifts recovery from 18.8 % to 22.1 %, eliminates the no‑patch failure mode (≈8 → 0 per run), and increases edit iterations from 4.4 to 7.8 per task.

**Figure 7.** Pass rate on SWE-Bench-Lite stratified by gold-patch shape, averaged over three evaluation runs per checkpoint. Lite contains no multi-file tasks; the multi-function single-file bucket is small (n=54) and shows no gain on this slice, with the +4.0 pp end-task improvement coming entirely from the single-function single-file bucket (n=246, +4.9 pp).

Overall, the mid‑training regime does not endow the agent with a brand‑new capability; rather, it refines the stopping policy, allowing the model to persist through negative observations and converge on correct edits.

Corpus Details

Provides repository categories, license breakdown, and corpus statistics.

The corpus comprises 968 public GitHub repositories, whose licenses are overwhelmingly permissive (MIT, Apache‑2.0, BSD‑3) with the remainder split among copyleft and Creative Commons families. All licenses allow at least non‑commercial research use, and the few minor categories are grouped under “Other research‑permissive licenses.” The per‑repository license file is released alongside the repository list.

Figure 3 visualizes the repository distribution across ten topic categories. The bulk of the corpus lies in reference implementations, scientific computing, and small frameworks, while compiler and networking/security categories form thin tails kept intentionally for diversity. This balance ensures coverage of both common and niche code bases.

**Figure 3.** Distribution of the 968 source repositories across ten topic categories. The corpus is dominated by reference implementations, scientific computing, and small frameworks; compiler and networking/security tails are kept by design to maintain coverage diversity.

Table 5 aggregates the key corpus statistics after decontamination and quality filtering. The dataset contains roughly 78 K self‑contained Python files, yielding about 400 K Fill‑in‑the‑Middle samples and a mid‑training token budget of 2.6 B tokens. Single‑function targets account for ~320 K samples, while multi‑function targets with k=2 and k=3 contribute ~60 K and ~20 K samples respectively, with an average target length of 34 lines of code and 100 % of targets accompanied by Gemini‑3 chain‑of‑thought annotations.

Training Hyperparameters

Lists all hyperparameter settings and compute resources for training and post‑training pipelines.

All experiments use AdamW with bf16 mixed precision and a cosine learning‑rate schedule; LLAMAFACTORY handles FIM mid‑training, while R2E‑Gym, SWE‑Lego, and TORCHTUNE handle the respective post‑training pipelines.

This table lists the hyperparameters used for training, including the optimizer (AdamW), learning rate ($1.0 \times 10^{-5}$), LR schedule (Cosine), warmup ratio (0.1), weight decay (0.05), epochs (1), per-device batch size (1), gradient accumulation (16), effective batch size (128), sequence length (32,768), and precision (bf16).

**Table 8.** Agentic post-training hyperparameters for the three pipelines. R2E-Gym and SWE-Smith follow their official released scripts; SWE-Lego follows the official recipe except for the epoch count.

Effective (global) batch sizes assume eight GPUs; for example, a per‑device batch size of 1 with gradient accumulation of 16 yields a global batch of 128 in Table 7.

NeurIPS Checklist

Compliance checklist for the submission’s claims, limitations, and reproducibility.

The checklist confirms that the abstract and introduction state the function‑call/agent‑step correspondence, the function‑aware FIM mid‑training method, and the three‑axis robustness claim, all of which are supported by the experiments in Sections 3.2 and 3.3.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers