Discrete Flow Matching

Question: In line 32, you mention one advantage of FM is its flexibility in handling non-Gaussian target distributions. Have you demonstrated this case? Yes, in fact all the probability paths we use in this paper are non-Gaussian. We worked with mask source distribution which corresponds to a delta function concentrated on a special “mask” token.

Comment: The method's description is unclear and confusing. I cannot distinguish which part pertains to Campbell's methods and which part is yours. Thank you for the comment. We build upon Campbell’s work [1] but generalize it in ways that allow significant improvement in performance. In particular, we offer the following contributions: (i) we consider arbitrary data coupling $(X_0,X_1)$ and use it for conditioning; (ii) we offer a novel family of probability paths (equation 8) that generalizes the paths used in Campbell as particular cases; (iii) in particular, we show that incorporating polynomial schedulers $\kappa_t$ considerably improves performance; (iv) we provide a unified and closed-form formula for marginal probability velocity (rate) in equations 14, 16, 17 and show it has the exact same form as in continuous Flow Matching, see Table 1. Campbell provided this rate as an expectation and resorted to compute it individually for the masked and uniform noise cases; (v) we develop a general yet closed form formula for corrector sampling with arbitrary schedulers (equation 23). This generalizes Campbell’s stochastic sampling constant $\eta$ ($\alpha_t = 1 + t\eta$ and $\beta_t = \alpha_t - 1$), and we also note that Campbell’s stochastic sampling (Proposition 3.3 and equation 9 in [1]) incorporates the detailed balanced matrix in an implicit way and therefore requires a particular solution in each case; (vi) we show that particular polynomial correctors schedulers provide a further significant boost in results.

[1] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. "Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design."

Question: Why are the results from Austin et al. [2021a] poor? We use the Pytroch implementation from https://github.com/cloneofsimo/d3pm, with the same architecture (DiT), tokenizer (GPT2), and data (Open Web Text) as we use for our model. We will add these details to the revised version.

Comment: In Line 214, I'm not entirely sure what the difference is. As mentioned above (see (v),(vi) in the above answer) compared to [1] we provide more general (arbitrary schedulers) and closed form corrector steps. The discrete and continuous diffusion works [2, 3] define only corrector iterations and not sampling (i.e., do not progress in time, similar to case (ii) described in line 280 in our submission) and perform corrector iterations by incorporating equal forward and reverse rates for the diffusion probability paths they consider. In our case we develop a closed form corrector step incorporating arbitrary schedulers (see $\alpha_t,\beta_t$ in equation (23)) that includes both corrector sampling and corrector iterations as special cases.

[2] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. "A continuous time framework for discrete denoising models."

[3] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. “Score-based generative modeling through stochastic differential equations.”

Comment: In Figure 3, the Inception Score (IS) needs to be included to demonstrate the diversity of image generation. Per the reviewer’s request, we added a corresponding inception score graph in the PDF attached to the main response of rebuttal. Note that we observed a similar trend to the FID demonstrated in the original submission.

Question: Except the perplexity, how about the performance on BLEU, bertscore? BLEU and BERT-score are traditionally used for the “data-data” case (e.g., translation), that is when both source and target samples are data, and then generated samples are compared to target (test) data. We use “noise-data” (except the conditioning part) and therefore we couldn't find a reasonable way to apply these metrics in our case.

Comment: Diversity is a significant concern for flow-based models. What about the result on diversity-related metrics? We are not aware that flow-based models have a diversity concern in general, maybe the reviewer refers to distilled flow models? In any case, we didn’t encounter particular diversity issues with our model, and in Tables 2, 3, and Figure 7, with additional details provided in Appendix F.1, we present the entropy of tokens within generated sequences to illustrate the diversity of the model's predictions. Furthermore, our modeling allows for control over the diversity by adjusting the temperature during sampling from $p_{1|t}$.

To further address the reviewer’s concerns, we attach in the rebuttal’s main reply several uncurated generations, sampled with the same prompt. These uncurated generations demonstrate the diversity of our model's output.

Comment: In F.1, in the context of conditional generation, it's unclear what the source (src) and target (tgt) samples are. Are they conditional prompts, target prompts, or a combination of Gaussian noise and target prompts? In the context of conditional generation (Appendix F.1), the source-target pairs $(X_0,X_1)$ are as described in Equation 5 in the paper, i.e., $(X_0,X_1) = (\mathbb{I} \odot X_1 + (\mathbf{1}-\mathbb{I})\odot (\mathbb{m},\ldots,\mathbb{m}) , X_1)$, where $\mathbb{I} \in \set{0,1}^{N}$ is indicating the conditioning mask. There are no target prompts, nor a combination of Gaussian noise/target prompts. We did our best to understand the reviewer’s question, if we got it wrong, please clarify.

Comment: In section H, there is no qualitative analysis of unconditional generation. Per the reviewer's suggestion, we added unconditional samples, generated by our model (attached in the pdf in the main response). We will add these samples to the revised version.

Commnet: In Section H, I did not see the color red. I only observed the background colors: gray and yellow. This is a mistake, prompts are marked in gray, we will fix it in the revised version.

Comment: It will be great that the code can be provided. We are planning on releasing the code not much after the paper is published.

Comment: Related works are missing. We thank the reviewer for the suggestion, we will add the relevant related works to the revised version.

