Neural Clustering Process (NCP)

• Neural Clustering Process is an amortized inference framework that uses deep neural networks to rapidly generate approximate posterior cluster assignments for nonparametric Bayesian models.
• It employs two complementary architectures—pointwise (O(N)) and clusterwise (O(K))—to balance computational efficiency and accurate uncertainty quantification.
• Demonstrated in high-dimensional applications like neural spike sorting, NCP efficiently handles unbounded clusters while preserving permutation invariance.

The Neural Clustering Process (NCP) is an amortized-inference framework for probabilistic clustering based on nonparametric Bayesian mixture models. It leverages deep network architectures to learn fast, approximate posterior samplers for cluster-label assignments, accommodating datasets of arbitrary size and an unbounded number of clusters. The NCP approach departs from traditional MCMC- and variational-inference methods by training on labeled samples drawn from a generative mixture model, making test-time inference highly efficient and fully parallelizable. Two complementary amortized architectures, requiring either $O(N)$ or $O(K)$ neural network forward passes per clustering sample (with $N$ the dataset cardinality and $K$ the number of clusters), enable tractable, symmetry-preserving computation of approximate posterior labelings and deliver exchangeable samples for downstream uncertainty quantification. The NCP has demonstrated strong empirical performance for high-dimensional scientific applications such as neural spike sorting (Pakman et al., 2018).

1. Model Foundations and Objectives

The NCP framework is designed to address computational limitations and inaccuracy inherent in posterior inference for mixtures—especially nonparametric mixtures where the label space is combinatorially large and the number of mixture components is not fixed a priori. In canonical Bayesian mixture models (including Dirichlet-process (DP) mixtures), each instance $x_i$ is associated with a latent cluster assignment $c_i \in \mathbb{N}$, and inference seeks the posterior $p(c_{1:N}|x_{1:N})$. Conventional posterior inference methods—Gibbs sampling, split-merge MCMC, and variational approximations—require significant computation per sample and may suffer from missed modes or slow mixing.

NCP addresses these issues by:

• Training a neural sampler $q_{\theta}(c|x)$ on synthetic, fully labeled datasets $(x,c)$ from the tractable generative model $p(x,c)$.
• Ensuring that the learned sampler outputs independent, exchangeable samples from an approximate posterior for any new test dataset.
• Allowing for efficient, GPU-parallel sample generation of hundreds to thousands of clusterings in milliseconds—bypassing burn-in and autocorrelation penalties.
• Supporting nonparametric posteriors: the sampler does not limit the number of possible clusters, naturally handling variable and unbounded $K$.

2. Generative Model and Clustering Priors

NCP is instantiated on standard nonparametric Bayesian mixture models parameterized as follows:

• Hyperparameters $\alpha_1,\alpha_2 \sim p(\alpha)$.
• Cluster labels $c_{1:N} \sim p(c_{1:N}|\alpha_1)$, e.g., a Chinese Restaurant Process (CRP) prior.
• Cluster parameters $\mu_{1:K}|c_{1:N} \sim p(\mu_{1:K}|\alpha_2)$, for $K=\max c_{1:N}$.
• Data $x_i|c_i,\mu \sim p(x_i|\mu_{c_i})$, independently for $i=1,\ldots,N$.

The joint density decomposes as:

$$p(x|c,\mu) = \prod_{i=1}^N p(x_i|\mu_{c_i}),\;\; p(c) = p(c_{1:N}|\alpha_1),\;\; p(\mu) = p(\mu_{1:K}|\alpha_2).$$

This generalization readily supports infinite mixtures. When $p(c_{1:N}|\alpha_1)$ is a CRP, $K$ can be random and unbounded, which is central to the nonparametric Bayesian paradigm.

3. Amortized Inference Architectures

NCP specifies two neural architectures for amortizing posterior sampling, both encoding exchangeability and permutation symmetry.

3.1 O(N) "Pointwise" NCP

This approach exploits the sequential factorization:

$p(c_{1:N}|x) = \prod_{n=1}^N p(c_n|c_{1:n-1}, x).$

At each $n$, there are $K_n+1$ candidate clusterings (joining an existing $K_n$ cluster or starting a new one). The sampler approximates each conditional:

$$q_{\theta}(c_n|c_{1:n-1}, x) \approx p(c_n|c_{1:n-1}, x),$$

via a permutation-invariant neural network using the following summary statistics:

• Within-cluster sum: $H_k = \sum_{i:c_i=k} h(x_i)$,
• Between-cluster sum: $G = \sum_{k=1}^{K_n} g(H_k)$,
• Unassigned sum: $U = \sum_{i=n+1}^N u(x_i)$.

Assigning $x_n$ to cluster $k$ updates $H_k$ and, consequently, $G \mapsto G_k$. Label probabilities are computed as a softmax of neural logits $f(G_k,U)$. The networks $h$, $u$, $g$, and $f$ (typically MLPs or convolutional nets) enforce permutation-invariance. Forward-pass computation is $O(NK)$ arithmetic, but parallelization enables $O(N)$ passes on GPU.

3.2 O(K) "Clusterwise" CCP

