Neural Clustering Processes

TL;DR

This paper introduces neural network architectures that efficiently generate approximate posterior samples for clustering tasks, improving accuracy and speed over traditional methods, and applies this to neural spike sorting.

Contribution

The authors develop deep learning methods for fast, approximate Bayesian posterior inference in clustering, capable of handling nonparametric models and arbitrary dataset sizes.

Findings

Networks produce accurate posterior samples for clustering.

Methods outperform traditional inference in speed and accuracy.

Applied successfully to neural spike sorting.

Abstract

Probabilistic clustering models (or equivalently, mixture models) are basic building blocks in countless statistical models and involve latent random variables over discrete spaces. For these models, posterior inference methods can be inaccurate and/or very slow. In this work we introduce deep network architectures trained with labeled samples from any generative model of clustered datasets. At test time, the networks generate approximate posterior samples of cluster labels for any new dataset of arbitrary size. We develop two complementary approaches to this task, requiring either O(N) or O(K) network forward passes per dataset, where N is the dataset size and K the number of clusters. Unlike previous approaches, our methods sample the labels of all the data points from a well-defined posterior, and can learn nonparametric Bayesian posteriors since they do not limit the number of mixture components. As a scientific application, we present a novel approach to neural spike sorting for high-density multielectrode arrays.