AgentFold: Closed-Loop Agentic Search for Protein Folding Model Design

TL;DR

AgentFold turns protein folding model design into an executable search problem over ESMFold code variants. A multi-agent loop proposes, implements, debugs, trains, and evaluates changes, while an MCTS-style controller stores successful and failed interventions and re-scores the search tree every ten iterations. On a compact engineering-scale codebase of more than 2,000 lines, the system explores about 80 variants with roughly 5,000 GPU hours and about 170 million LLM tokens, and at a matched 36 evaluation budget it reaches best lDDT 0.285 versus 0.265 for an independent Codex proposal controller. The most important secondary finding is empirical: early soft learnable priors and gated refinement tend to yield stable local gains, whereas direct geometric forcing and geometry-conditioned feedback often collapse training.

Background positioning

This is not a new top-tier folding architecture benchmarked against AlphaFold3 or RoseTTAFold2. It is a systems-and-evidence paper. It uses a compact ESMFold-derived substrate to make repeated agentic model development affordable. The claim to place it in the literature is that it instantiates general autonomous discovery methods such as FunSearch, AlphaEvolve, and MCTS-based agentic heuristic design inside a tightly coupled scientific ML codebase. The paper's value is therefore threefold: a closed-loop engineering search protocol, a matched-budget comparison against proposal-only and random-search controls, and trace-based evidence that some architectural motifs are more robust than others.

Problem and Motivation

Scientific LLM agents are often evaluated on literature reasoning, tool use, code snippets, or experiment plans. That is not enough for scientific ML model improvement. A plausible architectural idea must survive three separate failures. First, it must be written into a real codebase, where interfaces between pair representations, invariant point attention, rigid frame updates, torsion prediction, and recycling can break. Second, it must train without divergence or runtime errors. Third, it must be judged by expensive, noisy, multi-objective structural metrics. The paper's folding example exposes all three at once: local distance accuracy, backbone accuracy, oligomeric quality, RMSD, TM-score, contacts, and stereochemistry can move in different directions.

The authors' working intuition is executable. Treat a candidate model as a concrete code snapshot, not as a text summary. Then use tree search to allocate expensive validation toward promising branches, and use structured memory to make failures informative. This intuition is tested rather than assumed: Table 1 and the ablations show that the integrated controller beats both Codex-only proposals and random selection under matched evaluation count. The paper also supports a weaker design intuition: architectural priors placed before coordinates are formed tend to be safer than perturbations imposed after geometry is already in the update path.

Core Section: From Textual Hypotheses to Executable Folding Model Search

The substrate and the scoring contract

AgentFold starts from a compact ESMFold-derived substrate. Appendix A.1 explains that the folding trunk is reduced to a single trunk block unless otherwise noted, so that each variant can be trained and evaluated under a fixed compute budget. This choice preserves the coupled sequence, pair, and structure modules, but it also means that every quantitative conclusion must be read against a simplified architecture. The evaluation protocol uses CAMEO2022 as a development benchmark and OpenStructure metrics: backbone lDDT, lDDT, oligomeric GDT-TS, RMSD, and TM-score. To rank variants for search, the paper defines the Normalized Weighted Relative Score.

Here $\mathcal{M}$ contains ten aggregated metrics: mean and median values of bb_lddt, lddt, oligo_gdtts, rmsd, and tm_score. $x_m$ is a model's value, $b_m$ is the fixed ESMFold baseline value, $w_m$ is a nonnegative weight set to uniform 0.1, and $r_m$ is the baseline-relative score; for metrics where larger is better, $r_m=x_m/b_m$, while for RMSD, where lower is better, $r_m=b_m/x_m$. NWRS therefore maps an exact baseline replica to 0.5 and saturates at 1.0. Its role in the search chain is not to prove final scientific superiority, but to allocate compute by giving one scalar to compare expensive candidates. The boundary is important: NWRS can hide metric conflicts. Section 4.2.3 explicitly notes that some high NWRS variants improve lDDT while leaving TM-score nearly unchanged and producing mixed RMSD changes. Thus NWRS is a ranking instrument, not a complete structural assessment.

