Invertible Attention

TL;DR

This paper introduces a novel invertible attention mechanism that can be integrated into invertible neural networks, enabling long-range dependency modeling while maintaining invertibility, validated on image reconstruction and dense prediction tasks.

Contribution

It proposes a mathematically grounded invertible attention mechanism by constraining the Lipschitz constant, compatible with invertible models.

Findings

Invertible attention achieves comparable performance to standard attention on dense prediction tasks.

The proposed method enables invertibility in attention modules, validated on CIFAR-10, SVHN, and CelebA.

Mathematically proven invertibility through Lipschitz constant constraints.

Abstract

Attention has been proved to be an efficient mechanism to capture long-range dependencies. However, so far it has not been deployed in invertible networks. This is due to the fact that in order to make a network invertible, every component within the network needs to be a bijective transformation, but a normal attention block is not. In this paper, we propose invertible attention that can be plugged into existing invertible models. We mathematically and experimentally prove that the invertibility of an attention model can be achieved by carefully constraining its Lipschitz constant. We validate the invertibility of our invertible attention on image reconstruction task with 3 popular datasets: CIFAR-10, SVHN, and CelebA. We also show that our invertible attention achieves similar performance in comparison with normal non-invertible attention on dense prediction tasks. The code is available at https://github.com/Schwartz-Zha/InvertibleAttention