Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: The paper proposes two new generative models for non-ordinal discrete random variables. 

First, the paper proposes Argmax Flows, which apply argmax operation on continuous random variables, modeled by flow-based models. Note that estimating the likelihood of a model with a surjection requires a valid stochastic inverse, which should satisfy right-inverse (argmax constraint in the paper) and absolute continuity. For stochastic inverses of the argmax operation, the authors propose three novel methods: (1) thresholding, (2) Gumbel, and (3) Gumbel thresholding. The proposed stochastic inverses and the Argmax Flows are jointly trained by maximizing the ELBO.

Second, the authors propose a diffusion-based generative model for non-ordinal discrete data. Consider a pre-defined Markov chain, by which a data distribution is transformed to a prior distribution. The model learns a reverse Markov chain from the prior to the data. In this context, the paper proposes to diffuse non-ordinal discrete data to a uniform categorical distribution. For each transition, with a small probability, every value is independently decided to be either resampled or remaining the same; when we resample, a new value will be drawn from a uniform categorical distribution. For the reverse Markov chain, the authors acknowledge that a closed-form (conditional) reverse transition exists for the proposed diffusion, and they follow categorical distributions. Based on that, the paper proposes an efficient parameterization of the reverse transitions, which reduces the variance of the ELBO estimation.

In the experiments, the paper demonstrates that the proposed methods improve modeling language and segmentation maps (of cityscapes images) compared to previous generative models designed for ordinal data.
main_review: (Originality)  
The originality of the paper is clear:  
First, the paper proposes novel dequantization methods for non-ordinate discrete data, which are stochastic inverses of the argmax operation—noting that it is nontrivial to satisfy the necessary conditions to design valid stochastic inverses of a given surjection. The authors discuss how to achieve the constraints for a given surjection (aka argmax) to design its probabilistic inverse.  

Second, the paper also introduces a diffusion-based model for non-ordinate discrete data.  

Finally, the paper demonstrates the practicalities of the proposed methods via language modeling experiments and density estimation tasks on segmentation maps.  

(Significance & Quality)  
In general, the results of the paper are important. However, the followings aspects of the paper can be improved.

First, while one of the main contributions is the novel stochastic inverses of argmax operation, the discussions about their properties are limited. Aside from the proposed inverses satisfy the necessary constraints, I consider it important to discuss the expressivity of the proposed methods or potential failure cases. For instance, due to the truncations in the variational distributions, the proposed methods may not be optimal when prior $p(u_x) > 0$ for $u_x > T, x \\in \\{1, \dots, K\\}$ (considering $T > 0$). What can we do instead? In other instances, it would be helpful to provide a simple example of how the proposed models represent dequantizations for uniform categorial distribution or other toy cases.

Second, the connection between the Argmax Flows and diffusion-based models is limited. I believe that they are both important contributions. Based on the current version, however, I found that they are somewhat different topics. I believe that clarifying their connection improves the significance of the paper.

(Clarity)  
The paper is self-contained and has a well-organized structure to present relevant topics. In particular, the paper concisely raises the difficulty of previous methods regarding modeling non-ordinal data and explains how the proposed models address the issue.
limitations_and_societal_impact: (Limitations)  
Please find comments on (Significance & Quality) section in the main review. 

(Social Impact)  
N/A
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes two methods for generative modeling of categorical data, one based on normalizing flows and the other on diffusions. The first method models observed data as the result of a continuous normalizing flow, followed by an argmax operation rather than a rounding one sometimes used for ordinal data. The model is then trained by introducing a carefully constructed approximate posterior (of the continuous variable given the categorical observation), and maximizing the ELBO. The second method is a fairly straightforward extension of commonly used diffusions to the categorical setting.
main_review: Overall I quite liked the paper, it's very clearly written, well motivated, and showed good results. I found the posterior constructions from section 3.1 ensuring that the approximate posterior has the correct support to be elegant solutions; and the proposed argmax flows do beat competing dequantization schemes empirically, as well as VAE-based approaches. While argmax flows are still outperformed by purely autoregressive approaches, I think this is OK, and is actually a similar behaviour as what is observed with images.

While I found showing that categorical diffusion models can get reasonable performance interesting (Figure 5 is nice), the results for diffusions are weaker than for argmax flows, and the paper reads like two papers pasted together. Aside from the application, there really isn't that much in common between the proposed approaches and although the paper remains clear, I think this weakens the narrative of the paper.

