Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective

The results on ImageNet 256 fall short of state-of-the-art and even popular baselines, including both the autoregressive baseline and the proposed method. The network structure follows MaskGit; however, MaskGit reports an FID of 6, which is only half of the number reported in Table 3. What accounts for this discrepancy?

We thank the reviewer for raising this point. We have further investigated the gap in FID between our trained models and previous works [1] and identified the underlying causes:

1. Sub-optimal data preprocessing.
2. Sub-optimal architecture, which also included one dimensional positional encoding as opposed to two dimensional rotary encodings.

To mitigate these differences in the revision we follow the setup of [1] for data preprocessing, architecture of both the VQVAE and the model, and training hyper-parameters, for exact details see updated Appendix E.4 in the revision. After retraining our model and baselines we are able to achieve SOTA results on tokenized ImageNet 256, see Table 3 in the revision.

Why did the authors train the VQ-VAE for ImageNet rather than using off-the-shelf pretrained weights?

We trained our own VQVAE because we used the face-blurred variant of the ImageNet dataset in our experiment. As outlined in the paper introducing this variant (https://www.image-net.org/face-obfuscation/), the face-blurred variant has enhanced privacy. Note that our VQVAE reached the same rFID on face-blurred ImageNet as reported by Sun et. al. on unblurred ImageNet (rFID=2.20). We believe the differences between the datasets are rather minor and that we have properly reproduced prior results.

[1] Sun, Peize, et al. "Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation." arXiv preprint arXiv:2406.06525 (2024).

[2] Chang, Huiwen, et al. "Maskgit: Masked generative image transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

Is the use of the concept flux in this context novel, or already described in literature? If so, please provide a citation.

The concept of flux is quite prevalent in physics and dynamical systems. In the context of discrete generative models, it was previously used in [1], where also the relation to flux in the context of continuous flows is explained in figure 2 of [1]. However, we are the first to formulate a kinetic energy optimization problem for the probability velocity as in equation 18. Considering the flux as the problem’s unknown leads to a convex problem which we are able to solve in closed form for certain cases (i.e., symmetric weights).

line 174 'abusing a bit the previous ..'

We thank the reviewer for his comment, we fixed the typo.

lines 202 - 206: it would be nice if the black triangles where vertically aligned

We thank the reviewer for his comment, we fixed the alignment of the triangles in Eq. (18).

[1] Gat, Itai, et al. "Discrete flow matching." arXiv preprint arXiv:2407.15595 (2024).

The velocity $u_t(x,z)$ is only defined for $x \ne z$ in Eq. (5). The authors need to explicitly define $u_t(x,z)$ for the case where $x=z$, or explain why this case does not need to be defined separately.

The probability velocity is first defined in equations 4 and 5 where we note that the probability velocity must satisfy two constraints called Rate Conditions:

1. For any $x \ne z$, $u_t(x,z) \ge 0$.
2. For any $z$, $\sum_x u_t(x,z)=0$.

Hence, for any probability velocity $u_t(z,z)=-\sum_{x\ne z}u_t(x,z)$. This uniquely sets $u_t(z,z)$. We have made this clear also in the revision.

What is $\bar{i}$ in Line 117? Define this when it is first introduced in the paper.

We thank the reviewer for this comment, indeed we did not define the $\bar{i}$ notation properly. We now removed this notation and provided an explicit expression instead (see equation 7 in the revision).

Line 227: Explicitly explain the relationship between the symmetric condition mentioned and the detailed balance equation in Markov Chains, if any. This would help clarify the theoretical foundations of the proposed method.

We impose this condition in order to relax the problem statement (19), specifically, this condition allows us to drop the non-negative flux condition (19c) and arrive at a closed-form solution later on. Detailed balance usually shows up in deriving stationary processes. Here, we derive velocities specifically for non-stationary processes. It looks graphically similar to detailed balance but they are not related.

The authors are requested to provide a brief comparison of the key mathematical or conceptual differences between their kinetic optimal velocity formulation and the approach used in SEDD [1].

Following this request we now added Appendix F in the revised paper, where we outline in detail the relations between our method and SEDD [1]. We agree this is important to clarify the landscape of discrete diffusion methods and thank the reviewer for this suggestion!

In particular, we expand upon: (i) the fact that our work allows using a strict superset of the probability paths possible in SEDD, introducing new and important design options for discrete diffusion/flow; (ii) the exact conversion between our velocity and SEDD score when this conversion exists; and (iii) how their loss function is related to ours. We have further updated our experiments where we re-implement SEDD for apples-to-apples comparison; please see general response for details.

[1] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).

While the kinetic energy framework (Eq. 18) allows injecting problem dependent weighting $w_t(x,z)$ in theory, solving for arbitrary weights requires numerical approximation in practice, which can be challenging as mentioned by the authors. This can limit the practical potential of the design space.

