Mutual Information Constraints for Monte-Carlo Objectives

TL;DR

This paper introduces a method to better estimate mutual information in Monte-Carlo objectives for variational autoencoders, improving latent variable usage and reducing posterior collapse.

Contribution

It develops estimators for the true posterior's KL divergence from the prior using sample recycling, enhancing latent variable training in Monte-Carlo objectives.

Findings

Improved rate-distortion performance with better mutual information control.

Reduced posterior collapse in models with continuous and discrete latents.

Encouraged evaluation of inference methods across different mutual information levels.

Abstract

A common failure mode of density models trained as variational autoencoders is to model the data without relying on their latent variables, rendering these variables useless. Two contributing factors, the underspecification of the model and the looseness of the variational lower bound, have been studied separately in the literature. We weave these two strands of research together, specifically the tighter bounds of Monte-Carlo objectives and constraints on the mutual information between the observable and the latent variables. Estimating the mutual information as the average Kullback-Leibler divergence between the easily available variational posterior $q(z|x)$ and the prior does not work with Monte-Carlo objectives because $q(z|x)$ is no longer a direct approximation to the model's true posterior $p(z|x)$. Hence, we construct estimators of the Kullback-Leibler divergence of the true posterior from the prior by recycling samples used in the objective, with which we train models of continuous and discrete latents at much improved rate-distortion and no posterior collapse. While alleviated, the tradeoff between modelling the data and using the latents still remains, and we urge for evaluating inference methods across a range of mutual information values.