PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Objective Protein Optimization via Modular Diffusion Steering

Meir Roketlishvili, Felix Huber, Wenda Chu, Yisong Yue, Raul Astudillo

arXiv:2610.28943v1cs.LG

TL;DR

Protein engineering seeks useful tradeoffs among competing properties, but large sequence spaces and limited experimental batches make them difficult to discover. MODS steers a fixed diffusion model with randomized scalar rewards and updates its surrogates as feedback arrives. Across six tasks, it performs strongly with both diffusion models, with results ranging from the highest mean final hypervolume to performance similar to evaluated baselines.

  • Problem

    Protein engineering must find useful tradeoffs among competing properties despite an exponentially growing sequence space and experimental campaigns limited to a few batched measurements.

    Source: §1 · p. 1

  • Method

    MODS samples augmented Chebyshev weights and a fitness-surrogate ensemble member to create scalar rewards that steer a fixed diffusion model toward different objective tradeoffs.

    Source: §1 · p. 2

  • Results

    Across six protein optimization tasks, MODS performed strongly with both diffusion models, achieving either the highest mean final hypervolume or performance similar to evaluated baselines for each model.

    Source: §6 · p. 9

  • Takeaways & Limitations

    MODS supports different diffusion-model and compatible steering combinations while updating fitness surrogates from new observations without retraining the generative model.

    Source: §6 · p. 9

    Laboratory studies remain necessary to determine whether the observed gains translate to experimentally measured fitness, and the evaluation covers a limited set of tasks, diffusion priors, and steering procedures.

    Source: §6 · p. 9

Abstract

from arXiv · show

Protein engineering often requires improving several properties simultaneously, motivating the search for high-performing sequences that represent useful fitness tradeoffs. Discovering these sequences is challenging because the design space is combinatorially large and experimental campaigns often permit only a few rounds of batched measurements. We introduce \textit{MODS}, a modular framework for multi-objective protein optimization via inference-time diffusion steering. For each proposal, MODS samples augmented Chebyshev scalarization weights and a fitness-surrogate ensemble member to define a scalar steering reward. Varying the weights directs generation toward different objective tradeoffs, while varying ensemble members diversifies the predictions used for guidance. Because MODS expresses multi-objective acquisition through scalar rewards, it can be paired with different diffusion models and compatible steering procedures. After each evaluation batch, MODS updates the surrogates without retraining the generative model. We evaluate MODS on six protein optimization tasks using predictive oracles trained on experimental assay data, with two diffusion models and two steering procedures. MODS achieves competitive final hypervolume across all six tasks with both diffusion models, and substantially outperforms the best evaluated baseline on several tasks. Our results demonstrate a flexible and effective approach to multi-objective protein optimization with limited feedback.

1 INTRODUCTION

Protein optimization must balance competing objectives under limited experimental feedback. MODS uses randomized scalar rewards to steer a fixed diffusion model toward different tradeoffs, and achieves substantial final-hypervolume gains over evaluated baselines across six assay-derived benchmarks.

1 INTRODUCTION

  • Competing protein properties, exponentially growing sequence spaces, and costly experiments make useful tradeoffs difficult to find within a few batched rounds.The goal is a set of high-performing sequences representing useful fitness tradeoffs, rather than necessarily one best sequence.Source: §1 · p. 1
  • Prior multi-objective methods combine learned representations with Bayesian optimization, hypervolume acquisition, predictor gradients, or preference-conditioned generators, but often bind acquisition to a particular generator or procedure.This motivates a modular approach compatible with different generative models and steering mechanisms.Source: §1 · p. 2
  • MODS samples augmented Chebyshev weights and a surrogate-ensemble member per proposal, varying both the targeted objective tradeoff and predicted fitness landscape.Each pair defines a scalar reward, separating multi-objective acquisition from generation.Source: §1 · p. 2
  • After each evaluation batch, MODS updates its surrogates and steers the fixed generative model with new rewards, without retraining that model.This feedback loop progressively refines task-specific steering.Source: §1 · p. 2
  • Across six assay-derived oracle benchmarks, MODS consistently delivers substantial gains in final hypervolume over evaluated baselines.Experiments also examine two diffusion models, two steering procedures, and steering strength.Source: §1 · p. 3

2 RELATED WORK

