Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity
This addresses a bottleneck in training models with discrete latent variables for researchers and practitioners in machine learning, offering an incremental improvement over existing approximation methods.
The paper tackles the computational challenge of marginalizing over discrete or structured latent variables in neural network training by introducing a method that uses sparse mappings to enable exact and efficient marginalization, achieving good performance across three tasks with the practicality of sampling-based approximations.
Training neural network models with discrete (categorical or structured) latent variables can be computationally challenging, due to the need for marginalization over large or combinatorial sets. To circumvent this issue, one typically resorts to sampling-based approximations of the true marginal, requiring noisy gradient estimators (e.g., score function estimator) or continuous relaxations with lower-variance reparameterized gradients (e.g., Gumbel-Softmax). In this paper, we propose a new training strategy which replaces these estimators by an exact yet efficient marginalization. To achieve this, we parameterize discrete distributions over latent assignments using differentiable sparse mappings: sparsemax and its structured counterparts. In effect, the support of these distributions is greatly reduced, which enables efficient marginalization. We report successful results in three tasks covering a range of latent variable modeling applications: a semisupervised deep generative model, a latent communication game, and a generative model with a bit-vector latent representation. In all cases, we obtain good performance while still achieving the practicality of sampling-based approximations.