# Discrete Variational Autoencoding via Policy Search

arXiv: https://arxiv.org/abs/2509.24716 · HTML: https://mlsanity.com/paper/2509.24716/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 38 references that are in the map.
- First posted on arXiv 29 Sep 2025.

## Builds on (most cited first, 19 of 19)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control](https://mlsanity.com/paper/1909.12238/) (2019)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [PoseGPT: Quantization-based 3D Human Motion Generation and Forecasting](https://mlsanity.com/paper/2210.10542/) (2022)
- [MuProp: Unbiased Backpropagation for Stochastic Neural Networks](https://mlsanity.com/paper/1511.05176/) (2015)
- [Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator](https://mlsanity.com/paper/2010.04838/) (2020)
- [LocoMuJoCo: A Comprehensive Imitation Learning Benchmark for Locomotion](https://mlsanity.com/paper/2311.02496/) (2023)
- [Effective Sample Size, Dimensionality, and Generalization in Covariate Shift Adaptation](https://mlsanity.com/paper/2010.01184/) (2020)
- [Supervised Policy Update for Deep Reinforcement Learning](https://mlsanity.com/paper/1805.11706/) (2018)
- [Inferring Smooth Control: Monte Carlo Posterior Policy Iteration with Gaussian Processes](https://mlsanity.com/paper/2210.03512/) (2022)

## Abstract (arXiv)

Discrete latent bottlenecks in variational autoencoders (VAEs) offer high bit efficiency and can be modeled with autoregressive discrete distributions, enabling parameter-efficient multimodal search with transformers. However, discrete random variables do not allow for exact differentiable parameterization; therefore, discrete VAEs typically rely on approximations, such as Gumbel-Softmax reparameterization or straight-through gradient estimates, or employ high-variance gradient-free methods such as REINFORCE that have had limited success on high-dimensional tasks such as image reconstruction. Inspired by popular techniques in policy search, we propose a training framework for discrete VAEs that leverages the natural gradient of a non-parametric encoder to update the parametric encoder without requiring reparameterization. Our method, combined with automatic step size adaptation and a transformer-based encoder, scales to challenging datasets such as ImageNet and outperforms both approximate reparameterization methods and quantization-based discrete autoencoders in reconstructing high-dimensional data from compact latent spaces.