Matching Observations to Distributions: Efficient Estimation via Sparsified Hungarian Algorithm

TL;DR

This paper presents a fast, randomized algorithm for matching observations to known distributions using a sparsified Hungarian method, with theoretical guarantees and statistical analysis showing the MLE's effectiveness in high-dimensional settings.

Contribution

It introduces a novel sparsified Hungarian algorithm that reduces runtime for maximum likelihood estimation in distribution matching problems, with proven statistical bounds.

Findings

The new algorithm achieves $ ilde{O}(n^2)$ runtime, improving over the traditional $ ext{O}(n^3)$.

Separation of $ ext{log }k$ between Gaussian means suffices for perfect matching.

Expected mismatch rate decreases at rate $ ext{O}(( ext{log }k)^2/ ext{distance}^2)$.

Abstract

Suppose we are given observations, where each observation is drawn independently from one of $k$ known distributions. The goal is to match each observation to the distribution from which it was drawn. We observe that the maximum likelihood estimator (MLE) for this problem can be computed using weighted bipartite matching, even when $n$, the number of observations per distribution, exceeds one. This is achieved by instantiating $n$ duplicates of each distribution node. However, in the regime where the number of observations per distribution is much larger than the number of distributions, the Hungarian matching algorithm for computing the weighted bipartite matching requires $\mathcal O(n^3)$ time. We introduce a novel randomized matching algorithm that reduces the runtime to $\tilde{\mathcal O}(n^2)$ by sparsifying the original graph, returning the exact MLE with high probability. Next, we give statistical justification for using the MLE by bounding the excess risk of the MLE, where the loss is defined as the negative log-likelihood. We test these bounds for the case of isotropic Gaussians with equal covariances and whose means are separated by a distance $\eta$, and find (1) that $\gg \log k$ separation suffices to drive the proportion of mismatches of the MLE to 0, and (2) that the expected fraction of mismatched observations goes to zero at rate $\mathcal O({(\log k)}^2/\eta^2)$.