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

Models with discrete latent variables drive extensive research in machine learning applications, such as language classification and generation [34, 7], molecular synthesis [12], or game solving [18]. Compared to their continuous counterparts, discrete latent variable models can decrease the computational complexity of inference calculations, for instance, by discarding alternatives in hard attention models [13], they can improve interpretability by illustrating which terms contributed to the solution [21, 34], and they can facilitate the encoding of inductive biases in the learning process, such as images consisting of a small number of objects [2] or tasks requiring intermediate alignments [18]. Finally, in some cases, discrete latent variables are natural choices, including when modeling datasets with discrete classes [26]. Nonetheless, models involving discrete latent variables are hard to train, with the key issue being to estimate the gradients of the resulting non-differentiable objectives. While one can use unbiased estimators, such as REINFORCE [33], their variance is typically high [22, 19, 30, 4, 20, 31]. In variational auto-encoders (VAEs) with continuous latent variables, the reparameterization trick provides a successful 1 alternative [10, 25].