Robust Model-Based Optimization for Challenging Fitness Landscapes
Authors: Saba Ghaffari
Ehsan Saleh
Alex Schwing
Yu-Xiong Wang
Martin D. Burke
Saurabh Sinha
Venue: ICLR 2024 poster (2024)
Keywords: Model based optimization
protein engineering
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Protein design, a grand challenge of the day, involves optimization on a fitness landscape, and leading methods adopt a model-based approach where a model is trained on a training set (protein sequences and fitness) and proposes candidates to explore next. These methods are challenged by sparsity of high-fitness samples in the training set, a problem that has been in the literature. A less recognized but equally important problem stems from the distribution of training samples in the design space: leading methods are not designed for scenarios where the desired optimum is in a region that is not only poorly represented in training data, but also relatively far from the highly represented low-fitness regions. We show that this problem of “separation” in the design space is a significant bottleneck in existing model-based optimization tools and propose a new approach that uses a novel VAE as its search model to overcome the problem. We demonstrate its advantage over prior methods in robustly finding improved samples, regardless of the imbalance and separation between low- and high-fitness samples. Our comprehensive benchmark on real and semi-synthetic protein datasets as well as solution design for physics-informed neural networks, showcases the generality of our approach in discrete and continuous design spaces. Our implementation is available at https://github.com/sabagh1994/PGVAE.

Published as a conference paper at ICLR 2024
ROBUST MODEL-BASED OPTIMIZATION FOR CHALLENG- ING FITNESS LANDSCAPES
Saba Ghaffari' Ehsan Saleh ! Alexander G. Schwing!
Yu-Xiong Wang'
Martin D. Burke' 1
Saurabh Sinha?
'University of Minoii Jrbana-Champaa "Georgia Institute of Technology
(sabag2, ehsans2, aschwing, yxw, mdburke] @ sillincis.edee saurabh. s sinha@bme, gatech.edu.
ABSTRACT
Protein design, grand challenge of the day, involves optimization on fitness landscape, and leading methods adopt a model-based approach where a model is trained on training set (protein sequences and fitness) and proposes candidates to explore next. These methods are challenged by sparsity of high-fitness samples in the training set, a problem that has been in the literature. A less recognized but equally important problem stems from the distribution of training samples in the design space: leading methods are not designed for scenarios where the desired optimum is in region that is not only poorly represented in training data, but also relatively far from the highly represented low-fitness regions. We show that this problem of "separation" in the design space is a significant bottleneck in existing model- based optimization tools and propose a new approach that uses a novel VAE as its search model to overcome the problem. We demonstrate its advantage over prior methods in robustly finding improved samples, regardless of the imbalance and separation between low- and high-fitness training samples. Our comprchensive benchmark on real and semi-synthetic protein datasets as well as solution design for physics-informe neural networks, showcases the generality of our approach in discrete and continuous design spaces. Our implementation is available at
INTRODUCTION
Protein engineering is the problem of designing novel protein sequences with desired quantifiable properties, e.g., enzymatic activity, fluorescence intensity, for variety of applications in chemistry and bioengineering (Fox et al.) 2007) et al., 2017 Biswas et al., (02T). Protein engineering approached by optimization over the protein fitness landscape which specifies the mapping between protein sequences and their measurable property, i.e., fitness. landscape is extremely sparse, i.e., only a minuscule fraction of sequences have non-zero fitness, and rugged, i.c.. peaks of "fit" sequences are narrow and separated from each other by deep valleys (Romero & Arnold 2009), which greatly complicates the problem of protein design. Directed evolution is the most widely adopted technique for sequence design in laboratory environment (Arnold 1998). In this greedy local search approach, first set of variants of naturally occurring ("wild type") sequence are tested for the desired property, then the variants with improved property form the starting points of the next round of mutations (selected uniformly random) and thus the next round of sequences to be tested. This process is repeated until an adequately high level of desired property achieved. Despite advances, this strategy remains costly and laborious, prompting the development of model- -uidel searching schemes that support more efficient exploration of the sequence space (Biswas et al. 2018 Brookes & Listgarten 2018 G6mez--Bobann
believed that the protein fitness
et al., 2018 Brookes et al.) 2019: Angermueller et al., 2019; [Sinai et al. 2020 Ren et al.] 2022). In particular, there is emerging agreement that optimization schemes that utilize ML models of the sequence-fitnen relationship, learned from training sets that grow in size the optimization progresses, can furnish better candidates for the next round of testing, and thus accelerate optimization, as compared to model-free approaches such Bayesian optimization (Mockus) 2012 Sinai et al. 2020). Our work belongs to this genre of model- based optimization for sequence-function function landscapes.
Botth authors contributed equally.

---

