Learning Design-Score Manifold to Guide Diffusion Models for Offline Optimization

TL;DR

ManGO is introduced, a diffusion-based framework that learns the design-score manifold, capturing the design-score interdependencies holistically, and outperforms 24 single- and 10 multi-objective optimization methods across diverse domains.

Abstract

Optimizing complex systems, from discovering therapeutic drugs to designing high-performance materials, remains a fundamental challenge across science and engineering, as the underlying rules are often unknown and costly to evaluate. Offline optimization aims to optimize designs for target scores using pre-collected datasets without system interaction. However, conventional approaches may fail beyond training data, predicting inaccurate scores and generating inferior designs. This paper introduces ManGO, a diffusion-based framework that learns the design-score manifold, capturing the design-score interdependencies holistically. Unlike existing methods that treat design and score spaces in isolation, ManGO unifies forward prediction and backward generation, attaining generalization beyond training data. Key to this is its derivative-free guidance for conditional generation, coupled with adaptive inference-time scaling that dynamically optimizes denoising paths. Extensive evaluations demonstrate that ManGO outperforms 24 single- and 10 multi-objective optimization methods across diverse domains, including synthetic tasks, robot control, material design, DNA sequence, and real-world engineering optimization.

Figures (5)

• Figure 1: Overview of the ManGO framework for offline optimization. (a) Illustration of offline optimization: it identifies optimal designs for an unknown black-box function using an offline dataset (no environment interaction), where designs represent function inputs and scores correspond to outputs. (b) Training a diffusion model on score-augmented data to learn the joint design-score manifold. (c) Fidelity estimation via unconditional samples generated by the trained ManGO model: the fidelity metric determines whether to activate inference-time scaling during conditional generation. (d) Bidirectional conditional generation: it leverages preferred-score or preferred-design conditions to generate corresponding designs or scores, illustrated via the self-supervised importance sampling (self-IS) method at denoising timestep $t$ for sample $i$. (e) Conceptual illustration of ManGO: it learns on the design-score manifold to enhance out-of-distribution generation (OOG) capability, contrasted with design-space learning that struggles with OOG issues under unseen conditions YangGXZWCW24. (f) Case study on superconductor's temperature optimizationsuperconductor: it demonstrates the superior OOG performance via ManGO versus the design-space approach (i.e., DDOM) across varying ratios of top data removal.
• Figure 2: Visualization of manifold learning, trajectory generation, and generation capabilities of ManGO. Note that unconditional and conditional samples are generated via ManGO without guidance and with preferred-score guidance, respectively. (a-b) Manifold and trajectory comparisons for the Branin (SOO) and OmniTest (MOO) tasks. The generated manifold is constructed via ManGO’s design-to-score prediction within the feasible region of designs. Close alignment between the ManGO-generated and original manifold, confirming the model’s proficiency in learning complex design-score relationships. Generated trajectories visualize ManGO’s score-to-design mapping under minimal score and design constraints, highlighting its capacity to perform targeted denoising toward desired regions. (c) Branin task: Unconditional samples (green) match preferred scores from the training dataset, while conditional samples (blue) extrapolate beyond the training minimum (grey dashed line). (d) OmniTest task: Conditional samples better approximate preferred scores and Pareto-dominate the training data (grey) compared to unconditional samples. These results indicate that ManGO effectively reconstructs in-distribution samples during unconditional generation—reflecting well-learned manifold structure—while enabling OOG of superior samples through conditional guidance, demonstrating robust extrapolation based on the learned manifold.
• Figure 3: Pareto front generation under different guidance conditions. Across all subfigures of (a) RE21, (b) ZDT3, (c) DTLZ7, and (d) RE41, columns from left to right respectively show the results of no guidance, standard guidance, and the proposed self-IS-based guidance. The progressive improvement in generation quality highlights ManGO’s capability in OOG under conditional guidance. Furthermore, the enhanced performance with self-IS-based guidance illustrates the ManGO’s feasibility for more delicate guidance mechanisms.
• Figure 4: Comprehensive ablation studies validating key components of the ManGO framework. (a-b) Flexibility to guidance specification: Performance sensitivity to deviations between guided scores and true optima on Ant (a) and Superconductor (b) tasks, demonstrating ManGO's stability under suboptimal guidance conditions. (c-d) Fidelity-adaptive scaling: Performance gains relative to baseline fidelity thresholds in offline SOO (c) and MOO (d) tasks, with dashed lines indicating empirically optimal thresholds ($\tau_{\text{opt}} = 0.827$ for SOO; $\tau_{\text{opt}} = 0.87$ for MOO). (e-h) Inference-time scaling efficiency: HV (e,g) and IGD (f,h) versus number of function evaluations (NFE) for ZDT3 (e-f) and RE21 (g-h) task for three approaches, comparing standard denoising, self-IS scaling, and FKS scaling methods. Consistent performance improvements are achieved through adaptive noise-space exploration.
• Figure 5: Ablation on duplication size of Self-IS-based ManGO.