Inverse Design for Conditional Distribution Matching

2026-05-10

AI Summary

Traditional inverse design methods are limited to point-wise target outputs and struggle to accommodate design requirements expressed as target distributions. This work formalizes, for the first time, the distribution-level inverse design problem and introduces a new paradigm termed Conditional Distribution Matching (CDM), defining two task variants: CDMS and CDMO. The authors propose MLGD-F, a plug-and-play inference algorithm that efficiently solves these tasks without additional training. MLGD-F leverages a pre-trained score-based diffusion model combined with a single-step conditional sampler, using a matching loss to guide gradient updates. The method successfully recovers inputs whose outputs align with complex target distributions—including discrete mixtures and continuous low-rank supports—demonstrating effectiveness across synthetic data, structured image transformation, and generative editing tasks.

Abstract

Generative models are powerful tools for sampling from a learned distribution $\mathcal{P}(Y \mid X)$, and inverse-design methods invert this map to find an input $x$ that produces a desired point output $y^*$. However, many design goals are naturally distributional rather than pointwise, incorporating the inherent uncertainty of $Y$ and targeting a specific form for it, a task not addressed by standard inverse design. To address this issue we introduce Conditional Distribution Matching (CDM), a new inverse-design problem class in generative modeling: given a joint distribution $\mathcal{P}(X, Y)$ and a target distribution $\mathcal{G}(Y)$, find an input $x^*$ whose induced conditional distribution $\mathcal{P}(Y \mid X = x^*)$ matches $\mathcal{G}$. We formally define two variants: Conditional Distribution Matching Sampling (CDMS) and Conditional Distribution Matching Optimization (CDMO). To solve these problems, we propose MLGD-F (Matching-Loss Guided Diffusion with a Fast inner sampler), a plug-and-play inference-time algorithm that combines a pretrained score-based diffusion model with a pretrained fast conditional sampler, requiring no additional training or fine-tuning. By leveraging single-step conditional sampling, MLGD-F enables tractable gradient computation, making the estimation of $\mathcal{P}(Y \mid X)$ both memory-efficient and computationally lightweight. We validate MLGD-F on synthetic benchmarks, structured image transformations, and generative editing optimization, demonstrating reliable recovery of inputs whose conditional distributions match diverse user-specified targets, including discrete mixtures and continuous low-rank supports.

Problem

• Inverse Design
• Conditional Distribution Matching
• Generative Modeling
• Distributional Target
• Score-based Diffusion

Innovation

• Conditional Distribution Matching
• Inverse Design
• Score-based Diffusion
• Distributional Optimization
• Plug-and-play Inference