Enhancing Layer Interaction Using Key-Correlated Layer Attention

Technical Summary: Enhancing Layer Interaction Using Key-Correlated Layer Attention

Problem Statement
Recent advancements in deep neural network (DNN) design have utilized Layer Attention to enhance inter-layer interactions by treating layer outputs as sequence tokens. However, standard layer attention suffers from quadratic computational complexity ($O(D^2)$) relative to network depth $D$, alongside significant memory growth due to the need to cache outputs from all preceding layers.

To address these issues, prior works introduced Recurrent Layer Attention (RLA) and linear approximations like MRLA-L. While RLA reduces computational overhead by caching Key ($K$) and Value ($V$) representations, it functions as a static mechanism where early layer features remain fixed once generated, limiting dynamic information updates. Furthermore, linear approximations such as MRLA-L attempt to reduce complexity by linearly approximating the attention output. The authors observe that this approach causes the contributions of earlier layers to be exponentially attenuated (as shown in their analysis of the decay coefficients $\beta_l$), rendering it difficult to establish effective long-range cross-layer dependencies. Additionally, existing linearization methods often fail to address the memory bottleneck of caching previous layer outputs or the specific challenges of applying linear attention to hierarchical DNN structures.

Methodology: Key-Correlated Layer Attention (KCLA)
The proposed Key-Correlated Layer Attention (KCLA) is a novel linearized layer attention method designed to achieve linear computational complexity while preserving dynamic updates and long-range dependency modeling. The methodology is built upon three core insights and mechanisms:

1. Key Correlation Observation: Empirical analysis reveals that Key representations across different layers in standard layer attention exhibit high cosine similarity. Leveraging this, KCLA approximates the standard layer attention output by factoring out the linear transformation of the Value ($V$) and approximating the Key representations as being directionally correlated. This allows the aggregation of layer outputs into a recurrent state ($U^t$) without storing the full history of $K$ and $V$ matrices.
2. Dynamic Softmax Restoration: A critical challenge in linearizing attention is the loss of the softmax operation, which is essential for normalizing attention scores and enabling dynamic distribution adjustments. KCLA introduces a modified multi-head attention mechanism to restore this capability:
   • It employs multiple heads with different scaling factors (temperatures) derived from a logarithmically uniform sequence.
   • It utilizes a learnable linear transformation $\Phi$ to map the norm of the Query ($\|Q^t\|$) and the correlation angle ($\cos \theta$) into head-specific weights.
   • This allows the model to dynamically adjust the attention distribution at each layer, effectively simulating the temperature parameter's role in softmax without computing the full $N \times N$ matrix.
3. Cross-Stage Adaptation: To accommodate DNNs with varying feature map sizes and channel dimensions across stages, KCLA implements a specific adaptation procedure (Algorithm 1). This involves normalizing accumulated features, applying a $1 \times 1$ convolution with stride for downsampling, and recomputing normalization factors to bridge dimensional disparities between stages.

Key Contributions

• Observation and Proof: The authors demonstrate that Key representations across different layers in standard layer attention are highly correlated in direction, providing the theoretical basis for linear approximation.
• Novel Linearization: KCLA is proposed as a linearized layer attention method that avoids the static nature of RLA and the exponential attenuation of MRLA-L. It achieves linear computational complexity and fixed spatial complexity independent of network depth.
• Dynamic Softmax Approximation: A new method is introduced to approximate the softmax operation using multi-head scaling and learnable transformations, enabling dynamic and adaptive adjustment of attention distributions at each layer.

Experimental Results
The authors evaluated KCLA on three fundamental vision tasks: image classification, object detection, and medical image segmentation.

• Image Classification: On CIFAR-100, KCLA achieved the best performance across ResNet depths of 20, 56, and 110, outperforming MRLA-L, DLA-L, and standard layer attention. On ImageNet-1K, KCLA integrated into ResNet-50 and ResNet-101 improved Top-1 accuracy by +1.7% and +1.4% over the baseline, respectively. It maintained a slight but consistent advantage over MRLA-L while using significantly fewer parameters than the dynamic DLA-L method.
• Object Detection: On the COCO2017 dataset using Faster R-CNN and Mask R-CNN, KCLA demonstrated competitive performance. With ResNet-101, it surpassed the more complex DLA-L architecture on several metrics (AP$_{75}$, AP$_S$, AP$_M$) while using 2.2M fewer parameters.
• Medical Image Segmentation: On ISIC2017, ISIC2018, and Kvasir-SEG datasets, KCLA consistently outperformed MRLA-L and the baseline ResNet-UNet, achieving superior mIoU and DSC scores with minimal parameter overhead.
• Complexity Analysis: Unlike MRLA-B, whose memory and FLOPs grow quadratically with depth, KCLA exhibits linear growth in both computational load and memory cost.

Significance and Claims
The paper claims that KCLA successfully bridges the gap between the efficiency of linear attention and the representational power of standard attention. By directly approximating the standard layer attention mechanism based on Key correlations, KCLA avoids the static feature representation issues of RLA and the long-range dependency degradation of MRLA-L.

The authors emphasize that their method maintains the ability to model long-range cross-layer dependencies efficiently, offering a desirable balance between efficiency and performance. They note that the proposed method for restoring the softmax function is a critical factor in achieving these results. While the current implementation relies on pre-set temperature parameters, the authors identify the design of a fully self-adaptive, dynamic temperature multi-head linear layer attention as a direction for future optimization. The code is publicly available to facilitate reproducibility.