Diversity By Design: Leveraging Distribution Matching for Offline Model-Based Optimization

We thank Reviewer sKFn for their helpful and constructive feedback on our manuscript, and appreciate their efforts in helping us improve our work. Please find our responses below - we would be happy to discuss further and answer any follow-up questions as necessary. Thank you!

1. (Assumption in Lemma 3.1) Thank you for this comment. The inequality in Equation (16) follows directly from monotonicity of the integral and holds even if $q^\pi(x)$ is non-zero in a measure-zero set. If we have misunderstood your concern, please let us know.

2. (Reward Term in Eq. (5)) Thank you for this comment. We include a pure reward term in (5) because the forward surrogate model $r_\theta(x)$ provides an important (albeit sometimes inaccurate) signal of design quality outside the dataset $\mathcal{D}$. This strategy builds off prior work in offline MBO (e.g., Trabucco et al. ICML (2021)., Yu et al. NeurIPS (2021)., Yao et al. NeurIPS (2024).). Without this reward term, the Shannon entropy term in Lemma 3.3 would not appear in our derivation and so the generative policy is not explicitly encouraged to find designs that maximize against the $r_\theta(x)$ surrogate signal.

3. (Clarification of "Baseline") Thank you for this question. In general, there are two "optimizers" used in our experiments: (1) a model-training optimizer used to train a surrogate reward function on the offline dataset; and (2) a generative optimizer used to generate new designs by optimizing against the trained surrogate function from step (1). "Baseline" only refers to the latter generative optimizer.

The reward functions used for all the experiments in Table 1 were trained to approximate the oracle reward function using Adam as the model-training optimizer.

We have better clarified this point in our revised manuscript.

4. (Limitations of DynAMO) We thank the Reviewer for this comment. You are correct - our method does indeed trade reward for diversity. In many applications of offline MBO, we argue it is most important to obtain a good maximum score over a good median score. We highlight this potential limitation of DynAMO in our revised Limitations section, included at this link.

5. (Algorithm 1 Presentation) Thank you for this comment. We have revised Algorithm 1 here to improve its legibility. Regarding the while loops, the outer while loop terminates when the optimizer $a^b$ no longer improves, and is implemented as lines 281-308 of our main.py source code file in our Supplementary Material. The inner while loop terminates when the optimal Lagrange multiplier is found, and is implemented in lines 226-256 in the src/dogambo/models/mlp.py file (which is invoked in line 304 of our main.py source code file).

We would be happy to make any additional changes to the presentation of Algorithm 1 to improve its legibility.

6. (Table 1 Formatting) We appreciate your careful eye to detail! For the TFBind8 optimization task using the Adam optimizer, the top performing method in terms of design quality was RoMA+ with a score of $96.5\pm 0.0$ (lower bound on 95% CI is 96.5). However, ROMO scores $95.6\pm 0.0$ (upper bound on 95% CI is 95.6) and GAMBO scores $94.0\pm 2.2$ (upper bound on 95% CI is 96.2). Because the 95% CI intervals of ROMO and RoMA+ (and also GAMBO and RoMA+) are non-intersecting, ROMO and GAMBO are correctly left unbolded.

7. (Meaning of "Myopic") We follow Deisenroth et al., Ngo, and others to refer to policies that are trained according a value function with an discount factor of $\gamma=0$ in reinforcement learning as "myopic". Intuitively, this means that the surrogate reward function used by the REINFORCE algorithm only considers the current input design in computing the reward, and not future possible states like in traditional RL applications. We have better clarified this point in our revised manuscript.

We thank Reviewer DFtV for their thoughtful feedback and careful consideration of our work. Please find our responses to your comments below.

1. (Combination of KL Divergence and Adversarial Constraints) Thank you for this comment - indeed, we believe that it is a strength of our method that DynAMO is able to leverage existing individual techniques--such as KL divergence penalization from imitation learning and adversarial constraints from the GAN literature--and combine them together in a unique way to solve a new problem (offline MBO). Namely, leveraging KL divergence penalization has not been shown in prior MBO literature, and leveraging adversarial regularization has only recently been explored by Yao et al. NeurIPS (2024). Combining them together in a meaningful and principled way to solve the problem of offline MBO is the main novelty and contribution of our work.

Regarding the unique value of combining both techniques, we refer the Reviewer to Supp. Tables A8-A10 in Appendix E.2, which shows that using both KL-divergence and adversarial regularization together in combination significantly improves the performance of DynAMO than when using either technique alone.

