Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models

We thank the reviewer for their detailed review and constructive feedback, as well as the positive comments on the exposition and evaluations.

To address the questions:

Efficiency of training a model for each task

It is true that the outsourced diffusion sampler must be trained separately for each likelihood constraint $p(y\mid x)$. However, for the diffusion sampler, the inference-time cost per sample is much lower than for MCMC (i.e., the sampling cost is amortized). In fact, the total cost for generating samples for evaluation was much lower using our method (including training time) compared to MCMC, for instance:

• The CIFAR-10 evaluation required generating 1000 samples, which took 10 hours with Hamiltonian Monte Carlo (using multiple chains in parallel), while the outsourced diffusion sampler took 5 hours for both training and sampling (of which sampling was a neglible fraction).
• A similar pattern held for MCMC in the protein experiments: the entire experiment took 8 hours for MCMC and 4 total hours for our method.

Additional details on the timing comparisons between MCMC baselines and our method are available in Appendix B.

Regarding the need to retrain for new constraints, we point out that the regular functionality of generative models remains intact when using our method to sample the latent space. Namely, we can still apply approximate training-free methods for guidance to solve intractable sampling problems (e.g., for inpainting or linear inverse problems), with the only difference being that the noise components are obtained from our diffusion sampler. Our method offers a computational trade-off whereby additional compute can be spent to sample more accurately according to a given posterior distribution, compared to the generative model with standard noise.

Size of example images: Can we use larger models?

The largest latent space we explore in our experiments is with Stable Diffusion 3, featuring a latent resolution of $16 \times 64 \times 64$, corresponding to an image resolution of $3 \times 512 \times 512$. We consider this to be a reasonably high resolution for practical applications: most image generative models, save those with excessive hardware and sampling time requirements, use similar latent dimensions. (Although diffusion models trained in pixel space with higher dimension exist, they typically require far more sampling steps than ones that work in a latent space.) Generally, larger latent spaces do increase training times, but our experiments demonstrate that our method scales effectively even with high-fidelity models like SD3.

We note that for many generative models like GANs, the noise space is of much lower dimension than the data space (e.g., 512 vs. $3\times256\times256$ for the FFHQ StyleGAN3 we consider), which is further motivation to perform outsourced sampling.

Finally, thank you for your suggestions on reorganizing the figures. We will incorporate this feedback into our final submission.

Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.

We appreciate your good questions and constructive feedback. We have responded to your points below.

Algorithm details

We agree that the addition of an algorithm block would help readers understand how to implement the core training loop, and added this to Algorithm 1 in page 2 of the linked pdf.

FID scores not state-of-the-art

While FID scores are indeed worse than adjoint matching on CIFAR-10, we highlight that ODS is a general approach applicable to any generative model, while adjoint matching is specific to flow models (and not all flow models, see end of §3.1) and requires differentiating the reward.

The FID scores reported in [Venkatraman et al.] are an order of magnitude lower than those presented in our paper. On examining the underlying codebases we discovered that the discrepancy arises from the FID packages used. Specifically, [Venkatraman et al.] uses this codebase, while we use this one. The distinction lies in the handling of image resolution: CIFAR images are smaller than ImageNet images, and only the latter implementation (which we use) accounts for this difference appropriately.

Efficiency

Our efficiency claim refers to lower training time of ODS relative to finetuning methods, while having much cheaper inference than MCMC. There is a small (25 step) inference overhead relative to Adj. matching (for reducing this, see Appendix C.1).

Which part of the proposed method is new?

Trajectory balance has previously been employed for training diffusion samplers from unnormalized densities (notably in [Lahlou et al.] and [Sendera et al.]), and a variant of it, relative TB, is a diffusion-specific fine-tuning method in [Venkatraman et al.]. Our novel contribution is the general capability to train a sampler (using TB) of the Bayesian posterior in noise space (§3.2) for any generative model, providing a general-purpose framework for amortized posterior inference. The exact way in which ODS generalizes relative TB is detailed in Appendix A.1.

Training-free baselines and classifier guidance

Unbiased classifier guidance requires a time-dependent likelihood gradient: $\nabla \log p(y \mid x_t) = \nabla \log \mathbb{E}_{p(x_0 \mid x_t)}[p(y \mid x_0)]$

