Variational Inference Optimized Using the Curved Geometry of Coupled Free Energy

TL;DR

This paper proposes a novel variational inference framework using coupled free energy to better handle heavy-tailed distributions, improving robustness and stability in training models like CVAEs, especially against outliers.

Contribution

It introduces a coupled free energy approach that extends variational inference to heavy-tailed distributions, enhancing robustness and training stability.

Findings

3% improvement in FID score over VAE on CelebA after 5 epochs

Robust training against severe outliers

Stable training process with heavy-tailed distributions

Abstract

We introduce an optimization framework for variational inference based on the coupled free energy, extending variational inference techniques to account for the curved geometry of the coupled exponential family. This family includes important heavy-tailed distributions such as the generalized Pareto and the Student's t. By leveraging the coupled free energy, which is equal to the coupled evidence lower bound (ELBO) of the inverted probabilities, we improve the accuracy and robustness of the learned model. The coupled generalization of Fisher Information metric and the affine connection. The method is applied to the design of a coupled variational autoencoder (CVAE). By using the coupling for both the distributions and cost functions, the reconstruction metric is derived to still be the mean-square average loss with modified constants. The novelty comes from sampling the heavy-tailed latent distribution with its associated coupled probability, which has faster decaying tails. The result is the ability to train a model robust against severe outliers, while assuring that the training process is stable. The Wasserstein-2 or Fr\'echet Inception Distance of the reconstructed CelebA images shows the CVAE has a 3\% improvement over the VAE after 5 epochs of training.