Direct Optimization through $\arg \max$ for Discrete Variational Auto-Encoder

TL;DR

This paper introduces a direct optimization method for discrete variational auto-encoders that bypasses softmax relaxations, enabling effective training of models with structured discrete latent variables.

Contribution

It proposes a novel direct loss minimization approach for the $ ext{arg} ext{max}$ objective in discrete VAEs, extending to structured latent variables.

Findings

Effective training of discrete VAEs demonstrated

Outperforms softmax relaxation methods

Applicable to structured discrete latent models

Abstract

Reparameterization of variational auto-encoders with continuous random variables is an effective method for reducing the variance of their gradient estimates. In the discrete case, one can perform reparametrization using the Gumbel-Max trick, but the resulting objective relies on an $\arg \max$ operation and is non-differentiable. In contrast to previous works which resort to softmax-based relaxations, we propose to optimize it directly by applying the direct loss minimization approach. Our proposal extends naturally to structured discrete latent variable models when evaluating the $\arg \max$ operation is tractable. We demonstrate empirically the effectiveness of the direct loss minimization technique in variational autoencoders with both unstructured and structured discrete latent variables.