Design Editing for Offline Model-based Optimization

TL;DR

Offline model-based optimization (MBO) seeks to maximize a black-box objective $f(x)$ using only an offline dataset, yet gradient-based surrogates often overestimate due to out-of-distribution extrapolation. DEMO addresses this by first generating pseudo-designs via surrogate-guided gradient ascent and then editing them with a diffusion prior learned from the offline data to align with the valid design manifold. This two-phase approach decouples high-performance search from realism regularization, showing competitive performance across seven tasks that span continuous and discrete domains. The method offers practical impact for design tasks in robotics, materials, and biology by reducing spurious optima and improving reliability under offline constraints.

Abstract

Offline model-based optimization (MBO) aims to maximize a black-box objective function using only an offline dataset of designs and scores. These tasks span various domains, such as robotics, material design, and protein and molecular engineering. A common approach involves training a surrogate model using existing designs and their corresponding scores, and then generating new designs through gradient-based updates with respect to the surrogate model. This method suffers from the out-of-distribution issue, where the surrogate model may erroneously predict high scores for unseen designs. To address this challenge, we introduce a novel method, Design Editing for Offline Model-based Optimization (DEMO), which leverages a diffusion prior to calibrate overly optimized designs. DEMO first generates pseudo design candidates by performing gradient ascent with respect to a surrogate model. While these pseudo design candidates contain information beyond the offline dataset, they might be invalid or have erroneously high predicted scores. Therefore, to address this challenge while utilizing the information provided by pseudo design candidates, we propose an editing process to refine these pseudo design candidates. We introduce noise to the pseudo design candidates and subsequently denoise them with a diffusion prior trained on the offline dataset, ensuring they align with the distribution of valid designs. Empirical evaluations on seven offline MBO tasks show that, with properly tuned hyperparameters, DEMOs score is competitive with the best previously reported scores in the literature.

Figures (6)

• Figure 1: Illustration of DEMO: A diffusion model, acting as the prior distribution, is trained on the offline dataset. Pseudo design candidates are acquired by performing gradient ascent with respect to a learned surrogate model. New designs are generated by modifying pseudo design candidates toward the valid distribution captured by the diffusion prior.
• Figure 2:
• Figure 3: Computational efficiency analysis.
• Figure 4: Selecting $m$ near $0$ results in generated designs that retains most properties of pseudo design candidates. Conversely, setting $m$ near $1000$ generates designs that align closely with the distribution of existing designs. Optimal designs are achieved by choosing $m$ in the mid-range, effectively utilizing information from both the pseudo design candidates and the diffusion prior.
• Figure 5: The proportion is calculated as the number of new designs which surpass $\mathcal{D}(\textbf{best})$ divided by the budget $128$, indicating the reliability to consistently generate new higher-scoring designs. This figure demonstrates that DEMO is more reliable than Diffusion-only in all tasks.