This necessitates training a classifier on noised inputs, as using one trained on clean data renders the gradient intractable. Such methods therefore make stronger assumptions than our setting, which assumes a black-box classifier on clean data.

Training-free methods like DPS are biased approximations to classifier guidance but work well in practice. DPS uses the approximation $\nabla \log p(y \mid x_t) \approx \nabla \log p(y \mid \mathbb{E}[x_0 \mid x_t])$, and requires the same assumptions as our proposed algorithm making it an appropriate baseline.

We’ve added results for DPS (Diffusion Posterior Sampling, [Chung et al.]), a strong training-free baseline for approximate posterior sampling with diffusion and flow models using the I-CFM prior on CIFAR-10. Results are in Table 1 of the response to Reviewer qmBQ and also here. DPS achieves high reward but poor FID, reflecting highly biased sampling.

Additional experiments with SD1.5

We conducted additional experiments with Stable Diffusion 1.5 to compare against DDPO [Black et al.], DPOK [Fan et al.], and RTB [Venkatraman et al.], as requested. We also included RTB with the I-CFM prior as a baseline for the CIFAR-10 class-conditional sampling task. Results are shown in Tables 2 and 1 of our response to reviewer qmBQ and are available here. On SD1.5, ODS achieves the a strong balance between reward and sample diversity.

Stability of training

ODS uses TB, an off-policy RL objective that does not require backpropagating through the full sampling chain. Instead, it only needs local gradients of log-likelihoods for each sampling step (see [Venkatraman et al.], Appendix H.1)

Such off-policy optimization has three benefits, according to prior work (see, e.g., [Nüsken & Richter], [Sendera et al.]):

• The reward function need not be differentiable, since TB avoids using log-reward gradients;
• TB avoids instability and mode collapse seen in methods that backprop through the full SDE trajectory (e.g., PIS [Zhang & Chen]);
• Memory usage is minimal thanks to a gradient accumulation strategy (included in our code), avoiding storage of full computation graph

In our experiments, we found training to be very stable for all tasks except for the protein diversity task, where policy collapse sometimes occurs.

Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.

Thank you for the detailed review and constructive feedback, as well as pointing out that the paper is very well written.

You raise some important concerns, which have helped us improve the paper and which we hope to address:

Posterior evaluation

We agree that more comprehensive posterior evaluations are important for a thorough assessment.

First, we emphasize that the FID scores reported in our original CIFAR-10 experiments capture both reward and coverage/diversity, and serves as a proxy metric for closeness to the true posterior. FID remains a widely accepted and robust metric for evaluating the sample quality of image generative models.

For the CIFAR-10 experiments (comparing our amortized sampler to ground truth posterior samples), we tried posterior evaluation with TARP [Lemos et al., 2023] and PQMass [Lemos et al., 2024], an unconditional two-sample variant of TARP, but were not able to obtain results that seemed meaningful. For example, running TARP frequently showed prior samples as perfectly calibrated and unbiased relative the conditional dataset, which is obviously false -- but the discrepancy is captured correctly in the FID scores. (Published code from those papers was used; to ensure the results do not contain an error, we checked that our code could reproduce values in line with those reported on CIFAR-10 in [Lemos et al., 2024].)

For other experiments, ground truth samples from the posterior are not available, complicating evaluation (for instance, FID cannot be computed), which is the reason we resort to other metrics. For prompt-conditioned sampling we instead report the average CLIP feature distance, following the approach used by [Venkatraman et al.]. In our Stable Diffusion 1.5 experiments, the CLIP feature diversity scores exhibit a consistent and interpretable trend: the prior achieves the highest diversity, while DDPO, which performs greedy reward maximization without KL regularization, yields the lowest diversity. These trends support the reliability of our evaluation metrics in capturing meaningful aspects of sample diversity.

At the reviewer's suggestion, we also report the ELBO (i.e., the lower bound on the log-normalizing constant $\log Z$, where the ELBO-to-log-likelihood gap is the reverse KL divergence). This is a widely used metric in the diffusion samplers literature, including [Sendera et al.], and is maximized by a perfect sampler of the target density. We note however, that this metric cannot be used to evaluate inference-time baselines or MCMC.