2. (Theoretical Guarantee in Section D.4) To clarify, our primary contribution of this work is the DynAMO algorithm. The goal of the theoretical bound presented in Section D.4 is to use standard, validated techniques from the literature to show theoretically what we already observe in our presented experiments: DynAMO can improve the diversity of designs obtained using different generative optimization methods. Indeed, we use fairly standard methods and techniques to arrive at our result in Section D.4 - for example, see Ma et al. Proc NeurIPS (2022), Huang et al. (2024), and Liang (2016) for additional background.

3. (Method Limitations) Thank you for this comment. We highlight that many of our discrete MBO tasks (i.e., ChEMBL, Molecule) are problems over high-dimensional design spaces as discussed in prior work (Maus et al. NeurIPS (2022), Yao et al. NeurIPS (2024), Trabucco et al. ICML (2021)). Furthermore, the DKitty task in our work includes a sparse and highly-sensitive oracle reward model as discussed in Trabucco et al. ICML (2022). We have better highlighted these points in our revised manuscript, and have also included an additional paragraph explicitly discussing the challenges with optimization in these settings in alignment with the Reviewer's feedback.

4. (Setting Hyperparameters) Thank you for this comment. We highlight that all experimental results presented in the main text use the same hyperparameter values (e.g., $\tau=1.0$ and $\beta=1.0$) across all optimizers and MBO tasks. Furthermore, Supp. Figures A4-A6 show that the performance of our method is robust to the choice of hyperparameters. In short, while we domain expertise and tuning of hyperparameters might improve our method further, it is by no means necessary to achieve good experimental results using DynAMO.

We thank Reviewer kDhU for their thoughtful comments and feedback on our proposed work, which we believe have substantially improved the quality of our manuscript. Please find our responses to your comments below.

1. (Performance According to the Median@128 Metric) Thank you for this comment. DynAMO does indeed trade reward (according to the Median@128 metric) for diversity. Diversity enables us to find a greater subset of optimal and sub-optimal designs than baseline methods; naturally, this may result in a lower median score since DynAMO cannot simply return points near the global maximum (e.g., see Supp. Fig. A2). In many applications of offline MBO, it is more important to obtain a good maximum score than a good median score, and we therefore prioritize the Best@128 and pairwise diversity metrics in our evaluation of DynAMO. We highlight this potential limitation of DynAMO in our revised Limitations section, included at this link.

