A Lightweight Transformer Network for Hyperspectral Image Classification

Abstract

Transformer is a powerful tool for capturing long-range dependencies and has shown impressive performance in hyperspectral image (HSI) classification. However, such power comes with a heavy memory footprint and huge computation burden. In this paper, we propose two types of lightweight self-attention modules (a channel lightweight multi-head self-attention module and a position lightweight multi-head self-attention module) to reduce both memory and computation while associating each pixel or channel with global information. Moreover, we discover that transformers are ineffective in explicitly extracting local and multi-scale features due to the fixed input size and tend to overfit when dealing with a small number of training samples. Therefore, a lightweight transformer (LiT) network, built with the proposed lightweight self-attention modules, is presented. LiT adopts convolutional blocks to explicitly extract local information in early layers and employs transformers to capture long-range dependencies in deep layers. Furthermore, we design a controlled multi-class stratified sampling strategy to generate appropriately sized input data, ensure balanced sampling, and reduce the overlap of feature extraction regions between training and test samples. With appropriate training data, convolutional tokenization, and lightweight transformers, LiT mitigates overfitting and enjoys both high computational efficiency and good performance. Experimental results on several HSI datasets verify the effectiveness of our design.