Paper:
Mixture models, a basic building block in countless statistical models, involve latent random variables over discrete spaces, and existing posterior inference methods can be inaccurate and/or very slow. In this work we introduce a novel deep learning architecture for efficient amortized Bayesian inference over mixture models. While previous approaches to amortized clustering assumed a fixed or maximum number of mixture components and only amortized over the continuous parameters of each mixture component, our method amortizes over the local discrete labels of all the data points, and performs inference over an unbounded number of mixture components. The latter property makes our method natural for the challenging case of nonparametric Bayesian models, where the number of mixture components grows with the dataset. Our approach exploits the exchangeability of the generative models and is based on mapping distributed, permutation-invariant representations of discrete arrangements into varying-size multinomial conditional probabilities. The resulting algorithm parallelizes easily, yields iid samples from the approximate posteriors along with a normalized probability estimate of each sample (a quantity generally unavailable using Markov Chain Monte Carlo) and can easily be applied to both conjugate and non-conjugate models, as training only requires samples from the generative model. We also present an extension of the method to models of random communities (such as infinite relational or stochastic block models). As a scientific application, we present a novel approach to neural spike sorting for high-density multielectrode arrays.

Review:
Summary:
This paper introduces a novel deep learning architecture for efficient amortized Bayesian inference over mixture models. Unlike previous approaches to amortized clustering, the proposed method allows us to treat local discrete labels of data points and infer the unbounded number of mixture components, making it more flexible as in the case of Bayesian nonparametrics. It is shown that the resulting algorithm can be parallelized and applied to both conjugate and non-conjugate models. The authors also suggest an extension to models of random communities and a novel approach to neural spike sorting for high-density multielectrode arrays based on the proposed method.
Strengths:
The paper is generally well written and the relationship to previous works is well described. Empirical results seem quite convincing, for example, the clustering results presented in Fig. 2 and Fig. 3 clearly show not only the inferred number of clusters, but also the posterior probability which indicates that reasonable samples are assigned higher probability.
Weaknesses:
- Overall, the idea looks very original and promising, but I find some technical details are not easy to understand under the current form, especially for non-experts in this domain. I would recommend the authors to elaborate a bit more on the proposed architecture and the variable-input soft-max function in Sect. 2.1.
- On page 8, the authors mention that the NCP is much more efficient compared with MCMC, for example, in the Gaussian 2D example. However, regarding the DPMM clustering model, it is known that MCMC methods are generally slower compared with variational inference, which is computationally faster. I think it would be interesting to add a discussion or comparison with variational inference in terms of computational efficiency.
- If I understand correctly, the NCP is essentially based on a sequential sampling procedure. The authors claim that the proposed method is easily parallelized using a GPU, but there does not seem to be sufficient details on the GPU-parallelization of sequential sampling.
Minor comments:
The size of some figures appears too small, for example Fig. 6 and Fig. 10, which may hinder readability.