Designing a Conditional Prior Distribution for Flow-Based Generative Models

1. The proposed solution appears overly simplistic and may be ineffective for text-to-image generation. As acknowledged in the paper, text space is continuous, making it questionable to model the prior distribution with a GMM. This issue likely explains why, although the method improves CLIP Score on ImageNet, it underperforms compared to baseline Stable Diffusion on MS-COCO, as shown in Figure 5. Even it is possible to decrease FID while increasing the CLIP Score as shown in Table 2, the performance gain is too minimal.
2. Moreover, MS-COCO is too simple a benchmark for evaluating text-to-image performance, given its low text diversity. Performance evaluation on more challenging benchmarks, such as PickScore or DrawBench datasets, would strengthen the claim.
3. Even with the use of a pretrained auto-encoder to convert images to latent space, it’s challenging to justify that the transformed latents adhere to a GMM prior. While the distance reduction relative to BatchOT and CondOT is noted, the GMM prior itself remains unconvincing.
4. Minor issue: According to ICLR formatting guidelines, table captions should be placed above the tables.

The main concern is that a GMM-based prior may be too simplistic to represent text-to-image space effectively. Demonstrating that FID and KID scores improve while maintaining or increasing text-image alignment on more complex benchmarks would improve the evaluation. I will increase my score if my concerns are addressed.

• The reviewer is unfortunately a little skeptical regarding the extent of the novelty of the research, as its novelty winds down to, if the reviewer is not mistaken, the proposition of choosing Gaussian Mixture Distribution as the source distribution. The reviewer would like to note that the generative scheme proposed by Atong et al (Conditional Flow Matching) can be applied to "the generation of samples from $c$-conditioned distribution $\mu^c$ by training a "c-parametrized" vector field $v(t, x | c)$ through the loss $E_{c, x_1^c, x_0^c} [ \| v(t, \psi_c(t) | c) - (x_1^c - x_0^c) \|^2]$ with $x_1^c \sim \mu^c_1$, $x_0^c \sim \mu_0^c, \psi_c(t) = t x_1^c - (1-t) x_0^c$, and that it is customary to choose $\mu_0^c = N(E[x|c], Var[x|c])$ in practice, instead of a non-informative, c-independent prior. At the time of the generation, one may just integrate the ODE $\dot{x}^c(t) = v(t, x^c(t) | c), ~~~~ x^c(0) \sim \mu^c_0$ forward in time. This scheme is mathematically consistent as well, because it simply amounts to simultaneously solving $c$-parameterized set of continuity equations for $\{(\mu^c_t, v(t, x_t | c) ) \}.$ Such a scheme appears, for example, in Isobe et al (Extended Flow matching) as well, and the reviewer feels that this has been explored elsewhere as well. Note that this scheme only differs from the proposed method only in that the latter chooses $\mu_0^c= \mu_0 = Mixture({N(E[x|c_i], Var[x|c_i]})$.
  In such a scheme, it is also critical that the regressors $m: c \to E[x|c]$, $v: c \to Var[x|c]$ to be used in the way of $\mu_0^c= N(m(c), v(c))$ is trained so that they can well inter/extrapolate $E[x|c], Var[x|c]$ for $c$ not in training (corresponds to (21) in this research). The reviewer acknowledges that the paper partially mentions this matter, but the reviewer also believes that it is particularly a nontrivial problem when $c$ is "not" dense everywhere. That being said, the reviewer would like to know if the followings have been considered:

  • How does the current GMM construction of the prior compare against the usage of c-parametrized prior $\mu_0^c = N(E[x|c], Var[x|c])$ for the generation of $\mu^c$ with $c$-parametrized vector field $v(t, x | c)$? Is there any merit in choosing common $\mu_0 = Mixture({N(E[x|c_i], Var[x|c_i]})$ for the generation of $\mu^c$ for each $c$ ?
    • if the choosing of common $\mu_0$ is essential, what would be the mechanism behind it? What are the situations in which the choosing of common $\mu_0$ would be beneficial?
  • How does the current regression scheme ($P_\theta$) fares in various situations, such as with the presence of
    • condition c that lies far from the bulk of the conditions (outlying c)
    • condition c with small number of samples for $\mu^c$ (rare condition)?
  • The Design of Source distribution as the one the reviewer presented above as well as the design of GMM distribution presented in this paper seem to come from the intuition that Flow matching performs better when the Wasserstein distance between the source distribution and the target distribution is smaller. Has any theoretical investigation been done in this direction? How would the choice of isotropic Gaussian in GMM empirically fare against non-isotropic counterpart when each conditional distribution is highly non gaussian?

• While the reviewer very much values the amount of experiments done, for the reason the reviewer outlayed above, the reviewer feels that more theoretical / more ablative studies are required to substantiate the contribution.

Weaknesses

1. One of the key ideas of the paper is that shorter average distances yield straighter trajectories (e.g. L273), which is already discussed previously in the work of Pooladian et al. 2023 [w-a]. However, none of the toy examples (Fig 2, 3, 4) show the trajectories taken by the proposed approach. I suggest presenting some trajectory diagrams for these toy experiments (for instance, like the trajectories shown in the minibatch OT paper [w-b]).
2. One of the implicit assumptions of the paper is that in a class conditional dataset, each class is a mode, and is sufficiently disentangled from the other classes. Additionally, there is also an assumption of homogeneity in the data of each class, ignoring the possibility of internal modes in a class. It would be useful for the paper if analysis was provided on different types of conditional data distributions.
   • For example, consider a dataset from the Datasaurus Dozen, like VLines. Suppose line 1 and 3 (the odd lines) in VLines are one class (A), while the even 2 and 4 are another class (B). The mean of class A falls on samples from class B, and vice versa. Showing that the proposed approach still works better in this case than a standard normal would add to the strength of the paper. Or if not, it would be useful to know what properties of the target dataset are ill-suited for the proposed approach.
3. Rather than the irrelevant DDPM, the comparison should include RectifiedFlow, (e.g. 2-RF) [w-c] that is well established to have straight paths and low NFEs.
4. An important concern I have is whether classifier-free guidance can still be applied on a flow trained with a disentangled source. Does it work out of the box? Or does some adjustment have to be made? (Such as all classes sampling from a common N(0, I) with some probability p.)
   • How does the proposed method compare with existing approaches when CFG is applied? Since many state-of-the-art results with flow models are achieved by applying guidance.

I am highly amenable to improving my score if my concerns are addressed.

[w-a] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. 2023.

[w-b] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport, 2023.

[w-c] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. 2022.

• The novelty is rather limited, since the proposed conditional prior is assumed to be a Gaussian. As a result, it is not clear if the proposed method can be effective for more complex image categories, or as a general conditional PDF estimator for other modalities. (i.e., it is easy to construct a data distribution where the class conditional Gaussians are identical for multiple classes). Can you discuss potential extensions to handling more complex conditional priors, or add an experiment showing the performance on a case where priors overlap?
• The training details are missing form the paper, and so do the details of the pre-trained models that were used (i.e., VQ-VAE).