Related work spans surrogate-guided optimization, generative steering, and multi-objective acquisition. MODS builds on inference-time steering and scalarization while emphasizing an acquisition strategy that can pair with different generators and compatible steering procedures.

Protein generative models and steering

  • Protein generative models provide plausible-sequence priors; MODS follows inference-time steering by updating task-specific surrogates as feedback accumulates while keeping the generator fixed.Other approaches adapt the generative model itself using learned rewards.Source: §2 · p. 3

Protein optimization with limited feedback

  • Limited-feedback protein optimization often uses localized mutation spaces, while generative models can search larger combinatorial spaces; prior ensemble-guided diffusion work focused on single-objective design.MODS addresses the additional challenge of competing objectives.Source: §2 · p. 3

Multi-objective acquisition strategies

  • Hypervolume acquisition targets objective values and tradeoff coverage, while randomized scalarization converts objectives into scalar problems; Chebyshev scalarizations can recover nonconvex-front tradeoffs that weighted sums cannot.In MODS, weights select tradeoffs and sampled surrogate members select predicted fitness landscapes.Source: §2 · p. 3

Multi-objective generative protein optimization

  • Prior multi-objective protein methods use Bayesian optimization, hypervolume acquisition, predictor gradients, reward-conditioned generation, or offline guided flows; MODS adaptively updates surrogates after batches and pairs acquisition with different compatible generators and steering mechanisms.This modular pairing is identified as a key contribution.Source: §2 · p. 3

3 PROBLEM SETUP

The setup considers fixed-length protein sequences with multiple unknown objectives and limited batched evaluations. The goal is to identify evaluated sequences near the Pareto front while retaining diverse objective tradeoffs.

3 PROBLEM SETUP

  • The search space consists of fixed-length sequences over an amino-acid alphabet, each with M ≥ 2 unknown fitness objectives to maximize.The objective vector is denoted f1(x), . . . , fM(x).Source: §3 · p. 4
  • Because fitness evaluations are expensive, only a limited number of sequences can be queried in a small number of batched rounds.Source: §3 · p. 4
  • Conflicting objectives motivate Pareto-optimal tradeoffs: a sequence is dominated if another is at least as good on every objective and strictly better on one.The Pareto front maps the nondominated sequence set into objective space.Source: §3 · p. 4
  • A pretrained diffusion model, trained independently on unlabeled sequences, remains fixed while the optimizer uses accumulated observations to generate batches.After T rounds, the evaluation budget is |D0| + TB; the goal is to approach the Pareto front with diverse tradeoffs.Source: §3 · p. 4

4 THE MODS FRAMEWORK

MODS steers a fixed diffusion model for multi-objective protein optimization using scalar rewards built from surrogate predictions and sampled tradeoff weights. Its modular workflow generates and evaluates batches, then updates the surrogates without retraining the generative model.

4 THE MODS FRAMEWORK

  • MODS combines randomized augmented Chebyshev scalarizations with surrogate-ensemble sampling to target different objective tradeoffs and predicted fitness landscapes.Each weight vector specifies a tradeoff; each sampled ensemble member supplies a plausible predicted fitness landscape.Source: §4 · p. 4, §4.1 · p. 5

4.1 SURROGATES AND SCALARIZED REWARDS

  • Augmented Chebyshev scalarization converts normalized vector predictions into scalar rewards, with weights changing the favored tradeoff and an augmentation term rewarding improvements in remaining objectives.Unlike linear scalarizations, this construction can recover optimal tradeoffs on nonconvex Pareto fronts.Source: §4.1 · p. 5
  • Scalar rewards let MODS reuse a diffusion model and compatible steering procedure, while changing weights redirects generation toward different tradeoffs.Figure 2 illustrates three steered distributions from the sixth GB1 round using the same surrogate but different scalarization weights.Source: Figure 2 · p. 4, §4.1 · p. 5

4.2 GENERATING CANDIDATES WITH DIFFUSION STEERING

  • For each sampled surrogate-and-weight pair, MODS approximately samples a reward-tilted distribution, balancing predicted reward against closeness to the pretrained diffusion model.Larger α emphasizes predicted reward; smaller α keeps the distribution closer to the diffusion model.Source: §4.2 · p. 5