Comment: The paper lacks some experiments on common benchmarks used in existing discrete diffusion models, e.g., LM1B and OWT. First, please note that our experimental setup already includes the OWT dataset (see Table 2 and lines 268-273). Second, per the reviewer’s suggestion, we trained and evaluated our and baselines models on the LM1B dataset, please see Table 1 in the main rebuttal reply.

Experimental setup. For the reviewer’s convenience, we would like to detail all the evaluations done in this work. We experiment with two data modalities, text and images on small and large scale datasets.

Comment: I didn't find limitations specifically discussed in the paper. In the conclusions (Section 5), we mention the following limitations: (1) Discrete flow matching requires a higher number of function evaluations compared to its (deterministic) continuous counterpart. We attribute this to its stochastic sampling, similar to sampling by approximating a solution of an SDE which typically possesses a lower convergence order than their ODE solvers counterparts (see e.g., [1]). (2) There remains a performance gap between autoregressive modeling and our proposed approach.

[1] Sauer, T., 2011. Numerical solution of stochastic differential equations in finance. In Handbook of computational finance (pp. 529-550). Berlin, Heidelberg: Springer Berlin Heidelberg.

Question: Why do posteriors in Equation 15 define a proper distribution? Equation 15 is a proper distribution as follows: $\sum_{x^i} \hat{w}^j_t(x^i|X_t) = \sum_{x_0,x_1} \overbrace{\big( \sum_{x^i} w^j(x^i|x_0,x_1)\big )}^{=1} p_t(x_0,x_1|X_t) = \sum_{x_0,x_1} p_t(x_0,x_1|X_t) = 1,$ where in the first equality we change the summation order.

Comment: Descriptions of the schedulers on lines 96 and 98 are not complete. The conditions in lines 96 and 98 should be understood as additional conditions to the general conditions presented in Line 91, i.e., $\sum_{j} \kappa_{t}^{i,j}=1$ and $\kappa_{t}^{i,j} \ge 0$. Combining the conditions in Line 91 with the ones in Lines 96 and 98 guarantees proper distributions. We realize now this is confusing and will clarify this in the revised version of the paper - thank you.

Comment: Figure 2. The plots presented in the figure have $d=2$ unlike what is stated in the caption. This is indeed a typo, but should be $d=4$ (and $N=2$), that is, the figure depicts the state space of two tokens $x=(x^1,x^2)$ where each token can get a value in a vocabulary of size $d=4$.

Comment: Random variable notation creates confusion. We agree with the reviewer that this notation is confusing and we will use a lower case letter, e.g., $z$, instead in the revised paper.

Comments: Typos Lines 160, 163, and 582. Thanks, we will fix it in the revised version.

Comment: Summation on page 19. Please note that in the case of $j=\ell$ the term in the summation equals zero, thus, summation over $j \ne \ell$ will yield the same result.

Question: Does one have to define the length of the generated sequence before generation? If so how this can be decided? No, the length does not have to be defined before generation. We train the model using flattened data samples which are separated by an end-of-text (EOT) token. The model can predict EOT token in any location in the generated sequence $x^1,x^2,\ldots,x^N$ and the EOT will indicate the end of the generated text (similar to autoregressive modeling).

Question: What is the computational cost if one continues generating the sequence in the autoregressive way? The model generates a sequence of length $\leq N$, where $<N$ length can happen if the EOT token is predicted. If it does not (some texts are longer than $N$ tokens) one can potentially continue to generate by conditioning on the last $K<N$ tokens and predict the next $N-K$ tokens, where $K$ is user defined: e.g., large $K$ will provide more context but shorter extension. The cost of predicting these $N-K$ new tokens will be equivalent to a full sequence generation with our model. It is very interesting to develop methods to accelerate the generation, but we leave it to future research.

Comment: I am a bit confused with Eq 10 especially the scheduler terms. Equation 10 proposes a second instantiation of a conditional probability path $p_t(\cdot|x_0,x_1)$ where given some pair $(x^i_0,x^i_1)$ of source and target tokens, the token at time $t$ is: $x^i_1$ with probability $\kappa^1_t$; $x^i_0$ with probability $\kappa^3_t$; and a uniformly distributed random token with probability $\kappa^2_t$. This probability path is reminiscent of the brownian bridge in the continuous diffusion world. Lastly, note that the conditions on the schedulers $\kappa^j_t$ that guarantee this conditional path interpolates between $\delta_{x_0}(x^i)$ and $\delta_{x_1}(x^i)$ and stays a proper PMF for all $t\in[0,1]$ are the ones described in lines 98 and 91. We will clarify this in the revised version.

Question: What are the main reasons for discrete flow matching to require a significantly large number of evaluations? Discrete flow matching requires a higher number of function evaluations compared to its (deterministic) continuous counterpart. We attribute this to its stochastic sampling, similar to sampling by approximating a solution of an SDE which typically possesses a lower strong convergence order than their ODE solvers counterparts. In particular, for the Euler sampling case the deterministic Euler method has strong convergence order of $1$ while the non-deterministic Euler (Euler-Maruyama) has strong convergence order of $\frac{1}{2}$, see e.g., [1]. While this gives some intuition we do agree an analysis of global convergence of discrete sampling is interesting and defer it to future work; in this paper we only show local convergence error (see $o(h)=O(h^2)$ term in equation (22)).

[1] Sauer, T., 2011. Numerical solution of stochastic differential equations in finance. In Handbook of computational finance (pp. 529-550). Berlin, Heidelberg: Springer Berlin Heidelberg.