The mechanism: MCTS over snapshots, agentic evolution, and periodic re-scoring

Figure 1 presents the architecture as a dual loop. The inner loop samples variants from the Experience Pool, evolves executable code, runs experiments, and analyzes outcomes. The outer loop periodically re-scores candidates every ten iterations, because a single rollout is an expensive training and evaluation job rather than a cheap MCTS simulation. Each tree node is a concrete code snapshot, and top-k sampling selects high-scoring parents together with diverse reference nodes. Sampling siblings from the same parent is a near-controlled comparison: most code is held fixed, so gains or collapses can be attributed more cleanly to the typed intervention.

The search policy is hybrid. It combines objective parsing with an agent critic. Equation in Section 3.5 gives the score for an intervention $e$.

$S_L$ is extracted from training loss signals, $S_B$ is the benchmark score, and $S_A$ is a Critic Agent score reflecting intervention rationale, coherence with prior evidence, and likely destabilizing risk. The sum's function is prioritization: it tells the sampler which branches deserve the next expensive training job. The paper is careful to state that $S_A$ is used only to prioritize experiments and not to claim final improvements. If the Critic term is ignored, the controller becomes purely objective-driven; if it is over-trusted, the tree may chase architecturally plausible but empirically weak edits. The authors mitigate this by re-scoring periodically from stored metrics and attribution reports, not by a single terminal ranking.

The agent pipeline itself is a small software engineering organization. A Deduplicator screens proposed motivations against historical records. A Unified Planner then owns both design and implementation, rather than passing hypotheses to a separate coder. The paper argues that this reduces interface mismatch between what is proposed and what is actually coded. A Debugger observes the training environment and asks the Planner to fix syntax or runtime errors until training launches. After execution, an Analyst writes structured attributions linking metric deltas to intervention types, and a Searcher attaches relevant literature. This design choice is intuitive and important for engineering, but the paper does not provide a dedicated ablation for the unified planner, the deduplicator, or the literature Searcher. The matched-budget comparison tests the integrated controller, not each component separately.

Which interventions survived: bias before geometry and gated refinement

The strongest architecture, esmfold_struct_enhanced_v4, is introduced in the case study and architecture appendix as a parameter-light modification that leaves the core IPA, frame, and FAPE loop intact. Figure 7 shows the substrate and variant: ESMFold stacks invariant point attention blocks over single and pair representations, then uses geometric heads to update backbone frames and torsions. The enhanced variant inserts learned bias terms before or alongside IPA, including a residue-index-conditioned bias MLP, while keeping downstream geometry heads unchanged.

The enhanced attention logits add dynamic pair bias, sequence separation bias, and structure-aware bias to the standard IPA terms.

Here $a_{h,i,j}$ is the attention logit for head $h$ at residue pair $(i,j)$, $q_{h,i}$ and $k_{h,j}$ are query and key features, $\alpha$ scales content similarity, $z_{i,j}$ is the pair representation, $\Delta {\mathrm{residx}}_{i,j}$ is sequence separation, $s_i$ and $s_j$ are structural states, $p_{h,i,j}$ denotes the original IPA point geometry term, and $m_{i,j}$ denotes masking. Compared with vanilla IPA, the change is not a new global folding algorithm; it is a soft prior placed inside the attention channel before coordinates are strongly coupled. This matches empirical pattern P1: bias before geometry. If the same signal were imposed as a fixed frame offset after refinement, the paper's traces suggest it can collide with the rigid update process. The formulation is useful because it explains why small added parameters can matter: they modulate how attention distributes residue-pair relationships without replacing the geometry heads.

The second stable motif is gated refinement. The paper describes a learned gate on the backbone update vector.