4.3 THE COMPLETE MODS PROCEDURE

  • Each round fits surrogates on accumulated observations, generates a steered batch, evaluates its objectives, and adds the results to the dataset; the final output is the observed nondominated sequences.The surrogate member and weights are held fixed through each candidate’s denoising trajectory.Source: §4.3 · p. 6, §4.2 · p. 6

5 EXPERIMENTS

MODS is evaluated across six protein tasks using fixed assay-trained predictive oracles, multiple diffusion and steering configurations, and a shared evaluation budget. DAPS variants perform strongly across tasks, while the best steering strength depends on the model, procedure, and task.

5 EXPERIMENTS

  • The experiments compare MODS with three iterative protein-optimization baselines and examine sensitivity to diffusion model, steering procedure, and steering strength.The Table 1 caption identifies hypervolume results across six tasks, reporting means and standard errors over five repetitions.Source: §5 · p. 6, Table 1 · p. 7

5.1 EXPERIMENTAL SETUP

  • The six tasks cover five two-objective problems and one three-objective problem, evaluated with fixed predictive oracles trained on experimental assay data.Each oracle averages CNN-ensemble predictions and applies a task-scaled penalty beyond the labeled mutational neighborhood.Source: §5.1 · p. 6
  • MODS pairs MDLM and D3PM with classifier guidance or DAPS, training campaign-specific surrogate ensembles from observations without oracle-training data.The three baselines are ARLM + DPO, LaMBO-2, and SAGE-Prot; unlike them, MODS keeps its generative model fixed and decouples surrogates from it.Source: §5.1 · p. 6
  • Every method receives a budget of 900 sequences: 100 initial samples followed by eight rounds of 100 evaluations, repeated five times.Performance is measured by hypervolume over all evaluated sequences up to each round, using a reference point formed from the component-wise minima across all methods' evaluated sequences.Source: §5.1 · p. 7
  • For each diffusion model, the selected DAPS setting has higher final mean hypervolume than the selected CG setting on every task.The main results show MDLM + DAPS at α = 212 and D3PM + DAPS at α = 214, settings selected using separate random seeds.Source: §5.1 · p. 7

5.2 RESULTS

  • On five tasks, both DAPS variants achieve higher final mean hypervolume than ARLM + DPO, LaMBO-2, and SAGE-Prot; on Src, D3PM + DAPS achieves the highest final mean hypervolume.MDLM + DAPS reaches 66.4 ± 2.8 on Src, compared with 67.5 ± 2.4 for ARLM + DPO, 66.2 ± 3.2 for SAGE-Prot, and 45.7 ± 1.2 for LaMBO-2; each DAPS variant also improves over unsteered sampling from its diffusion model.Source: §5.2 · p. 7, Table 1 · p. 7
  • For GB1, ERBIN-PDZ1, KRAS, and Src, the highest-hypervolume Pareto fronts extend farther along complementary objective directions.Figure 4 omits the three-objective NHERF3-PDZ1 task because its Pareto front cannot be represented in a single two-dimensional plot.Source: §5.2 · p. 7, Figure 4 · p. 9
  • Steering strength affects performance differently across diffusion model–task combinations, with intermediate strength performing best in several MDLM settings.For MDLM, performance peaks at intermediate steering strength on KRAS and Src with DAPS, and NHERF3-PDZ2 and KRAS with CG; further strengthening reduces final hypervolume in those settings.Source: §5.2 · pp. 7–8

6 CONCLUSION

MODS uses sampled scalarization weights and surrogate ensemble members to steer fixed generative models toward different objective tradeoffs, updating surrogates as observations accumulate. Across six tasks, it performs strongly with both diffusion models, but the evaluation relies on predictive oracles and a limited set of tasks and model–steering choices.

6 CONCLUSION

  • MODS samples augmented Chebyshev scalarization weights and surrogate ensemble members to build rewards for different objective tradeoffs, while keeping the generative model fixed as surrogates update.The reward interface supports compatible inference-time steering procedures without retraining the generative model.Source: §6 · p. 9
  • Across six tasks, MODS either achieves the highest mean final hypervolume or performs similarly to evaluated baselines for each diffusion model.Performance depends on the choice of generative model and steering procedure.Source: §6 · p. 9
  • The evaluation uses predictive oracles trained and validated on experimental assay data, so laboratory studies remain necessary to test whether gains translate to experimentally measured fitness.The oracles apply a distance-based penalty beyond labeled mutational neighborhoods, where prediction accuracy is uncertain.Source: §6 · p. 9
  • The authors identify evaluation on a limited set of protein tasks, diffusion priors, and steering procedures as a boundary on establishing MODS's generality.They suggest studies across more protein families, generative models, and steering procedures.Source: §6 · p. 9

