Across scientific domains, generating new models or optimizing existing ones while meeting specific criteria is crucial. Traditional machine learning frameworks for guided design use a generative model and a surrogate model (discriminator), requiring large datasets. However, real-world scientific applications often have limited data and complex landscapes, making data-hungry models inefficient or impractical. We propose a new framework, PropEn, inspired by ``matching'', which enables implicit guidance without training a discriminator. By matching each sample with a similar one that has a better property value, we create a larger training dataset that inherently indicates the direction of improvement. Matching, combined with an encoder-decoder architecture, forms a domain-agnostic generative framework for property enhancement. We show that training with a matched dataset approximates the gradient of the property of interest while remaining within the data distribution, allowing efficient design optimization. Extensive evaluations in toy problems and scientific applications, such as therapeutic protein design and airfoil optimization, demonstrate PropEn's advantages over common baselines. Notably, the protein design results are validated with wet lab experiments, confirming the competitiveness and effectiveness of our approach.
Paper
Similar papers
Peer review
Summary
The paper develops a simple but novel deep learning framework to implicitly optimize a seed input design to iteratively improve upon a particular property $$ g(\cdot) $$, until it reaches a fixed-point. In this way, the paper develops a novel framework for design optimization and provides a theoretical derivation to showcase that this implicit setup leads to an optimal learned neural network function $$ f^*(\cdot) $$ whose iterative refinement of a seed design , approximates the direction of gradient of the property function of interest. The framework has been evaluated on three diverse design tasks and compared with several state-of-the-art design methods.
Strengths
- The proposed method is well motivated and novel for the problem of design optimization with neural networks without adversarial classification (which are inherently hard to train). - The various experiments on toy as well as realistic design tasks like airfoil design optimization, antibody binding affinity design optimization demonstrate that the proposed method based on a simple `pair-wise matching` criterion is able to achieve better results than state-of-the-art models previously used for the same design task. - The insight and the proof demonstrating that the constructed optimal function $f^*(\cdot)$ approximates the direction of the gradient of a property function $g(\cdot) $ even though it is not explicitly trained to do so, is interesting and will benefit the space of design optimization.
Weaknesses
- Further analysis about the effect of $ \Delta_x $ and $\Delta_y $, their selection methodology and their sensitivity in various design contexts is necessary as the matching step critically depends on these two parameters. - The method has been demonstrated only in the context of designs that improve a single property $ g(\cdot) $. However in reality designs usually need to consider multiple properties simultaneously. This brings into question the applicability of the current method to practical problems. - A rigorous analysis of the "goodness" of the seed design and the effect of said goodness on the convergence speed / accuracy is necessary.
Questions
1. How are $ \Delta_x $ and $\Delta_y $ set / tuned? 2. How does property under consideration $ g(\cdot) $ influence the calibration / selection of $ \Delta_y $? 3. How can the current model be augmented to function with multiple properties that might need to be simultaneously improved? 4. How are seed designs selected for initial inputs? i.e, what is considered a minimum viable seed design to be improved upon / refined until convergence? How does the model react to bad (i.e., flawed / noisy) designs?
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
A limitation of the current work is its inability to incorporate more than a single property of interest. However, the authors have explicitly mentioned this limitation and the reviewer believes that despite this, the proposed method is still valuable and can serve as a good foundation for future research in design optimization.
Summary
The paper address design optimization, which is the process of optimizing over a "design" parameter space to optimize over one or more observable outcomes in many scientific and engineering problems. The proposed framework "PropEn" uses a three step process, first identifying a "matched dataset" that pairs every sample with another with a better outcome; this is followed by training a model that takes a sample and predicts, essentially creating an autoregressive sampler. The technique is evaluated on multiple science and engineering benchmarks.
Strengths
* The matched objective is an interesting application of regressive sampling, and a different view on augmenting a small dataset, which is a common issue in many science problems. * The paper goes into theoretical justifications as to why a matched reconstruction objective is useful, suggesting that it approximates the gradient of an explicitly trained property predictor. * Its also shown that the generated samples have a high likelihood close to the training samples, indicating that they are indeed sampling from the training distribution.
Weaknesses
* PropEn is an interesting idea for design optimization but it needs to be more thoroughly benchmarked against existing methods. For example, there is a host of methods that do Bayesian optimization, and variants thereof, that can work with as few as 10 samples (albeit in lower dimensions) which the paper does not consider. I think it sheds light into the behavior of PropEn in lower dimensions. More recent methods that use DNN based UQ estimates are able to go beyond 5-10D datasets effectively as well. * At a high level, matching and matched reconstruction appear to be very closely related to diffusion models. In my understanding, the diffusion process can be framed as matched reconstruction, where the property to maximize is the likelihood of the sample. Framing matched reconstruction through the lens of diffusion might not only strengthen the paper's formulation but also provide other theoretical insights. * Its unclear to me why neighborhoods in the design space are constructed using L2-balls, why is this the best way to identify similarity in x? How sensitive is the optimization to the size of the neighborhood? Or in another sense, a related question is -- how sensitive is matched reconstruction to the choice of sampling? for e.g. if you start with an LHC sampling (or any other variants) vs other types. * On small amounts of data, what kinds of regularizations keep the main model from overfitting?
Questions
Please see above
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
Yes, it appears so.
Summary
This work proposes the method PropEn, which is inspired by the concept of matching techniques in econometrics. Using PropEn (specifically in scenarios with a lack of large datasets), the authors can expand the dataset, which will inherently help in design improvement, etc. To do this, they train a network to learn a mapping from an initial sample to another sample with an improved target attribute selected during the dataset curation phase. This model can further be sampled auto-regressively until it converges to a final candidate design.
Strengths
- The paper is very well written, and the motivation, methods, and results are clearly presented and easy to understand. - The proposed method is simple, generalizable, and appears effective. - The experiments and corresponding results are well-defined and support the claims made in the paper.
Weaknesses
- Additional experiments that would add substantial value to the work would be to extend the data matching method to other models, rather than a simple encoder-decoder model with a reconstruction loss. As an example, provide experiments for diffusion models where the conditioning signal is the x_0 seed (or x_i in autoregressive sampling), and the generated design is x_{i+1}, which would solidify your claims about the generalizability of the data matching method. - Have the authors considered active learning approaches? This work is very close to active learning ideas, applied in the context of low data regimes. - The aspect of evaluating the target properties on the neighboring samples is not clearly explained. - the examples shown are simple and not really huge.
Questions
See above in weakness
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
2
Limitations
No limitations are mentioned by the authors.
I have read the comments and am satisfied with the clarifications and additional results. I will raise my score to borderline accept
Official Comment by Authors
Dear Reviewer, Thank you for taking the time to provide your valuable feedback. We are glad to hear that the clarifications and additional results we provided addressed your concerns and we are grateful for your decision to raise the score. Best regards.
Summary
The paper presents a generative framework for property enhancement. The proposed framework consists of only a generative model, and it's missing a discriminator that is usually found in other frameworks for guided design. This is achieved by training the generative model on a "matched" dataset that consists of paired examples $(x, x')$ where $x'$ is an "enhanced" version of $x$ from its immediate proximity ("enhanced" meaning it has a higher value of the property that is sought to be maximized in the guided design process, and "proximity" depends on the problem domain). A mathematical analysis of the proposed method shows that the learnt generative function converges to the gradient of the function computing the property and that examples sampled by the function are as likely as the training set. These analysis motivate the iterative application of the generative function to converge to a stationary point with enhanced property. The authors claim that this framework is particularly well suited for applications with scarce data, and apply it to a toy problem and airfoil optimization, as well as protein optimization. In the latter application the method shows superior performance compared to other methods, as measured by wetlab experiments.
Strengths
1. The paper presents a simple method with mathematical justifications why it should work. 1. Additionally to the toy examples, the method is applied to airfoil enhancement and protein optimization. The predictions of the latter are confirmed by wetlab experiments, and the reported results are better than state-of-the art methods.
Weaknesses
1. The paper is missing several important implementation details. I would be especially interested in knowing the details of the training datasets and the constructed matched datasets, which I could not find anywhere in the paper or the appendix. Section B.1 in the appendix should be expanded to explain in detail what models were used. For example, line 571 states that a "ResNet with 3 blocks" was used to generate one-hot encoded AhO aligned sequences. But it's not clear to my why a 2D convolutional architecture would be used for such a problem, how exactly the data was represented, and what motivates the use of that precise model. 1. The airfoil optimization is missing a strong baseline. From the presented results I find it hard to judge whether the method indeed works well across domains. I'm not familiar with the protein optimization experiments, so I cannot judge the power of the method by the presented results, especially since important implementation details are missing (see above weakness). Since all the findings in Section 3.2 are based on very few datapoints (per optimization target) and there is a wide intra-method variability, I would like to see some statistical analysis for presented data.
Questions
1. Section 2.3 lines 126-127 "As a property to optimize, we choose the log-likelihood of the data as estimated by a KDE with Gaussian kernel" – it seems that optimizing for likely examples would both optimize the property as well as the likelihood of the examples. But theorem 1 is about increasing the property value and theorem 2 is about sampling likely examples. So I would have thought it more natural to choose an illustrative example that decouples the property to be optimized from the likelihood of the examples. 1. I couldn't find the numbers reported in Table 3 in referenced papers [38, 39]. If it is a reproduction by the authors, then I would like to see that clearly stated in the main text or the table caption. 1. The assumption of a "sufficiently overparametrized model" (line 472) seems a pretty strong one to me. How realistic is this assumption in the presented settings, and what would be the consequence for Theorem 1, and PropEn applicability in general, if the assumption does not hold? Various nits, questions, comments: 1. line 33: it should say "even for deep neural networks" (remove "with") 1. caption Figure: missing space "value.Bottom" 1. line 39: missing a word at "[6] the" 1. line 48: "have only be used provide" 1. line 57: "PropEn", not "Propen" 1. line 113: "from the its" 1. line 125: please add a reference to the "well-konwn 2d pinwheel dataset" (I don't know which dataset is meant exactly); similarly, the "8-Gaussian" dataset (line 157) is also missing a reference 1. line 154: why is the regularized variant called "mixup/mix"? I find this confusing, as there is a well-known mixup technique in ML that does not seem related (https://arxiv.org/abs/1710.09412) 1. line 165: how are metrics "uniqueness" and "novelty" defined? 1. line 185: reference for "NACA airfoils" missing 1. Section 3.1.2: how were the initial airfoil designs chosen? 1. lines 216-218 and Figure 4: when varying one threshold, what is the value of the other threshold? a 2D density plot might point a more complete picture of the measured improvement when varying both thresholds. 1. lines 219-222 and Figure 4: we can see that in a single example the $C_l/C_d$ ratio increased monotonically; but it's quite a stretch from that to claim that "all designs ... are deemed plausible" and that there is (always) a "consistent enhancement"; Figure 7 in the appendix (which is not mentioned in the main paper) does not add substantially to these claims 1. line 225: it would be nice to reference figures/tables to show that "mixup variant of PropEn may require longer training" 1. line 234: "referred", not "refereed" 1. Figure 5 (left): why do the points have different sizes? 1. Figure 8: what is the take-away from this figure?
Rating
6
Confidence
2
Soundness
2
Presentation
3
Contribution
3
Limitations
The "NeurIPS Paper Checklist" in the appendix still needs to be filled in (currently it only shows the instructions).
Official Comment by Authors
Dear Reviewer, Thank you for your valuable feedback and support. As the author-reviewer discussion period deadline approaches (Aug 13), we kindly remind you to provide any additional feedback. We hope our responses have addressed your concerns, and if so, would greatly appreciate a score update. If there are remaining issues, please let us know, and we will promptly address them. Best regards.
Thank you for attentive replies to my original review. First some minor replies: [Mixup/Mix] Yes indeed, I still find this confusing: in the cited Mixup paper the regularization works by creating convex combinations from pairs of input data and labels. In the case of PropEn the regularized variant works by incorporating a reconstruction regularizer $ℓ(f_θ(x), x)$. It is not quite clear to me how this regularized PropEn variant "interpolates between examples", as you claim in your response. [Details on Models] Even with the additional piece of information that the ResNet is 1D I'm not quite sure about the exact models used. For example, the toy dataset model (Section B.1, lines 556-562) mentions a 2 layer MLP with 100 units per layer, and the main text mentions (Section 3.1.1, lines 162-163) that the explicit baseline uses an "auto-encoder of the same architecture as PropEn augmented with a discriminator for guidance". I would like to see more details about the discriminator and the encoder-decoder. More broadly, it would be nice to have a discussion about the models used and how they were selected. Especially since the authors choose their own baselines, care should be taken that not only the proposed method yields good results, but also that the baselines are properly optimized wrt their architecture and hyper parameters. [Concerning all comments in the response] Thank you for the clarifications, I assume you would add them also to the paper? (e.g. the response currently says "We add the following table in the supplement", but it doesn't mention whether the submission would be updated to any of the other replies). Then, some more comments on the main weakness raised in my initial comment: **[Missing comparison to strong baselines]** After reading the paper and the author's reply to my review, I'm still not convinced that the method really performs better than other available methods. The main problem in my opinion is that the authors implement their own baselines and the paper lacks details that convince that the methods are applied in such a way that they really yield best possible results. One way to mediate this would be to take strong results from existing papers and then show that PropEn performs better on the exact problems evaluated by these papers. Another way would be to show convincingly in the paper that the methods were tuned in such a way that they yield best possible results. As for the added Airfoil baseline, I'm not convinced that these are really strong baselines, as [Diffusion-based 2DAirfoil Generation](https://github.com/tonyzyl/Diffusion-based-2DAirfoil-Generation/blob/main/models/airfoil_MLP.py) is a Github repo with a single star, and [tiny-diffusion](https://github.com/tanelp/tiny-diffusion) is "A minimal PyTorch implementation of probabilistic diffusion models for 2D datasets", where Table 1 states that the Airfoil dataset has 400 dimensions. Overall, I appreciate the author's efforts in their reply, but I don't think that they sufficiently addressed the main weaknesses mentioned in my initial review, so I remain with my initial rating of 4 (borderline reject).
Thank you for reviewing our rebuttal and further clarifying your remarks. **[Mixup]** We will change the name of this variant to **Cross xx2xx**, and we are open to suggestions from the reviewer. By interpolation, we refer to the behaviour observed during training in this mode. Specifically, the network learns to reconstruct (a control sample) `x` into (a treatment sample) `x'`, `x''`, `x'''`, etc. This means the network never learns a perfect reconstruction of one specific treatment, but instead generates variants in-between the two. **[Details on Models]** The rebuttal response was limited to 6,000 characters, so we couldn't include all the details. Please note that in Section B.1, lines 556-562, we specify 30 neurons per layer, not 100 as mentioned in the reviewer's comment. To avoid confusion, please find below the networks used for the toy and engineering experiments, which we will add to the supplementary material of the paper. For WJS, we used a 1D Conv architecture (please see A.1 in [38]), and for Lambo, a bert-small transformer (please see C.1 in [39]). Please note that the code for all of the experiments will be released. For optimizing the parameters of the baselines in the toy and engineering experiments, we conducted a grid search over the learning rate (`[1e-2, 1e-5]`), weight decay (`[1e-2, 1e-5]`), number of epochs (`[300, 1000, 5000]`), batch size (`32, 64, 128`), and number of neurons per layer (`[30, 50, 100]`). For the antibody experiments, each baseline was implemented and trained as advised by the authors of the original publications. **Toy example architectures:** AE( (criterion): MSELoss() (encoder): Sequential( (0): Linear(in_features=2, out_features=30, bias=True) (1): ReLU() (2): Linear(in_features=30, out_features=30, bias=True) (3): ReLU() (4): Linear(in_features=30, out_features=30, bias=True) (5): ReLU() (6): Linear(in_features=30, out_features=15, bias=True) ) The decoder is analogous to the encoder, starts with the compressed representation (15) and progressively reconstructs the original input size (2 features). (Discriminator): Sequential( (0): Linear(in_features=2, out_features=30, bias=True) (1): ReLU() (2): Linear(in_features=30, out_features=30, bias=True) (3): ReLU() (4): Linear(in_features=30, out_features=30, bias=True) (5): ReLU() (6): Linear(in_features=30, out_features=1, bias=True) ) ) Diffusion: MLP( (time_mlp): PositionalEmbedding( (layer): SinusoidalEmbedding() ) (input_mlp1): PositionalEmbedding( (layer): SinusoidalEmbedding() ) (input_mlp2): PositionalEmbedding( (layer): SinusoidalEmbedding() ) (joint_mlp): Sequential( (0): Linear(in_features=2, out_features=128, bias=True) (1): GELU(approximate='none') (2): Block( (ff): Linear(in_features=128, out_features=128, bias=True) (act): GELU(approximate='none') ) (3): Block( (ff): Linear(in_features=128, out_features=128, bias=True) (act): GELU(approximate='none') ) (4): Block( (ff): Linear(in_features=128, out_features=128, bias=True) (act): GELU(approximate='none') (5): Linear(in_features=128, out_features=2, bias=True)) **[Concerning all comments in the response]** Yes, we will add all modifications from the rebuttal to the manuscript. We appreciate all remarks and believe that by incorporating these modifications, we will improve our submission. **[Missing comparison to strong baselines]** We disagree with the reviewer regarding the lack of strong baselines: * In the toy example, AE with guidance is the most standard baseline from the literature. As is common in all ML publications, to ensure a fair comparison, we compare the same architectures for the implicit and explicit baselines (after hyperparameter tuning). * In the engineering example, tiny diffusion is the most appropriate baseline to compare to given the MLP architecture used for PropEn. The diffusion model was developed to generate hundreds of points in 2D, which is exactly the same as the airfoil setup (so the reviewer's concern about dimensionality mismatch and the inappropriateness of the baseline is not grounded). We believe this is the fairest baseline we could add for comparison. * In the antibody design experiment, we compare against **three** SOTA methods. Finally, we would like to remind the reviewer that we never claim in the manuscript or the rebuttals that PropEn is the best method across the board. We believe the true advantage of PropEn is its **applicability and efficiency across many domains**, which is what we aim to confirm in our experimental evaluation. All other baselines are application-specific, requiring numerous modifications and inductive biases included in the architecture and training. This is also why it is difficult to have the same baselines across all applications, as most methods focus on a single application, while PropEn is general.
Thank you for the quick and exhaustive reply to my latest comments. [Mixup] I think using any other name, such as "Cross xx2xx" would be clearer. Thank you for the update. [Details on Models] Thank you for the additional information about the used models. I think these details should also be briefly mentioned in the appendix. **[Missing comparisons to strong baselines]** Since I am not familiar with the guided design literature and airfoil design / protein optimization, I cannot get a solid impression of the performance of the proposed algorithm with the data presented in the paper. If the method were compared directly to results from the literature, that would certainly give a better grounded impression of the performance. But I broadly agree with the rebuttal from the authors, and since no other reviewer commented on this point, I have increased my rating accordingly.
We are pleased to hear that the reviewer agrees with our rebuttal, and we sincerely appreciate all the comments and remarks provided during the discussion period. We will incorporate the additional results and clarifications into our manuscript. Best regards
Official Comment by Authors
Dear Reviewer, Thank you for your valuable feedback and support. As the author-reviewer discussion period deadline approaches (Aug 13), we kindly remind you to provide any additional feedback. We hope our responses have addressed your concerns, and if so, would greatly appreciate a score update. If there are remaining issues, please let us know, and we will promptly address them. Best regards.
Official Comment by Authors
Dear Reviewer, Thank you for your valuable feedback and support. As the author-reviewer discussion period deadline approaches (Aug 13), we kindly remind you to provide any additional feedback. We hope our responses have addressed your concerns, and if so, would greatly appreciate a score update. If there are remaining issues, please let us know, and we will promptly address them. Best regards.
Thank you for the response and overall comments. I do appreciate the authors on actually implementing the designs into a wet lab framework, as a test for real world applicability -- a high bar for evaluating the impact of the work. A note on BO/AL, i think the distinction the rebuttal makes is subtle and it would still be useful to compare against BO/AL on known benchmarks to show the benefits, considering the broad applicability of BO/AL methods. There have been works in the past that use BO in conjunction with a generative models, this may be an informative experiment. In light of these comments, I will raise my score to 6.
Thank you for reconsidering your score recommendation. In addition to the comparison with Lambo, which uses BO to guide design optimization, we could include an experiment in the supplementary of the final manuscript that evaluates PropEn in BO on standard test functions like Branin-Currin or DTLZ, where we have ground truth solutions for comparison. This would involve a different problem setup, considering multiple rounds, and would require adapting PropEn to handle design selection independently. We will also include an extended discussion on this point. Best regards
Decision
Accept (poster)