Distribution Compression in Near-Linear Time

Abhishek Shetty, Raaz Dwivedi, Lester Mackey

Keywords: Distribution compression, linear time, thinning, i.i.d. sampling, Markov chain Monte Carlo, maximum mean discrepancy, reproducing kernel Hilbert space

Abstract: In distribution compression, one aims to accurately summarize a probability distribution $\mathbb{P}$ using a small number of representative points. Near-optimal thinning procedures achieve this goal by sampling $n$ points from a Markov chain and identifying $\sqrt{n}$ points with $\widetilde{\mathcal{O}}(1/\sqrt{n})$ discrepancy to $\mathbb{P}$. Unfortunately, these algorithms suffer from quadratic or super-quadratic runtime in the sample size $n$. To address this deficiency, we introduce Compress++, a simple meta-procedure for speeding up any thinning algorithm while suffering at most a factor of $4$ in error. When combined with the quadratic-time kernel halving and kernel thinning algorithms of Dwivedi and Mackey (2021), Compress++ delivers $\sqrt{n}$ points with $\mathcal{O}(\sqrt{\log n/n})$ integration error and better-than-Monte-Carlo maximum mean discrepancy in $\mathcal{O}(n \log^3 n)$ time and $\mathcal{O}( \sqrt{n} \log^2 n )$ space. Moreover, Compress++ enjoys the same near-linear runtime given any quadratic-time input and reduces the runtime of super-quadratic algorithms by a square-root factor. In our benchmarks with high-dimensional Monte Carlo samples and Markov chains targeting challenging differential equation posteriors, Compress++ matches or nearly matches the accuracy of its input algorithm in orders of magnitude less time.

One-sentence Summary: We introduce a simple algorithm for compressing an $n$-point summary of a probability distribution into a $\sqrt{n}$-point summary of comparable quality in $O(n \log^2 n)$ time.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2111.07941/code)