Boosting the Generalization Ability for Hyperspectral Image Classification Using Spectral-Spatial Axial Aggregation Transformer

Abstract

In the hyperspectral image classification (HSIC) task, the most commonly used model validation paradigm is partitioning the training-test dataset through pixelwise random sampling. By training on a small amount of data, the deep learning model can achieve almost perfect accuracy. However, in our experiments, we found that the high accuracy was reached because the training and test datasets share a lot of information. On nonoverlapping dataset partitions, well-performing models suffer significant performance degradation. To this end, we propose a spectral-spatial axial aggregation transformer model, namely, SaaFormer, which preserves generalization across dataset partitions. SaaFormer applies a multilevel spectral extraction structure to segment the spectrum into multiple spectrum clips such that the wavelength continuity of the spectrum across the channel is preserved. For each spectrum clip, the axial aggregation attention mechanism, which integrates spatial features along multiple spectral axes, is applied to mine the spectral characteristic. The multilevel spectral extraction and the axial aggregation attention emphasize spectral characteristics to improve the model generalization. The experimental results on five publicly available datasets demonstrate that our model exhibits comparable performance on the random partition while significantly outperforming other methods on nonoverlapping partitions. Moreover, SaaFormer shows excellent performance on background classification.