Invertible Attention

Jiajun Zha,Yiran Zhong,Jing Zhang,Richard Hartley,Liang Zheng

DOI: https://doi.org/10.48550/arXiv.2106.09003

2021-06-27

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 <a class="link-external link-https" href="https://github.com/Schwartz-Zha/InvertibleAttention" rel="external noopener nofollow">this https URL</a>

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem that this paper attempts to solve is to introduce the attention mechanism into the reversible network, enabling it to capture long - distance dependencies. Specifically:

1. **Background problem**: The traditional attention mechanism has been proven to be an effective method for capturing long - distance dependencies in deep learning, but it has not been applied to reversible networks yet. This is because for a network to be reversible, each component in the network must be a bijective transformation (i.e., satisfy one - to - one mapping), and the standard attention module is not bijective.

2. **Core challenge**: How to design a reversible attention mechanism that can both maintain the original functions and performance and ensure the reversibility in the network.

3. **Solution**: The author proposes a new method - Invertible Attention. By carefully constraining its Lipschitz constant, the attention module can meet the requirements of reversibility while maintaining the original functions.

4. **Verification and application**:
- The author verifies four different types of reversible attention mechanisms (Gaussian, Embedded Gaussian, Dot - Product, and Concatenation) on the image reconstruction task.
- The proposed reversible attention module is embedded into i - ResNet, demonstrating its improvement on the image generation task.
- On the camouflaged object detection task, the performance of reversible attention and non - reversible attention is compared, and the results show that they perform equally well.

### Formula summary

To ensure the reversibility of the attention module, the author constrains its Lipschitz constant through the following formula:

\[ L(A) = k(\mu_R c_F + c_R \mu_F) \]

where:
- \( A(x) = R(x)F(x) \)
- \( F \) and \( R \) are the feature mapping and response function respectively, and they each have Lipschitz constants \( c_F \) and \( c_R \) respectively.
- \( \mu_F \) and \( \mu_R \) are the norm upper limits of \( F(x) \) and \( R(x) \) respectively.
- \( k \) is the constant boundary between the \( L_1 \) norm and the \( L_2 \) norm.

In addition, the author also implements the final reversible attention module through the following formula:

\[ f(x) = x + W_L A(x) \]

where \( W_L \) is the weight matrix of the last convolutional layer, and its Lipschitz constant is set to \( \frac{c}{k(\mu_R c_F + c_R \mu_F)} \) to ensure the reversibility of the entire module.

### Summary

Through these methods, the author successfully introduces the attention mechanism into the reversible network, solves the problem of the irreversibility of the traditional attention mechanism, and verifies its effectiveness and practicality on multiple tasks.