Addition of the ELBO metrics for CIFAR-10 and SD3 are included in Tables 1 and 3 repectively, within the linked PDF.

Additional baselines

We include experimental results for the RTB baseline proposed by [Venkatraman et al.] for flow models trained with independent coupling, which can subsequently be fine-tuned as diffusion models. This evaluation is done for RTB using the I-CFM prior on CIFAR-10; we also include new results comparing our method against RTB, DPOK [Fan et al.] and DDPO [Black et al.] with Stable Diffusion 1.5 as a prior with the same setup used by [Venkatraman et al.]. These are included in Tables 1 and 2 of the response to Reviewer qmBQ, and also in the linked pdf.

We note that for CIFAR-10, the RTB baseline is significantly more unstable to train: requiring LoRA fine-tuning [Hu et al.], as well as lower learning rates to avoid a quick policy collapse. This aligns with design choices in the RTB paper [Venkatraman et al.]. Even with this, the training remains somewhat unstable for the flow model architecture, being very sensitive to training hyperparameters. Within the training time allocated to ODS, the model achieves a modest improvement in reward; however, the generated samples often fail to consistently belong to the target class and frequently exhibit visual artifacts. This degradation in quality is further reflected in poorer posterior metrics (FID, ELBO). We argue the training instability is caused by the flow prior, which uses a different noise schedule as well as much fewer inference steps compared to experiments in [Venkatraman et al.]. This instability could lead to "reward-hacking" behavior where reward improves at the cost of sample quality.

By contrast, ODS improves largely within the first 2-3 GPU hours, and only marginally afterwards. The general difficulty of training RTB with the flow prior reinforces the claim that the latent posterior $p(z|y)$ is a simpler target than the image posterior $p(x|y)$.

Finally, thank you for your helpful suggestions on improving the notation and writing clarity, as well as for pointing out the typos. We will implement this feedback for our final submission.

Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.

Thank you for the helpful review. We address the weaknesses and questions raised by you in our response below.

Comparison with model-specific approaches

We thank the reviewer for pointing out the additional model-specific baselines with which to compare out general approach.

• We note that for GANs, we compare to a latent space exploration baseline similar to PULSE, Hamiltonian Monte Carlo (HMC).
• We have added the DPS baseline to the CIFAR-10 experiments.
• We have added RTB as a diffusion-specific baseline for CIFAR-10 and Stable Diffusion-1.5 experiments.

The additional results are included in Tables 1 and 2 below, and also in the linked pdf.

Table 1: CIFAR-10 posterior sampling results, averaged over 10 classes

In the CIFAR-10 results we note that DPS obtains a higher reward than Outsourced Diff., with a worse FID score (and ELBO -- see the response to Reviewer nZ5P). This points to possible reward hacking. Additionally, the added RTB baseline for this experiment proved much more unstable during training than our method, reflected in its worse FID (and ELBO) scores, while improving the reward relative to the prior. Some RTB runs underwent policy collapse, requiring early stopping for the reported results.

Table 2: SD 1.5 fine-tuning results. DDPO, DPOK and RTB results taken from [Venkatraman et al.]

For the SD1.5 experiment, we see that our method achieves a slightly lower reward than RTB, while obtaining a higher diversity over all baselines (other than the prior).

Relationship to non-amortized methods (golden noise and test-time scaling)

We agree that incorporating a discussion comparing our proposed amortized sampler with these inference-time optimization methods would strengthen the paper. Notably, the latent space HMC baseline used in our CIFAR-10 and FFHQ experiments shares conceptual similarities with these inference-time noise optimization approaches. While the methods in prior work do not asymptotically sample from the true Bayesian posterior over the noise space, they effectively bias the sampling process toward high-reward regions, often producing qualitatively similar results. In contrast, our method frames prior-regularized reward fine-tuning as posterior inference, whereas these inference-time techniques are more aligned with direct reward maximization. We will integrate this discussion into the appendix of the final version.

Organization of results

We thank the reviewer for their feedback on the organization of results. We will merge some of the results that share similar metrics and baselines (such as CIFAR and FFHQ), and reorganize the tables to make the presentation more clear for the final submission.

Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.