2. (ChEMBL and UTR Tasks) We would be happy to move the results associated with these two tasks to the Appendix in accordance with this feedback. (Notably, DynAMO's average optimality gap consistently increases and its average ranking stays the same or improves after excluding the ChEMBL and UTR tasks.)

3. (Normalization of Results) All of the quality metrics reported have already been max-min normalized (0,1) according to the equation $\hat{y}=\frac{y-y_{\text{min}}}{y_{\text{max}}-y_{\text{min}}}$ in alignment with prior work in the offline MBO literature. We multiplied the normalized metrics $\hat{y}$ by 100 to improve the legibility of our results, as highlighted in the table captions.

4. (Good and Bad Examples in Offline RL) We note that KL constraints are also commonly used in offline RL, where the dataset can include non-expert trajectories. Intuitively, the direction of the KL term keeps the generated samples inside the support of the offline dataset, with failing to match the offline dataset outside its support being penalized less. We can further tune the tradeoff by using the temperature hyperparameter $\tau$ defined in Equation (6). Using a $\tau>0$, designs that are good are upweighted in the reference distribution $p^\tau(x)$ compared to bad designs - see Fig. A1 for examples. As a result, we are able to use the standard intuition and techniques derived from imitation learning, and show in our work that our approach works both theoretically and empirically.

5. (Figure 1) We thank the Reviewer for this comment. Indeed, DynAMO generates a diversity of designs, many of which are not concentrated around the global maxima. However, the point of DynAMO (and Figure 1) is that at least a few of the proposed designs are located around the different global maxima. This is in contrast with the baseline method, which proposed a batch of designs that might all be near-optimal, but are all clustered around the same optima and region of the input space. Our primary motivation of DynAMO is that obtaining a small number of diverse and high-quality designs is better than obtaining a large number of high-quality but near-identical designs. If a particular application of offline MBO does not suit this use case, then DynAMO may not be well-suited for the application at hand. We have better clarified this point in Figure 1 of our revised manuscript.

6. (Diversity for Secondary Objectives) To demonstrate how diversity can enable better downstream exploration of secondary objectives, we compare DynAMO-enhanced optimization methods against their baseline counterparts and report the results here. We find that if an optimization method achieves a higher pairwise diversity score, it also consistently achieves a larger variance in the values of the secondary objectives, supporting the idea that secondary objectives can benefit from diversity.

To be clear, we do not claim that DynAMO will discover designs that maximize multiple objectives while only optimizating against a single objective. The goal is that by increasing diversity, we can better interrogate other design properties, which will hopefully be different from one another.

7. (Additional Baselines) We thank the Reviewer for sharing these additional baselines with us. We have included all 5 referenced methods in our updated design quality and diversity results. Our results do not significantly change with the inclusion of these additional baselines.

8. (Choice of Terminology) Thank you for this comment. We initially chose to leverage terminology from both the offline MBO and RL literature to more explicitly highlight how DynAMO adapts techniques inspired from offline RL to the MBO setting. We have revised our manuscript to remove references to reward functions and other RL-centric terminology to better align with existing offline MBO literature.

We appreciate Reviewer scBu for their insightful feedback and thorough evaluation of our work. We have provided our responses to your comments below and would be happy to answer any follow-up questions. Thank you!

1. (Presentation of Results) To improve the legibility of our results, we summarize the main findings on the task-averaged Rank and Optimality Gap metrics in the main table in our revised manuscript, and move the detailed breakdown of algorithm performance by task to the Appendix for the interested reader.

2. (Diversity versus Quality) To be clear, we do not claim that DynAMO will propose designs of higher quality than other baseline methods across all settings. (In fact, no baseline method is state-of-the-art across all tasks.) The goal of DynAMO is to increase the diversity of designs without significantly reducing the quality of best design proposed. Indeed, across multiple offline MBO tasks evaluated, there is often no statistically significant difference between the quality of designs proposed by DynAMO and other baseline methods according to the Best@128 evaluation metric. However, DynAMO is consistantly state-of-the-art in producing diverse sets of designs across all settings. We have better clarified this primary goal and use case of DynAMO in our manuscript.

3. (Equation (5) for Diversity) To clarify, we do not want to sample designs that are truly OOD, as we cannot guarantee the correctness of these designs using an arbitrary surrogate $r_\theta(x)$. Ideally, it would be great to discover designs that are both novel and significantly OOD, but this is not possible while guaranteeing correctness in the offline setting. Consistent with prior work (e.g., Yao et al. NeurIPS (2024).), DynAMO seeks to generate interpolated, in-distribution points to balance diversity with naive reward maximization. Through the regularization in Eq. (5), we aim to learn a generative policy that proposes designs with coverage of the search space similar to the offline dataset $\mathcal{D}$. If $\mathcal{D}$ is too small or has poor coverage, then our method would likely be less effective - we discuss this in further detail in our Limitations section of our revised manuscript.

4. (Source Critic Constraint) The primary goal of the source critic is to penalize designs that are OOD with regards to the offline dataset $\mathcal{D}$. If $c^*(x)$ is large, then $x$ is likely from the same distribution as $\mathcal{D}$ and so we can trust that the prediction $r_\theta(x)$ is a good estimate of $r(x)$ with reasonable confidence. If $c^*(x)$ is small, then $x$ is likely not from $\mathcal{D}$ and so we penalize the reward associated with $x$ to avoid choosing the design as a proposed candidate. Note that this source critic is different from the forward surrogate $r_\theta$ (approximating oracle $r$), and so we cannot replace $c^*$ with the forward surrogate.

We note that our use of $c^*(x)$ builds on prior literature from Yao et al. Proc NeurIPS (2024) and Arjovsky et al. Proc ICML (2017) - we refer the Reviewer to these prior works for additional discussion.

5. (Fairness in Comparisons) To ensure fairness in comparisons, we tuned the hyperparameters of DynAMO according to only on its performance using the Grad. optimizer, following the same paradigm as COMs, ROMO, and GAMBO. All hyperparameter values for all methods (including DynAMO) were then held constant for all of the other optimizers. While each of the methods (including DynAMO) might benefit from further optimizer-specific and task-specific hyperparameter fine-tuning, such an approach is not feasible in real-world applications. All of the quality metrics reported have already been max-min normalized to a (0, 1) scale according to the equation $\hat{y}=\frac{y-y_{\text{min}}}{y_{\text{max}}-y_{\text{min}}}$ in alignment with prior work. However, we then multiplied the normalized metrics $\hat{y}$ by 100 to improve the legibility of our results, as highlighted in the table captions where relevant.

6. (Optimization Initialization Strategy and Distribution of Oracle Objective Values) We thank the Reviewer for suggesting an alternative strategy to improve the diversity of designs through random initialization. We compared top-$k$ and random-$k$ initialization strategies for first-order optimization methods in this table. For most methods, neither initialization strategy results in consistently more diverse designs (except for DynAMO, which benefits from random-$k$ initialization). Furthermore, both COMs and RoMA suffer from lower quality of designs proposed when a random-$k$ initialization strategy is used. These results highlight the significance of DynAMO as a novel algorithm to consistently increase the diversity of final candidates.