Mixed Variational Flows for Discrete Variables

TL;DR

This paper introduces a novel variational flow method for discrete distributions that avoids continuous embeddings, providing more reliable and faster approximations with minimal tuning.

Contribution

Develops a measure-preserving invertible map for discrete distributions and a mixed variational flow that directly models discrete data without embedding.

Findings

MAD Mix outperforms continuous-embedding flows in accuracy.

MAD Mix is significantly faster to train.

Provides reliable i.i.d. sampling and density evaluation for discrete data.

Abstract

Variational flows allow practitioners to learn complex continuous distributions, but approximating discrete distributions remains a challenge. Current methodologies typically embed the discrete target in a continuous space - usually via continuous relaxation or dequantization - and then apply a continuous flow. These approaches involve a surrogate target that may not capture the original discrete target, might have biased or unstable gradients, and can create a difficult optimization problem. In this work, we develop a variational flow family for discrete distributions without any continuous embedding. First, we develop a measure-preserving and discrete (MAD) invertible map that leaves the discrete target invariant, and then create a mixed variational flow (MAD Mix) based on that map. Our family provides access to i.i.d. sampling and density evaluation with virtually no tuning effort. We also develop an extension to MAD Mix that handles joint discrete and continuous models. Our experiments suggest that MAD Mix produces more reliable approximations than continuous-embedding flows while being significantly faster to train.