📊 statistics

Inverse Design for Conditional Distribution Matching

This paper introduces Conditional Distribution Matching (CDM), a new inverse-design problem class for finding inputs that induce a specific target conditional distribution rather than a single point, and proposes MLGD-F, a training-free inference-time algorithm that combines a pretrained diffusion model with a fast conditional sampler to efficiently solve this problem.

Original authors: Ori Meidler, Shaul Tolkovsky, Or Zuk

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ori Meidler, Shaul Tolkovsky, Or Zuk

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: From "One Perfect Shot" to "The Perfect Mix"

Imagine you are an architect working with a magical, pre-trained AI builder. This builder is amazing at creating houses, but you can't change how it thinks or learns; it's "frozen." You can only give it a blueprint (an input) and it builds a house (an output).

The Old Way (Standard Inverse Design):
Usually, if you wanted a house, you would tell the builder: "Build me a house that looks exactly like this specific photo." The builder would try to find a blueprint that results in that one perfect house.

The Problem: What if you don't want just one house? What if you want a blueprint that, when used by this builder, produces a mix of houses? Maybe you want 50% modern cottages and 50% Victorian mansions, or a smooth blend of styles. The old method can't do this because it's obsessed with hitting a single target point, not a whole distribution of possibilities.

The New Way (This Paper's Solution):
The authors introduce a new method called Conditional Distribution Matching (CDM). Instead of asking for one specific house, they ask: "Find a blueprint that makes the builder produce exactly the mix of houses I want."

The Goal: You specify the desired "flavor profile" of the output (e.g., "I want 75% female portraits and 25% male portraits"). The system finds the input that, when fed into the frozen AI, generates that exact statistical mix.

The Challenge: The "Black Box" and the "Slow Motion"

The paper faces two main hurdles:

The Frozen Builder: You can't retrain the AI. You have to work with what you have.
The Speed Trap: To figure out if a blueprint is good, the AI has to generate many sample houses to see if they match your desired mix.
- The Analogy: Imagine the builder is a slow-motion camera. To check one blueprint, it has to take 30 slow-motion frames (steps) to build the house. If you need to check 100 variations to get the math right, and you have to do this 100 times during the search, the process takes forever and requires a supercomputer that doesn't exist (it runs out of memory).

The Solution: MLGD-F (The "Fast-Forward" Guide)

The authors created an algorithm called MLGD-F (Matching-Loss Guided Diffusion with a Fast inner sampler). Here is how it works using an analogy:

1. The "Fast-Forward" Sampler (The Inner Loop)
Instead of asking the builder to take the full 30 slow-motion steps to check a blueprint, the authors use a "distilled" version of the builder.

The Metaphor: Think of the original builder as a master chef who takes 30 minutes to bake a cake. The "distilled" builder is a sous-chef who has memorized the recipe and can bake the same cake in one step (or very few steps).
Why it matters: Because this "sous-chef" is so fast, the system can instantly generate hundreds of sample houses to check if they match your desired mix. This makes the math possible without crashing the computer's memory.

2. The "Loss-Guided" Search (The Outer Loop)
The system starts with a random blueprint. It asks the "sous-chef" to generate a batch of houses. It compares this batch to your target mix (e.g., "You gave me too many Victorian mansions").

It calculates a "score" (how far off the mix is).
It uses that score to nudge the blueprint slightly in the right direction.
It repeats this process, slowly refining the blueprint until the builder's output perfectly matches your desired distribution.

What They Proved (The Experiments)

The team tested this on three levels, like training for a marathon:

The Practice Run (Synthetic Data): They used simple math shapes (Gaussian mixtures).
- Result: MLGD-F found the perfect input 11 times faster than the slow method, with just as much accuracy.
The Intermediate Test (MNIST Digits): They used images of handwritten numbers.
- The Task: "Find a number image that, when rotated, looks like a mix of 0s, 1s, and 8s."
- Result: The system found specific digit shapes (like the circle of a '0') that naturally satisfied the rotation requirements, proving it could handle complex image spaces.
The Grand Challenge (Stable Diffusion): They used a massive, real-world AI that generates high-quality portraits.
- The Task: "Start with a sketch of a man. Find a modified sketch that, when fed to the AI, produces a 50/50 mix of men and women, or a smooth gradient of ages from 40 to 79."
- Result: The system successfully tweaked the original sketch (changing just a few lines around the eyes and hair) to shift the AI's output distribution exactly as requested.
- Key Finding: Without the "Fast-Forward" (distilled) sampler, this task would require 375 GB of computer memory (impossible on standard hardware). With their method, it only needed 43 GB.

The Takeaway

This paper solves a specific problem: How do you control a frozen AI to produce a specific variety of outputs, rather than just one specific output?

They did this by combining a "slow, perfect" AI (the frozen model) with a "fast, approximate" AI (the distilled sampler) to guide the search. This allows users to define complex goals—like "make the output diverse" or "balance the demographics"—and have the system find the input that achieves it, all without needing to retrain the massive AI models.

In short: They figured out how to tune a radio (the input) so that the static (the output) creates a perfect playlist of songs, rather than just playing one song on repeat. And they did it using a remote control that works 15 times faster than the old one.

Technical Summary: Inverse Design for Conditional Distribution Matching

1. Problem Formulation

Generative models have established themselves as powerful tools for sampling from learned distributions $P(Y | X)$ and for inverse design, where the goal is to find an input $x$ that produces a specific point output $y^*$ . However, many design goals are inherently distributional rather than pointwise. For instance, a designer may seek an amino acid sequence $x$ such that the resulting protein structure distribution $P(\text{structure} | X=x)$ matches a specific mixture of folds, or a generative pipeline that produces outputs balanced across demographic groups.

Standard inverse design methods are limited to point targets and cannot express goals where the desired outcome is a distribution over outputs. This paper introduces Conditional Distribution Matching (CDM), a new problem class in generative modeling. Given a joint distribution $P(X, Y)$ and a target distribution $G(Y)$ , the objective is to find an input $x^*$ such that the induced conditional distribution $P(Y | X = x^*)$ matches $G$ .

The authors formalize two variants of this problem:

Conditional Distribution Matching Sampling (CDMS): Finding a sample from a tilted distribution $Q_\beta(x) \propto P(x) e^{-\beta L(x)}$ , where $L(x) = \|P(Y | X=x) - G(Y)\|$ is the distributional distance.
Conditional Distribution Matching Optimization (CDMO): Finding the input $x^*$ that minimizes the distance $L(x)$ (the limit as $\beta \to \infty$ ).

Crucially, this setting assumes the generative pipeline (the conditional model $P(Y|X)$ ) is frozen and cannot be modified; the only lever available is the input $x$ .

2. Methodology: MLGD-F

To solve these problems, the authors propose MLGD-F (Matching-Loss Guided Diffusion with a Fast inner sampler). This is a plug-and-play, inference-time algorithm that requires no additional training or fine-tuning of the generative models.

Core Architecture

MLGD-F combines two pretrained components:

A Score-Based Diffusion Model ( $s_\theta$ ): Models the marginal prior $P(X)$ .
A Fast Conditional Sampler ( $f_\phi$ ): Models the conditional distribution $P(Y | X)$ . Crucially, this sampler must be a "few-step" model (e.g., a consistency model or a distilled diffusion model like SDXL-Turbo) capable of generating samples in a single step or very few steps.

Algorithmic Workflow

The method operates as a Loss-Guided Diffusion (LGD) process with a specialized inner estimator:

Outer Loop: Performs reverse diffusion steps to generate $x$ . At each step $t$ , it computes a Tweedie estimate $\hat{x}_0$ of the clean sample.
Inner Estimator (The Novelty): Instead of using a fixed loss function, the algorithm evaluates the distributional loss $L(x) = \|P(Y | X=x) - G(Y)\|$ $L (x) = ∥ P (Y ∣ X = x) - G (Y) ∥$ at each outer step.
- For a query $x$ , the fast sampler $f_\phi$ generates $n_{cond}$ samples from $P(Y | X=x)$ .
- These samples are compared against a fixed set of target samples $S_G \sim G$ using a distributional distance metric (e.g., Maximum Mean Discrepancy (MMD) or Sliced Wasserstein Distance (SWD)).
- The gradient of this loss with respect to $x$ is computed via automatic differentiation through $f_\phi$ .
Guidance: The computed gradient $\nabla_x L(\hat{x}_0)$ is added to the denoising update, guiding the trajectory toward inputs that induce the target distribution.

Why Fast Samplers?

The use of a single-step (or few-step) conditional sampler is essential for tractability. Standard diffusion models require unrolling $K^*$ steps (e.g., 30–1000) for sampling. Backpropagating through such a chain at every outer diffusion step would result in prohibitive memory costs ( $\Theta(K^*)$ ) and computational overhead. By using a distilled, single-step sampler ( $K_s=1$ ), the memory footprint is reduced by a factor of $K^*/K_s$ , making gradient computation feasible on standard hardware.

3. Key Contributions

Problem Definition: The paper defines CDM, moving beyond characterizing samples by their own properties to seeking inputs whose induced conditional distributions match user-specified targets.
Algorithm (MLGD-F): A novel inference-time algorithm that enables tractable gradient-based optimization for distributional matching by combining a pretrained diffusion prior with a fast, differentiable conditional sampler.
Theoretical Analysis: The authors provide an analysis of the plug-in loss estimator and the gradient estimator. They establish that the error decomposes into finite-sample variance, sampler output fidelity ( $\epsilon_s$ ), and Jacobian fidelity ( $\epsilon_g$ ). They formalize the memory-fidelity trade-off, showing that reducing the number of inner steps reduces memory linearly while introducing a bounded gradient discrepancy.
Empirical Validation: The method is validated across three scales:
- Synthetic: Mixture of Gaussians, demonstrating near-optimal recovery and speedups over standard diffusion-based inner loops.
- MNIST: A rotation task where the goal is to find digit images whose rotation angle distribution matches specific targets (uniform, bimodal, unimodal).
- Stable Diffusion: Large-scale image editing using SDXL and SDXL-Turbo. The method successfully edits scribbles to produce images with balanced gender distributions, skewed gender ratios, and continuous interpolations of gender and age.

4. Results

Efficiency: In synthetic benchmarks, MLGD-F achieved an 11× speedup over standard MLGD (using a slow diffusion sampler) while recovering near-optimal solutions. In high-dimensional settings, the single-step model's cleaner gradient signal allowed it to outperform the more accurate but variance-heavy multi-step sampler.
Memory Tractability: In the Stable Diffusion experiments, using a standard 30-step sampler (SDXL-Base) as the inner loop would require approximately 375 GB of VRAM (infeasible), whereas the distilled 2-step sampler (SDXL-Turbo) required only 43 GB.
Distributional Control: MLGD-F successfully recovered inputs inducing conditional distributions matching diverse targets:
- Discrete Mixtures: Achieved balanced (50/50) and skewed (25/75) gender distributions in generated portraits, outperforming baselines like "average scribbles" (which produce ghosted artifacts) and unguided SDEdit.
- Continuous Submanifolds: Successfully navigated continuous low-rank supports, such as interpolating between feminine and masculine features or varying age, while remaining on the data manifold.

5. Significance and Claims

The paper claims that CDM opens a new direction for optimizing over induced conditional distributions without retraining generative models.

Plug-and-Play Nature: A central advantage is that MLGD-F inherits quality improvements from the underlying pretrained models automatically. As better diffusion priors and few-step samplers become available, the optimization quality improves without algorithmic changes.
Feasibility at Scale: The authors emphasize that the few-step model is not merely a speedup but a necessity for tractability. Replacing the single-step sampler with a standard multi-step model renders the method infeasible at scale due to memory constraints.
Limitations: The authors are modest about the method's limitations. Optimization quality is bounded by the fidelity of the distilled conditional sampler; if the sampler poorly approximates the true conditional, the guidance gradient will be biased. Furthermore, the stochastic nature of the inner-loop estimation often requires multiple independent runs to find a good solution, multiplying compute costs. The method currently requires the conditional model to be differentiable end-to-end.
Broader Impact: The paper acknowledges that while distributional control can reduce bias (e.g., enforcing demographic balance), the same mechanism could theoretically be misused to steer frozen pipelines toward harmful distributions. However, since the method operates on frozen models, existing safety mechanisms remain active.

In summary, this work provides a framework for "inverse design" where the target is not a single point but a distribution, enabling users to control the statistical properties of generative model outputs through input optimization alone.