Published as a conference paper at ICLR 2024
Fitness
Figure 1: Challenges of imbalance and separation in fitness landscape. Each plot shows sequence space (x-y plane) and fitness landscape whitebbie- gradient), along with training data composition (white circles and stars). (A-C, left to right) In each of these hypothetical scenarios, sparsity of high- training samples (white stars) relative to low-fitness samples (white circles), also called "imbalance" presents challenge for MBO. Moreover, panel C shows greater degree of separation between low- and high-fitness samples, compared to B and A, presenting significant additional challenge for MBO, above and beyond that due imbalance. The rightmost panel the schematic representation of real-world dataset of enzyme variants designed for an unnatural substrate (xyz) distinct from the substrate of the wild-type enzyme (xyz). The dataset comprises few non-zero fitness variants (stars) that are far from the bulk of training samples, which have zero fitness (white circles). Hypothetical peaks have been drawn the rare non-zero fitness variants, to illustrate that the fitness landscape presents the twin challenges of imbalance and separation, similar to that in panel C.
Intuitively, the success of fitness optimization depends on the extent to which functional proteins are represented in the experimentally derived data (training set) so that the characteristics of desired sequences can be inferred from them. Prior work has examined this challenge of "sparsity" in fitness landscapes, proposing methods that use a combination of "exploration" " and "exploitation" to search in regions of the space less represented in training data (Romero al.) [013]] Gonzalez al. 2015 Yang et al. 2019; Hie & Yang. (022). Optimization success also depends on the distribution of training samples in the sequence space, in particular
whether the desired functional sequences
are proximal to and easily reachable from the frequent but low-fitness training samples. This second challenge of "separation" (between the optima and training samples) in fitness landscape is relatively unexplored in the literature. In particular, it is not known how robust current search methods are when the optimum located in a region that is poorly represented in the training set and is located relatively far (or separated due to rugged landscape) from the highly represented regions Figure [ A real-world example of this is the problem of designing an enzyme for an unnatural target substrate, starting from the wild-type enzyme for related natural substrate. Most variants of the wild type enzyme are not functional for the target substrate, thus the training set is sparse in sequences with non-zero fitness; furthermore, the rare variants that do have non-zero activity (fitness) for the target substrate are located relatively far from the wild-type and its immediate neighborhood that forms the bulk of the training set Figure [ (rightmost panel).
We study the robustness of model-guided search schemes to the twin challenges of imbalance and separation in fitness landscape. We explore for the first time how search algorithms behave when training samples of high fitness are rare and separated from the more common, low-fitness training samples. (Here, separation is in the design or sequence space, not the fitness space.) Furthermore, given a fixed degree of separation, we investigate how the imbalance between the low- and high-fitness samples the training set affect the performance of current methods. A robust algorithm should have consistent performance under varying separation and imbalance.
To this end, we propose a new model-based optimization (MBO) approach that uses VAE (Kingma & Welling 2014) as its search model. The latent space of our VAE is explicitly structured by property (fitness) values of the samples (sequences) such that more desired samples are prioritized over the less desired ones and have higher probability of generation. This allows robust exploration of the regions containing more desired samples, regardless of the extent of their representation in the train set and of the extent of separation between low- and high- samples in the train set. We refer to the proposed approach as a "Property-Prionte Prioritized Generative Variational Auto- Encoder" (PPGVAE).
Our approach is designed with the goal of obtaining improved samples in less number of MBO steps (less sampling budget), as desired in the sequence design problem. Methods that rely on systematic exploration techniques such Gaussian processes (66mez-Bombarns not converge in small number of rounds (Srinivas et al. 2009);: a problem that is exacerbated by higher dimensionality of the search space (Frazier 2018
et al.) (018)) may
Djolonga et al. 2013). In general, optimization
2

---

Published as a conference paper at ICLR 2024
VAE
PPGVAE (Hard Constraint)
PPGVAE (Soft Constraint)
2
-4
10
Z1
Z1
Z1
Figure 2: Latent space of our PPGVAE vs Vanilla VAE. PPGVAE and vanilla VAE were trained on a toy MNIST-derived dataset where property values decrease monotonically with digit value (zero has highest property value). Vanilla VAE (Left) scatters the rare samples of digit zero (blue) and samples of next-highest property value (digit one, orange) in the latent space, whereas PPGVAE (Middle and Right) maps digits with higher property values closer to the origin. This results in the classes with greatest property values having higher probability of generation. PPGVAE was run two modes, where the relationship loss was enforeed in strong (Middle) or soft (Right) manner (see text).
with fewer MBO steps can be achieved by either 1) bringing more desired (higher fitness) samples closer together and prioritizing their exploration over the rest, as done our approach, or 2) using higher weights for more desired samples in a weighted optimization setting (Brookes & Listgarten 2018 Brookes et al., 2019; Gupta & Zou 2019). Neither of these can be achieved by methods that the the property values (Kang & Cho (2018) or encode the properties as separate latent variables along with the samples (Guo et al. 20000 Chan et al. 2021). This is the key methodological gap in the sate-of-the-arr that is addressed by new VAE technique
condition the generation of samples
for masel-based optimization.
Through extensive benchmarking on real and semi-synthetic protein datasets we demonstrate that MBO with PPGVAE is superior to prior methods in robustly finding improved samples regardless of 1) the imbalance between low- and high-fitness training samples, and 2) the extent of their separation in the design space. Our approach is general and not limited to protein sequences, i.c., discrete design spaces. We further investigate MBO with PPGVAE on continuous designs spaces. In an application to physics-informed neural networks (PINN) (Raissi et al., 2019). we showcase that our method can consistently find improved high quality solutions, given PINN-derived solution sets overpopulated with low quality solutions separated from rare higher quality solutions. In section 2 MBO is reviewed.
PPGVAE explained in section followed by experiments in section
2
BACKGROUND
Model Based Optimization. Given (z. y) pairs as the data points, e.g., protein sequence x and its associated property (e.g.. pKa value), the goal of MBO is to find æ E X' that satisfy an objective S regarding its property with high probability. This objective can be defined as maximizing the property value y, i.e., S (vly > ym } where ym is some threshold. Representing the search model with (2) (with parameters 0), and the property oracle as pa(y(z) (with parameters (3). MBO is commonly performed via an iterative process which consists of the following three steps at iteration t (Fannjiang & Listgarten 2020
1. Taking K samples from the search model, Vi E {I...., xf ~ Por
2. Computing sample-specific weights using monotonic function f which method-speeiife
w; := fp E S)z)));
(1)
3. Updating the search model parameters via weighted maximum likelihood estimation (MLE):
K
g(+1 = arg max
wy log(po(z)))
(2)
The last step optimizes for search model that assigns higher probability to the data points satisfying the property objective SS, ie., where pa(y E Sl(r) is high. Prior work by (Brookes & Listgarten
3

---

Published as a conference paper at ICLR 2024
2.5
AN
0.3
0.1
MYmar
1.5
0.03
0.01
0.0
Lowest
Highest
Au
Step
Separation
Low Separation
Medium Separation
High Separation
AYmax
0.03
0.01
0.1
0.3
1
0.1
0.3
0.1
0.3
1
PPGVAE
CEM-PI
RWR
Figure 3: Robustness to imbalance and separation in MBO for GMM. A bimodal GMM is used as the property oracle (top Left), i.e., fitness (Y) landscape on one-dimensional sequence space (X). Separation is defined as the distance between the means of the two modes (Ap). Higher values of Aμ are associated with higher separation. Train sets were generated by taking N samples from the less desired mode 141 and pN (imbalance ratio p S 1) samples from the For a fixed separation, PPGVAE achieves robust relative improvement of the highest property sample generated (AY max). regardless of the imbalance ratio (Bottom panels). Performance of PPGVAE. aggregated over all imbalance ratios, stays robust to increasing separation (top Right). PPGVAE converges in less number of MBO steps (top Middle).
desired mode μl2-
2018) (DbAS) and (Brookes et al. 2019) (CbAS) have explored variants of weighting schemes for the optimization in the second step.
weighted regression (RWR) (Peters & Schaal 2007)) and CEM-PI (Snoek et al.) 2012) have additionally been benchmarked by CbAS, each providing different weighting scheme. RWR has been used for policy learning Reinforeement Learning (RL)
and CEM-PI maximizes the probability of improvement over the best current value using the cross entropy method (Rubinstein] 1997[1999)
Common to all these methods is that weighted MI_E could suffer from reduced effective sample size. In contrast, our PPGVAE does not use weights. Instead, it assigns higher probability to the high fitness (more desired) data points by restructuring the latent space. Thus, allowing for the utilization of all samples in training the generative model (see Appendix(A[C]
Exploration for Sequence Design. In addition weighting based generative methods, model-based RL (Angermueller et al.) 2019) (Dyna PPO) and evolutionary greedy approaches (Sinai et al. 2020) (AdaLead) have been developed to perform search in the sequence space for improving fitness. More recently, (Ren et al. 2022) (PEX) proposed an evolutionary search that prioritizes variants with improved property which fall closer to the wild type sequence.
PROPERTY-PRIORITIE GENERATIVE VARIATIONAL AUTO-ENCODER
To prioritize exploration and generation of rare, high-fitness samples, our PPGVAE uses property (fitness) values to restructure the latent space. The restructuring enforces samples with higher property to lie closer to the origin than the ones with lower property. As the samples with higher property lie closer to the origin, their probability of generation higher under the VAE prior distribution (0, I). Representing the encoder and its parameters with Q and 0, the structural constraint on N samples is imposed by
((bg)1b)) i, E (1,.., log(Pr((()) log (og(P(()6)) T(u -y)),
4

---

Published as a conference paper at ICLR 2024
Low Separation
High Separation
3
3
1.0
2
2
2
0.8
1
1
1
0.6
0
0
0.4
x
-2
-2
-2
2
4
2
4
2
4
PC1
PC1
PC1
200
100
50
Low Separation
0.3
0.3
0.1
4a
0.03
Smaller
0.01
0.01
0.03
0.1
0.3
0.01
0.03
0.1
0.3
Lowest
Medium
Highest
p
Separation
PPGVAE
AdaLead
CbAS
CEM-PI
RWR
Figure 4: Robustness imbalance and separation in MBO for AAV dataset. PCA plot for protein sequences in the dataset, colored with their property values (top Left). Blue and red color spectrum are used for less and more desired samples, respectively. Top middle and right panels show train sets with low and high separation, respectively, between the abundant less-desired and rare more-desired samples. PPGVAE achieves robust relative improvements (shown here for the low separation scenario), regardless of the imbalance ratio (bottom Middle). Its performance also stays robust increasing separation (bottom Right). PPGVAE performance is only slightly affected by reducing its sampling budget MBO step (N,) (bottom Left).
where 2á (2;) and are the latent space representation and property value of sample x, respectively. The probability of the encoded representation Pr(u)) is computed w.r.t. the VAE prior distribution (o, I) over the latent space, i.e., Pr(u() exp( -4
Intuitively, if higher values property y are desired, then y/ S y: results in sample getting mapped closer to the origin. This results in a higher probability of generating sample i than sample j. The extent prioritization between each pair of samples is controlled by the hyper-paramedt т, often referred as the temperature. The structural constraint is incorporated into the objective of vanilla VAE as relationship loss that should be minimized. This loss is defined as,
Cr
Σ
(I(1og(Pr((E)) log(Pr((()) T(y: -
Combined with the vanilla VAE, the final objective of PPGVAE to be maximized is,
DKL (((( - -
(5)
where Xy is hyper-parameet controlling the extent which the relationship constraint is enforced.
Here, we abused the notation and wrote Q(=(x) (x) for Pr(=)Q(z)). To understand the impact of the structural constraint on the mapping of samples in the latent space, we define 9; := log(Pr((()) TWh- -
Also, assume that the samples are independent and identically distributed (i..d.). Then minimizing the relationship loss can be rewritten
C, x Eg..g ((g: 9;) = E (((g: - E(g:)) (aj E(g))))
(6)
Using the i.i.d. assumption this is simplified to,
C, x 2Var(g)) = (Va((lo((((H)) -y:)) = 2Var( Pg T
pbg
Ty:)
(7)

---

Published as a conference paper at ICLR 2024
Low Separation
30
1.0
0.8
0.3
1SNE2
0.1
0
0.4
0.03
-10
0.2
0.0
-20
20
0.01
0.03
0.1
0.3
Lowest
Highest
tSNE1
Separation
Low Separation
30
ON pcl
1.0
20
0.8
0.3
10
0.1
ESNE2
0
0.03
Smaller er N
0.4
0.01
-10
0.2
-20
0.0
0.01
0.03
0.1
0.3
0
8
12 step
16
20
-20
20
ISNE1
PPOVAE
AdaLead
CbAS
CEM-PI
RWR
Figure 5: Robustness to imbalance and separation in MBO for GB1 dataset. The 1SNE plot for the appended sequences of semi-synthetic GB1 dataset (top Left). Bottom left panel represents an example of train set for low separation between less and more desired samples, i.e., appended sequence of length three (see Figure A3 for an example of high separation). For a fixed separation level, PPGVAE provides robust improvements, regardless of the imbalance ratio (top Middle). is also robust to the degree of separation, measured by aggregated performance over all imbalance ratios (top Right). PPGVAE has faster convergence (bottom Right) and achieves similar improvements with less sampling budget per MBO step (N,) (bottom Middle).
Therefore, minimizing C, is equivalent to minimizing the variance. This is equivalent to setting the random variable RHS of Equation to a constant value C: Vi: 11 Hb = 2(C -ry)). This implies the distribution of samples with the same property value to be the same sphere. The sphere lies closer to the origin for the samples with higher property values. This ensures that higher property samples have higher probability of generation under the VAE prior N N(0, I), while allowing for all samples to fully contribute to the optimization.
To demonstrate the impact of relationship loss the latent space, PPGVAE, with two-dimensional latent space, was trained on toy MNIST dataset (Deng) 2012). The dataset contains synthetic property values that monotonically decreasing with the digit class. Also, samples from the digit zero with the highest property are rarely represented (see Appendix [D3).. Strong enforcement of the relationship loss, using relatively large constant A,, aligns samples from each class on circles whose radius increases as the sample property decreases. Soft enforcement of the relationship loss, by gradually decreasing \y, makes the samples more dispersed, while roughly preserving their relative distance to the origin (Figure
EXPERIMENTS
We will compare the performance of MBO with PPGVAE to the baseline algorithms that use generative models optimized with weighted MLE (CbAS, RWR, CEM-PI (Brookes et al. 2019)) for search in discrete (c.g., protein sequence) and continuous design spaces. For sequence design experiments, we additionally included AdaLead as a baseline.
In all optimization tasks, we 1) provide definition of separation between less and more desired samples in the design space X, 2) study the impact of varying imbalance between the representation of low- and high-fitness (property value) samples in the training set, given fixed separation degree, and 3) study the impact increasing separation. The ground truth property oracle was used in all experiments. The performance is measured by AY representing the relative improvement of the

---

Published as a conference paper at ICLR 2024
Low Separation
1.0
0.8
20
0.3
ISNE2
0.6
&Ymat
0.1
0.4
0.03
-20
0.2
0.01
0.0
-20
0.03
0.3
Highest
ISNE1
Separation
Low Separation
+
200
+**
100
-#-
50
1
ON
0.8
0.3
20
0.1
ISNE2
0.6
0.4
Smaller
0
0.0
0.01
0.03
D.1
0.3
0
16
20
-20
20
step
ISNE1
PGVAE
AdaLead
CbAS
CEM-P
RWR
Figure 6: Robustness to imbalance and separation in MBO for PhoQ dataset. The panels and
their corresponding observation have the same semantics as in Figure
highest property found by the model the highest property in the train set, i.e., initial set the beginning of MBO (see Appendix D for further details).
4.1 GAUSSIAN MIXTURE MODEL.
We bimodal Gaussian mixture model (GMM) the property oracle, in which the relationship between r and property y is defined as, We choose the second Gaussian as the more desired mode by setting 02 > an (Figure 3}. Here, separation is defined as the distance between the means of the two Gaussian modes (Ayx). Larger values of Ap are associated with higher separation. For each separation level, the train sets were generated by taking N samples from the less desired mode, and taking pN (imbalance ratio p S 1) samples from the more desired mode (see Appendix D.3).
exp(-(x M1)2 /2/ ) a2 exp(-(< (pp--(x- H2) 22222).
For fixed separation level, compared the performance of PPGVAE and baseline methods for varying imbalance ratios. The relative improvements achieved by PPGVAE are consistently higher than all other methods and are robust to the imbalance ratio. Other methods achieve similar improvements when high-fitness samples constitute a larger portion the train set (larger p). This happens smaller for lower separation levels (smaller Aμr,, indicating that the impact of imbalance is offset by a smaller separation between the optimum and the dominant region training data (see Figure] Figure A4).
We then compared the relative improvement aggregated over all imbalance ratios, as the separation level increases. All methods perform well for low separation levels. PPGVAE stays robust the degree of separation, whereas the performance of others drops by increasing separation (see Figure top right). The difficulties encountered by other generative methods higher separation levels is due to the difference between the reconstruction of low- and high-fitness samples. As Aμ increases, the more desired (high fitness) samples get mapped to a farther locality than the less desired ones. This makes the generation of more desired samples less likely under the VAE prior distribution A/(0, 7). This is exacerbated when more desired samples are rarely represented (small imbalance ratio) in the train set. For smaller Aju, more desired samples get mapped to a similar locality as less desired ones, regardless of their extent of representation in the train set. PPGVAE stays robust to the the imbalance ratio and extent of separation, desired samples over the less desired ones by mapping them closer to the latent space origin. Similar explanation be provided for the rest of optimization tasks benchmarked in this study.
it always prioritizes generation and exploration of more
7

---

Published as a conference paper at ICLR 2024
Low Separation
High Separation
50
50
50
tSNE2
0
0
-50
-50
-50
-50
50
-50
0 ISNEl
50
-50
0 tSNE1
50
tSNE1
Low Separation
1
0.3
0.3
AV mom
0.1
0.1
0.03
0.03
0.01
0.01
0.1
0.3
Lowest
Medium
Highest
12
16
20
P
Separation
MBO Step
PPGVAE
CbAS
CEM-PI
RWR
Figure 7: Robustness to imbalance and separation in MBO for PINN. The tSNE plot for the PINN-derived solutions colored with their property values (top Left). Blue and red color spectrum used for lower and higher quality solutions, respectively. Top middle and right panels show train sets with low and high separation between the abundant less-desired and rare more-desired samples. PPGVAE achieves robust relative improvements, regardless the imbalance ratio p (bottom Left). Its performance stays robust to increasing separation (bottom Middle). It also converges in less number of MBO steps compared other methods (bottom Right).
4.2 REAL PROTEIN DATASET
To study the impact of separation, real protein datasets with property (activity) measurements more broadly dispersed in the sequence space are needed. Such datasets are rare among the experimental studies of protein fitness landscape (Johnston et al.) 2023), past efforts have mostly been limited to optimization around wild type sequence. We chose the popular AAV (Adeno-ar Aadeno-sssociated virus) dataset (Bryant al. 2021) in ML-guided design (Sinai et al. 202 Mikos al, 2021]. This dataset consists of virus viability (property) measurements for variants of AAV capsid protein, covering the span of single 28-site mutations, thus presenting wide distribution in sequence space.
The property values were normalized to [0, 1] range. Threshold of 0.5 was used to define the set of low (y < 0.5) and high (y > 0.5) fitness (less and more desired respectively) mutants in the library. To define a proxy measure of separation, we considered the minimum number of mutated sites in the low-fitness training samples. larger value of possible only when the low-fitness samples are farther from the high-fitness region, i.e., at higher separation (see Figure D. For a fixed separation train sets were generated by taking N samples from the low-fitness mutants containing at least number of mutations, and pN (p 1) samples from the high-fitness mutants (regardless of the number mutations) (see Appendix D).
Given a fixed separation level, PPGVAE provides robust relative improvements, regardless of the imbalance ratio (Figure 4 Figure A5). CEM-PI is the most competitive with PPGVAE for low separation levels, however its performance decays for small imbalance ratios separation increases (see Figure A5). The performance other methods improve as the high-fitness samples represent a higher portion of the train set.
Next, we compared the performance of methods, aggregated over all imbalance ratios, as the separation level increases. PPGVAF, is the most robust method to varying separation. CEM-PI is
8

---

Published as a conference paper at ICLR 2024
the most robust weighting- based generative method. Its performance is similar to PPGVAE for low separation levels and degrades as separation increases.
As it is desired to achieve improved samples with less sampling budget, also studied the impact of varying sample generation budget per MBO step, given a fixed separation level. As expected, the performance increases by increasing the sampling budget for both CEM-PI and PPGVAE. Further- more, on average PPGVAE performs better than CEM-PI for all sampling budgets (Figure I This is the direct result of prioritizing the generation and exploration of high-fitness samples relative the low-fitness ones in the latent space of PPGVAE.
4.3 SEMI-SYNTHETIC PROTEIN DATASETS
Next, used two popular protein datasets GB1 (Wu et al. 2016 and PhoQ (Podgornaia & Laub 2015) with nearly complete fitness measurements on variants of four sites. However, these data sets exhibit narrow distribution in sequence space (at most mutations), and are not ideal study the effect of separation. We thus transformed the landscape of these proteins to generate a more dispersed dataset of mutants on which we could control for the separation of less and more desired variants. In this transformation, first certain threshold on the property was used to split the dataset into two sets of low and high fitness mutants (see Appendix (D). Second, specific sequence of length I, was appended to high fitness mutants, while a random sequence of the same length was appended to the low fitness ones. Here, the length of the appended sequence determines the extent of separation. Higher separation is achieved by larger values of L and makes the optimization more challenging. Note that the separation is achieved by changing the distribution of samples in the design space X, while keeping the property values unchanged. For each separation level, train sets were generated by taking N samples from the low fitness mutants and pN samples from the high fitness mutants (see Figures and Figure A3).
For a fixed separation level, the performance of all methods improve as high fitness samples constitute higher portion of the train set, i.e., higher imbalance ratio (Figure 5 PPGVAE is more robust to the variation of imbalance ratio and it is significantly better than others when high fitness samples are very rare. Same observations hold for all separation levels studied (see Figures A6 A7). Similar as before, CEM-PI is the most competitive with PPGVAE for low separation; however as the separation increases its performance decays. Furthermore, the reduction of sampling budget does not affect the performance of PPGVAR as much as CEM-PI (Figures 5).
4.4 IMPROVING PINN-DERIVED SOLUTIONS TO THE POISSON EQUATION
Our method can also be used in continuous design spaces. We define the task of design as finding improved solutions given a training set overpopulated with low quality PINN-derived solutions. Details on train set generation and separation definition are covered in Appendix [D.3] Similar conclusions hold for the robustness of PPGVAE to the separation level and imbalance ratio (Figure 7.
Common in all optimization tasks, PPGVAE achieves maximum relative improvement in less number of MBO steps than others (see Figures A9 A11 Characteristies of the latent space and sample generation have been studied for PPGVAE, and contrasted with prior generative approaches including (Gômez-Bombarell et al. 2018) in Appendix A Sensitivity of PPGVAE performance to the temperature discussed in Appendix E. (see Figure AI3
5 CONCLUSION
We proposed a robust approach for design problems, in which more desired regions are rarely explored and separated from less desired regions with abundant representation to varying degrees. Our method is inherently designed to prioritize the generation and interpolation of rare more-desired samples which allows it to achieve improvements in less number of MBO steps and less sampling budget. As it stands, our approach does not use additional exploratory mechanism to achieve improvements, however it could become more stronger by incorporating them. It is also important to develop variants of our approach that are robust to oracle uncertainties. and study the extent to which imposing the structural constraint can be restrictive in some design problems.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This research was funded by Molecule Maker Lab Institute: an AI research institute program supported by National Science Foundation under award No. 2019897. This work utilized resources supported by 1) the National Science Foundation's 's Major Research Instrumentation program, grant No. 1725729 (Kindratenko al., 2020), and 2) the Delta advanced computing and data resource which is supported by the National Science Foundation (award OAC 2005572) and the State of Illinois.
REFERENCES
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International conference on learning representations, 2019.
Frances H Amold. Design by directed evolution. Accounts of chemical research, 31(3):125-131, 1998.
Surojit Biswas, Gleb Kuznetsov, Pierce J Ogden, Nicholas J Conway, Ryan P Adams, and George M Church. Toward machine- machine-guided guided design of proteins. BioRxiv, pp. 337154, 2018.
Surojit Biswas, Grigory Khimulya, Ethan Alley, Kevin M Esvelt, and George M Church. Low-n protein engineering with data-efficient deep learning. Nature methods, 18(4):389-396, 2021.
David Brookes, Hahnbeom Park, and Jennifer Listgarten, Conditioning by adaptive sampling for robust design. In International conference on machine learning, pp- 773-782, PMLR, 2019.
David H Brookes and Jennifer Listgarten. Design by adaptive sampling. arXiv preprint arXiv:1810.03714, 2018.
Drew H Bryant, Ali Bashir, Sam Sinai, Nina K Jain, Pierce J Ogden, Patrick Riley, George M Church, Lucy Colwell, and Eric D Kelsic. Deep diversification of an aav capsid protein by machine learning. Nature Biotechnology, 39(6):691669) 2021.
Alvin Chan, Ali Madani, Ben Krause, and Nikhil Naik. Deep extrapolation for generation. Advances in Neural Information Processing Systems, 34: 34:14084-14096 2021.
Christian Dallago, Jody Mou, Kadina Johnston, Bruce J Wittmann, Nicholas Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks in fitness landscape inference for proteins. bioRxiv, pp. 2021-11, 2021.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 1916):44 2012. 2012.
Josip Djolonga, Andreas Krause, and Volkan Cevher. High- figh-diimessioal gaussian process bandits. Advances in neural information processing systems, 26, 2013.
Clara Fannjiang and Jennifer Listgarten. Autofocused oracles for model-based design. Advances in Neural Information Processing Systems, 33: 13:15295-12956, 2020.
David Firth. Bias reduction of maximum likelihood estimates. Biometrika, 80(1):27-38, 1993.
Richard J Fox, S Christopher Davis, Emily C Mundorff, Lisa M Newman, Vesna Gavrilovic, Steven K Ma, Loleta M Chung. Charlene Ching, Sarena Tam, Sheela Muley, et al. Improving catalytic function by prosar-driven enzyme evolution. Nature biotechnology, 25(3):338- 2007.
Peter Frazier, A tutorial on bayesian optimization. arXiv preprint 1:XXi:1807.028111 2018.
Saba Ghaffari, Ehsan Salch, David Forsyth, and Yu-Xions Wang. On the importance firth bias reduction in few-shot classification. International Conference Learning Representations, 2022.
10

---

Published as a conference paper at ICLR 2024
Rafael Jomeez-oombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez- Sánchez-Lengelinn Lengeling, Dennis Sheberla, Jorge Aguilera-Ipanan Iparraguirre Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik, Automatic chemical design using data-driven continuous continuuus representation of molecules. ACS central science, 4(2):268-27) 222267, 2018.
Javier Gonzalez, Joseph Longworth, David C James, and Neil D Lawrence. Bayesian optimization for synthetic gene design. arXiv preprint arXiv:1505.000 0067770...7.... 2015.
Xiaojie Guo, Yuanqi Du, and Liang Zhao. Property controllable variational autoencoder via invertible mutual dependence. In International Conference on Learning Representations, 2020.
Anvita Gupta and James Zou. Feedback gan for dna optimizes protein functions. Nature Machine Intelligence, 1(2):105-111) 2019.
Brian L Hie and Kevin K Yang. Adaptive machine learning for protein engineering. Current opinion in structural biology, 72:145-152, 2022.
Kadina E Johnston, Clara Fannjiang, Bruce Wittmann, Brian L Hie, Kevin K Yang, and Zachary Wu. Machine learning for protein engineering. arXiv preprint arXiv:2305.16634 2023.
Seokho Kang and Kyunghyun Cho. Conditional molecular design with deep generative models. Journal of chemical information and modeling, 59(1):43-555 -2, 2018.
Volodymyr Kindratenko, Dawei Mu, Yan Zhan, John Maloney, Sayed Hadi Hashemi, Benjamin Rabe, Ke Xu, Roy Campbell, Jian Peng, and William Gropp. Hal: Computer system for scalable deep learning. Practice and Experience in Advanced Research Computing, PEARC '20, pp. 41-48, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 978 450366892. doi: 10.1145/331 11114//331990.3966690 1790.339649..
org ortp:///////.......1..1....1 111.:11..11111111111 239.8111999933969669 3396649
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations, 2014.
Leslie Kish and Martin Richard Frankel. Inference from complex samples. Journal of the Royal Statistical Society: Series B (Methodological), 36(1):1-22, 1974.
Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Advances in Neural Information Processing Systems,
HA Daniel Lagassé, Aikaterini Alexaki, Vijaya Simhadri, Nobuko Katagiri, Wojciech Jankowski, Zuben E Sauna, and Chava Kimchi-Sarfate Recent advances in (therapeutic protein) drug develop- ment. F1000Research, 6, 2017.
Georgios Mikos, Weitong Chen, and Junghae Suh. Machine learning identification of capsid mutations to improve aav production fitness. bioRxiv, pp. 2021-06, 2021.
J. Mockus. Bayeisan approach to global optimization: theory and applications. Springer Science & Business Media, volume 37, 2012.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
Jan Peters and Stefan Schaal. Reinforcement learning by reward- rewghe-weighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning, pp-
745-750, 2007.
Anna I Podgornaia and Michael T Laub. Pervasive degeneracy and epistasis in a protein protein-protein interface. Science, 347(6222) 2015.
Yuchi Qiu and Guo-Wei Wei. Clade 2.0: Exolution-driven cluster leassing-sssisted directed evolution. 2022.
Journal of Chemical Information and Modeling,
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-infomme neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.

---

Published as a conference paper at ICLR 2024
Zhizhou Ren, Jiahan Li, Fan Ding, Yuan Zhou, Jianzhu Ma, and Jian Peng. Proximal exploration for goodel-guieed protein sequence design. In International Conference on Machine Learning, pp- 18530. PMLR, 2022.
Philip A Romero and Frances H Amold. Exploring protein fitness landscapes by directed evolution. Nature reviews Molecular cell biology, 10(12):866887, 2009.
Philip A Romero, Andreas Krause, and Frances H Arnold. Navigating the protein fitness landscape with gaussian processes. Proceedings of the National Academy of Sciences, 110()):193-E) EE01,
2013.
Reuven Rubinstein.
The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 1:127-190, 1999.
Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1):89-112 1997.
Ehsan Saleh, Saba Ghaffari, Timothy Bretl, Luke Olson, and Matthew West. Learning from integral losses in physics informed neural networks. arXiv preprini arXiv:2300. 17387, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Sam Sinai, Richard Wang, Alexander Whatley Stewart Slocum, Elina Locane, and Eric D Kelsic. Adalead: A simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint arXiv:2010.02 141, 2020.
Sam Sinai, Nina Jain, George M Church, and Eric D Kelsic. Generative aav capsid diversification by latent interpolation. bioRxiv, pp. 2021-04, 2021.
Jasper Snock, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25, 2012.
Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias Seeger. Gaussian process opti- mization in the bandit setting: No regret and experimental design. arXiv preprint arXiv:0912.3995, 2009.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. In International Conference on Machine Learning, pp. 21658-21676. PMLR, 2022.
Nicholas Wu, Lei Dai, C Anders Olson, James O Llihh-Smith, and Ren Sun. Adaptation in protein fitness landscapes is facilitated by indirect paths. Elife, 5:e16965, 2016.
Kevin K Yang, Zachary Wu, and Frances H Arnold. Machine-learning-guiee directed evolution for protein engineering, Nature methods, 16(8):687 -644, 2019.
12

---

Published as a conference paper at ICLR 2024
VAE
PPGVAE (Soft Constraint)
CbAS
2
2.5
1
0.0
N
-.5
-2
-5.0
-2
RWR
CEM-PI
Bombarelli
ceweneeeeeee Seardil
1-
N
oyt ditttiiaaaen
0-
5
-1
-2.5
0.0
2.5
5.0
-2.5
0.0
2.5
5.0
-2.5
0.0
2.5
5.0
Z1
Z1
Z1
0
1
2
3
7
8
Figure A1: Two-dimensional latent space of our PPGVAE and other methods trained on the toy
MNIST dataset.
A CHARACTERISTICS OF THE LATENT SPACE AND SAMPLE GENERATION
To contrast PPOVAE with prior generative based methods, all methods were trained the toy In weighting based methods, there is less distinction among the mapping
MNIST dataset of Figure 2 regions of different classes, the weighting becomes extreme (Figure AT). Extreme weighting makes less samples contribute to the optimization, thus affecting the diversity in sample generation.
CEM-PI, CbAS and RWR have the most to least extreme weighting. In contrast weighting based approaches, PPGVAE and Bombarelli have more distinct mapping of classes. Thus, sample generation using PPGVAE has more diversity. By increasing the standard deviation (a,,) of the sampling prior (O, I).). PPGVAE is capable of generating all types of digits, whereas the least extreme weighting method (RWR) generates three types of digits only (Figure A2).
is easily seen that the organization of points in the latent space of PPGVAE is significantly different from Bombarelli's s. PPGVAE enforces the samples with higher properties to lie closer to the origin, whereas they are just separated by digit and ordered by property in Bombarelli's. By design, PPGVAE generates improved samples with higher probability, while exploration mechanisms need to be carefully designed for Bombarelli's Finally, by sampling from VAE prior NO,,, 7), i.e., o, = 1, PPGVAE generates the top two highest property classes as well as their interpolations, more frequently than others. Optimization over the latent space of Bombarelli's 's barely produces samples from the top two highest property classes.
B RELATED WORK
Machine learning approaches have been developed improve the design of DNA, molecules and proteins with certain properties. Early work in ML-assisted design (Gómez- Bombarelli et al.) 2018). structures the samples the latent space by joint training of property predictor on the latent space. However, further systematic exploration mechanisms the latent space are required to find improved samples. Variants weighting based MBO techniques have also been proposed facilitate the problem of protein design (Brookes & Listgarten 2018: Gupta & Zou 2019 Brookes al. 2019). CbAS Brookke et al. 2019) proposed a weighting scheme which prevents the search model from exploring regions of the design space for which property oracle predictions are unreliable. CbAS (Brookes et al., 2019) was built on an adaptive sampling approach (Brookes Listgarten 2018) that leverages uncertainty in the property oracle when computing the weights for MBO. Reward
13

---

Published as a conference paper at ICLR 2024
O
1 i 1 2
0: I
0, I 0
PPGVAE Constraint)
/ 0 0 3
PPGVAE (Hard Constraint)
/
PPGVAE (Hard Constraint)
4
3
0
0
0
5 O 0 1'
5'
7
O 1 0 L
1
3
2
/ 1
1 5 2 6
2
2 233 /
/
1
1
7
1
/
1
3
3
1
0
O
0
a.
1
C
/
LA
5 2 1 5
2
O
1
1
/
3
2
3
6
3
10
O
O
1 2 1 1
1
1
5
1 2 6 2
1
I
O
2
A i
1 /
4
/
O
4
2
2
5
3
8 /
0
0
/
2
O
1
3
8
2
===1 1 /
a,
0,
0
O I
2
0
O
I
I
I
/
/ /
/ /
/
O
/
I
I
0
/
0 /
O
/
I
O
0
/
0 O / /
O
1
/
I
/
1
0
1
CbAS
0
1
1
CbAs
0
(
/
CbAS
0
í / í
1
0
1 0 0
/
/
1000
/
1
100/11
0 0
1000
1
IIIIO
0
1
Ø
1.
O₂ 1
O, |
O, =
0
3
2
O
2
0
/
0 O
4
1
I
I
I I
X
I I
/
2 /
2
/
0
/ 1 1
1
2
/////
/
0 1
/
2 / 2 ///11
O
/ 0 R / 1 1 / 0 / 0 / /////
2
/
21
1
/
2 2
2 1
1 /
\
2
0
1200
2
1
O
RWR
0
( 1
1
RWR
0
1
2 2
/
0
0
0
2
0
0
/
102/1
0 0 1
I / &
1
0 0 1
1
1000
100
3
1
/
2
011000
0
21100 0,
0 7
O₂ 1
0:
CEM-PI
CEM-PI
CEM-PI
9
9
9
9
9
9
9
Bombarelli
7
9
5
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9 9
9
9
9
9
9
9
9 9
9
9 7
9
7
Figure A2: Samples generated from different methods trained on the toy MNIST dataset (Figure A1 For all methods except Bombarelli Hoanez-Bombarelli et al., 2018), normal sampling distribution (/ o with varying standard deviation [1, 2, was used. For Bombarelli, the optimization on the latent space was performed with different starting point for each generated sample.
14

---

Published as a conference paper at ICLR 2024
Notation
Description
X'
The design space
22
Sample from the design space (x, E X)
Vi
Property value associated with 2;
Weight associated with s;
Qa
Encoder with parameters 0
Encoded representation of x;
Pr(p))
Probability of the encoded representation w.r.t the VAE pior
C
Relationship loss
Coefficient of the relationship constraint
Temperature of relationship constraint
po(u(æ)
Property oracle
po(x)
Search model with parameters 0
N.
Number of samples generated per MBO step
P
Relative proportion of more-desired to less-desired samples in train set
Netr
Effective sample size in weighted MLE
Table Al: Mathematical notations used paper.
(benchmarked by ChAS) uses weights that do not take into account oracle uncertainty. Not requiring a differentiable oracle is common to these methods. On the other hand, simple baseline proposed by (Trabucco et al.) 2022) requires learning differentiable property oracle that is used to find design samples by gradient ascent. In an effort to avoid failures due imperfect oracles, (Kumar & Levine (220)) proposed learning an inverse map from property space to design space and search for optimal property during optimization instead. Protein design has also been formulated as sequential decision-making problem (Angermueller et al.) 2019), in which proximal policy optimization (PPO) (Schulman et al. 2017) has been used to search for design sequences. Recently, an evolutionary greedy approach has been shown to be competitive with prior algorithms (Brookes & Listgarten 2018 Brookes al. 2019: Angermueller et al. (019). The latest work on protein design (Ren et al. 2022) (PEX), focuses on exploring the rugged landscape close to the wild-type sequence, to find high fitness sequences with low mutation counts.
Weighted Regression (RWR) (Peters & Schaal 2007)
In our benchmark for sequence design, we only included AdaLead as a baseline, selected over alternatives has been demonstrated to have superior performance Dyna PPO (Angermueller et al. 2019). PEX (Ren et al., 2022) was not included, as is specifically designed to perform local search starting from the wild type, while our interest the more challenging scenario of optima being distal to the wild type.
C EFFECTIVE SAMPLE SIZE IN WEIGHTED MLE
Common to all weighting based generative methods is that weighted MLE could suffer from reduced effective sample size, which leads to estimation of parameters (++1 with higher variance and higher bias in extreme cases. This is inevitable in train sets with severe imbalance between the abundant undesired (e.g., zero property) and rare desired (e.g., nonzero property) samples. In contrast, PPGVAE does not weights. Instead, it assigns higher probability to the more desired data points by restructuring the latent space. Thus, allowing for the utilization of all samples training the generative model, regardless of the extent of imbalance between less and more desired samples.
15

---

Published as a conference paper at ICLR 2024
Dataset
Encoder Architecture
Protein
indim - 64 LReLU 20
PINN
indim - 64 LReLU - 10
MNIST
indim - 512 - ReLU + 256 ReLU +
GMM
indim -> 64 LReLU 64 LReLU -
Table A2: VAE architecture used for each benchmark dataset. Encoder and decoder have symmetrical architecture. All generative based methods share the same architecture
In imbalanced datasets, sample weights (w;) are typically uneven. As example, assume dataset of size 100 with five positive (desired property values), and 95 negative (undesired) samples. One weighting scheme can assign non-zero weights the five desired samples and zero weights otherwise. Therefore, only five out of 100 samples contribute to the objective in Equation [ This leads to higher bias and variance in the maximum likelihood estimator (1+1 (MLE). Next, present the same argument mathematically.
Representing the log-likelihodd log(po(z:)) of the generative model with 4, Equation can be rewritten as maximizing
K
wyli
(A8)
Assuming Σ wy and i.i.d. samples, the effective sample size New (Kish & Frankel 1974 can be defined such that
Var(I) New
Var(4)).
(A9)
Net = (EK 1= wf)- It can be proved that 1 S Nett S K where the equality for the lower and upper bound holds
According to Equation A8
Nen K, Vi, K'
and
Nor = 1, wy w 0.
(A10)
As mentioned earlier, uneven weights are expected in imbalanced datasets. As the weights become more uneven, Net approaches its lower bound. Therefore, with imbalanced datasets, Ner tends to drop and Var(I) increases (Owen 2013). This turn increases the estimation bias and variance of of the MLE g+++1
(Firth 1993 Ghaffari al. 2022). Both estimation bias and variance are ONN=1)
Our search model does not require weighting of the samples to prioritize the generation of some over the others. Instead, it directly uses the property values to restructure the latent space such that samples with better properties have higher chance of being generated and interpolated. For this reason PPGVAE has Netr = K, which makes it robust to the issues associated with parameter estimation in weighting-basee MBO techniques.
D DETAILS OF THE EXPERIMENTS
D.1 MBO SETTINGS AND IMPLEMENTATION DETAILS
We performed 10 rounds MBO on the GMM benchmark and 20 rounds MBO on the rest of the benchmark datasets. In all experiments temperature (r) was set to five for PPGVAE with no further tuning. We used the implementation and hyper-parameters provided by (Brookes et al. 2201)). for CbAS, Bombarelli, RWR, and CEM-PI methods. The architecture of VAE the same for all methods (Table The formula to compute the weights for each weighting baasee based method are included in the Appendix of (Brookes et al. 2019)
16

---

Published as a conference paper at ICLR 2024
Dataset
Measured Property
Separation Criterion
Separation Quantities
Imbalance Ratio (p)
GMM
Bimodal Gaussian function
Distance between the two modes (Aμ)
4,6,8,10,12
0.05, 0.1, 0.2, 0.4, 0.8,
PINN
-log(WMSE)
Property percentile for more desired samples
(30, 40), (40, 50), (50, 60), (60, 70), (70, 80)
0.05, 0.1, 0.2, 0.4, 0.8
GBI
Folded protein enrichment
Length of the appended sequence
3,4,5,8
0.0125, 0.025,
0.05, 0.1, 0.2, 0.4, 0.8
PhoQ
Yellow fluorescent protein level
Length of the appended scquence
3,4.6,8
0.0125, 0.025, 0.05, 0.1, 0.2, 0.4, 0.8
AAV
Capsid viability
Minimum number of mutations for samples with less desired property (y < 0.5)
6,8,10,12,15
0.0125, 0.025, 0.05, 0.1, 0.2, 0.4, 0.8
Table A3:
Settings used in train set generation for each benchmark dataset.
The ground truth property oracle was used in all experiments, which is equivalent to limiting the search space the sequences with measured properties in the protein datasets. The performance AYmax is reported with 95% bootstrap confidence interval.
In AdaLead, The oracle query size and the experiment batch size were both the number of samples generated per MBO step (N.). This was to run AdaLead in a comparable setting to other weighting- weighting-based MBO approaches. Ground-truth oracle was used in all experiments, i.c., the search space was limited to the samples existing in the dataset.
For each setting imbalance ratio and separation level, optimization was performed with 20 different random seeds. In AdaLead, the starting sequence was randomly selected from the initial train set for each seed. We noticed that AdaLead performance is highly dependant on where search starts in the sequence space. This justifies the high variability in its performance. This is a common problem to all approaches based on evolutionary search, as it may not be known priori where to start the search.
D.2 PROTEIN DATASETS
AAV: The dataset contains variants of the protein located in the capsid of AAV virus along with their viability measurements. It consists of roughly 284K variants with single to 28 sites mutations. The properties were normalized into range [0, 1]. The dataset was obtained from (Dallago et al. 2021).
GB1: It contains the empirical fitness landscape for the binding of protein G domain BI to an antibody. Enrichment of the folded protein bound to the antibody was defined as fitness. Fitness was measured for 149,361 sequences for variants of amino acids four sites. The fitness was normalized to range [0, 1] this paper. The dataset overpopulated with low or zero-fitness sequences. The sequences along with their properties were obtained from Qiu & Weil 2022
17

---

Published as a conference paper at ICLR 2024
30
1.0
20
10
10
04
0.4
-10
-20
0.0
0.0
-20
20
-20
20
ISNE1
ISNE1
Low Separation
High Separation
30
20
0.8
20
0.8
10
N
0.4
0.4
-10
-20
-20
0.0
ISNE1
ISNE1
Figure A3: GB1 transformed landscapes and train set examples. Transformed landscapes with appended sequence of length three and six are shown in top left and top right panels, respectively. Examples of train sets taken from each landscape are shown in bottom left (low separation) and bottom right (high separation) panels, respectively.
is a combinatorial library of four amino acid sites located at the interface between PhoQ kinase and sensor domains. Signaling of two-component regulatory system was measured by yellow fluorescent protein (YFP) levels, i.e., fitness. The fitness landscape contains 140,517 measured sequences. The dataset is overpopulated with low zero-fitness sequences. The fitness was normalized to range [0, 1] in this paper. The sequences along with their properties were obtained
PhoQ:
from (Qiu & Wei 2022).
D.3 GENERATION OF TRAIN SETS WITH VARYING IMBALANCE AND SEPARATION
In each dataset, we specify the definition of less and more desired samples as well as the separation criterion in the design space X. Table A3 includes the settings used to generate train sets in all benchmark datasets.
GMM: The property oracle is bimodal GMM in which the second mode has higher peak (higher mode) than the first mode (lower mode). Each mode is represented by triplet (Mi o, a;) where a; determines the peak height and the mode index. The triplet was set to (0, 0.25, 1) and (u2. 2.5) for the lower and higher modes, respectively. Mean of the second mode is variable as it determines separation level.
Train sets consisted of less and more desired samples, which were taken from the two modes. N samples were taken from the lower mode using (0, 0.6) distribution. pN samples representing the higher mode were taken uniformly from range (uz + 2 μ2 + 02)-
The separation the design space was defined as the difference between the means of two modes Ap. Higher Áu is associated with higher separation.
AAV: The entire AAV dataset consists of 284K samples (Bryant et al. 2021). which are split into two sets of "sammledd" and "designed" variants (Dallago et al.] (021). We only used the sampled variants in this study. Threshold 0.5 on the viability scores was used to define the set of more desired (> 0.5) and less desired (< 0.5) variants.
Train sets consisting of less and more-desired samples were generated by, taking N samples from the less desired variants, whose mutation count more than specified threshold, and taking pN samples from the more-desired variants, whose properties fall in the (5,10) percentile of the property values.
18

---

Published as a conference paper at ICLR 2024
The separation is determined by the minimum number of mutations present in the less dess-desired samples. The separation of less and more-desired samples in the train set increases as the minimum number of mutations increases.
GB1: To study the impact separation, the sequences GBI dataset were transformed to have more dispersed coverage of the landscape in the design space. First, the property values were normalized into range [0, 1]. Threshold of 0.001 on the property was used to define the set of less desired (< 0.001) and more-desired variants (> 0.001). Then, the sequences of less-desired variants were appended with random sequence of length L, whereas more-desired variants were appended with specific sequence of the same length. Here, the length of the appended sequence specifies the degree of separation. Larger length is associated with higher separation.
Train sets were generated by sampling N samples from the less-desired and pNV samples from the more-desired variants. The transformed landscape of GB1 associated with low and high separation, along with examples of its train sets are shown Figure A3
PhoQ: The same procedure as GB1 I was used to generate the train sets.
PINN: We used dataset of PINN-derived solutions the Poisson equation for the potential of three point charges located diagonally on 2D plane (Saleh et al. 2023). The solutions were pooled from batch PINNs trained with different seeds different training epochs. Negative log weighted MSE (wMSE) between the PINN-derived solution and the analytical solution was used as the property value. Note that in practice, the exact average loss can be a proxy of the solution quality without the analytical solution. Higher quality solutions have higher properties.
In generating the train sets, low fitness samples were taken from (0, 15) percentile of the property values, whereas pN high fitness samples were taken from a specified range of property percentiles (P1, Pb) Pi 2 30. Higher values P1 are associated with higher separation levels. Note that in this case, separation of samples by property values is concordant with their separation in the design space X. See Table A3 for the values of (P1, P2) that have been tested.
In practice, accurate solutions be sporadic with stochastic optimizers, PPGVAE can interpolate higher-quality solutions given imbalanced sets of solutions.
Toy MNIST: This dataset used demonstrate the latent space of PPGVAE and other methods, as well as their sampling generation characteristics. The dataset has rare representation of zero digits relative the other digits (imbalance ratio p 0.01). Samples belonging to the digit class C' have property values distributed as N N(10 0.001).
19

---

Published as a conference paper at ICLR 2024
E ADDITIONAL PLOTS
E.1 STUDYING THE IMPACT OF IMBALANCE RATIO PER SEPARATION LEVEL
For given separation, the performance improves as the imbalance ratio increases. This is expected as the more-desired samples constitute a higher portion of the train set. As separation level increases, the task of optimization becomes harder. Therefore, other methods become mostly ineffective for low imbalance ratios, and only show improvements for higher imbalance ratios. This behavior is consistently seen all benchmark tasks.
Lowest Separation
Low Separation
Medium Separation
0.3
0.1
ma
0.03
0.01
High Separation
Highest Separation
0.1
0.3
1
0.3
PPGVAE
CbAS
0.1
mama
CEM-PI
0.03
RWR
0.01
0
0.1
0.3
1
0.1
0.3
Figure A4: Performance imbalance ratio for each separation level in the GMM benchmark.
Lowest Separation
Low Separation
Medium Separation
0.3
max
0.1
0.03
0.01
High Separation
Highest Separation
PPGVAE
0.3
AdaLead
CbAS
AYaax
0.1
CEM-PI
0.03
RWR
0.01
0.01
0.03
0.1
0.3
1
0.01
0.03
0.1
0.3
1
Figure A5: Performance ys imbalance ratio for each separation level in the AAV benchmark.
20

---

Published as a conference paper at ICLR 2024
PPGVAE
-
AdaLead
ChAS
CEN-PI
RWR
Lowest Separation
Low Separation
High Separation
Highest Separation
0.03
0.01
0.01
0.03
0.1
0.3
0.01
0.03
0.1
0.3
0.01
0.03
0.3
0.01
0.03
0.1
0.3
Figure A6: Performance imbalance ratio for each separation level in the GB1 benchmark.
PPGVAE
-
AdaLead
-
CbAS
CEM
RWR
Lowest Separation
Low Separation
High Separation
Highest Separation
0.03
0.01
0.01
0.03
0.1
0.3
0.01
0.03
0.1
0.3
0.01
0.03
0.1
0.3
0.01
0.03
0.1
0.3
Figure A7: Performance imbalance ratio for each separation level in the PhoQ benchmark.
Lowest Separation
Low Separation
Medium Separation
man
0.1
0.03
0.01
High Separation
Highest Separation
0.3
PPGVAE
CbAS
mam
CEM-PI
RWR
0.03
0.01
0.1
0.3
0.1
0.3
Figure A8: Performance vs imbalance ratio for each separation level in the PINN benchmark.
21

---

Published as a conference paper at ICLR 2024
E.2 STUDYING THE CONVERGENCE RATE FOR DIFFERENT IMBALANCE RATIOS
For each benchmark task, we looked at the progress of each method as number of MBO steps increases. For higher imbalance ratios, all methods have faster convergence relative the low imbalance ratios. Furthermore, our PPGVAE consistently has the highest convergence rate to the highest improvement. This expected, as PPGVAE prioritizes the interpolation and generation of more desired samples, regardless of the extent of their representation in the train set.
PPGVAE
CbAS
CEM-PI
RWR
0.05
0.1
p= 0.2
0.1
AYmax
0.03
0.01
p 0.4
p= 0.8
p 1.0
AYmax
0.03
0.01
8
10
0
2
10
8
10
MBO Step
MBO Step
MBO Step
Figure A9: Performance vs the number of MBO steps for different imbalance ratios in GMM benchmark. Separation level is set to medium.
p=
p=
0.05
p=
0.01
0.2
p= 0.4
p= 0.8
PPOVIE
0.3
CbAS
RWR
0.03
0.01
16
12
16
8
12
20
MBO step
step
MBO step
Figure A10: Performance vs the number of MBO steps for diffferet imbalance ratios AAV benchmark. Separation level is set to low.
22

---

Published as a conference paper at ICLR 2024
0.0125
0.025
0.05
0.1
11
0.3
0.03
0.01
p 0.2
p= 0.4
p
PPOVAE
CbAS
CEM-PI
0.03
0.01
16
20
12
16
0
12
20
MBO
MB0
MBO
Figure A11: Performance the number of MBO steps for different imbalance ratios in GB1 benchmark. Separation level is set low, CEM-PI is the most competitive with PPGVAE for higher imbalance ratios.
p= 0.0125
p= 0.025
p = 0.05
p= 0.1
1
0.3
0.1
0.83
0.01
0
== 0.2
p=
0=0.8
CbAS
CEM
0.03
0.01
16
12
16
12
20
MBO
MB0 step
MBO step
Figure A12: Performance the number of MBO steps for different imbalance ratios in PhoQ benchmark. Separation level set low, CEM-PI is the most competitive with PPGVAE for higher imbalance ratios.
23

---

Published as a conference paper at ICLR 2024
E.3 TEMPERATURE SENSITIVITY AND TRAINING ON HIGH FITNESS SAMPLES
To study the sensitivity of PPGVAE performance to the temperature, the GMM experiments with the lowest imbalance ratio and the highest separation (most challenging scenario) were performed with varying temperatures (r) Figure A13 The performance is almost the same for log1o(7) E [- Also, the sensitivity to temperature decreases the number of MBO steps increases.
We also repeated the AAV experiments for CEM-PI in which only high fitness samples were used as the initial train set (CEM-PI/High) Figure A13 Aggregated performance over all imbalance ratios for PPGVAE is better than both CEM-PI and CEM-PI/High. This demonstrates the importance of including all samples the training using PPGVAE. Furthermore, CEM-PI/High has better performance than CEM-PI for higher separation, showing that filtering the samples might be beneficial for weighting based approaches as the optimization task gets harder by separation criterion.
Temperature Sensitivity
All High Fitness
0.3
(t)
0.3
PPGVAE
CEM-PI
0.1
2.00
CEM-PI/High
&Ynor
-1.0
AYmor
0.1
1.0
0.03
0.01
0.01
0
2
4
10
Lowest
Medium
Highest
MBO Step
Separation
Figure A13: Impact of varying temperature in GMM benchmark (Left) and performance comparison between training on all vs high fitness samples for CEM-PI (Right).
24


Reviews:
[{"review_id": "UhnzZO5IAk", "reviewer": "Reviewer_e46j", "summary": "The paper addresses a crucial challenge in protein design, which involves optimization on a fitness landscape. The main concern in current model-based optimization methods is the sparsity of high-fitness samples in training datasets and the separation problem - wherein the desired optimum is situated in a region that is poorly represented and far from low-fitness areas. This paper pinpoints that existing tools do not efficiently handle this separation problem in the design space.\n\nThe authors introduce a new method using Property-Prioritized Generative Variational Auto-Encoder (PPGVAE). This VAE's latent space is structured by the fitness values of the samples, ensuring higher prioritization and generation probability for more desired sequences. This new method aims for better results with fewer optimization steps, which is particularly valuable for sequence design problems.\nA comparative advantage of this approach over prior methods is demonstrated via extensive benchmarks on real and semi-synthetic protein datasets. \n\nThe PPGVAE proves to be superior in robustly finding improved samples, regardless of the imbalance between low- and high-fitness samples and the degree of their separation in the design space. The authors further extend the versatility of their method by testing it on continuous design spaces, showcasing its efficacy on physics-informed neural networks (PINN).", "questions": "1. How easily can PPGVAE be extended to prioritize or balance multiple properties simultaneously? Would this require a significant alteration to the existing framework? \n2. Could you provide insights into the sensitivity of the model's performance to changes in the temperature in the relationship loss?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper recognizes the less-explored challenge of \"separation\" in protein design space, which is a significant departure from recent studies that have mostly focused on the sparsity of high-fitness samples. And The PPGVAE proposed in the paper is an effective approach to tackling the separation issue in model-based optimization.\n\n2. The paper does not just present a theoretical model but comprehensively demonstrates its effectiveness through extensive benchmarking on real and semi-synthetic protein datasets. Beyond just protein datasets, the paper further validates the model on continuous design spaces, exemplified with physics-informed neural networks (PINN).", "weaknesses": "1. Limitation on the conducting experiments exclusively on the real protein dataset of AAV. While this might result in high accuracy and performance metrics within this context, the method may not readily translate to other proteins or protein datasets, especially if they possess distinct characteristics or functionalities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nW184lyYZ4", "reviewer": "Reviewer_soEo", "summary": "Note: I am not familiar with the field of protein design, so this review is largely from the perspective of generative modelling.\n\n**Problem Setting**\n\nWe are considering the problem of protein design, in which data consists of sparse examples of high-fitness designs. Due to the nature of the optimization problem, it is desirable to learn a model that can judge the fitness of specific sequences.\n\n**Novel Idea**\n\nThis work proposes the learning of a weighted VAE such that high-fitness samples are weighted in the probability distribution. The intuition is that exploration in VAE space is more likely to result in succesful generations with this bias in place. However, all samples can still be used in training the model.\n\nThe PPGVAE is structured so that designs with higher fitness are generated closer to the origin. This means they will be more likely to be sampled when using a normal distribution prior. This is implemented in the form of a constraint between the ratio of log probabilities vs. fitness, and this constraint is relaxed to a weighted penalty.\n\n**Findings**\n\nExperiments showcase that the PPGVAE model can correctly separate between two modes. The modes are defined as the probabilities represented by a two-mode Gaussian mixture model. On protein datasets, the method showcases that the PPGVAE can separate between low and high fitness examples, and that separated models lead to faster convergence when using MBO.", "questions": "See above for a list of questions.\n\nIs the intent of this work to be applicable to other domains, or only protein design? If the answer yes, as implied by the title, it would be greatly strengthening to apply this method on more classical generative modelling tasks larger than MNIST.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper presents a useful method for learning an exploration prior for model based optimization. In the area of protein design, there are sparse examples of high fitness designs, and many low fitness designs. A good prior should bias towards the space of high-fitness designs, while still making use of all examples. This work presents a simple and clean objective that learns a variational auto-encoder. The methodology is clear and theoretical justification is provided. It may provide significance in the specific problem domain, although the method itself is domain-agnostic.", "weaknesses": "The figures of the paper were confusing in terms of what message they were trying to convey. A more informative caption describing why the results are important would improve the clarity here.\n\nThe experiments do showcase the the proposed architecture helps in terms of separating high-fitness examples. It would be insightful to include some experiments on what happens if the VAE is simply trained on only the high-fitness examples and the low-fitness examples are dropped altogether. \n\nIt is unclear the difference between the hard constraint and the soft constraint. Are these just different hyperparameters on the penalty? Or is it a true constraint using i.e. a Lagrange multiplier?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "P03iqY1rPr", "reviewer": "Reviewer_R8Dk", "summary": "This paper identifies the scenarios where the desired optimum is in a region that is not only poorly represented in training data but also relatively far from the highly represented low-fitness regions for model-based optimization for sequence-function landscapes, named separation, and proposes a new method using a variational auto-encoder to explicitly structure the latent space by property values of the sequences such that more desired samples are prioritized over the less desired ones and have higher probability of generation. Empirical studies on three real and semi-synthetic protein datasets show the robustness of the proposed method.", "questions": "What is the maximum dimensionality examined in the experiments? How will the proposed method perform when increasing the dimensionality of an optimization problem?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThis paper identifies the problem of separation for model-based optimization for sequence-function landscapes.\n2.\tA robust method based on a variational auto-encoder is proposed to solve the separation and sparsity problems.\n3.\tEmpirical studies are conducted on three real and semi-synthetic protein datasets to study the effectiveness of the proposed method.", "weaknesses": "1.\tA clear description of the robustness concept needs to be provided for friendly reading.\n2.\tA description of the organization of this paper as well as a Conclusion session should be added.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eXjKOEi8iU", "reviewer": "Reviewer_7H3r", "summary": "Authors propose to use a Variational Autoencoder (VAE) to model a protein engineering task. Obtained model would allow scientists from the area to better know where to look for promising compounds in such difficult search space. They argue that their approach is more effective and efficient than other normally attempted approaches, such as RL guided searches or evolutionary optimization. Authors provide some experimental assessments that include experiments with artificial datasets to properly evaluate their approach.", "questions": "I'm leaning a bit towards rejection of the paper, just for the reason I mentioned in the 'weaknesses' section, about novelty and contribution. \nBut if other reviewers go with acceptance, I won't argue against.\nAC has been alerted to seek another opinion if needed.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Unfortunately, Bioinformatics is far outside my scope areas. I can't really assess quality of the paper. For the untrained eye, everything seems to fit. The proposed approach seems correct for the task at hand (as far as I could understand it), but I'm completely unfamiliar with the related work, can't really say if this has been attempted before for this particular domain, or to what extent. All I can say is that presentation of the paper is good, and the artificial dataset experiments somewhat seem to validate what authors attempted to do.", "weaknesses": "As I already mentioned, I can't really assess the draft.\nThe only one thing that I'd like to bring forward, it's that for the given bioinformatics task described -as far as I could understand it- VAEs seem like a natural thing to try. I'm surprised nobody has done it before; but since I'm unfamiliar with the Related Work, surely there are things that I'm missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)