Guided Trajectory Generation with Diffusion Models for Offline Model-based Optimization

Thank you for the positive feedback and further suggestions that could enhance our manuscript.

(Weakness 1) The authors have adopted the Diffusion model for offline model-based optimization, which is a current trending approach. It is recommended that the authors emphasize their unique contributions and clarify the distinctions between their method and previous approaches.

We acknowledge that several concurrent works adopt diffusion models for offline model-based optimization. DDOM [1] is the pioneering work utilizing diffusion models for offline MBO. Following DDOM, DEMO [2] trains a diffusion model to match a pseudo-target distribution constructed by gradient ascent and uses the model to edit designs in the offline dataset. Diff-Opt [3] considers a constrained optimization setting and introduces a two-stage framework that begins with a guided diffusion process for warm-up, followed by a Langevin dynamics stage for further correction. Diff-BBO [4] measures the uncertainty of generated designs to select the optimal target value for conditioning the diffusion model.

Our method uniquely distinguishes itself from prior and concurrent works by using diffusion models to generate trajectories toward high-scoring regions, learning to improve solutions from the dataset. While we mention this part in the related work section, we will modify the manuscript to emphasize our unique contributions in the introduction.

(Weakness 2) The model filters candidates through a trained proxy model. The authors should elaborate on how data handling is managed during experiments to ensure that there is no label leakage and discuss whether the comparisons with other models are conducted fairly.

In the filtering stage, we train our proxy model with only offline datasets for fair comparison.

(Weakness 3) The authors should provide a more detailed description and discussion of the experimental results.

Sorry for missing detailed description and discussion on the experiment results. We will provide more thorough analysis on our main experiment in the manuscript as written below.

Our experiment results generally surpass forward approaches, which struggle to fall into OOD designs, especially in high-dimensional settings. We also observe that our method outperforms inverse approaches, including DDOM, which utilizes the diffusion model. It demonstrates that generating trajectories towards high-scoring regions can be more effective than generating a single design, as we can distill the knowledge of the landscape of the target function into the generator. Our method achieves higher performance compared to BONET, which also generates trajectories. It indicates that our novel trajectory construction strategy effectively guides the diffusion model to explore diverse paths toward high-scoring regions.

[1] Krishnamoorthy, Siddarth, Satvik Mehul Mashkaria, and Aditya Grover. "Diffusion models for black-box optimization." International Conference on Machine Learning. PMLR, 2023.

[2] Yuan, Ye, et al. "Design Editing for Offline Model-based Optimization." arXiv preprint arXiv:2405.13964 (2024).

[3] Kong, Lingkai, et al. "Diffusion models as constrained samplers for optimization with unknown constraints." arXiv preprint arXiv:2402.18012 (2024).

[4] Wu, Dongxia, et al. "Diff-BBO: Diffusion-Based Inverse Modeling for Black-Box Optimization." arXiv preprint arXiv:2407.00610 (2024).

Thank you for the critical reviews, insightful feedback, and further suggestions to enhance our manuscript.

(Weakness 1) More varied discrete benchmarks and high variance in TFBind10 results

We excluded NAS as it takes too long to evaluate. Instead, to verify the effectiveness of our method on various discrete tasks, we conduct additional experiments on UTR. We use the reported scores in PGS for other baselines. As shown in Table 19 attached in PDF, we achieve the second-best results in the UTR task.

For high variance in TFBind10, we observe that it is due to the variability of the seeds. When we conduct experiments on more seeds (16 seeds), we find that our method still showed high performance with lower variance.

Performance of GTG on TFBind10 with more random seeds.

(Weakness 2) It’s not clear how this approach deals with out-of-distribution (OOD) problems.

Thank you for your interest in how our method addresses out-of-distribution (OOD) problems. We introduce classifier-free guidance to guide the diffusion model towards exploring high-scoring regions. To mitigate the OOD issue, we introduce $\alpha$, which controls the exploration level of the generated trajectories. We find that setting $\alpha=0.8$ generally exhibits good performance while high $\alpha$ results in sub-optimal OOD designs. Please refer to Section 5 and Appendix D.2 for a more discussion on the impact of $\alpha$.

(Weakness 3) Using COMS or ROMA models as a proxy

As you mentioned, the proxy function can be fragile when estimating OOD designs. To mitigate this issue, we use ensembles of MLP as a proxy with rank-based reweighting during training to focus on high-scoring regions as written in Appendix B.2. Nevertheless, using COMs and ROMA is a promising approach to enhance the robustness against OOD designs.

To this end, we replace proxy with COMs and ROMA and evaluate the performance of our method on TFBind8 and Dkitty tasks. As shown in Table 19 attached in PDF, we find no significant difference in performance based on the choice of different proxy functions.

(Weakness 4) Absence of principled mathematical explanation

While there is no principled mathematical explanation, our method shows promising results across various real-world tasks. We also conduct a thorough analysis of our method using Toy 2D experiment visualizations in Figure 2. Additionally, we perform experiments on practical settings such as sparse and noisy datasets, demonstrating the versatility of our method.

(Weakness 5) Absence of median performance

We apologize for missing the median performance results. We will update the manuscript with median performance results. As shown in the table 20 attached in the PDF, we also achieve the highest mean rank (4.0) in terms of the median performance.

(Question 1) For the discrete tasks, Is a latent space used, or is everything kept discrete during diffusion and proxy training?

As mentioned in Appendix A.2, we convert discrete input into a continuous vector and adopt a continuous diffusion and proxy model.

