[Paper Review] Large Separable Kernel Attention: Rethinking the Large Kernel Attention Design in CNN

Kin Wai Lau, Lai-Man Po+1 more|arXiv (Cornell University)|Sep 4, 2023

Advanced Neural Network ApplicationsComputer Science3 citations

TL;DR

This paper proposes Large Separable Kernel Attention (LSKA), a computationally efficient alternative to Large Kernel Attention (LKA) in CNNs, by decomposing 2D depth-wise convolutions into cascaded 1D horizontal and vertical kernels. LSKA reduces FLOPs and memory usage by up to 50% with similar performance to LKA, while enhancing shape bias and robustness on corrupted ImageNet data.

ABSTRACT

Visual Attention Networks (VAN) with Large Kernel Attention (LKA) modules have been shown to provide remarkable performance, that surpasses Vision Transformers (ViTs), on a range of vision-based tasks. However, the depth-wise convolutional layer in these LKA modules incurs a quadratic increase in the computational and memory footprints with increasing convolutional kernel size. To mitigate these problems and to enable the use of extremely large convolutional kernels in the attention modules of VAN, we propose a family of Large Separable Kernel Attention modules, termed LSKA. LSKA decomposes the 2D convolutional kernel of the depth-wise convolutional layer into cascaded horizontal and vertical 1-D kernels. In contrast to the standard LKA design, the proposed decomposition enables the direct use of the depth-wise convolutional layer with large kernels in the attention module, without requiring any extra blocks. We demonstrate that the proposed LSKA module in VAN can achieve comparable performance with the standard LKA module and incur lower computational complexity and memory footprints. We also find that the proposed LSKA design biases the VAN more toward the shape of the object than the texture with increasing kernel size. Additionally, we benchmark the robustness of the LKA and LSKA in VAN, ViTs, and the recent ConvNeXt on the five corrupted versions of the ImageNet dataset that are largely unexplored in the previous works. Our extensive experimental results show that the proposed LSKA module in VAN provides a significant reduction in computational complexity and memory footprints with increasing kernel size while outperforming ViTs, ConvNeXt, and providing similar performance compared to the LKA module in VAN on object recognition, object detection, semantic segmentation, and robustness tests.

Motivation & Objective

• To address the quadratic growth in FLOPs and memory usage caused by large kernel depth-wise convolutions in Large Kernel Attention (LKA) modules.
• To enable the use of extremely large kernels in CNN attention modules without increasing computational complexity.
• To improve model robustness and shape bias by modifying the inductive bias of attention mechanisms.
• To benchmark LSKA's performance against LKA, ViTs, and ConvNeXt on corrupted ImageNet and downstream vision tasks.
• To provide a more efficient and scalable attention mechanism for vision transformers and CNNs with large receptive fields.

Proposed method

• Decompose the 2D depth-wise convolution kernel in LKA into cascaded 1D horizontal and vertical convolutions to reduce parameter and FLOP growth.
• Apply the same attention mechanism as in LKA but with separable kernels, maintaining the same feature refinement process via 1×1 convolutions.
• Use a two-stage design: first local feature extraction with small kernels, then long-range dependency modeling with dilated depth-wise separable convolutions.
• Introduce LSKA-trivial and LSKA variants to evaluate the impact of kernel decomposition on performance and efficiency.
• Leverage the same training and inference pipeline as LKA, ensuring fair comparison across models.
• Use shape and texture dimensionality estimation via mutual information to analyze inductive bias shifts induced by kernel size and decomposition.

Experimental results

Research questions

• RQ1 Can separable kernel decomposition reduce the computational cost of large kernel attention in CNNs without sacrificing performance?
• RQ2 How does the LSKA design affect the model's inductive bias toward shape versus texture in feature representations?
• RQ3 How does LSKA compare to LKA, ViTs, and ConvNeXt in terms of robustness on corrupted ImageNet benchmarks?
• RQ4 What is the impact of increasing kernel size on the effective receptive field and feature representation in LSKA compared to LKA?
• RQ5 Does the separable kernel design lead to better generalization and efficiency in vision tasks like classification, detection, and segmentation?

Key findings

• LSKA reduces FLOPs by up to 50% compared to LKA at kernel size 65, with only a 0.3% drop in top-1 accuracy on ImageNet.
• LSKA maintains comparable performance to LKA across image classification, object detection, and semantic segmentation tasks, while significantly reducing memory and computation.
• The LSKA design increases the model's bias toward shape representation, especially with larger kernels, as shown by shape-texture dimensionality estimation.
• LSKA outperforms ViTs and ConvNeXt on all five corrupted ImageNet benchmarks, demonstrating superior robustness.
• The effective receptive field of LSKA saturates at kernel sizes of 35, 53, and 65, similar to LKA, indicating stable long-range modeling.
• LSKA achieves 74.8% top-1 accuracy on ImageNet with a 65×65 kernel using only 4.07M parameters and 0.85 GFLOPs, while LKA with the same kernel uses 4.73M parameters and 1.12 GFLOPs.