# Self-Improvement Imitation with Biologically Guided Search for Protein Design Under Oracle Budgets

arXiv: https://arxiv.org/abs/2605.26690 · HTML: https://mlsanity.com/paper/2605.26690/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 50 references that are in the map.
- First posted on arXiv 26 May 2026.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Deep Reinforcement Learning that Matters](https://mlsanity.com/paper/1709.06560/) (2017)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacement](https://mlsanity.com/paper/1903.06059/) (2019)
- [Biological Sequence Design with GFlowNets](https://mlsanity.com/paper/2203.04115/) (2022)
- [Conditioning by adaptive sampling for robust design](https://mlsanity.com/paper/1901.10060/) (2019)
- [Conservative Objective Models for Effective Offline Model-Based Optimization](https://mlsanity.com/paper/2107.06882/) (2021)
- [AdaLead: A simple and robust adaptive greedy search algorithm for sequence design](https://mlsanity.com/paper/2010.02141/) (2020)
- [Self-Improvement for Neural Combinatorial Optimization: Sample without Replacement, but Improvement](https://mlsanity.com/paper/2403.15180/) (2024)
- [Design by adaptive sampling](https://mlsanity.com/paper/1810.03714/) (2018)
- [Self-Improved Learning for Scalable Neural Combinatorial Optimization](https://mlsanity.com/paper/2403.19561/) (2024)
- [Self-Labeling the Job Shop Scheduling Problem](https://mlsanity.com/paper/2401.11849/) (2024)
- [Robust Optimization in Protein Fitness Landscapes Using Reinforcement Learning in Latent Space](https://mlsanity.com/paper/2405.18986/) (2024)
- [Improved Off-policy Reinforcement Learning in Biological Sequence Design](https://mlsanity.com/paper/2410.04461/) (2024)
- [GraphXForm: Graph transformer for computer-aided molecular design](https://mlsanity.com/paper/2411.01667/) (2024)
- [ProSpero: Active Learning for Robust Protein Design Beyond Wild-Type Neighborhoods](https://mlsanity.com/paper/2505.22494/) (2025)

## Abstract (arXiv)

Protein sequence optimization under tight oracle budgets requires methods that explore vast combinatorial spaces while making each evaluation informative. Existing reinforcement learning and off-policy generative approaches often degrade under surrogate noise, and position-agnostic mutation proposals risk disrupting functionally critical residues. We introduce SILO, a trajectory-level self-improvement imitation framework for oracle-budgeted protein design. SILO uses a hierarchical edit policy that decomposes each mutation into a position choice followed by a residue choice. In each active-learning round, the policy samples candidate trajectories via incremental stochastic beam search without replacement (SBS), and a UCB-based proxy ensemble, combined with an alanine-scan fitness score (AFS), selects candidates with functionally relevant edits for in silico oracle evaluation. The policy is then updated by next-action cross-entropy imitation on the round's best oracle-labeled trajectories, avoiding value-function estimation. Across eight reproduced protein fitness landscapes and five strong baselines from prior work, SILO achieves the highest maximum and top-100 mean fitness on 8 of 8 landscapes within our evaluations, often exhibiting faster early-stage improvement. In low-data and noisy-proxy stress tests on two landscapes per setting, SILO remains competitive or best when several baselines degrade. Ablations show that SBS with AFS account for much of the gains, with iterative imitation providing additional improvement. Code is available at: https://github.com/grimmlab/SILO.git