Inverse Design for Conditional Distribution Matching

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.

• The paper introduces the Conditional Distribution Matching framework, generalizing inverse design by matching full output distributions to user-specified targets.
• It presents MLGD-F, a plug-and-play inference-time optimization method that combines fast conditional samplers with diffusion models.
• Empirical results on synthetic tasks, MNIST, and Stable Diffusion demonstrate notable speedups and improved distributional fidelity.

Conditional Distribution Matching: A Generalization of Inverse Design

Problem Formulation and Conceptual Innovation

The paper "^^^^0^^^^" (2605.09439) introduces Conditional Distribution Matching (CDM) as a formalization of distributional inverse design. Unlike traditional inverse design—where the goal is to find input $x$ such that the output $y = f(x)$ matches a point target $y^*$—CDM seeks an input $x^*$ whose induced conditional distribution $P(Y | X = x^*)$ matches a user-specified target $G(Y)$. Two variants are defined: Conditional Distribution Matching Sampling (CDMS), which aims to sample inputs whose induced conditional match $G(Y)$, and Conditional Distribution Matching Optimization (CDMO), which minimizes the distributional distance between $P(Y | X = x)$ and $G(Y)$.

This problem generalizes loss-guided generation and pointwise inverse design by targeting distributions rather than specific values. The formulation is relevant in scenarios where outcomes are inherently uncertain or multimodal—such as protein sequence design with structural ambiguity, or generative pipelines requiring demographic balancing.

Methodological Contributions

MLGD-F: Plug-and-Play Inference-Time Optimization

The authors propose MLGD-F (Matching-Loss Guided Diffusion with a Fast inner sampler), which integrates a pretrained score-based diffusion model with a pretrained, differentiable, few-step conditional sampler. This enables efficient gradient-based optimization at inference time, without retraining or fine-tuning the generative models. MLGD-F operates as follows:

• The outer loop performs loss-guided reverse diffusion, where the score network is augmented by the gradient of a distance metric between $P(Y | X = x)$ and $y = f(x)$0.
• The inner loop leverages a fast conditional sampler (consistency model or distilled diffusion model), allowing tractable and memory-efficient estimation and differentiation of the distributional loss.

The proposed pipeline exploits single-step or few-step sampling in the conditional model, which avoids the prohibitive memory and computational cost of backpropagating through extended multi-step diffusion chains.

Theoretical Analysis

The paper provides a rigorous analysis of MLGD-F’s estimator for the distributional loss and its gradient. The error decomposition is explicit: finite-sample variance, sampler-output bias, and Jacobian bias. Bounds are stated in terms of the fidelity of the distilled sampler to the true conditional, and the propagation of memory savings versus gradient discrepancy is analytically quantified. Notably, replacing a $y = f(x)$1-step teacher diffusion with a $y = f(x)$2-step student conditional sampler reduces memory usage and runtime by a factor of $y = f(x)$3, but introduces gradient errors proportional to the distillation fidelity gap.

Empirical Validation

Synthetic Benchmarks

On synthetic Mixture of Gaussians tasks across increasing dimensionalities, MLGD-F attains near-optimal distribution matching with a speedup of $y = f(x)$4 (2D) to $y = f(x)$5 (10D) compared to slow loss-guided diffusion. The variance and bias trade-off between sampler fidelity and gradient variance is characterized: in higher dimensions, the cleaner gradient provided by few-step samplers outweighs fidelity loss, making MLGD-F superior in scalability.

MNIST Rotational Ambiguity

For the MNIST rotated-digits task, MLGD-F optimizes over image space to recover digits whose induced distribution over rotation angles matches diverse targets (uniform, bimodal, unimodal). Without explicit digit-class supervision, the method produces semantically meaningful digits—circular "0" for uniform rotation, symmetric digits for bimodal targets, and canonical upright digits for unimodal targets.

Stable Diffusion Image Editing

MLGD-F is applied to large-scale distribution-guided image generation with SDXL and SDXL-Turbo. The method successfully edits scribble images to ensure the downstream conditional distribution (e.g., over gender, age) matches user-specified targets in CLIP embedding space. Quantitatively, MLGD-F achieves 22–33% improvement in relative MMD scores compared to baseline methods, and recovers distributional proportions closely tracking targets (e.g., 47.4% male for a 50% target, 26.4% for 25%). The approach handles cases from discrete mixtures to continuous low-rank manifolds robustly.

Memory profiling demonstrates that MLGD-F’s fast sampler (Ks=2) consumes 43 GB VRAM, while multi-step diffusion (K+=30) would require 375 GB, making the latter infeasible for research and deployment.

Practical and Theoretical Implications

Plug-and-Play Tractability and Inheritance of Model Advances

MLGD-F’s design allows optimization on fixed, frozen generative pipelines, enabling direct inheritance of advances in pretrained diffusion or conditional models. This is distinguished from adversarial distribution matching approaches that require retraining per target.

Few-step samplers, by virtue of their shallow computation graph, enable scalable optimization—especially for high-dimensional and complex output spaces. The single-step property not only provides tractable gradient computation but may also yield cleaner signals in the gradient oracle, as observed in high-dimensional synthetic experiments.

Limitations and Directions for Future Work

Optimization is bottlenecked by the fidelity of the distilled conditional sampler—if it poorly approximates the true conditional, bias is irreducible. The stochastic nature of the inner estimator typically necessitates multiple independent runs, increasing compute demand and runtime. The method currently requires differentiable conditional models; extension to black-box or non-differentiable settings would require gradient-free optimization.

The extension of CDM to discrete input spaces, integration of exact sampling correctors (e.g., twisted SMC, MCMC), and theoretical guarantees for Jacobian fidelity in distilled models are natural future directions.

Broader Impact and Risks

Distributional control offers applications in bias correction, fairness, and targeted diversity. However, unconstrained specification of target distributions could enable deliberate steering of output distributions in potentially harmful ways, emphasizing the need for robust safety and interpretability mechanisms.

Conclusion

Conditional Distribution Matching, as formalized by this work, expands the scope of inverse design by allowing optimization for user-specified distributions over outputs rather than pointwise targets. MLGD-F operationalizes this at inference time, leveraging the tractability of few-step conditional samplers to achieve scalable distributional inverse design. Empirical and theoretical analysis demonstrate its efficacy across synthetic and real-world generative pipelines. The approach lays groundwork for gradient-based optimization over induced conditional distributions in generative modeling, with significant practical and theoretical implications for controlled generation and distributional fairness in AI systems.