AI USE STATEMENT

  • The study reports using generative AI for research-methodology or experiment feedback, translation, software-code editing, information sourcing, and paper readability edits.The authors state they reviewed and verified AI-assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

  • The study uses publicly available datasets and reports no human subjects, personally identifiable information, sensitive personal data, or private or proprietary user data.The authors considered biosafety and dual-use concerns and state they do not intentionally generate or release sequences with known harmful function.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • At submission, the planned release of code, processed data, checkpoints, configurations, and analysis materials was not yet publicly available.An anonymized reproducibility package was to be made available to reviewers during rebuttal.Source: REPRODUCIBILITY STATEMENT · p. 10

A DATASETS AND PREPROCESSING

Each task uses separate unlabeled homolog sequences to train the diffusion model and labeled mutational sequence–fitness data to train predictive oracles. The homolog data are processed into fixed-length sequences with redundancy reduction, while the labeled data provide fitness targets.

A DATASETS AND PREPROCESSING

  • Diffusion training uses unlabeled UniRef90 homologs, while labeled mutational sequence–fitness pairs train the oracles as proxies for wet-lab measurements.The two data sources remain separate: fitness labels are not used for diffusion training, and homolog data are not treated as fitness labels.Source: A DATASETS AND PREPROCESSING · p. 15
  • Homologs are retained only with at least 75% gap-free parent-sequence coverage, then processed into fixed-length pseudo-natural sequences.Insertions are removed and gaps are replaced with parent residues.Source: A DATASETS AND PREPROCESSING · p. 15
  • To reduce redundancy, homologs are clustered at 80% identity, 5% of clusters are held out, and remaining sequences are resampled with cluster-size weighting.For GB1, all single amino acid substitutions augment the relatively small homolog set.Source: A DATASETS AND PREPROCESSING · p. 15
  • Fitness datasets come from mutational studies; GB1 stability labels combine published single-mutant estimates with double-mutant values formed by summing corresponding single-mutant effects.For other proteins, fitness values are taken directly from the corresponding mutational studies.Source: A DATASETS AND PREPROCESSING · p. 15

B DIFFUSION MODELS AND SAMPLE QUALITY

The learned diffusion models produce sequences with perplexities that differ from random sequences and better approximate the training-data distribution. Across all six protein groups, generated-sequence diversity and mean predicted fitness trade off.

B DIFFUSION MODELS AND SAMPLE QUALITY

  • Figure B.1 compares training-data and generated-sequence perplexities with random and training sequences, using ProGen2-base likelihood as the reference.The 764 M-parameter ProGen2-base model supplies the perplexity calculation.Source: B DIFFUSION MODELS AND SAMPLE QUALITY · p. 16
  • Across all datasets, learned-model perplexities differ clearly from random sequences and better approximate the training-data distribution.Random samples consistently have the highest perplexity, while natural MSA sequences used for training have much lower perplexity.Source: B DIFFUSION MODELS AND SAMPLE QUALITY · p. 16
  • Across all six protein groups, random samples have high diversity but the lowest mean predicted fitness, while target sequences occupy the high-fitness, low-diversity corner.Figure B.1’s scatter plots compare mean oracle fitness with sequence diversity, defined as average per-position amino-acid Shannon entropy.Source: B DIFFUSION MODELS AND SAMPLE QUALITY · p. 16

C FITNESS ORACLE EVALUATION AND VALIDATION

The fitness oracles average predictions from ten CNNs trained on sequence–fitness data. Their predictions are penalized beyond a reference-sequence distance because oracle accuracy becomes less reliable farther from the labeled neighborhood.