Minor notational issue: It looks from the unnumbered equation between equations 1 and 2 that "P" is reserved for probabilities, and "p" for continuous densities, but for diffusion models "P(x_0)" is used, but so are "p(x_{t-1}|x_t)" and "p(x_T)".

======================================================================================================

EDIT 1 AFTER REBUTTAL

I have read the other reviews as well as the author's rebuttal. Although I still think the two proposed methods are mostly disconnected, and find the contributions for argmax flows more exciting and novel than the ones for multinomial diffusions, I also still think the proposed argmax flow model is an elegant approach with good empirical results deserving publication, and thus have kept my original score.

======================================================================================================
limitations_and_societal_impact: Limitations and potential negative societal impacts are adequately discussed.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This paper proposes two distinct methods for learning categorical distributions. The first is based on an extension to variational dequantization with argmax operators which allows normalizing flow models to learn categorical distributions. The second is based on an extension to Bernoulli diffusion probabilistic models so that it can handle larger than 2 categories. Proof-of-concept experiments demonstrate that the proposed methods can successfully handle high dimensional categorical distributions.
main_review: **Significance**: Normalizing flow models and diffusion models are mostly designed for continuous distributions. Extending them to the realm of  discrete distributions has important practical values and is highly relevant to the NeurIPS community. This paper proposes new methods to achieve this.

**Clarity**: The submission is clearly written and very easy to follow.

**Quality and originality**: Here is a list of concerns:
1. Argmax flows and multinomial diffusions are two completely different model families. The ways of extending them to categorical distributions are also very different. It is therefore unclear why authors are combining these two hardly related ideas into a single paper. Authors should either separate these ideas into two different piece of publications, or have a more unified story.

2. I am not impressed by the novelty of argmax flows. It is a simple extension to variational dequantization. The argmax flow layers cannot be composed to form deeper architectures. In fact, it is a simple "preprocessing" layer attached to standard normalizing flows, same as variational dequantization. 

3. It is unclear why argmax flows can perform better than VAEs when priors are learnable and parameterized by normalizing flows. Authors argue that argmax flows are better than variational dequantization because the continuous variables are not restricted to the unit interval [0, 1). However, if this is the right intuition for why argmax flows are better, we should instead use VAEs and train a normalizing flow model on the latent space as the prior. The latent space of VAEs is even more flexible than argmax flows, since it is not restricted by the argmax constraint. Although authors have compared VAE + IAF with their argmax flows in one experiment, I find it unsatisfying as the difference in model architectures is unclear, and comparison in other experimental settings is also needed.

4. I am not impressed by the novelty of multinomial diffusions. It is a straightforward generalization to the Bernoulli diffusion models in [Sohl-Dickstein et al. 2015], and as the author correctly noted, [Song et al. 2020] has already included a same formulation in the appendix. Moreover, in the language data experiment, authors use binary Cartesian products to represent multiple categories. Doesn't this reduce the problem to modeling binary data? If so, isn't multinomial diffusion the same as Bernoulli diffusion in the original 2015 paper?

**Post rebuttal update**:
Thanks for the response. Authors have addressed some of my questions, so I'm happy to increase my score. However, my major concerns still remain valid:

First, as noted by other reviewers as well, the connection between argmax flows and multinomial diffusion models are underwhelming. Authors responded that the biggest connection of them is that they both model discrete distributions. This is not convincing. Modeling discrete distributions is not an unsolved problem—autoregressive models, VAEs, energy-based models, and sometimes even GANs can model discrete distributions. We can't just publish paper by putting together separate ideas on VAEs, EBMs, etc, if they both model discrete distributions.

Second, I'm still not convinced that the proposed argmax flows are useful. The proposed method is a special case of a VAE. The paper should focus more on why a special variant of a VAE, with a hand-designed decoder, is necessary for good performance on discrete data. In the paper, authors motivate argmax flows as a more flexible variant to variational dequantization. However, VAEs are even more flexible and should be better from that reasoning. Authors also mentioned in the response that argmax flows are easier to optimize due to the fixed decoder, which seems to say that limiting the flexibility can contribute to the performance. This contradicts with the motivation from variational dequantization to argmax flows.

limitations_and_societal_impact: No obvious concerns.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The paper proposed two methods to learn categorical distributions. 

