🧬 biology

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

The paper introduces SILO, a trajectory-level self-improvement imitation framework that combines hierarchical edit policies, stochastic beam search, and an alanine-scan fitness score to achieve superior protein sequence optimization under tight oracle budgets compared to existing reinforcement learning and generative baselines.

Original authors: Ashima Khanna, Dominik Grimm

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ashima Khanna, Dominik Grimm

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are a master chef trying to create the perfect new recipe for a dish. You have a basic recipe (the "wild-type" protein), and your goal is to tweak the ingredients to make it taste even better (higher "fitness"). However, you have a strict rule: you can only taste-test your creations a limited number of times before you run out of money or ingredients. This is the challenge of protein design under an "oracle budget."

In the scientific world, the "taste test" is a computer simulation or a lab experiment called an oracle. It tells you if a protein works, but it's expensive and slow. The problem is that there are more possible ingredient combinations than there are stars in the sky. If you just guess randomly, you'll waste your budget. If you rely too much on a cheap, imperfect prediction model (a "surrogate"), you might get tricked into thinking a bad recipe is good.

The paper introduces a new method called SILO (Self-Improvement Imitation with Biologically Guided Search) to solve this puzzle. Here is how it works, broken down into simple concepts:

1. The Problem with Old Methods

Previous attempts to solve this were like two different types of chefs:

The Reinforcement Learning Chef: This chef tries to learn a "value" for every possible dish. But if the taste-test machine is noisy (gives inconsistent results), the chef gets confused, makes bad guesses, and the whole strategy falls apart.
The Random Mutation Chef: This chef changes ingredients without thinking about where they are. They might accidentally swap a crucial spice that holds the whole dish together, ruining the recipe.

2. The SILO Solution: A Smart Apprentice

SILO acts like a smart apprentice who learns by watching the best versions of the dish they have made so far. Instead of trying to guess the "value" of every possible future dish, it simply copies the specific steps that led to the best results.

Here are the three secret ingredients of SILO:

A. The "Step-by-Step" Editing (Hierarchical Policy)

Instead of changing the whole recipe at once, SILO breaks the process into two small steps for every change:

Pick a spot: "Let's change the 5th ingredient."
Pick a new ingredient: "Let's swap it for salt."
This prevents the chef from making chaotic, all-at-once changes that break the dish. It keeps the search organized.

B. The "Biological Safety Net" (Alanine-Scan Fitness)

This is the most creative part. In biology, there's a technique called Alanine Scanning. Imagine you want to know if a specific ingredient is critical to the dish. You swap it with a neutral, bland ingredient (Alanine) to see if the flavor disappears.

How SILO uses it: Before spending its precious "taste-test" budget on a new recipe, SILO asks: "If I replaced these specific changes with bland ingredients, would the dish still taste good?"
If the answer is yes, it means those spots are flexible and safe to change.
If the answer is no, it means those spots are critical, and SILO avoids making risky changes there.
This acts as a filter, ensuring the chef doesn't waste time on recipes that are likely to fail because they broke a vital part of the structure.

C. The "Self-Improvement Loop" (Imitation Learning)

SILO runs in rounds:

Generate: It creates many new recipe variations using a smart search technique (Stochastic Beam Search) that explores many different paths without repeating itself.
Filter: It uses the "Safety Net" and a prediction model to pick the top candidates.
Taste-Test: It sends the best few to the expensive oracle for a real evaluation.
Learn: It looks at the winners. It doesn't try to calculate why they won; it simply memorizes the exact sequence of steps (the "trajectory") that turned the original recipe into the winner.
Imitate: It updates its "apprentice brain" to be more likely to make those same steps next time.

3. The Results: Winning the Cooking Contest

The authors tested SILO on eight different "cooking challenges" (protein datasets) and compared it against five other top methods.

The Outcome: SILO consistently found the highest-scoring recipes (maximum fitness) and the best average group of recipes (top-100 mean fitness) across all eight challenges.
Speed: It found good recipes faster than the others, meaning it used its limited "taste-test" budget more efficiently.
Resilience: When the prediction model was noisy (like a broken taste-test machine) or when the chef started with very little data, SILO kept performing well while other methods struggled or failed.

Summary Analogy

Think of protein design as navigating a massive, foggy maze to find the highest peak.

Old methods were like trying to map the whole mountain at once (often getting lost in the fog) or just taking random steps.
SILO is like a climber who:
1. Takes small, deliberate steps (hierarchical editing).
2. Checks the ground stability before stepping (biological safety net).
3. Looks at the path the previous successful climber took and decides, "I'll follow those exact footsteps," rather than trying to guess the terrain's value from scratch.

The paper concludes that combining smart searching, biological safety checks, and learning from past successes is a powerful way to design better proteins when you can't afford to test every single possibility.

Technical Summary: Self-Improvement Imitation with Biologically Guided Search for Protein Design Under Oracle Budgets

Problem Formulation

Protein engineering is framed as a sequential decision-making optimization problem aimed at maximizing a "fitness" function (e.g., thermostability, binding affinity) under a strict evaluation budget. The core challenge lies in the vast combinatorial size of the sequence space ( $20^L$ ) and the prevalence of epistasis, which creates rugged, sparse, and non-convex fitness landscapes. Existing approaches face three persistent limitations:

Instability in Value-Based Optimization: Reinforcement learning (RL) methods relying on value functions are sensitive to hyperparameters and prone to error propagation when trained on noisy surrogate rewards.
Position-Agnostic Mutations: Many methods propose mutations without regard for residue conservation, risking the disruption of functionally critical sites.
Inefficient Exploration: Under tight oracle budgets, methods often fail to prioritize high-value regions or avoid redundant evaluations.

