Function-Aware Fill-in-the-Middle as Mid-Training for Coding Agent Foundation Models¶
🕒 Published (v1): 2026-07-14 07:44 UTC · Source: Arxiv · link
Why this paper was selected
Fill-in-the-Middle mid-training teaches models to handle tool returns; novel foundation for coding agents
Ask a follow-up
Open an assistant pre-loaded with this paper's context.
💬 Ask ChatGPT✦ Ask Claude
TL;DR¶
The paper proposes a self-supervised mid-training stage that exploits the structural isomorphism between function call sites (caller → callee → return → continuation) and coding-agent steps (history → action → observation → continuation). By masking functions selected via program dependency graph analysis with a complexity–inferability criterion and embedding chain-of-thought rationales, they improve SWE-Bench-Verified by up to +5.3 points while also recovering capability erosion that agentic post-training otherwise causes on non-agent benchmarks.
Problem¶
Standard left-to-right pretraining under-exposes models to the bidirectional conditioning structure that coding agents require at inference time: integrating externally-computed tool returns into ongoing reasoning. Random-span FIM used in existing code LLMs addresses this only incidentally—span boundaries are syntactically arbitrary, there is no reasoning supervision, and the objective is diluted across trillions of pretraining tokens. A dedicated mid-training stage aligned with agent-relevant inductive biases is missing between pretraining and agentic post-training.
Method¶
The method exploits the observation that the agent loop \((h_t \to a_t \to o_{t+1} \to \text{continuation})\) is structurally identical to a function call site (pre-call context → call → return value → downstream code). The training pipeline has five components:
-
Corpus: 968 GitHub Python repositories (≈400K FIM samples, 2.6B tokens) decontaminated against SWE-Bench source repositories by name and commit timestamp.
-
Program dependency graph (PDG): ASTs are parsed to extract call edges \(E_{\text{call}}\) (caller→callee) and sibling edges \(E_{\text{sib}}\) (same-class methods sharing instance state).
-
Function selection via double criterion: Each function \(v\) receives a complexity score $\(\hat{H}(v) = w_\ell \phi(\text{LoC}(v), c_\ell) + w_c \phi(\text{CC}(v), c_c) + w_d \phi(D(v), c_d)\)$ and an inferability score $\(\hat{I}(v) = \alpha C_{\text{caller}} + \beta C_{\text{callee}} + \gamma C_{\text{sig}} + \delta C_{\text{doc}} + \varepsilon C_{\text{class}}\)$ Combined as a harmonic-mean-like joint score \(\text{FIM}(v) = \frac{\hat{H}\hat{I}}{\hat{H}+\hat{I}+\epsilon} \cdot \rho(\Delta(v))\), with a difficulty penalty \(\rho\) down-weighting unlearnable targets; functions scoring above threshold \(\tau=0.08\) are selected.
-
Multi-function group masking: Connected groups of \(k=2\) or \(k=3\) functions are jointly masked (8 topology patterns: caller-callee, co-callee, sibling, etc.), with the main recipe being an 80%/15%/5% single/pair/triple mixture.
-
CoT augmentation: Gemini-3-Flash generates step-by-step rationales and candidate bodies (with no access to ground truth); a separate Gemini-3-Flash judge filters to the top ≈400K samples. The retained rationale is placed before the body in the FIM middle span, teaching the model a think-then-act pattern.
Mid-training uses standard FIM loss on the middle span; the existing agentic post-training pipelines (R2E-Gym, SWE-Smith, SWE-Lego) are then applied without modification.
Key Contributions¶
- Isomorphism framing: Formalizes the structural equivalence between function call sites and agent action–observation–continuation loops, motivating internet-scale code as a source of agent-relevant inductive bias.
- Function-aware FIM pipeline: PDG-based target selection with the base-model-agnostic \(\hat{H}\)–\(\hat{I}\) double criterion and multi-function group masking.
- CoT integration: Chain-of-thought rationales embedded inside the FIM middle span, teaching reasoning-before-code consistent with agent think-then-act.
- Robustness validation: Gains confirmed across two model sizes (7B/14B), two post-training pipelines (R2E-Gym, SWE-Smith), and one alternative base (Qwen3-8B + SWE-Lego).
- Capability preservation: Mid-training mitigates the hidden capability erosion from agentic post-training on non-agent coding and cross-domain tool-use benchmarks.
- Open release: 968-repository decontaminated corpus, selection pipeline, and mid-training checkpoints.
Results¶
- SWE-Bench-Verified gains (vs. reproduced post-training-only baselines):
- Qwen2.5-Coder-7B + R2E-Gym: 17.80 vs. 15.00 → +2.80
- Qwen2.5-Coder-7B + SWE-Smith: 17.60 vs. 12.30 → +5.30
- Qwen2.5-Coder-14B + R2E-Gym: 29.20 vs. 26.20 → +3.00
- Qwen3-8B + SWE-Lego: 35.00 vs. 31.80 → +3.20
- SWE-Bench-Lite gains:
- 7B/R2E-Gym: +3.67; 7B/SWE-Smith: +0.50; 14B/R2E-Gym: +4.00; Qwen3-8B: +5.40
- Capability preservation (14B, R2E-Gym):
- LiveCodeBench: 24.10 (post-only) → 35.20 (+11.10, near Instruct ceiling of 37.20)
- BFCL: 15.80 → 18.20 (+2.40)
- τ-bench: 3.40 → 7.30 (+3.90)
- 6-benchmark average: 16.04 → 19.56 (+3.52)
- Ablations (7B, 200K budget):
- No CoT FIM vs. full Gemini-CoT FIM: +1.18 vs. +2.43 average over baseline (FIM structure contributes ~half the gain independently of distillation)
- Random masking vs. full PDG+\(\hat{H}\)+\(\hat{I}\): 13.95 vs. 15.60 average
- 80%/15%/5% mix: best at 16.10 average vs. 15.60 (single only)
Limitations¶
- Mid-training corpus is Python-only; cross-lingual generalization is untested.
- Qwen3-8B comparison jointly confounds base model and post-training pipeline (SWE-Lego), limiting cross-family conclusions.
- FIM-only checkpoints degrade instruction-following, so mid-training cannot be evaluated in isolation—all gains are measured only after post-training.
- CoT rationale generation requires frontier model inference (Gemini-3-Flash) at scale, adding pipeline cost; self-CoT closes most but not all of the gap.
- \(\hat{H}\) and \(\hat{I}\) are hand-designed proxies; a learned predictability score could improve selection but would couple the pipeline to a specific reference model.
- Ablations restricted to 7B due to compute constraints; behavior at larger scales or with different architectures is not tested.
- SWE-Smith pairing yields only +0.50 on Lite despite +5.30 on Verified, suggesting the benefit magnitude is sensitive to post-training data distribution in ways not fully explained.
Relevance to Agentic AI / LLM Agents¶
This work directly addresses the base-model alignment gap that agentic post-training pipelines currently skip: the structural mismatch between left-to-right pretraining and the bidirectional conditioning required when an agent must integrate tool returns into ongoing reasoning. The function-call/tool-call isomorphism is a theoretically grounded and empirically validated bridge, showing that internet-scale code already contains the inductive bias needed for agentic conditioning if exploited at the right granularity. The cross-domain transfer to τ-bench and BFCL—despite a Python-only mid-training corpus with no tool-use trajectories—provides concrete evidence that the inductive bias is structural rather than surface-level, with implications for any LLM agent framework where actions return observations that condition future steps. The demonstration that agentic post-training silently erodes general capabilities by double-digit margins (LiveCodeBench −13.1, BFCL −7.4) and that mid-training largely recovers this erosion is a significant practical finding for practitioners building coding agents on top of general-purpose base models.