Gaussian Mixture Flow Matching Models

Thank you for your thoughtful review. We have uploaded a revised manuscript and essential code in this anonymous link (full code will be released upon publication):

https://anonymous.4open.science/r/anonymous_gmflow-63FE
backup: https://limewire.com/d/CgAn9#jkBxDmC3qh

Weaknesses and Questions

1. Comparison with GMS.

We’ve made a direct comparison with GMS and reported the results on unconditional CIFAR-10 image generation in Appendix B.3 of the revised manuscript. We train GMFlow with $K=2$ from scratch using the same U-Net backbone as GMS. We choose CIFAR-10 because re-training an ImageNet model using the backbone of GMS is expensive and time-consuming.

Here are the FID results. GMFlow significantly outperforms GMS and other moment-matching methods in few-step sampling. The competitor results are from the original GMS paper.

2. Computational complexity of KR transport

We would like to point out that KR transport is highly efficient since it's basically 1D CDF mappings. In the paper, we stated that GMFlow incurs only 0.005 sec of overhead per step, which is minimal compared to the total inference time of 0.39 sec per step (most of which is spent on DiT). This includes the KR transport. The computation complexity of per-pixel KR transport is $O(K\cdot C^2)$ ($C$ is the channel size). We have added more details on spectral sampling in Figure 9 of the revised manuscript.

Section 3.2 needs to be expanded

Thank you for the suggestion. In the revised manuscript, we have fully rewritten and expanded Section 3.2 to make it more clear. We kindly refer the reviewer to the revised manuscript for more details.

3. Why construct a Gaussian mask using two surrogate Gaussians?

The reasons are twofold:

• The division of a GM PDF by another GM PDF doesn't have a closed-form solution, and would require slower numerical approximations.
• If a raw GM PDF is put into the denominator, it's likely to be unstable since the denominator can be very small in some regions.

We have already experimented with full GM formulations using adaptive importance sampling, and they are slow and unstable. In contrast, using Gaussians is simple and performs well in our experiments.

4. Intuition behind this design of guided Gaussian.

In the revised manuscript (Line 183–196), we have added two paragraphs explaining the intuition behind the design.

5. Why do you scale the shift $\Delta \mu_n$ by $s_c$?

This is to satisfy bias–variance decomposition, such that when $\tilde{w} = 1$, all the energy in the variance $s_c$ is converted into the bias $s_c \Delta \mu_n$ , making the hyper-parameter $\tilde{w}$ more meaningful. Please refer to the revised manuscript (Line 183–189) for more details.

6. How do you obtain $\mu_c$ and $\mu_u$?

In the revised manuscript (Line 174–175 and Appendix A.1), we added that we approximate the conditional and unconditional GM predictions as isotropic Gaussian surrogates by matching the mean and total variance of the GM. This is equivalent to minimizing the KL divergence.

7. Orthogonal projection trick

We apply this to $\mu_\text{c} - \mu_\text{u}$, prior to normalization.

8. Please provide the analytical solution of Equation (8).

Your idea is correct. The appendix has discussed conflation of two Gaussians with powers. A Gaussian PDF divided by a Gaussian PDF is basically multiplication with power 1 and -1, and the result is still a Gaussian. The code is simply implemented as

gm_output = gm_mul_gaussian(gm_cond, gaussian_mul_gaussian(gaussian_guided, gaussian_cond, 1, -1))

where all the operations are analytical.

Additional Questions

2. and 5. (Typos)

Thank you for pointing these out! These typos been fixed in the revised manuscript.

3. Please explain how you apply DDPM sampling scheme to flow models in details.

Added in the revised manuscript (Appendix A.5).

4. Is $\mathcal{L}_\text{spec}$ used for training even if GM-ODE will be used for sampling?

Yes, we use the same model for both ODE and SDE sampling. The spectrum MLP is completely isolated and does not back-propagate to the main model. Therefore, the loss weights do not matter.

6. Is there any typo in the definition of ${\mu_s}_k$?

We confirmed that it's correct. This is the same as converting the flow velocity to the score function, i.e., $s_t(x_t) = -\frac{1}{\sigma_t} ( x_t + \alpha_t u)$.

Is $s$ shared across all spatial locations? What does "the mean of per-pixel GM variances" refer to?

Yes, $s$ is shared. The GM variance refers to the GM’s total variance divided by $D$, which is also dependent on $\{\mu_k\}$. We have added more details in the revised manuscript (Fig. 9).