The objective is to design sequences with higher fitness than a starting sequence within a limited number of oracle queries ( $N$ rounds, $K$ queries per round), utilizing a proxy model to guide selection while reserving the expensive oracle for final evaluation.

Methodology: SILO

The authors introduce SILO (Self-Improvement Imitation with Biologically Guided Search), an active-learning framework that integrates structured sampling, biologically informed selection, and trajectory-level imitation learning. The process operates in an iterative loop:

1. Hierarchical Edit Policy and MDP Formulation

SILO casts candidate generation as a Markov Decision Process (MDP) over edit trajectories. Instead of generating full sequences at once, a transformer-based policy $\pi_\theta$ decomposes each mutation into two hierarchical stages:

Level-0: Selecting a position $i \in \{1, \dots, L\}$ to modify.
Level-1: Selecting a replacement amino acid $a \in \mathcal{V}$ conditioned on the chosen position.
The policy is trained on frozen ESM Cambrian embeddings, which encode evolutionary priors, using a lightweight transformer encoder with two policy heads.

2. Structured Sampling (Incremental SBS)

To generate diverse candidate trajectories, SILO employs Incremental Stochastic Beam Search (SBS). Unlike standard beam search, SBS samples trajectories without replacement by perturbing log-probabilities with Gumbel noise. After a trajectory is sampled, its probability mass is removed, and the distribution is renormalized. This ensures a diverse set of action trajectories is explored within the mutational budget ( $d_{max}$ ).

3. Biologically Informed Selection

Before oracle evaluation, candidates are scored using a proxy ensemble $f_\phi$ and a specific objective function $S(x)$ :
$S(x) = \text{UCB}_{\text{proxy}}(x) + \text{UCB}_{\text{AFS}}(x)$

$\text{UCB}_{\text{proxy}}(x)$ : An Upper Confidence Bound based on the proxy's mean prediction and uncertainty, balancing exploration and exploitation.
$\text{UCB}_{\text{AFS}}(x)$ : An Alanine-Scan Fitness Score. This term evaluates an alanine-substituted variant of the candidate (where mutated positions are replaced with alanine). If the alanine variant is predicted to have high fitness, it suggests the mutated positions are functionally relevant and tolerant to neutral changes. This acts as a ranking signal to down-weight candidates likely to disrupt protein function.

4. Trajectory-Level Self-Improvement Imitation (SIL)

Unlike RL methods that optimize value functions, SILO uses Self-Improvement Imitation Learning.

Data Collection: The top $K$ candidates from the SBS sampling are evaluated by the oracle.
Imitation Target: The top $P$ oracle-labeled sequences (BESTFOUND) and their corresponding action trajectories are selected.
Training: The policy $\pi_\theta$ is updated via cross-entropy minimization on history-next action pairs $(h(t), a(t))$ derived from the BESTFOUND trajectories. This allows the policy to learn directly from the actions of top performers without relying on noisy value estimates.

Key Contributions

Structured Sampling and Biologically Informed Selection: The combination of Incremental SBS for diverse trajectory generation and the Alanine-Scan Fitness Score (AFS) for selection drives the majority of performance gains. The AFS term specifically mitigates the risk of disrupting conserved residues.
Adaptation of SIL to Oracle-Budgeted Settings: The paper adapts the SIL paradigm (previously used in routing and molecular design with cheap evaluations) to the protein design context. By training on oracle-labeled trajectories rather than surrogate values, SILO avoids the instability associated with value-function estimation under noisy proxies.
Empirical Robustness: The framework demonstrates consistent performance across eight protein fitness landscapes, maintaining competitiveness even in low-data regimes and under noisy proxy conditions where other baselines degrade.

Experimental Results

The authors evaluated SILO against five strong baselines (PEX, AdaLead, GFN-AL- $\delta$ CS, MLDE, ProSpero) across eight benchmark datasets (AAV, AMIE, E4B, GFP, LGK, Pab1, TEM, UBE2I).

Performance: SILO achieved the highest maximum fitness and top-100 mean fitness on all 8 landscapes. It frequently discovered high-fitness sequences in earlier active-learning rounds, indicating more efficient use of the oracle budget.
Low-Data Regime: In stress tests with reduced initial labeled data (10%, 20%, 50% of $D_0$ ), SILO remained stable and competitive, whereas methods like GFN-AL- $\delta$ CS showed noticeable degradation.
Noisy Proxy: Under synthetic Gaussian noise added to the proxy model (SNR levels of -25, -15, -5), SILO maintained high performance, while baselines like ProSpero and PEX degraded significantly.
Ablation Studies:
- Removing the structured sampling (replacing SBS with standard Beam Search) reduced performance.
- Removing the AFS term reduced performance, confirming the value of biologically guided selection.
- A "frozen policy" variant (no imitation updates) remained competitive, suggesting that sampling and selection account for a substantial portion of the gains, but iterative imitation provided additional improvements, particularly on AAV and Pab1.

Significance and Claims

The paper claims that SILO represents an effective direction for oracle-budget protein sequence optimization by successfully integrating learning, sampling, and biologically guided selection.

Stability: By sidestepping value-function estimation and relying on cross-entropy imitation of oracle-labeled trajectories, SILO offers a more stable alternative to RL in settings with noisy surrogate rewards.
Efficiency: The method identifies high-fitness candidates faster than baselines, making it suitable for scenarios where experimental evaluations are costly.
Biological Relevance: The incorporation of the Alanine-Scan Fitness Score ensures that the search is guided by biological principles, reducing the likelihood of generating non-functional variants.

The authors conclude that while SILO excels in in silico benchmarks, future work should focus on evaluating the method under larger mutation budgets and validating its performance in wet-lab settings to address potential domain shifts between oracles and biological reality.