Generalization Gap in Amortized Inference
This work addresses overfitting issues in VAEs, which is important for applications like lossless compression, but it appears incremental as it builds on existing VAE frameworks.
The paper tackles the generalization gap in Variational Auto-Encoders (VAEs), identifying amortized inference as the main cause of overfitting, and proposes a new training objective that improves generalization, demonstrated through enhanced performance in image modeling and lossless compression.
The ability of likelihood-based probabilistic models to generalize to unseen data is central to many machine learning applications such as lossless compression. In this work, we study the generalization of a popular class of probabilistic model - the Variational Auto-Encoder (VAE). We discuss the two generalization gaps that affect VAEs and show that overfitting is usually dominated by amortized inference. Based on this observation, we propose a new training objective that improves the generalization of amortized inference. We demonstrate how our method can improve performance in the context of image modeling and lossless compression.