Gradients should stay on Path: Better Estimators of the Reverse- and Forward KL Divergence for Normalizing Flows

TL;DR

This paper introduces a new path-gradient estimator for normalizing flows that reduces variance, accelerates training convergence, and mitigates mode-collapse, improving variational inference performance.

Contribution

It presents a novel, easy-to-implement path-gradient estimator for both reverse and forward KL divergences in normalizing flows, enhancing training stability and accuracy.

Findings

Lower variance in gradient estimates

Faster convergence during training

Reduced susceptibility to mode-collapse

Abstract

We propose an algorithm to estimate the path-gradient of both the reverse and forward Kullback-Leibler divergence for an arbitrary manifestly invertible normalizing flow. The resulting path-gradient estimators are straightforward to implement, have lower variance, and lead not only to faster convergence of training but also to better overall approximation results compared to standard total gradient estimators. We also demonstrate that path-gradient training is less susceptible to mode-collapse. In light of our results, we expect that path-gradient estimators will become the new standard method to train normalizing flows for variational inference.