A Generic Shared Attention Mechanism for Various Backbone Neural Networks

TL;DR

This work identifies a strong cross-layer correlation in self-attention maps, revealing that per-block SAMs are parameter-inefficient when networks deepen. It proposes Dense-and-Implicit-Attention (DIA), a layerwise shared attention unit, and DIA-LSTM, which uses a GAP-augmented LSTM to bridge current and prior layer information, offering substantial parameter savings while improving performance across ResNet, Transformer, and UNet backbones on classification, generation, and detection tasks. DIA achieves dense, implicit cross-layer connections and acts as a regularizer that stabilizes training, with lightweight variants (DIA-LSTM Light) that retain most benefits while minimizing overhead. The approach demonstrates statistically significant gains over multiple SAMs and backbones on benchmarks like CIFAR, STL-10, ImageNet, MS COCO, and medical imaging, and provides practical guidance on deployment, including cost considerations and BN-sharing caveats. Overall, DIA presents a broadly applicable method to improve attention module efficiency and network performance in diverse vision problems.

Abstract

The self-attention mechanism has emerged as a critical component for improving the performance of various backbone neural networks. However, current mainstream approaches individually incorporate newly designed self-attention modules (SAMs) into each layer of the network for granted without fully exploiting their parameters' potential. This leads to suboptimal performance and increased parameter consumption as the network depth increases. To improve this paradigm, in this paper, we first present a counterintuitive but inherent phenomenon: SAMs tend to produce strongly correlated attention maps across different layers, with an average Pearson correlation coefficient of up to 0.85. Inspired by this inherent observation, we propose Dense-and-Implicit Attention (DIA), which directly shares SAMs across layers and employs a long short-term memory module to calibrate and bridge the highly correlated attention maps of different layers, thus improving the parameter utilization efficiency of SAMs. This design of DIA is also consistent with the neural network's dynamical system perspective. Through extensive experiments, we demonstrate that our simple yet effective DIA can consistently enhance various network backbones, including ResNet, Transformer, and UNet, across tasks such as image classification, object detection, and image generation using diffusion models.

Figures (15)

• Figure 1: The paradigm of self-attention module (SAM).
• Figure 2: Linear correlation of attention maps across different layers. Experiments on ImageNet and CIFAR100 are conducted using ResNet50 and ResNet164 as backbones, respectively. The attention modules "SE" and "ECA" refer to those from two well-known channel attention networks, SENet hu2018squeeze, and ECANet wang2020eca. (a) and (b) respectively represent the linear relationship of attention maps generated in two randomly sampled blocks in the first stage and the last stage. The visualization of individual input samples is provided by (a) and (b), while (c) displays the distribution of Pearson correlation coefficients across multiple input samples in different pair of blocks.
• Figure 3: The Residual neural network on dynamical system perspective. The forward process of ResNet can be modeled as a forward numerical method for a dynamical system. The feature map of any block in the same stage can be regarded as a point in the same vector space. For example, in Stage 1, "A1", "B1", and "C1" are the feature maps from three adjacent blocks, and these feature maps can be regarded as homogenous.
• Figure 4: The DIA architecture in the residual network. $F_{ext}$ means the operation for extracting different scales of features. $F_{emp}$ means the operation for emphasizing or recalibrating features.
• Figure 5: The showcase of DIA-LSTM in detail. In the LSTM cell, $c_t$ is the cell state vector and $h_t$ is the hidden state vector. GAP means global average pool over channels and $\otimes$ means channel-wise multiplication.