Thank you for your thoughtful review. We have uploaded a revised manuscript and essential code in this anonymous link (full code will be released upon publication):

https://anonymous.4open.science/r/anonymous_gmflow-63FE
backup: https://limewire.com/d/CgAn9#jkBxDmC3qh

The GM idea is adopted in [1]

Please note that [1] is essentially a moment matching method, which only converts its moment predictions into a bimodal GM during inference. This is fundamentally different from our GMFlow formulation, which directly learns GM parameters. In comparison, [1] employs three L2 losses for three moments, whereas we use a single loss for all Gaussian components.

Moreover, our formulation can generalize to more GM components, and we further propose few-step ODE sampling with analytical substeps and probabilistic guidance.

Theoretical analysis of using GM

Gaussian mixtures are well known to yield better approximation with more components [2, 4]. In diffusion SDE sampling, it is also widely recognized that improving the accuracy of the reverse transition distribution can greatly improve few-step sampling. Taken together, these observations already support the claim that using more Gaussian components should lead to reduced SDE sampling errors.

The practical performance ultimately depends on how the network learns the GM parameters, a process that is difficult to analyze theoretically. Therefore, we believe that our empirical validation through experiments is more important here.

• [2] Huix et al. Theoretical guarantees for variational inference with fixed-variance mixture of gaussians.
• [4] Bishop. Mixture Density Networks

Minimum of the cross-entropy training loss

Cross-entropy equals the sum of the data entropy and the KL divergence. Although the data entropy is intractable, it is solely determined by the data and independent of the model, thus it can be ignored in the loss function. Note that Eq. (6) omits all irrelevant constant terms.

Our goal is to minimize the KL divergence between the predicted GM and the ground truth distribution, which has a minimum of 0 when the GM perfectly fits the ground truth.

Flow matching loss in [3] achieves minimum 0 and reduces the error variance

In general, the flow matching loss does not have a tractable form. While [3] analyzes special cases with analytic velocity fields, our work focuses on practical applications.

Theoretically, the "error variance" of stochastic flow matching loss equals the sum of denoising distribution variance and the squared error of the velocity. The former is determined by the data, whereas the latter can be reduced to near 0 with enough capacity.

Different from standard flow matching, our goal is not just to learn accurate local velocity, but also to capture the underlying denoising distribution so that a global velocity field can be analytically derived for multi-substep GM-ODE sampling.

Spectral sampling for image generation is a detour

We want to clarify that spectral sampling is just an optional extension to standard diffusion sampling with minimal overhead.

• Optional: It's not used in our ODE solvers. In our SDE solvers, it's an optional trick to improve FID, whereas Precision is not affected, as shown in our ablation studies (Table 4).
• Extension: When setting $K=1$ and using a uniform spectrum, spectral sampling is equivalent to standard diffusion sampling. Therefore, we consider it a more general approach instead of a detour.
• Minimal overhead: The total overhead of GMFlow is less then 2% in inference time.

Choosing suitable K

With pixel-wise factorization, multimodality is indeed less apparent. From a practice point of view, we choose the best $K$ based on the evaluation metrics of interest. We also found that early-step training NLL can serve as a potential indicator.

Why the best Precision drops from K=8 to K=16

We believe this is due to numerical errors in spectral sampling, which grow larger with increasing $K$. We have tested the case without spectral sampling, where both $K=8$ and $K=16$ yield the same Precision of 0.946, with NFE=8.

Results of K=2

We have added an experiment using $K=2$ on the CIFAR-10 dataset (training on ImageNet is beyond the time frame of rebuttal), and our few-step results are much better than GMS. Please refer to our response to Reviewer A5dP.

Inference time

Sorry for the confusion. We have rewritten this paragraph for improved clarity: GMFlow adds only 0.005 sec of overhead per step (batch size 125, A100 GPU) compared to its flow-matching counterpart, which is minimal compared to the total inference time of 0.39 sec per step---most of which is spent on DiT.

Scale the experiment to large dimensions and preserving spatial correlations

This will be one of our future directions. We have recently conducted preliminary experiments using patch-wise factorization instead of pixel-wise factorization, which show promising results.

Thank you for your thoughtful review. We have uploaded a revised manuscript and essential code in this anonymous link (full code will be released upon publication):

https://anonymous.4open.science/r/anonymous_gmflow-63FE
backup: https://limewire.com/d/CgAn9#jkBxDmC3qh

Pre-defined number of mixtures can limit the types of distributions. Using an excessive number of mixtures leads to high costs and requires a large model