We agree; we have opened up the opportunity for many future works to explore this direction by proposing a new framework for constructing velocities given probability paths. We note that it is possible to enable a larger set of weights by using efficient numerical solvers, but we decided to focus the current study on analytical solutions.

We note that while the design space of the kinetic optimal velocity is limited, the design space of the probability path is not. Even having a single velocity formulation that can work with general discrete paths is already very useful, and converting to any other velocity for the same probability path can theoretically be done using the probability-preserving component proposed in Section 5 so the limitation is minimal.

I would be interested to know how do DFM approaches compare with FM approaches in the current state of the art. For example, with equal training time and computational resources, how does discrete flow matching compare with continuous flow matching in the task of pixel space image generation?

Following the reviewer question we train a continuous flow matching (FM)~[1] CIFAR10 model with the Cond-OT scheduler. Similar to our discrete CIFAR10 model, we use the U-Net architecture as in [2]. This results in an FID of 2.59, which outperforms our discrete metric-induced path model. We mainly focus our paper on comparison to discrete models, and we note that staying within discrete space provides additional benefits such as neural compression. In this paper we considerably extended the design space of discrete flows, and showcase that it allows us to improve discrete flows on that task of image generation. Closing the gap to continuous flows is an important future goal, and we believe our work is a significant step toward achieving it. Furthermore, note that we have now also reproduced and outperformed LlamaGen on ImageNet256, a state-of-the-art discrete model when scaled up [3].

[1] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).

[2] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in neural information processing systems 34 (2021): 8780-8794.

[3] Sun, Peize, et al. "Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation." arXiv preprint arXiv:2406.06525 (2024).

Is the construction in Section 4.2 unique, and is it general enough?

Yes, the construction in 4.2 is indeed unique: it is equivalent to the shortest path problem on the hyper-sphere. Note however, that 4.2 is constructed only for the weight choice as in equation 25. It is general in the sense it can be used to map any source distribution $p$ to any target $q$.

How does the proposed method relate to discrete rectified flow, and how do their equilibrium probability paths differ when using the same source and target domains?

We are not sure whether the reviewer's intention by “discrete rectified flow” is related to Discrete Flow Matching[1] or Rectified Flow[2,3] on a discrete target data (e.g., text), hence we will consider both interpretations.

Discrete flow matching

Discrete Flow Matching [1] focuses on conditional probability path called the mixture path, of the form

$p_t(x^i|x_1^i)=(1-\kappa_t)p(x^i) + \kappa_t\delta_{x_1^i} (x^i).$

In contrast, we considerably extend the design space of discrete flow matching to allow general conditional probability paths. For example, we introduce a metric induced probability path in equation 27,

$p_t(x^i| x_1^i) = \text{softmax}\left( -\beta_t d(x^i, x_1^i) \right),$

where $\beta_t$ is a monotonic scheduler and $d(x^i,x_1^i)$ is an arbitrary dissimilarity measure. We demonstrate the advantage of this extended family of probability paths in our Image generation experiments on CIFAR10 in Figure 2. and on ImageNet 256 in Table 3.

Another example is given in equation 2 utilizing a token dependent scheduler,

$p_t(x^i|x_1^i)=(1-\kappa_t(x_1^i))p(x^i) + \kappa_t(x_1^i)\delta_{x_1^i} (x^i),$

and later in section 4.2 we show it is kinetic optimal for a certain weight of the kinetic energy.

Rectified flow for discrete data

There exists a work called Language Rectified flow~[3], where the discrete data is embedded in a continuous space (e.g., $\mathbb{R}^D$) and a latent rectified flow model is defined in the continuous latent space. They do not define a Markov process that directly acts in discrete space.

Could you provide more background on the mixture path and explain how it connects to the claim in Section 4.2?

The mixture path is a particular choice of a conditional probability path that interpolates between source and target distributions using convex linear combinations, as introduced in section 2 equation 2,

$p_t(x^i|x_1^i)=(1-\kappa_t(x_1^i))p(x^i) + \kappa_t(x_1^i)\delta_{x_1^i} (x^i).$

Special cases of the mixture path were used in previous works [1,4,5,6,7], however no additional theoretical justification was given for it so far.

In section 4.2 we show that the mixture path minimizes the kinetic energy with the specific weighting $w_t(x,z)=\frac{1}{p_t(x)}$.

[1] Gat, Itai, et al. "Discrete flow matching." arXiv preprint arXiv:2407.15595 (2024).

[2] Liu, Xingchao, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow." arXiv preprint arXiv:2209.03003 (2022).

[3] Zhang, Shujian, et al. "Language rectified flow: Advancing diffusion language generation with probabilistic flows." arXiv preprint arXiv:2403.16995 (2024).

[4] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).

[5] Campbell, Andrew, et al. "Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design." arXiv preprint arXiv:2402.04997 (2024).

[6] Shi, Jiaxin, et al. "Simplified and Generalized Masked Diffusion for Discrete Data." arXiv preprint arXiv:2406.04329 (2024).

[7] Sahoo, Subham Sekhar, et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).