Attentive Contractive Flow with Lipschitz-constrained Self-Attention

TL;DR

This paper introduces Attentive Contractive Flow (ACF), a novel invertible model incorporating localized self-attention to enhance expressivity and training efficiency of normalizing flows, leading to more realistic samples and better noise resilience.

Contribution

The paper proposes ACF, a new invertible flow model using contractive flows and self-attention, improving expressivity, convergence speed, and robustness of normalizing flows.

Findings

ACF improves bits per dim metric across models.

ACF leads to faster training convergence.

ACF samples are more realistic and locally coherent.

Abstract

Normalizing flows provide an elegant method for obtaining tractable density estimates from distributions by using invertible transformations. The main challenge is to improve the expressivity of the models while keeping the invertibility constraints intact. We propose to do so via the incorporation of localized self-attention. However, conventional self-attention mechanisms don't satisfy the requirements to obtain invertible flows and can't be naively incorporated into normalizing flows. To address this, we introduce a novel approach called Attentive Contractive Flow (ACF) which utilizes a special category of flow-based generative models - contractive flows. We demonstrate that ACF can be introduced into a variety of state of the art flow models in a plug-and-play manner. This is demonstrated to not only improve the representation power of these models (improving on the bits per dim metric), but also to results in significantly faster convergence in training them. Qualitative results, including interpolations between test images, demonstrate that samples are more realistic and capture local correlations in the data well. We evaluate the results further by performing perturbation analysis using AWGN demonstrating that ACF models (especially the dot-product variant) show better and more consistent resilience to additive noise.