${\Delta }_{\mathrm{frame}}$ is the raw six-dimensional rigid update, $g$ is a learned parameter vector with the same dimensionality, $\sigma$ is the sigmoid function, and $\odot$ is elementwise multiplication. In the search argument, this gate implements pattern P2: multiplicative refinement. It can shrink uncertain updates instead of adding a fixed displacement. If the gate is disabled, the update becomes additive forcing; if the gate is learned poorly, it may suppress useful movement. The paper's ablation supports the first failure mode: Section 4.4 reports that removing BackboneUpdate gating lowers mean lDDT by 0.012 relative to esmfold_struct_enhanced_v4, and removing IPA bias lowers it by 0.017. The complementary drops indicate that both mechanisms contribute to the final variant, though the ablation is performed on a single strong branch rather than across the whole tree.

The negative side of the pattern set is equally informative. Pattern P3 says to avoid feeding immature geometry into attention or frame updates. Table 5 gives a failure contrast for esmfold_net_differential_geometry: the stored record reports backbone lDDT 0.015, all-atom lDDT 0.000, and TM-score 0.091, while the parameter count is 31.03M in Section 4.3.2. That example is large, but it does not dominate; instead it collapses. This supports the authors' claim that the gains are not simply capacity effects.

Experimental Evidence and Search Behavior

Table 1 reports the matched controller comparison after 36 evaluations each. AgentFold achieves best lDDT 0.285 and best NWRS 0.526. The independent Codex proposals baseline reaches 0.265 and 0.512, and the random controller reaches 0.260 and 0.510. The same table also reports Top-5 selections by NWRS, where AgentFold is ahead on both lDDT and NWRS.

The interpretation needs care. The absolute best-lDDT advantage over Codex is 0.020, which corresponds to the paper's headline 7.5% relative improvement. The Top-5 advantage over random search is more revealing of controller quality: 0.267 versus 0.242 in Top-5 lDDT. This comparison supports the integrated search mechanism, but the paper itself says it does not isolate individual components. It shares edit space, models, prompts, checks, training, and evaluator with Codex proposals, which is a strong matched-budget design. It does not share the tree, intervention history, and structured memory, so the result cannot separate MCTS from agentic memory or planner architecture.

Figure 2 visualizes why a tree controller matters. Darker nodes indicate higher average lDDT, while gray nodes mark collapsed variants with average lDDT below 0.1. The shape of the tree shows that the search is not a random scatter of high scores; it clusters productive modifications near certain ancestors.

Table 2 gives the global metric deltas for representative variants. The ESMFold row reports absolute mean values; other rows show mean deltas.

The most important pattern is heterogeneity. esmfold_struct_enhanced_v4 improves lDDT strongly but leaves backbone lDDT and oligomeric GDT-TS only modestly changed. esmfold_struct_distAware_v1 has a smaller lDDT gain but better backbone lDDT, GDT-TS, RMSD, and TM-score behavior according to the text. Several high NWRS variants improve local accuracy while leaving global fold metrics nearly unchanged or mixed. This is not a defect of AgentFold; it is a central scientific result. Closed-loop search found local structural refinements, not a global topology overhaul. It also explains why a single NWRS number cannot be taken as the whole story.

Table 3 sharpens the same conclusion through targeted loop and contact metrics.

The loop lDDT improvements are much larger than the global backbone lDDT improvements. For the top local-context variant, Table 3 shows a +0.063 loop lDDT change, while Table 2 shows only +0.049 all-atom lDDT and +0.002 backbone lDDT mean change. The esmfold_struct_enhanced_v4 row adds a +0.019 precision gain in the 12 to 24 sequence separation contact bin, which suggests the intervention affects medium-range pair relations, not only loop fitting. MolProbity improvements, notably -0.157 for esmfold_struct_enhanced_v4, also matter because physical plausibility is not optimized by NWRS alone. This supports the claim that stable gains arise from steering attention and update magnitude rather than forcing coordinates.

Parameter analysis adds a crucial confound check. Section 4.3.2 reports the ESMFold baseline at about 22.61M parameters. The highest NWRS variant #36, esmfold_struct_enhanced_v4, has 22.856M parameters, an increase of about 1.1%. Variant #47 adds approximately 0.015M parameters and achieves the second-highest NWRS in Table 2. Variant #28, esmfold_struct_distAware_v1, has 22.574019M parameters in Table 9, slightly below baseline while still improving NWRS in Table 2. Larger models are not reliably better: #24 and #40 have 28.46M parameters, #57 has 32.49M, and #60 has 31.03M yet collapses. These numbers weaken a naive capacity explanation and strengthen the placement explanation: biases before geometry and gates on updates are the decisive design variables in this search.

