Summary
This paper proposed a method to train a posterior $p^{post}(x)$ given a prior distribution $p(x)$ and some additional (possibly unnormalized) constraint function $r(x)$, when the prior is a pretrained diffusion model. Through the choice over $r(x)$, this setup can capture a wide range of tasks. The authors propose *relative trajectory balance* (RTB) constraint, which, if satisfied by the posterior under training, guarantees that $p^{post} \sim p(x)r(x)$. By training $p^{post}(x)$ to satisfy the RTB constraint, one can recover the desired posterior. Importantly, this training objective allows for efficient computation of the loss and off-policy training. The paper provides convincing empirical results on a variety of tasks ranging from conditional generation to offline RL.
Strengths
**Motivation and Theoretical Justification**
* The paper provides a theoretical justification (Proposition 1) for the proposed method, which is a natural extension of the trajectory balance constraint for the setting where we cannot easily sample from $p^{post}$. This is a realistic constraint that one often encounters when trying to perform posterior inference given a pretrained prior.
**Empirical Analysis**
* While it's clear that the formulation does capture a wide range of tasks, it is nice to see an experimental evidence of that. The generality of the proposed method is emphasized through the results on four different tasks.
* I appreciate that the authors took time to discuss some of the implementation details (Sec 2.3).
**Writing Quality**
* The presentation is very well-done. The problem being tackled is clear, and the technical aspects of the solution are intuitive and easy to follow.
Weaknesses
**Scalability**
* One main concern I have is how scalable the method is due to trajectory matching, especially as the prior becomes larger. The main benefit of this technique relies on the assumption that one aims to repurpose an existing diffusion prior, which is presumable a large foundation model -- so it's important that the method scales desirably to the prior's size, which is not clear from the results in the paper.
* Section H discusses some techniques employed to reduce memory usage, but even this relatively small-scale experiment (50 steps on latent diffusion with LoRA), the authors report that they could fit only 8 timesteps on an A100.
* It'd be very useful to get some idea of how sensitive the method is to these hyperparameters. One extreme case to consider would be to use single-step gradient subsampling for a much larger T (e.g. 1000 steps).
**Experimental Results**
* Sec 2.3 claims that the proposed method should easily generalize to conditional constraints (which I believe), but there doesn't seem to be any experiment that tests this claim. For example, it'd strengthen the paper to have a result where RTB training successfully extracts a class-conditional model (that takes a conditioning input) from an unconditional one (perhaps for the CIFAR experiment).
* The text-conditional image synthesis experiment seems to be done for only four hardcoded text prompts, which is pretty limited.
* While the text in-filling results are great to have, a more practical and interesting task would be fine tuning a pretrained LLM. As hinted in Sec. 2.4, I'm curious if RTB can be used to efficiently customize an LLM, where we deviate from the Gaussian transition kernel.
* Having a data-dependent reference as an upper bound on the performance of the proposed method would be useful (e.g. classifier guidance trained classifier trained on noisy data, or directly fine-tuning the prior itself on a specific task). This can tell us more about how much of the potential performance the model captures.
Questions
* For the MNIST even-digit experiment, why was the reward function chosen to be $max_{i \textrm{ even}} p(c=i \mid x)$? Wouldn't $\sum_{i \textrm{ even}} p(c=i \mid x)$ be a more reasonable reward? I'm curious if this particular choice of reward makes it more likely for certain methods to mode collapse (e.g. by maximizing the probability for a single even class), and whether the results in Table 2 would change for MNIST with a different $r(x)$.
* Could you share some more details of how LoRA was used for MNIST/CIFAR setup (e.g. which weights were LoRAed)? Also I'm curious what happens if you don't use LoRA and train the posterior on full set of weights (from prior weight or from scratch), as I LoRA could act as a regularizer.
* How long is the ratio between prior vs. posterior training in terms of flops or wall-clock time? How quickly the posterior can be fine-tuned could limit the applicability of this method on certain real-world tasks.
Limitations
Yes. For concerns and potential limitations, see the comments and questions above.