Amortizing intractable inference in diffusion models for vision, language, and control

Thank you for your feedback and positive assessment of the paper! We answer your question below:

Does Gaussian initial state violate assumptions?

Assuming the initial distribution is Gaussian does not violate the GFlowNet assumption, as it amounts to assuming the first generation step transitions from an abstract initial state to a sample from that Gaussian distribution. When training, e.g., with the trajectory balance objective, the likelihood of this first transition -- from the abstract initial state to the Gaussian noise sample -- appears in the decomposition of the sampling trajectory's forward likelihood. With RTB, the likelihood of this initial-to-noise step is the same for the prior and posterior models and thus cancels in the loss.

Typos

• The direction of the arrows in $x_1\leftarrow\dots\leftarrow x_0$ is intentional: the arrows always show the transitions that exist in the MDP, but the sampling of the noising trajectory happens in reverse (starting from data $x_1$ and ending at noise $x_0$). We can revise if you think this causes confusion.
• Equation (11): Yes, ${\rm NN}_2$ isn't taking $x_t$ as input. We will fix this.

Thanks for reading carefully!

Please let us know if there are any other questions we can answer during the discussion period.

Thank you for your comments. We are happy to see that you appreciated the breadth of applications presented in our paper and the value of the contribution.

Below we have attempted to answer your questions and to clarify a few misunderstandings in the review. We note that most of the listed weaknesses and questions pertain to specifics of individual experiments, not to the main contribution of the paper: a general-purpose method that can be applied to a wide range of tasks simultaneously, achieving results comparable to specialized methods that had been developed for each problem. We hope that our reponses will help to answer your questions and concerns.

Improvements in image quality

Measuring posterior sampling quality is difficult for text-to-image generation models, since we do not have ground truth samples from the posterior. We acknowledge that differences in visual quality compared to DPOK are not apparent, and the visual results (Figure 4) should be seen as complementing Figure 3 in showing comparable quality and diversity to baselines.

RTB training facilitates asymptotically unbiased posterior sampling. This may not necessarily translate to improvements in auxiliary fidelity metrics such as those we must resort to when ground truth posterior samples are unavailable (as in the text-to-image experiments). However, in the experiments where we do have access to unbiased samples from the target classes (MNIST/CIFAR-10), it does translate to better FID scores (Table 2).

In practice, text-to-image generation systems may use a combination of techniques to improve samples (e.g., collecting data from human raters, imposing auxiliary aesthetic constraints, etc.). Our work adds a novel, general, and probabilistically principled technique to this toolbox.

Fair comparisons

We answer a few questions related to comparisons between other methods and hours, hoping to clarify why we believe they are fair and informative.

Infilling baselines seeing the right context

Good question; no, it is not the case that baselines see only $x$ (left context) and not $y$ (right context). The fact that this was not clear suggests we should write it explicitly in the paper; thank you for directing our attention to it.

In fact, all the baselines with the exception of "Prompt (x)" do condition on both $x$ and $y$, making the methods comparable. For example, for the prompting and SFT autoregressive baselines, the model receives a prompt of the form "Beginning: {x}, End: {y}, Middle: ", following prior work (see Appendix F). The difference between "Prompt (x)" and "Prompt (x,y)" shows the importance of conditioning on $y$.

Algorithm components

To clarify, we apply the same collection of techniques with RTB and with methods from prior work wherever applicable. For example, all methods used the same noising schedules. Gradient clipping, fine-tuning the pretrained prior checkpoint, and parameter-efficient training were done for RL baselines as well.

On the other hand, some techniques used with RTB training are not applicable to other methods. Data collection methods used for off-policy training cannot be applied to compared to manifestly on-policy finetuning algorithms such as DPOK. The inference-time classifier guidance baselines, such as those in Table 2, do not require any training, so training hyperparameter choices do not apply to them.

Altogether, while we believe the comparisons we make are fair, we agree that ablations on the techniques used in the paper (while not being applicable to baseline methods) would add to the understanding of the proposed algorithm. We will include ablation studies about this in the final version.

Prior/posterior training data

It is important to note that the newly proposed algorithms and most baselines are "data-free" -- that is, the prior model is pretrained and we receive no ground truth samples from the posterior, but rather must discover its modes through exploration and can query for the reward of a generated sample.

How were samples for figures chosen?

In the image generation figures (Figures 3 and 4 and Appendix H.1), the samples are not cherry-picked. In Appendix H.1, in each row the images are generated from seeds 0-9.

RL baselines

The reporting of baseline numbers and highlighting in Table 4 exactly matches the form in a line of previous offline RL papers (see, e.g., [35,43,36]), where std is reported only for the newly proposed algorithm alongside baseline means. However, we will update the paper to include std for comparisons, and have attached an updated table with std bars from baselines which report this metric in the global rebuttal document (Table 4).

Line 112

In that line, $p^{\rm post}$ is a distribution over $x_1$, where $y$ is fixed. For a fixed $y$, the product is proportional as a function of $x_1$ to the posterior $p(x_1\mid y)$.

Thank you again for your feedback. We hope these answers and clarifications help; if they do, please consider updating your score. Let us know if there are any other questions we can answer during the discussion phase.

Dear reviewer,

Thanks for your response. We are happy to clarify these points:

1. The training free baselines (DPS, LGD-MC) cannot be used with discrete diffusion models for the text infilling task and are only derived for the particular case of inverse problems in continuous space. DPOK/DDPO are RL finetuning baselines which can in principle be used for all the tasks, but do require significant changes to be adopted for the case of discrete diffusion (which was not considered in the original work). Note that all the baselines are also biased samplers, and hence expected to give worse performance for tasks requiring unbiased posterior inference such as classifier guidance (as we show in the paper) and other important scientific tasks (for example [1]). In the context of offline RL, the D-QL baseline trains the policy with an objective very similar to DPOK - The policy is trained to maximize the Q function with a per-step KL regularization term with behavior policy. We can add some notes in the appendix to make these connections clearer in the final version.
2. We agree that the claim "matches state-of-the-art" is more accurate, and will update the abstract in the revision.