The ablation in Section 4.4 is more targeted. Removing the IPA bias lowers mean lDDT by 0.017 relative to esmfold_struct_enhanced_v4. Removing BackboneUpdate gating lowers mean lDDT by 0.012. The text says the ablated variants remain competitive but neither mechanism alone recovers the full improvement, indicating that the two changes are complementary rather than interchangeable. This is a meaningful causal signal because it tests the same strong branch. It is still limited: both ablations focus on the best branch and do not prove that bias before geometry is universally better than alternative geometric priors.

The case study reinforces the mechanism and its trade-off. Section 4.5 examines representative loop-region targets. For 7vd7_B, the paper reports loop backbone lDDT improving from 0.665 to 0.753 and loop RMSD decreasing from 8.298 to 3.436 Å, while loop lDDT slightly decreases from 0.225 to 0.209. That exact pattern is valuable because it prevents overclaiming. Better rigid backbone alignment can occur even when all-atom local contact accuracy does not uniformly rise. The authors interpret this as improved loop placement and backbone alignment, especially in flexible loop regions, rather than uniform improvement over every local score.

Overall evidence quality is mixed but useful. The paper's strongest evidence is the matched-budget controller comparison and the targeted metric analysis. The weakest evidence is the post-hoc pattern inference. The intervention traces are real and structured, and Table 5 supplies report-level evidence for P1 to P3, but these patterns are correlations across the search tree, not controlled causal laws. The system's computational scale is explicit: roughly 80 variants, about 5,000 GPU hours, and about 170M LLM tokens. However, all main results are on a one-block folding trunk, a 1,000-chain PDB subset, and CAMEO2022 development evaluation. The limitations section concedes that transfer to stronger folding systems and broader biological settings remains unverified.

Deep Insights and Conclusion

AgentFold's central contribution is not simply that an LLM improved ESMFold. It is an evidence-generation method for scientific ML engineering. By representing hypotheses as executable code snapshots, retaining failed runs in structured memory, and using tree search to allocate expensive validation, the system produces both improved models and reusable design signals. The improved models are modest in absolute terms on the full metric stack, but the 0.020 best-lDDT margin over Codex proposals in Table 1 is meaningful under a fixed 36-evaluation budget. The trace-based patterns are likely more reusable than the specific architecture names: soft pair or attention biases before coordinate instantiation and multiplicative gating on rigid updates appear to preserve training stability better than direct geometric perturbation.

The most concrete limitation is scope. A single trunk block is a deliberate computational shortcut, but it may underrepresent deeper recycling interactions where geometric errors can compound differently. The NWRS ranking metric may encourage local lDDT gains at the expense of RMSD and TM-score, as Table 2 already shows. The pattern set P is mined from stored agent reports and variant trees, so it remains descriptive. A stronger follow-up would test P1 to P3 by constructing paired branches: inject equivalent geometric strength after coordinates, before coordinates, and with gating versus forcing, under identical training data and multiple trunk depths. Without that, the paper should be read as a rigorous case study in agentic scientific engineering and empirical design heuristics, not as a universal theorem about folding architectures.

Future work should move from compact search to production science. The most promising extension is not merely more protein families, but model-specific edit interfaces that let the controller address AlphaFold3-like interaction modules, multi-chain interface losses, and recycling depth without breaking coupling. Another is budget-aware validation: use cheap structural surrogates or early-stopping signals to decide when to abort a branch, because 5,000 GPU hours is not negligible. A third is making trace mining causal, by using the same memory database to generate counterfactual probes rather than post-hoc reports. In its present form, AgentFold is a strong demonstration that closed-loop code search can produce measurable folding model improvements and expose design patterns that human or LLM readers might otherwise miss.