Consistent Amortized Clustering via Generative Flow Networks

TL;DR

This paper introduces GFNCP, a novel order-invariant amortized clustering framework using Generative Flow Networks, which improves clustering consistency and performance without sequential label assignment dependencies.

Contribution

The paper proposes GFNCP, a new order-invariant amortized clustering method based on Generative Flow Networks, ensuring consistent clustering posteriors and better performance.

Findings

GFNCP achieves order invariance in clustering.

GFNCP outperforms existing clustering methods on synthetic and real data.

Flow matching conditions ensure clustering posterior consistency.

Abstract

Neural models for amortized probabilistic clustering yield samples of cluster labels given a set-structured input, while avoiding lengthy Markov chain runs and the need for explicit data likelihoods. Existing methods which label each data point sequentially, like the Neural Clustering Process, often lead to cluster assignments highly dependent on the data order. Alternatively, methods that sequentially create full clusters, do not provide assignment probabilities. In this paper, we introduce GFNCP, a novel framework for amortized clustering. GFNCP is formulated as a Generative Flow Network with a shared energy-based parametrization of policy and reward. We show that the flow matching conditions are equivalent to consistency of the clustering posterior under marginalization, which in turn implies order invariance. GFNCP also outperforms existing methods in clustering performance on both synthetic and real-world data.