C FITNESS ORACLE EVALUATION AND VALIDATION

  • Each objective oracle averages predictions from an ensemble of ten one-dimensional CNNs trained on one-hot amino-acid sequences.The networks use two convolutional layers, batch normalization, GELU activations, and a final linear regression head.Source: C FITNESS ORACLE EVALUATION AND… · p. 16
  • The oracle penalty discourages optimization into distant sequence regions, where limited-neighborhood training data make predictions increasingly unreliable.The penalty is applied to raw oracle predictions to discourage exploration beyond the labeled neighborhood.Source: C FITNESS ORACLE EVALUATION AND… · p. 17
  • Predictions remain unchanged through the reference distance d0, after which a Hamming-distance penalty is controlled by λm.The experiments use d0 = 4 for the shorter-sequence tasks and d0 = 8 for KRAS and Src; λm is set to the training-label standard deviation for that objective.Source: C FITNESS ORACLE EVALUATION AND… · p. 17

D SURROGATE ARCHITECTURE AND TRAINING

MODS trains campaign-specific surrogate ensembles on accumulated observations, with different input handling for CG and DAPS. The appendix also reports how steering strength was selected and summarizes results and sensitivity analyses.

D SURROGATE ARCHITECTURE AND TRAINING

  • Each surrogate is an MLP that embeds one-hot amino-acid sequences into 64 dimensions and predicts the objectives through fully connected layers.For CGsteering, the input features are concatenated with a 64-dimensional timestep embedding.Source: D SURROGATE ARCHITECTURE AND TRAINING · p. 17
  • CG surrogates train on corrupted sequences and timesteps but target the clean sequence’s oracle objective vector; DAPS surrogates train on clean sequences.CG’s corruption process supplies examples across noise levels without additional oracle queries.Source: D SURROGATE ARCHITECTURE AND TRAINING · p. 17
  • Each optimization round trains ten newly initialized surrogate ensemble members from scratch on all observations collected during that campaign.Members use different random seeds and jointly minimize mean squared error over objectives; CG runs 1000 epochs and DAPS 200.Source: D SURROGATE ARCHITECTURE AND TRAINING · p. 17

E HYPERPARAMETER AND STEERING SENSITIVITY

  • Steering strength was selected using two-restart mean results by minimizing the largest hypervolume drop across tasks for each configuration.Table E.1 reports the results used for this selection.Source: E HYPERPARAMETER AND STEERING… · p. 17
  • The appendix reports five-restart MODS and baseline results, p-value comparisons, and hypervolume sensitivity to steering strength across rounds.Tables E.2 and E.3 report results and p-values; Figures E.1 and E.2 show steering-strength effects.Source: E HYPERPARAMETER AND STEERING… · p. 17

F EXTENDED OPTIMIZATION RESULTS

The extended results characterize MODS computational costs and sequence-distance behavior. The supplied tables report final hypervolume summaries and comparisons against baselines, but their captions do not provide the results’ cell values.

F EXTENDED OPTIMIZATION RESULTS

  • Surrogate training costs increase over optimization rounds, CG surrogate training is generally more expensive, and DAPS generation costs more than CG.CG surrogates are conditioned on the denoising step; DAPS generation incurs Metropolis-Hastings steps at each denoising step.Source: F EXTENDED OPTIMIZATION RESULTS · p. 17
  • Final hypervolume is reported as a five-repeat mean ± standard error, while hyperparameter-selection results use means over two repetitions.The tables mark the highest mean among a method’s parameters in bold and the highest mean across methods in green.Source: Table E.1 · p. 18, Table E.2 · p. 19
  • All baseline-comparison p-values are below 0.05 except for MODS (MDLM + DAPS) on Src.Table E.3 compares the Table 1 MODS configurations against each baseline for every problem.Source: Table E.3 · p. 20
  • The computational-cost table reports mean ± standard deviation in seconds on one NVIDIA RTX A6000, with generation cost independent of round and surrogate training growing with the labelled pool.Surrogate training grows as optimization accumulates labelled data.Source: F EXTENDED OPTIMIZATION RESULTS · p. 17, Table F.1 · p. 20

LaMBO-2

  • For all methods, many evaluated sequences exceed d0 = 8, so the distance penalty remains soft and permits evaluation beyond the oracle-training mutational neighborhood.The dotted line marks d0 = 8; the penalty applies above that threshold.Source: LaMBO-2 · p. 20, Figure F.1 · p. 26
  • MODS with MDLM produces closer but generally wider parent-sequence distances than MODS with D3PM, and distance from the parent has no direct relationship with hypervolume.The distance distributions vary across methods, generative models, and steering mechanisms.Source: LaMBO-2 · p. 20
Loading 2610.28943v1…