Publikation
Discrete Variational Autoencoding via Policy Search
Michael Drolet; Firas Al-Hafez; Aditya Bhatt; Jan Peters; Oleg Arenz
In: Computing Research Repository eprint Journal (CoRR), Vol. abs/2509.24716, Pages 1-28, arXiv, 2025.
Zusammenfassung
Discrete latent bottlenecks in variational autoencoders (VAEs) offer high bit ef-
ficiency and can be modeled with autoregressive discrete distributions, enabling
parameter-efficient multimodal search with transformers. However, discrete ran-
dom variables do not allow for exact differentiable parameterization; therefore,
discrete VAEs typically rely on approximations, such as Gumbel-Softmax repa-
rameterization 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 tech-
niques in policy search, we propose a training framework for discrete VAEs that
leverages the natural gradient of a non-parametric encoder to update the paramet-
ric encoder without requiring reparameterization. Our method, combined with
automatic step size adaptation and a transformer-based encoder, scales to chal-
lenging datasets such as ImageNet and outperforms both approximate reparame-
terization methods and quantization-based discrete autoencoders in reconstructing
high-dimensional data from compact latent spaces, achieving a 20% improvement
on FID Score for ImageNet 256.