[Paper Review] Amortized Bayesian inference for clustering models

Ari Pakman, Liam Paninski|arXiv (Cornell University)|Nov 24, 2018

Bayesian Methods and Mixture Models18 references3 citations

TL;DR

This paper introduces the Neural Clustering Process (NCP), an amortized Bayesian inference method for clustering models that uses neural networks to map symmetry-invariant representations of cluster arrangements into conditional cluster assignment probabilities. The approach enables fast, parallelizable, independent posterior sampling with computational efficiency comparable to a single Gibbs sweep, applicable to both conjugate and non-conjugate models.

ABSTRACT

We develop methods for efficient amortized approximate Bayesian inference over posterior distributions of probabilistic clustering models, such as Dirichlet process mixture models. The approach is based on mapping distributed, symmetry-invariant representations of cluster arrangements into conditional probabilities. The method parallelizes easily, yields iid samples from the approximate posterior of cluster assignments with the same computational cost of a single Gibbs sampler sweep, and can easily be applied to both conjugate and non-conjugate models, as training only requires samples from the generative model.

Motivation & Objective

• Address the computational inefficiency of MCMC methods in probabilistic clustering models, especially for non-conjugate models with high-dimensional parameter spaces.
• Overcome the lack of accuracy guarantees in variational inference by developing a method that combines the strengths of MCMC and deep learning.
• Enable fast, scalable, and parallelizable posterior sampling for cluster assignments by amortizing inference through neural network training on generative model samples.
• Preserve permutation symmetries in cluster assignments—within clusters, between clusters, and among unassigned points—through invariant representations.
• Develop a framework that generalizes across model types, including Dirichlet process mixture models, by relying only on samples from the generative model for training.

Proposed method

• Represent cluster configurations using distributed, symmetry-invariant features: within-cluster summaries $H_k = ext{sum}(h(x_i))$ for each cluster $k$, and global cluster summaries $G = ext{sum}(g(H_k))$.
• Model the conditional probability $p(c_n | c_{1:n-1}, extbf{x})$ using a neural network $f(G_k, Q, h_n)$, where $Q$ is the sum of embeddings for unassigned points and $h_n = h(x_n)$, with $G_k$ computed for each possible cluster assignment.
• Train the neural network parameters $\theta$ via stochastic gradient descent to minimize the expected negative log-likelihood of the conditional probabilities under the true posterior.
• Leverage permutation invariance in training by sampling permutations $\pi$ of data points, ensuring the model learns equivariant behavior under cluster relabeling.
• Apply Rao-Blackwellization to reduce variance in the training objective by computing exact conditional probabilities over $c_{n:N}$ for fixed $c_{1:n-1}$, improving sample efficiency.
• Use the trained network to generate independent, GPU-parallelizable posterior samples of cluster assignments for new data, bypassing sequential MCMC sampling.

Experimental results

Research questions

• RQ1 Can we develop an amortized inference method for clustering models that achieves both computational efficiency and accuracy comparable to MCMC, without requiring expensive Markov chain updates?
• RQ2 How can we design a neural network-based inference framework that respects the inherent permutation symmetries in cluster assignments—within clusters, between clusters, and among unassigned points?
• RQ3 To what extent can we generalize the method across conjugate and non-conjugate models using only samples from the generative model for training?
• RQ4 Can we achieve significant speedups over standard Gibbs sampling by enabling full parallelization of posterior samples while maintaining comparable variance?
• RQ5 Does the use of symmetry-invariant representations and Rao-Blackwellized training improve the convergence and stability of the neural network policy in cluster assignment prediction?

Key findings

• The Neural Clustering Process (NCP) achieves posterior sampling performance comparable to Gibbs sampling in terms of variance, with a significant reduction in wall-clock time.
• NCP enables massive parallelization on GPUs, reducing average runtime from 1969 seconds for Gibbs sampling (20,000 samples, 1000 burn-in) to 184 seconds for NCP.
• The method maintains symmetry invariance under permutations of cluster labels and data points, which is monitored during training and preserved in the model’s predictions.
• Rao-Blackwellization reduces the variance of the training objective by computing exact conditional probabilities over $c_{n:N}$ for fixed $c_{1:n-1}$, improving sample efficiency.
• The approach is applicable to both conjugate and non-conjugate models, as training only requires samples from the generative model, not explicit posterior computation.
• After training, NCP generates independent and identically distributed (i.i.d.) posterior samples of cluster assignments at the same computational cost as a single Gibbs sweep, enabling scalable inference.