We hope we've addressed your questions and concerns. Feel free to reach out with any additional questions.

[1] Adam, Alexandre, et al. "Posterior samples of source galaxies in strong gravitational lenses with score-based priors." arXiv preprint arXiv:2211.03812 (2022).

Thank you for your constructive feedback and good questions. We are glad you found the paper well-written and the empirical results compelling. We believe that the diversity of domains on which we showed the effectiveness of our approach makes this work both a valuable illustration of off-policy finetuning for intractable diffusion posteriors and a starting point for a number of possible domain-specific applications.

We have addressed each of your questions below and added three new experiments in response to your comments (please see the response to all reviewers).

Text infilling and LLM fine-tuning

We believe there may be a misunderstanding present. The infilling results are for the setting of fine-tuning a pretrained discrete diffusion language model, which does not use a Gaussian transition kernel. In discrete diffusion language models, the kernel is categorical at each position.

The text infilling experiment (Table 3 in the paper) thus supports two claims:

• that the proposed diffusion fine-tuning algorithm is applicable outside the Gaussian setting (and even outside the continuous setting);
• that fine-tuning of LLMs for intractable posteriors, such as infilling, using off-policy RL can be generalized beyond autoregressive language models (as done in [28]).

On scalability and cost

We answer a few questions related to scaling and fine-tuning cost comparisons.

Scalability

RTB actually scales surprisingly well to larger models and more diffusion steps (relative to simulation-based methods) despite being a trajectory-level objective.

An important observation about the RTB objective -- not included in our original submission -- is that computing the RTB gradient does not require storing the computation graph of all timesteps to be updated. We observe that the gradient of the RTB objective for a single trajectory is just the sum of per-step log-likelihood gradients scaled by the RTB residual: $\nabla_{\phi}L_{RTB} = 2\left(\log\frac{Z_{\phi}}{r(x_1)} + \sum_{i=1}^T\log\frac{p_{\phi}^{post}(x_{i\Delta t} \mid x_{(i-1)\Delta t})}{p_{\theta}(x_{i\Delta t} \mid x_{(i-1)\Delta t})}\right) \cdot \nabla_{\phi}\sum_{i=1}^T\log p_{\phi}^{post}(x_{i\Delta t} \mid x_{(i-1)\Delta t}).$ Because the likelihood gradients can be accumulated during the forward pass, this allows for a batched gradient accumulation version of the update. For trajectory length (number of diffusion steps) $T$ and accumulation batch size (number of time steps receiving a gradient signal in each backward pass) $B$, the number of batched forward passes required scales as $T/B$.

Only the accumulation batch size $B$, not the trajectory length $T$, is constrained by the memory budget. This means we can easily scale training with large number of diffusion steps $T$ without increasing the variance of the gradient through stochastic subsampling, with training time growing only linearly with number of time steps under a fixed memory budget.

The batched update implementation was not included in our original codebase, but we have performed experiments to validate it and will include them in the revision.

Finally, we highlight that this is in contrast to other diffusion samplers (e.g., PIS/DDS/DIS) which differentiate through the sampling (SDE integration) process and therefore need to store the entire computation graph. For these methods, memory necessarily scales with the length of the trajectory.

Cost of training posterior vs. prior

For our image experiments (Table 2 and Figure 2), a very small number of finetuning steps is necessary for convergence to finetune a posterior with RTB (usually between 100 and 300 are sufficient, although we run for 1500), where each step performs a gradient update on a batch of trajectories (from 16 to 32, depending on the experiment). In contrast, training a prior diffusion model even on MNIST requires tens of thousands of training iterations.

LoRA in MNIST/CIFAR setup

The experiments on discrete diffusion and offline RL (Sections 3.3 and 3.4) do not use LoRA, while the image experiments (Sections 3.1 and 3.2) do. Note that parameter-efficient fine-tuning is standard for large models such as Stable Diffusion and also adopted by baseline methods (e.g., DPOK).

For example, for MNIST and CIFAR we use LoRA (rank=32) on the key, query and value vectors of the denoiser’s attention mechanism of a UNet model. RTB can be applied without LoRA parameterization on these tasks too; please see the reponse to all reviewers and Table 1 in the attached pdf.

Conditional constraints

Using conditional constraints requires a more detailed study of architectures and conditioning strategies, since the unconditional prior cannot directly be finetuned. However, we have run a preliminary experiment on the MNIST even/odd digits task (Table 2); please see the response to all reviewers.

Data-dependent upper bound

Please see the response to all reviewers for such an upper bound on the text infilling task, where we have a larger ground truth dataset.

Max or sum in classifier reward

Thank you for the qustion. The choice of the reward function for the MNIST finetuning task (L220) was to discourage generation of ambiguous digits: for example, $\sum_{i\text{ even}} p(c=i\mid x)$ is high when $p(c=0\mid x)\approx p(c=2\mid x)\approx\dots\approx p(c=8\mid x)$. In the current formulation, high reward can only be achieved by generating an image whose likelihood of at least one class is high. We observed that replacing max by sum resulted in more such ambiguous digits, which would appear to be a weakness of the pretrained prior model.

Prompts for ImageReward task

For the results in Figures 3 and 4, we use a similar setup to DPOK [15]. The prompts are copied verbatim from that paper.

We hope we have answered your questions satisfactorily; if so, please consider updating your score. Let us know if there are any other questions we can answer during the discussion phase.