(Question 2) How many gradient steps are used for gradient ascent methods like COMs and RoMA? Is it 64? And is the trajectory length the same for BONET and PGS?

We strictly follow the original procedures of the baselines written in their papers and publicly available code. For BONET, while the original paper generates 4 trajectories with a prediction length of 64, we generate 4 trajectories with a prediction length of 32 to match the evaluation budget of 128.

(Question 3) Are designs picked at the end of the generation process, or can they be selected from the middle of the trajectory?

We aggregate all designs generated by trajectories and select candidates for evaluation with a filtering procedure.

(Question 4) Is the filtering applied to results generated by BONET or PGS? How does this affect their outcomes?

Both methods do not apply a filtering strategy after generation. While PGS generates synthetic trajectories, it ultimately selects final candidates using a policy trained with the trajectories. Therefore, we examine the effect of filtering only in BONET. As shown in the table, filtering can improve the performance of BONET. However, we find that our method still outperforms BONET in terms of performance.

Performance of BONET with filtering strategy.

(Question 5) How many designs are initially considered, and how many are selected at the end (128 out of how many)?

As mentioned in Section 4.4, we sample $N=128$ trajectories conditioning on $C=32$ context data points, so overall $4,096$ designs are initially considered. Then $Q=128$ designs are finally selected. As sampling procedure can be conducted in parallel, time complexity is not too problematic. For time complexity of sampling procedure, please refer to Appendix D.4.

(Question 6) Why is the perturbation epsilon different for DKitty?

We introduce perturbation epsilon ($\epsilon$) to prevent synthetic trajectories from converging into a single maxima in the offline dataset. We use small $\epsilon$ for Dkitty as there are relatively many samples with high scores in the offline dataset compared to other datasets.

We also observe that even with the same $\epsilon=0.05$ for the Dkitty task, there is no big difference in performance, and it still outperforms other baselines, as shown in the table. For analysis on $K$ and $\epsilon$ for constructing trajectories, please refer to Appendix D.1.

Performance of GTG on DKitty task with different perturbation epsilon ($\epsilon$).

Thank you for the valuable review and positive feedback!! As only two questions have been raised by the reviewer, we will address it in this response. If you have any additional questions, please do not hesitate to let us know!

(Weakness 1) Unreliable results of superconductor task.

We acknowledge the issues with Design-Bench, particularly in the Superconductor task. However, we have found that by fixing the seed for evaluation, we can obtain consistent outputs for multiple copies of the same inputs. Please note that we have reproduced baselines whose code is publicly available using the same evaluation procedure for a fair comparison.

(Weakness 2) The paper calls existing trajectory generation approaches (like sorted monotonic trajectories) as heuristics but the proposed approach is also a heuristic. Is there more principled/theoretical explanation for the proposed approach?

We apologize for overstating our proposed method. Our trajectory construction method can also be considered a heuristic since it does not include learning components. Nevertheless, we focus on constructing trajectories that give us more valuable information for learning to improve designs toward diverse high-scoring regions. We have empirically found that constructing trajectories with locality bias can guide a diffusion model to explore high-scoring regions in the Design-Bench and its variants. Furthermore, we also find that the computational time for constructing trajectories does not require a significantly large amount of time, as shown in Table 9 in Appendix B.1.

Thank you for the insightful review and positive feedback!

(Weakness 1 & Question 1) Clarification on the difference between the proposed method and Decision Diffuser.

We would like to highlight that our method aims to find a design that maximizes the target black-box function while Decision Diffuser is a planner to solve offline RL problems by generating high-rewarding trajectories that follow environment dynamics.

Instead of utilizing diffusion model to generate a single design, we train diffusion model with synthetic trajectories constructed from the dataset. While several prior works construct trajectories from datasets, we developed a novel method for constructing trajectories with locality bias to help generator better understand the landscape of the target function. As reviewer MZpe mentioned, how to construct trajectories is a interesting problem in the context of MBO. The ablation study summarized in Table 4 underscores that our method outperforms prior strategies across various tasks.

(Weakness 2) It is not clear that in loss function, only x is used to calculate loss or both x and y are used.

We apologize for the confusion in the notation. The trajectory $\tau$, which is a set of $(x, y)$ input pairs, should be used to calculate the loss in line 5 of Algorithm 2. We train diffusion model to generate trajectories, not a single design $x$. I will modify the manuscript to clarify this.

(Question 2) In Section 3.4 and Appendix D.3, the distribution of data used to train models must be close to the real distribution. Otherwise, out of distribution could occur, leading to an error evaluation of trajectories.

For both Section 3.4 and Appendix D.3, we train proxy model with only offline dataset. As our objective is filtering high-fildelity designs with the proxy, we introduce a rank-based reweighting suggested by [1] during training to make the proxy model focus on high-scoring regions. To prevent high error in evaluation, we train ensembles of MLP for robustness of the proxy model. For more details, please refer to Appendix B.2 and our code.

[1] Austin Tripp, Erik Daxberger, and José Miguel Hernández-Lobato. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Advances in Neural Information Processing Systems, 33:11259–11272, 2020.

(Weakness 3 & Question 3) Clarification of the concept of filtering in Section 3.4 and Section 5

We apologize for the confusion regarding the concept of filtering in the manuscript. Both Section 3.4 and Section 5 refer to the same concept. Among samples from multiple generated trajectories, we select candidates for evaluation by filtering with the proxy function.