Summary
The paper proposes an approach to generative modelling for discrete data, i.e. multidimensional distributions where the variable along every dimension can take value in a finite set. This is an alternative approach to autoregressive generative modelling for discrete data which is currently actively studied for language and code generation.
The philosophy of the proposed method is heavily inspired by the Flow Matching algorithm [1] and Continuous Time Markov Chains (CTMC), which were previously used in [2] to propose a similar model. That is, for a given CTMC, the authors define the vector field generating samples from this CTMC by local updates of the samples (independently along every dimension). Based on the PMF (Probability Mass Function) of the CTMC, the authors derive a formula for the vector field. Furthermore, they introduce the continuity equation analogous to the continuous case, which allows for an easy validation that the change of the density given by CTMC corresponds to the vector field.
The authors perform an ablation study of some of their design choices and extensive empirical studies for generation of code, language, and discrete-valued images. The proposed model outperforms the competitors bridging the gap between flow-based models and autoregressive models.
[1] Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. "Flow matching for generative modeling." *arXiv preprint arXiv:2210.02747* (2022).
[2] Campbell, Andrew, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. "Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design." arXiv preprint arXiv:2402.04997 (2024).
Strengths
The paper presents a complete study of the actively studied topic in the field. It is clearly written, the developments presented in the paper are novel and are properly studied empirically. The paper is of great interest to the NeurIPS community.
Weaknesses
The paper presents a complete study and its methodological part does not raise any major concerns. However, I would like to ask the authors to clarify the following question.
Why do posteriors in eq. 15 define a proper distribution? Indeed, according to eq. 8, the linear combination of the conditional distributions with the schedulers should define a correct distribution (sum up to $1$). If we simply try to sum over all possible $x^i$ in eq. 15, we won’t get $1$, hence eq. 15 does not define a posterior. This is an important detail when defining the objective in eq. 25, which is the cross entropy between distributions.
The minor
- Descriptions of the schedulers on lines 96 and 98 are not complete, i.e. there are no conditions that schedulers have to define correct probability distributions in eqs. (9,10) and the scheduler $\kappa^3$ is missing a description.
- Figure 2. The plots presented in the figure have $d = 2$ unlike what is stated in the caption.
- At the beginning of the paper, the authors introduce the notation of random variables as capital letters (e.g. $X_t$). This creates confusion in equations 13-17 because there the authors clearly mean $X_t$ to be the value of a random variable.
- There is a typo in line 160 when describing the panel of Fig. 2.
- There is a typo in line 163 when defining the vector field $v(x,z)$.
- There is a type in line 582. Index $\ell$ is not a function of $j$ according to eq. (36).
- I haven’t checked thoroughly, but I think equations on the top of page 19 should have a summation over $j\neq \ell$ instead of the summation over all possible $j$.
Questions
I would suggest adding a discussion of the variable-length generation. It is an important difference with autoregressive modelling and I’m wondering how the authors handle this issue given that their model has two important properties:
1. Conditional generation for partially masked sequences.
2. Independence of the conditional distributions between dimensions.
The questions I would like to have answered are:
1. Does one have to define the length of the generated sequence before generation? If so how this can be decided?
2. What’s the computational cost if one continues generating the sequence in the autoregressive way?
Limitations
NeurIPS Paper Checklist does not follow the required format.
The paper adequately discusses the limitations of the proposed approach. The only thing that is not cover enough in the paper is the variable-length generation.