ARM: Augment-REINFORCE-Merge Gradient for Stochastic Binary Networks

Abstract

To backpropagate the gradients through stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator that is unbiased, exhibits low variance, and has low computational complexity. Exploiting variable augmentation, REINFORCE, and reparameterization, the ARM estimator achieves adaptive variance reduction for Monte Carlo integration by merging two expectations via common random numbers. The variance-reduction mechanism of the ARM estimator can also be attributed to either antithetic sampling in an augmented space, or the use of an optimal anti-symmetric "self-control" baseline function together with the REINFORCE estimator in that augmented space. Experimental results show the ARM estimator provides state-of-the-art performance in auto-encoding variational inference and maximum likelihood estimation, for discrete latent variable models with one or multiple stochastic binary layers. Python code for reproducible research is publicly available.

arXiv commentICLR 2019

Cited by in corpus (20)

• Gradient Estimation with Stochastic Softmax Tricks
• Resource-Efficient Neural Networks for Embedded Systems
• $L_0$-ARM: Network Sparsification via Stochastic Binary Optimization
• Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity
• Not All Attention Is All You Need
• DisARM: An Antithetic Gradient Estimator for Binary Latent Variables
• Storchastic: A Framework for General Stochastic Automatic Differentiation
• Exploring the Back Alleys: Analysing The Robustness of Alternative Neural Network Architectures against Adversarial Attacks
• Straight-Through Estimator as Projected Wasserstein Gradient Flow
• Direct Optimization through $\arg \max$ for Discrete Variational Auto-Encoder
• Improved Gradient-Based Optimization Over Discrete Distributions
• Dichotomize and Generalize: PAC-Bayesian Binary Activated Deep Neural Networks
• Generative Semantic Hashing Enhanced via Boltzmann Machines
• Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks
• Coupled Gradient Estimators for Discrete Latent Variables
• Graph-Free Knowledge Distillation for Graph Neural Networks
• Adaptive Correlated Monte Carlo for Contextual Categorical Sequence Generation
• Dynamic Narrowing of VAE Bottlenecks Using GECO and L0 Regularization
• Neural Plasticity Networks
• Augment-Reinforce-Merge Policy Gradient for Binary Stochastic Policy