Method 1: ArgMax Flows. The sampling procedure is simply taking argmax of a continuous random variable. The learning procedure involves a variational distribution that satisfies the argmax constraint and optimizing ELBO. Several variation distributions are introduced. In addition, the Cartesian product of argmax flows is introduced to trade-off between symmetry and dimension.

Method 2: Multinomial DDPM. The diffusion process is defined by categorical distributions. The paper derived ELBO for these models, and introduced a special parameterization of the reverse process.

The paper evaluated the methods in several experiments including text generation and image segmentation. They show better performances than baseline methods.
main_review: The paper is original to my knowledge. It is well written and easy to follow. The research problem is significant because many tasks involve discrete data generation. 


Pros:
- The ArgMax Flow is generally an elegant framework. It is natural to optimize ELBO when likelihood cannot be computed. Then, the argmax condition for the variational distribution is a simple but smart solution that validates the ELBO objective. The several proposed variational distributions are elegant. 
- The Cartesian product is very useful especially when data is supported on a big, discrete domain. 
- The multinomial DDPM defines the diffusion process with categorical distributions, and therefore the model can be applied to categorical data. The closed-form ELBO is derived so we can directly optimize it. 
- The experiments are generally persuasive, and the character-level spell check is interesting. 


Cons / main concerns:
- Except that both ArgMax flows and multinomial DDPM are designed for categorical data, what are some inherent connections between these two types of models? Why do we want two models? Is there anything that one model can do while the other cannot? If one model is significantly more advantageous than the other in certain cases, is there any reason / recipe? 
- In eq(6), what does $v_{-x}$ mean? A clear definition is needed.
- Line 109, it seems the threshold function should be $\mathbb{R}^2\rightarrow(-\infty,T)$. This still looks strange: $T$ is an input; how can an input determine the range of the output? A re-formulation is needed here: e.g., $\mathrm{threshold}(\cdot;T)$ or $\mathrm{threshold}_T(\cdot)$.
- Regarding the Cartesian products, the authors had an additional experiment in the appendix, where the data domain size is 27. Theoretically, this method should be more useful when the data domain is very large. For instance, in 64*64 image generation, the data domain is [0,...,255] * [1,..,3] * [1,...,64] * [1,...,64]. In 16kHz 1-sec audio generation, the data domain is [-32768,...,32768] * [1,...,16000]. Therefore, it is more persuasive to show the effect of the Cartesian products if this method is evaluated on one of these tasks (although Table 1 indicates the rounding operation is more suitable). 
- The diffusion process defined by categorical distribution is a natural extension of the Bernoulli distribution in the original diffusion model. Although the paper mentioned it in the related work, I would strongly suggest the authors to write down the original Bernoulli version in the background section so that readers can have a clear sense of the background and the contribution of this paper. In addition, a clear definition for the $\mathcal{C}(\cdot|\cdot)$ in equation (11) is needed. There should be the unit vector $\textbf{1}$ after $\beta_t/K$. 
- Line 164-165: it is a good idea to parameterize $\mu(x_t,t)$, but why do you call it $\hat{x}_0$? Why can it approximate $x_0$? Either a theoretical or experimental validation is necessary. 
- Note: most of these concerns are about validation and explanation, and they do not harm the innovation or contribution of the paper. I would claim again that this paper is generally very elegant. 


Minor concerns:
- There should be commas or dots after each equation. There should be a number for the eq at line 56. 
- Some math symbols are bold somewhere but not bold elsewhere. Some symbols are confusing, e.g. in eq (4), I would prefer $v_{d,k}$ rather than $v_{dk}$.


Other questions just for curiosity:
- Is there a way to estimate the likelihood of ArgMax flows, even with a specially designed variational distribution?
- The spell check is interesting but seems a little too heuristic. If there are more mistakes then you probably need to start from an earlier step. Do you have a solution for that? E.g., one way is to learn a network that predicts the noise level of the input, and then determine which step to begin accordingly. 

limitations_and_societal_impact: Social impact is discussed in paper. Concerns on the limitations are in the main review.
review_title: null

Meta-review: This paper presents two approaches for learning expressive distributions over discrete variables. Although there are valid criticisms regarding the novelty (especially for the multinomial diffusion), the majority of the reviewers found the introduction of argmax flows an interesting and non-trivial extension of prior works. Given the community's interest in discrete variables, I recommend this paper for acceptance. For the final camera-ready version, I'd appreciate it if the authors could expand the discussion of limitations as suggested by 4fJX (in addition to other promised changes). 
Decision: Accept (Poster)