With pixel-wise factorization, a small number of Gaussians is enough to capture the distribution of a single latent pixel. The experiments also reveal that image generation metrics generally saturate at $K=8$.

On the other hand, GMFlow only expands the output channel of the final layer to accommodate the mixture components. With $K=8$, currently, the final layer has an output channel size of $(4 + 1) \times K = 40$, where 4 is the latent channel size. So the computational cost of increasing K is still minimal when compared to that of the remaining layers in DiT, especially considering the hidden dimensions of DiT are usually much larger.

typos

Thank you for pointing this out. We have fixed the typos in the revised manuscript.

Evaluating the proposed method across different model sizes and image sizes

We have added an experiment on unconditional CIFAR-10 image generation (please refer to our response to Reviewer A5dP), which employs different architecture (U-Net), model size (53M), and image resolution (32x32) from the DiT used in the paper. Unfortunately, training 512-res GMFlow and baseline methods are too computationally expensive and time-consuming, and is beyond the time frame of the rebuttal.

IS metric

Following previous work [3, 4], we employ Precision as the quality measurement instead of IS. This is because IS metric does not capture the image quality under a large guidance scale---in particular, IS is not sensitive to over-saturation. In fact, both our GMFlow model and flow matching baseline reach IS > 500 under the highest guidance scale evaluated, despite that the baseline has clearly worse visual quality.

• [3] Sadat et al. Eliminating Oversaturation and Artifacts of High Guidance Scales in Diffusion Models
• [4] Kynkäänniemi et al. Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models

Discussion of GIVT

Thank you for suggesting this interesting paper. In the revised manuscript, we have added a paragraph discussing GIVT and other generative models using GMs.

Thank you for your thoughtful review. We have uploaded a revised manuscript and essential code in this anonymous link (full code will be released upon publication):

https://anonymous.4open.science/r/anonymous_gmflow-63FE
backup: https://limewire.com/d/CgAn9#jkBxDmC3qh

Boundary of experiment in Table 2

Unfortunately due to limited computing resources, we were not able to conduct experiments for K values beyond 16. That being said, as shown in Fig. 5 (or Fig. 7 in the revised manuscript). we have shown that the FID and Precision metrics generally saturate when $K\ge 8$. The NLL values in Table 2 do not directly reflect image generation quality, and they are more relevant to some downstream applications (e.g., score distillation-like test-time optimization), which is not the main focus of this work. We are happy to add more discussions in the final draft.

Discussion on the combination of GM with other generative models such as VAE

Thank you for the suggestion. In the revised manuscript, we have added a paragraph discussing GM GANs and autoregressive Transformers. In the context of modern generative AI, however, VAEs alone are often regarded not as generative models but rather as representation compressors. Some GM VAE papers are focused on clustering instead of generation [1, 2]. We are happy to add more related works in the final draft.

• [1] Dilokthanakul et al. Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders.
• [2] Jiang et al. Variational Deep Embedding: AnUnsupervised and Generative Approach to Clustering

FID in Table 3

In Table 3, we aim to evaluate the saturation of different methods at their best precision, which typically requires large CFG values that are beyond reasonable ranges for FID comparison. The main comparisons on both FID and Precision are presented in Fig. 4 (or Fig. 6 in the revised manuscript).

In Line 70, $x_T \approx \epsilon$ does not take into account VE diffusion

Thank you for pointing this out. The manuscript states "A typical diffusion model", which does not cover all cases. In practice, mainstream diffusion models are trained with VP schedules, and they can be rescaled into VE diffusions during sampling (which is how the popular EDM Euler solver is implemented).

Will the results be better if $sI$ is replaced by a diagonal matrix?

We have tried predicting pixel-wise variances instead of the global $s$ in image generation. The cons outweigh the pros: it makes training less stable since it's more likely to have small variances in the denominator of the loss function, and the benefits for NLL can be equally achieved by increasing $K$.

Does pixel-wise factorization mean that each latent pixel is treated as a Gaussian component? How does this relate to K?

Yes. Each latent pixel (4 channels) is a 4-D GM of $K$ components. Please refer to the network architecture in Fig. 3 (or Fig. 4 in the revised manuscript).

Does pixel-wise factorization limit the method to a fixed resolution setting?

We don’t think so. We think the opposite might be true: without pixel-wise factorization, the entire latent grid would be treated as a high-dimensional GM, which would complicate diverse resolution generation due to varying dimensions of GM. We think pixel-wise factorization makes diverse resolution generation easier since the per-pixel dimensions are fixed.