Instead of sampling labels sequentially, CCP samples whole clusters:

$p(S_{1:K}|x) = \prod_{k=1}^{K} p(S_k|S_{<k},x),$

where $S_k$ are index sets for each cluster. Each factor is a mixture over a first index $d_k$ and a membership vector $b \in \{0,1\}^{m_k}$ (for the $m_k$ remaining points):

$$p(S_k|S_{<k},x) = \frac{1}{|I_k|} \sum_{d_k \in I_k} p(b|d_k, S_{<k}, x).$$

A conditional de Finetti representation enables modeling $p(b|d_k, S_{<k},x)$ by independent Bernoulli assignments with parameters from a neural context. The implementation employs a conditional VAE per cluster for both the continuous latent $z_k$ and the discrete $d_k$ (with Gumbel-Softmax relaxation). All $K$ clusters are processed in parallel, leading to $O(K)$ neural-network passes.

4. Training Objectives

NCP's objectives ensure the learned sampler matches the true posterior for samples from the generating model.

• For the pointwise architecture (NCP), the loss is:

$$L_{\mathrm{NCP}}(\theta) = -\mathbb{E}_{p(N)p(x,c)} \sum_{n=1}^N \log q_{\theta}(c_n|c_{1:n-1},x),$$

minimizing the averaged KL divergence $$\mathbb{E}_{p(N)p(x,c)}[KL(p(c|x) \Vert q_\theta(c|x))]$$. Gradients are backpropagated through the differentiable softmax; training does not require MCMC.

• For the clusterwise approach (CCP), the per-cluster ELBO is:

$$L_{\mathrm{CCP}}(\theta, \phi) = \mathbb{E}_{p(x, S_{1:K})} \sum_{k=1}^K \mathbb{E}_{q_{\phi}(z_k, d_k | ...)} \left[\log p_{\theta}(S_k, z_k, d_k | S_{<k}, x) - \log q_{\phi}(z_k, d_k|S_{\leq k}, x)\right].$$

Where $q_{\phi}$ acts as a variational posterior. Gumbel-Softmax and normal reparameterization are used for differentiability of discrete and continuous latents, respectively.

5. Sampling Algorithms and Computational Characteristics

Sampling procedures mirror the respective inference network structures:

• Pointwise NCP (O(N) passes): For each data point, condition on prior assignments, update cluster summaries, and sample the cluster label via softmax over the logits $f(G_k,U)$. Labels are assigned sequentially with dynamic cluster allocation and efficient updating of summary states.
• Clusterwise CCP (O(K) passes): At each step, uniformly select an unassigned index $d_k$, decode a membership vector $b$ via the VAE, and assign indices to the new cluster $S_k$. Iterate until all points are clustered.

Key computational distinctions:

• NCP: $O(NK)$ arithmetic per sample, but $O(N)$ GPU-forward passes due to parallelization over candidate clusters.
• CCP: $O(K)$ heavy forward passes, efficient for cases with small $K$ relative to $N$.

Memory usage scales as $O(N d_h)$ for NCP and $O(Kd)$ plus VAE overhead for CCP.

6. Diagnostics, Validation, and Empirical Performance

Diagnostic validation is performed via posterior-probability accuracy (e.g., matching $p(c_{41}|c_{1:40},x)$ in a 2D-Gaussian DP mixture), exchangeability tests (verifying negligible variance in NLL across data permutations), and Geweke-style tests (matching marginals between $q_{\theta}(\cdot|x)$ and $p(c)$).

Performance benchmarks:

• On high-density neural data, NCP produces thousands of iid approximate posterior samples in under one second on a single GPU, versus single correlated samples for collapsed Gibbs.
• In spike-sorting applications, NCP achieves or surpasses the clustering quality of KiloSort and variational MFM on real, synthetic, and hybrid datasets while providing uncertainty quantification.
• Empirical results indicate the approach maintains the fidelity of the underlying Bayesian posterior despite amortization.

7. Scientific Application: Neural Spike Sorting

A primary domain application is spike sorting from high-density multi-electrode array (MEA) data:

• Raw input: each spike is represented as a $7\times32$ spatiotemporal waveform.
• Direct clustering: NCP dispenses with manual feature construction; an encoder $h(\cdot)$, implemented as a ResNet-style 1D-convolution over time with 7 channels, outputs $h(x_i) \in \mathbb{R}^{256}$ for each spike.
• Training: Datasets are synthesized with ground-truth labels from a finite-mixture-of-finite-mixtures prior, using real spike templates augmented with structured noise.
• Test-time: NCP yields 150 high-likelihood clusterings for $N \approx 2000$ spikes in approximately 10 seconds via GPU-parallel sampling. The highest-probability clustering is used for generating spike templates, and the ensemble captures posterior uncertainty in ambiguous regions.
• The architecture's ability to sample from a well-defined posterior, handle unknown $K$, and obviate manual preprocessing distinguishes it from both heuristic and variational pipelines.

NCP exemplifies a general approach toward fast, scalable, nonparametric Bayesian clustering with full uncertainty quantification and has been validated across both simulated and challenging real-world tasks (Pakman et al., 2018).