Semi-Discrete Normalizing Flows through Differentiable Tessellation

Ricky T. Q. Chen, Brandon Amos, Maximilian Nickel

Keywords: deep probabilistic modeling, normalizing flows, dequantization, disjoint mixture modeling

TL;DR: We combine Voronoi tessellation with normalizing flows to construct a new invertible transformation that has learnable discrete structure. We construct new tessellation-based dequantization and disjoint mixture modeling approaches.

Abstract: Mapping between discrete and continuous distributions is a difficult task and many have had to resort to heuristical approaches. We propose a tessellation-based approach that directly learns quantization boundaries in a continuous space, complete with exact likelihood evaluations. This is done through constructing normalizing flows on convex polytopes parameterized using a simple homeomorphism with an efficient log determinant Jacobian. We explore this approach in two application settings, mapping from discrete to continuous and vice versa. Firstly, a Voronoi dequantization allows automatically learning quantization boundaries in a multidimensional space. The location of boundaries and distances between regions can encode useful structural relations between the quantized discrete values. Secondly, a Voronoi mixture model has near-constant computation cost for likelihood evaluation regardless of the number of mixture components. Empirically, we show improvements over existing methods across a range of structured data modalities.

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