# Enabling Automatic Repair of Source Code Vulnerabilities Using Data-Driven
Methods

Anastasiia Grishina  0000-0003-3139-0200 Simula Research LaboratoryKristian
Augusts gate 23OsloNorway0164<EMAIL_ADDRESS>
(2022)

###### Abstract.

Users around the world rely on software-intensive systems in their day-to-day
activities. These systems regularly contain bugs and security vulnerabilities.
To facilitate bug fixing, data-driven models of automatic program repair use
pairs of buggy and fixed code to learn transformations that fix errors in
code. However, automatic repair of security vulnerabilities remains under-
explored. In this work, we propose ways to improve code representations for
vulnerability repair from three perspectives: input data type, data-driven
models, and downstream tasks. The expected results of this work are improved
code representations for automatic program repair and, specifically, fixing
security vulnerabilities.

Automatic Program Repair, Static Analysis, Software Security, Natural Language
Processing, Graph-based Machine Learning, ML4Code

This is a preprint of the work published in the 44th International Conference
on Software Engineering (ICSE ’22) Doctoral Symposium, May 21–29, 2022,
Pittsburgh, PA, USA, https://doi.org/10.1145/3510454.3517063.
This version is licensed under a Creative Commons Attribution 4.0
International (CC BY 4.0) license

††journalyear: 2022††copyright: acmcopyright††conference: 44th International
Conference on Software Engineering Companion; May 21–29, 2022; Pittsburgh, PA,
USA††booktitle: 44th International Conference on Software Engineering
Companion (ICSE ’22 Companion), May 21–29, 2022, Pittsburgh, PA, USA††price:
15.00††doi: 10.1145/3510454.3517063††isbn:
978-1-4503-9223-5/22/05††submissionid: DS-11

##  1\. Introduction

Users around the world rely on various software-intensive systems, such as the
systems that assist in healthcare and transportation. These systems regularly
contain bugs and security vulnerabilities that can have various negative
effects, from economical to life-threatening (bbvaopenmind2015:most, 1, 2).
Software companies have been making a considerable effort to eliminate such
bugs and vulnerabilities. To support these activities, Automatic Program
Repair (APR) methods have received increasing attention in software
engineering research (monperrus2018:automatic, 3).

The rich body of literature on APR has been surveyed in various studies
(allamanis2018:survey, 4, 5, 3, 6, 7, 8). Following the terminology of the APR
surveys (monperrus2018:automatic, 3, 7), this PhD project focuses on
_learning-based methods_ of _static_ repair of _source code vulnerabilities_.
In _static_ analysis, we process code without executing it, as opposed to
dynamic analysis. When approached with _learning-based methods_ , APR requires
training data. We also refer to these methods as _data-driven methods_. Using
code organized as pairs of (buggy code; fixed code), data-driven approaches
employ deep learning techniques to predict a fix, also referred to as a patch.

The main objective of this work is to improve automatic vulnerability repair
in source code from the following three perspectives.

(a) _Input data_. Currently, code change history is the major source of input
data for training code representation models. In addition, forum discussions
are used jointly with code change history for code analysis and repair
(liu2018:mining, 9, 10, 11). This work will explore the effect of integrating
context from additional input modalities, such as programming tutorials, on
security-oriented APR.

(b) _Models_. Two characteristics of source code are used in literature: its
similarity to natural languages and the formal structure of code. In the
modelling phase, this work will focus on creating hybrid methods based on
these two approaches for code representation.

(c) _Downstream tasks_. To guarantee the security of given software, one
approach is to eliminate the known vulnerabilities. Instances of security
principles violation are published in Common Vulnerabilities and Exposures
(CVE) Records.111 https://cve.mitre.org/ However, the CVE list is not
exhaustive, and not all CVEs are bugs. For example, issues, such as storing
passwords without salt or using an unreliable hashing function, are commonly
not reported as errors by compilers, but they can compromise confidentiality
and integrity of an application. Moreover, repair of security vulnerabilities
is under-explored, compared to general bugs (sharma2021:survey, 8). Therefore,
we have chosen pattern mining for CVE types and vulnerability repair as
downstream tasks.

##  2\. Related Work

Two strands of research in data-driven APR methods form the basis for this
work. They consider different approaches to code representation that are used
for various downstream tasks, including code repair. One branch focuses on
similarities between natural language and code, while the other represents
code using its structure.

Allamanis et al. (allamanis2018:survey, 4) presents the hypothesis that
programming languages are forms of natural language. Therefore, Natural
Language Processing (NLP) can be reused for code analysis tasks. Examples of
code processing models are Tufano et al. (tufano2019:empirical, 12), SequenceR
(chen2019:sequencer, 13), CodeBERT (feng2020:codebert, 14), Codex
(chen2021:evaluating, 15). These models reuse RNNs (sutskever2014:sequence,
16), transformer architectures (devlin2019:bert, 17) with attention
(vaswani2017:attention, 18), and the copy mechanism (gu2016:incorporating,
19). Specifically, CodeBERT is a recent study on fine-tuning the BERT model on
a large varied code base to perform tasks related to translating from natural
language to programming language and the same in the reverse direction.
Another study that reused BERT is CuBERT (kanade2020:learning, 20). The
authors trained the BERT model on code and fine-tuned it on code understanding
and repair tasks. A recent advancement in code representation by GitHub and
OpenAI, Codex (chen2021:evaluating, 15) empowering Copilot, is built using the
GPT-3 model. Codex translates a human language description of a function to
code performing the described task.

One of the differences between source code and natural language is strict
syntactic and logical structure of code that can be used for code
representation. Code parsed according to its syntax forms an Abstract Syntax
Tree (AST) (long2017:automatic, 21). Similarly, data flow and control flow
graphs developed for code analysis earlier also represent code characteristics
(stevens1974:structured, 22, 23). For instance, code2vec (alon2019:code2vec,
24) uses both n-grams (n words considered as one token) and ASTs to
automatically predict function name from its code. GraphCodeBERT
(guo2021:graphcodebert, 25) obtains an AST and translates it to a data flow
graph to obtain one of three input representations of code that are tested on
four downstream tasks, including code repair. Although they have received less
attention than the mentioned graph types, knowledge graphs are also used in
code analysis and repair tasks (atzeni2017:codeontology, 26, 27, 28).

While the aforementioned research focuses on general program repair, we
concentrate on repair of security vulnerabilities. Compared to security
vulnerability detection and general bug repair, security vulnerability repair
has received less attention (shen2020:survey, 29, 8). For example, the
benchmark CodeXGlue (lu2021:codexglue, 30) has a setup for vulnerability
detection, but the repair (refinement) part contains only general bugs
following the dataset of Tufano et al. (tufano2019:empirical, 12). One study
on vulnerability repair uses generative adversarial networks to generate
security vulnerability patches in C/C++ (harer2018:learning, 31). Another
study focuses on repair of vulnerabilities of code generated by Codex
(pearce2021:can, 32). In this research, we aim at improving results for
vulnerability repair and generalize to other programming languages.

##  3\. Expected Contributions

The overview of this work is sketched in Figure 1. In the three phases of the
PhD project, we hypothesize that (a) using context from additional modalities
can improve APR performance; (b) when built upon NLP-based and graph-based
methods for code representation and generation, hybrid methods can capture
both natural language similarities and structural properties of code; (c)
mining patterns from code in the wild can help improve vulnerability repair
results. The overall objective is to automatically repair software
vulnerabilities in source code using data-driven methods.

Firstly, we aim to explore the impact of context from additional modalities on
code representation and bug repair compared to the mostly exclusive use of
code change history. Examples of input modalities are texts and code snippets
in tutorials, books on security in software engineering, and forum
discussions.

Secondly, we aim to analyze to what extent existing APR models can be adapted
to repair vulnerabilities. As baselines, we will use NLP-based and graph-based
methods. In this step, we will experiment with building hybrid methods using
baseline approaches. To the best of our knowledge, there is no work on
combination of these methods for vulnerability repair.

Finally, we will use code representations to mine patterns of errors and their
fixes within each CVE type. The pattern mining task is inspired by the study
that categorizes common mistakes in AI projects (islam2020:repairing, 11).
With patterns discovered for security vulnerability types, we will focus on
vulnerability repair in source code.

##  4\. Planned Evaluation of Results

In this project, we will reuse data gathered in CVEfixes
(bhandari2021:cvefixes, 33). CVEfixes provides both a dataset and a framework
for collecting data from GitHub, GitLab, and Bitbucket. The already extracted
dataset refers to CVE Records in the U.S. National Vulnerability Database to
classify code snippets by CVE type. In this dataset, code changes are
extracted on the file and method level. In downstream tasks, CVE type-specific
pattern mining will use vulnerable code classified by CVE type, while code
pairs of (vulnerable code; fixed code) will be used for training and
evaluation of vulnerability repair.

The target for the modelling phase (b) is to learn vector representations of
code as sequences of tokens or as graphs, so that code snippets of the same
CVE type in the same programming language are close to each other in the
embedding space. To measure the quality of code representations, we will use
metrics for clustering in machine learning and metrics in NLP: Silhouette
Coefficient, Normalized Mutual Information, Completeness, and V-measure for
clusters; perplexity, cosine and other similarity metrics from NLP.

Figure 1. Organization of the work: we focus on (a) different input
modalities; (b) model types; (c) two downstream tasks.

Alternatively, since APR is analogical to the NLP task of neural machine
translation, it can be evaluated with the Recall-Oriented Understudy for
Gisting Evaluation (ROUGE) and Bilingual Evaluation Understudy (BLEU) NLP
metrics (feng2020:codebert, 14, 34), and their extensions (ren2020:codebleu,
35, 30). In the context of vulnerability repair, ROUGE scores evaluate the
patch based on the number of occurrences of n-grams from the known repaired
code (reference sequence) in the patch (generated sequence). By contrast, BLEU
shows n-gram precision of the sequence generation. In detail, it is measured
based on the number n-gram occurrences from the patch (generated sequence)
that are also found in the known repaired code (reference sequence).

##  5\. Conclusion

In conclusion, this project focuses on automatic repair of source code
vulnerabilities. We hypothesize that repair performance can be improved with
the context extracted from different modalities in addition to widely used
code change history. Furthermore, the repair performance can be influenced by
combining models that build on naturalness hypothesis and formal
characteristics of code, such as syntax, structure, and dependencies. We aim
to test methods of pattern mining for each vulnerability type, and its repair.
Training will be done using the CVEfixes dataset, and the evaluation step will
reuse clustering and NLP metrics. The expected results of this work are
improved code representations for APR and, specifically, fixing security
vulnerabilities.

###### Acknowledgements.

This work is supported by the Research Council of Norway through the secureIT
project (IKTPLUSS #288787).

## References

  * (1) BBVA OpenMind  “The 5 Most Infamous Software Bugs in History”, https://www.bbvaopenmind.com/en/technology/innovation/the-5-most-infamous-software-bugs-in-history/, 2015
  * (2) Andi Wilson, Ross Schulman, Kevin Bankston and Trey Herr  “Bugs in the System: A Primer on the Software Vulnerability Ecosystem and Its Policy Implications”, 2016
  * (3) Martin Monperrus  “Automatic Software Repair: A Bibliography”  In _ACM Computing Surveys_ 51.1, 2018, pp. 1–24  DOI: 10.1145/3105906
  * (4) Miltiadis Allamanis, Earl T. Barr, Premkumar Devanbu and Charles Sutton  “A Survey of Machine Learning for Big Code and Naturalness”  In _ACM Computing Surveys_ 51.4, 2018, pp. 81:1–81:37  DOI: 10.1145/3212695
  * (5) Yuzhen Liu, Long Zhang and Zhenyu Zhang  “A Survey of Test Based Automatic Program Repair”  In _Journal of Software_ 13.8, 2018, pp. 437–452  DOI: doi: 10.17706/jsw.13.8.437-452
  * (6) Luca Gazzola, Daniela Micucci and Leonardo Mariani  “Automatic Software Repair: A Survey”  In _IEEE Transactions on Software Engineering_ 45.1, 2019, pp. 34–67  DOI: 10.1109/tse.2017.2755013
  * (7) Claire Le Goues, Michael Pradel and Abhik Roychoudhury  “Automated Program Repair”  In _Communications of the ACM_ 62.12, 2019, pp. 56–65  DOI: 10/gkgf29
  * (8) Tushar Sharma, Maria Kechagia, Stefanos Georgiou, Rohit Tiwari and Federica Sarro  “A Survey on Machine Learning Techniques for Source Code Analysis”  In _arXiv:2110.09610 [cs]_ , 2021  arXiv:2110.09610 [cs]
  * (9) Xuliang Liu and Hao Zhong  “Mining Stackoverflow for Program Repair”  In _International Conference on Software Analysis, Evolution and Reengineering (SANER)_ Campobasso: IEEE, 2018, pp. 118–129  DOI: 10.1109/SANER.2018.8330202
  * (10) Sonal Mahajan, Negarsadat Abolhassani and Mukul R. Prasad  “Recommending Stack Overflow Posts for Fixing Runtime Exceptions Using Failure Scenario Matching”  In _Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering (ESEC/FSE)_ , 2020, pp. 1052–1064  DOI: 10.1145/3368089.3409764
  * (11) Md Johirul Islam, Rangeet Pan, Giang Nguyen and Hridesh Rajan  “Repairing Deep Neural Networks: Fix Patterns and Challenges”  In _International Conference on Software Engineering (ICSE)_ Seoul South Korea: ACM, 2020, pp. 1135–1146  DOI: 10.1145/3377811.3380378
  * (12) Michele Tufano, Cody Watson, Gabriele Bavota, Massimiliano Di Penta, Martin White and Denys Poshyvanyk  “An Empirical Study on Learning Bug-Fixing Patches in the Wild via Neural Machine Translation”  In _ACM Transactions on Software Engineering and Methodology_ 28.4, 2019, pp. 19:1–19:29  DOI: 10.1145/3340544
  * (13) Zimin Chen, Steve James Kommrusch, Michele Tufano, Louis-Noël Pouchet, Denys Poshyvanyk and Martin Monperrus  “SEQUENCER: Sequence-to-Sequence Learning for End-to-End Program Repair”  In _IEEE Transactions on Software Engineering_ , 2019  DOI: 10.1109/tse.2019.2940179
  * (14) Zhangyin Feng et al.  “CodeBERT: A Pre-Trained Model for Programming and Natural Languages”  In _Findings of the Association for Computational Linguistics: EMNLP 2020_ , 2020, pp. 1536–1547  DOI: 10.18653/v1/2020.findings-emnlp.139
  * (15) Mark Chen et al.  “Evaluating Large Language Models Trained on Code”, 2021  arXiv:2107.03374
  * (16) Ilya Sutskever, Oriol Vinyals and Quoc V. Le  “Sequence to Sequence Learning with Neural Networks”, 2014  arXiv:1409.3215
  * (17) Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova  “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”  In _Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL, Vol 1)_ Association for Computational Linguistics, 2019, pp. 4171–4186  DOI: 10.18653/v1/n19-1423
  * (18) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser and Illia Polosukhin  “Attention Is All You Need”  In _Advances in Neural Information Processing Systems 30_ Curran Associates, Inc., 2017, pp. 5998–6008
  * (19) Jiatao Gu, Zhengdong Lu, Hang Li and Victor O.. Li  “Incorporating Copying Mechanism in Sequence-to-Sequence Learning”, 2016  arXiv:1603.06393
  * (20) Aditya Kanade, Petros Maniatis, Gogul Balakrishnan and Kensen Shi  “Learning and Evaluating Contextual Embedding of Source Code”, 2020  arXiv:2001.00059
  * (21) Fan Long, Peter Amidon and Martin Rinard  “Automatic Inference of Code Transforms for Patch Generation”  In _Proceedings of the 2017 11th Joint Meeting on Foundations of Software Engineering_ Paderborn Germany: ACM, 2017, pp. 727–739  DOI: 10.1145/3106237.3106253
  * (22) W.. Stevens, G.. Myers and L.. Constantine  “Structured Design”  In _IBM Systems Journal_ 13.2, 1974, pp. 115–139  DOI: 10.1147/sj.132.0115
  * (23) Frances E. Allen  “Control Flow Analysis”  In _ACM SIGPLAN Notices_ 5.7, 1970, pp. 1–19  DOI: 10.1145/390013.808479
  * (24) Uri Alon, Meital Zilberstein, Omer Levy and Eran Yahav  “Code2vec: Learning Distributed Representations of Code”  In _Principles of Programming Languages (POPL)_ ACM, 2019, pp. 1–29  DOI: 10.1145/3290353
  * (25) Daya Guo et al.  “GraphCodeBERT: Pre-training Code Representations with Data Flow”  In _9th International Conference on Learning Representations, ICLR 2021_ Virtual Event, Austria: OpenReview.net, 2021, pp. 1–18
  * (26) Mattia Atzeni and Maurizio Atzori  “CodeOntology: RDF-ization of Source Code”  In _The Semantic Web – ISWC 2017_ 10588 Cham: Springer International Publishing, 2017, pp. 20–28  DOI: 10.1007/978-3-319-68204-4˙2
  * (27) Hoan Anh Nguyen, Tien N. Nguyen, Danny Dig, Son Nguyen, Hieu Tran and Michael Hilton  “Graph-Based Mining of In-the-Wild, Fine-Grained, Semantic Code Change Patterns”  In _International Conference on Software Engineering (ICSE)_ , 2019, pp. 819–830  DOI: 10.1109/icse.2019.00089
  * (28) Ahmed El Amine Djebri, Antonia Ettorre and Johann Mortara  “Towards a Linked Open Code”  In _The Semantic Web_ 12731 Cham: Springer International Publishing, 2021, pp. 497–505  DOI: 10.1007/978-3-030-77385-4˙29
  * (29) Zhidong Shen and Si Chen  “A Survey of Automatic Software Vulnerability Detection, Program Repair, and Defect Prediction Techniques”  In _Security and Communication Networks_ 2020, 2020, pp. 1–16  DOI: 10.1155/2020/8858010
  * (30) Shuai Lu et al.  “CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation”, 2021  arXiv:2102.04664
  * (31) Jacob Harer, Onur Ozdemir, Tomo Lazovich, Christopher Reale, Rebecca Russell, Louis Kim and peter chin  “Learning to Repair Software Vulnerabilities with Generative Adversarial Networks”  In _Neural Information Processing Systems Conference (NeurIPS)_ Curran Associates, Inc., 2018, pp. 7933–7943
  * (32) Hammond Pearce, Benjamin Tan, Baleegh Ahmad, Ramesh Karri and Brendan Dolan-Gavitt  “Can OpenAI Codex and Other Large Language Models Help Us Fix Security Bugs?”  In _arXiv:2112.02125 [cs]_ , 2021  arXiv:2112.02125 [cs]
  * (33) Guru Bhandari, Amara Naseer and Leon Moonen  “CVEfixes: Automated Collection of Vulnerabilities and Their Fixes from Open-Source Software”  In _International Conference on Predictive Models and Data Analytics in Software Engineering (PROMISE)_ ACM, 2021, pp. 30–39  DOI: 10.1145/3475960.3475985
  * (34) Colin B. Clement, Dawn Drain, Jonathan Timcheck, Alexey Svyatkovskiy and Neel Sundaresan  “PyMT5: Multi-Mode Translation of Natural Language and Python Code with Transformers”, 2020  arXiv:2010.03150
  * (35) Shuo Ren et al.  “CodeBLEU: A Method for Automatic Evaluation of Code Synthesis”  In _arXiv:2009.10297 [cs]_ , 2020  arXiv:2009.10297 [cs]


# SaaFormer: Spectral-spatial Axial Aggregation Transformer for Hyperspectral
Image Classification

Enzhe Zhao, Zhichang Guo, Yao Li, Dazhi Zhang Coresponding author: Yao Li
(E-mail: yaoli0508@hit.edu.cn)Enzhe Zhao, Zhichang Guo, Yao Li and Dazhi Zhang
are with the Department of Computational Mathematics, School of Mathematics,
Harbin Institute of Techology, Harbin 150001, China.

###### Abstract

Hyperspectral images (HSI) captured from earth observing satellites and
aircraft is becoming increasingly important for applications in agriculture,
environmental monitoring, mining, etc. Due to the limited available
hyperspectral datasets, the pixel-wise random sampling is the most commonly
used training-test dataset partition approach, which has significant overlap
between samples in training and test datasets. Furthermore, our experimental
observations indicates that regions with larger overlap often exhibit higher
classification accuracy. Consequently, the pixel-wise random sampling approach
poses a risk of data leakage. Thus, we propose a block-wise sampling method to
minimize the potential for data leakage. Our experimental findings also
confirm the presence of data leakage in models such as 2DCNN. Further, We
propose a spectral-spatial axial aggregation transformer model, namely
SaaFormer, to address the challenges associated with hyperspectral image
classifier that considers HSI as long sequential three-dimensional images. The
model comprises two primary components: axial aggregation attention and multi-
level spectral-spatial extraction. The axial aggregation attention mechanism
effectively exploits the continuity and correlation among spectral bands at
each pixel position in hyperspectral images, while aggregating spatial
dimension features. This enables SaaFormer to maintain high precision even
under block-wise sampling. The multi-level spectral-spatial extraction
structure is designed to capture the sensitivity of different material
components to specific spectral bands, allowing the model to focus on a
broader range of spectral details. The results on six publicly available
datasets demonstrate that our model exhibits comparable performance when using
random sampling, while significantly outperforming other methods when
employing block-wise sampling partition.

###### Index Terms:

Hyperspectral (HS) image classification, random sampling, long sequence
information, block-wise sampling, Transformer.

##  I Introduction

Hyperspectral remote sensing images are a type of satellite or airborne image
that captures data across a wide range of wavelengths. Unlike traditional
remote sensing images that capture data across just a few wavelengths (e.g.,
red, green, and blue), hyperspectral images capture data across hundreds of
narrow spectral bands. Compared with other types of remote sensing images.
This allows for a more detailed and precise analysis of the surface of the
earth, as different materials and surfaces reflect light differently across
the electromagnetic spectrum[1]. Therefore, HSI has been applied in various
fields, such as astronomy, agriculture, and mineralogy. Consequently, HSI
classification has emerged as a highly active research area within the field
of hyperspectral research.In the context of a high spectral image
classification system, a typical sequential workflow encompasses the following
stages: image restoration, encompassing tasks like denoising and missing data
recovery[2], [3], dimensionality reduction[4], spectral unmixing[5, 6, 7, 8],
and feature extraction[9, 10]. These stages are executed in a consecutive
manner, forming the overall processing pipeline.

During the early stages of hyperspectral image analysis, researchers explored
and applied existing classification techniques from the fields of remote
sensing and image processing. The Maximum Likelihood Classifier(MLC)[11, 12]
and the Spectral Angle Mapper(SAM)[13] were among the earliest methods used
for hyperspectral image classification. MLC assumes that pixel spectra in each
class follow a specific probability distribution, typically a multivariate
Gaussian distribution. However, this assumption may not hold true for complex
hyperspectral data, leading to reduced classification accuracy. SAM is a
geometric method for hyperspectral image classification that measures spectral
similarity by calculating the angle between pixel spectra and reference
spectra. However, it heavily relies on spectral angles and disregards spectral
variations caused by factors such as illumination changes or surface
conditions, greatly limiting its capability to accurately classify similar
spectra. Furthermore, these methods may be affected by the curse of
dimensionality, especially when the number of training samples is limited
compared to the number of spectral bands. Feature extraction and selection
play a vital role in reducing dimensionality and capturing pertinent
information, leveraging technological advancements and enhanced computational
capabilities. Researchers have developed techniques such as Principal
Component Analysis, Linear Discriminant Analysis, and Wavelet Transform to
address the challenges posed by the curse of dimensionality. However, these
methods are often sensitive to outliers leading to misleading representations
of the data. Support Vector Machines[15, 16] and Random Forests[17, 18], to
better handle the complexity and high dimensionality of HSI. However, they
primarily focus on pixel-based classification without considering spatial
information. Hyperspectral images often contain valuable spatial and
contextual information, and disregarding the spatial relationships among
neighboring pixels may result in suboptimal outcomes.

Figure 1: Comparison of the accuracy of the overlap rate of training samples
and test samples for different models.

Inspired by the great success of deep learning[19], which is capable of
finding out intrinsic and potentially valuable knowledge from the vast amounts
of pluralistic data [20], convolutional neural network (CNN) has emerged as a
popular technique for hyperspectral image classification due to their ability
to automatically learn hierarchical feature representations from the spectral
bands of the input data. Taking into account the fact that hyperspectral data
is comprised of a multitude of successive and interconnected spectral bands,
Hu et al.[22] proposed 1DCNN Classifier to capture and assimilate localized
dependencies that arise between adjacent spectral bands, thus enriching the
overall understanding and interpretation of HSI. Concurrently, Mou et al.[23]
effectively leverage the temporal aspect, treating it as a time series, and
employ RNN framework as a viable approach to model the inherent dependencies
and extract intricate spectral information. According to [24], [25], spatial
context information plays an important role in the classification of HSI.
2DCNN has been introduced into this domain, due to the powerful feature
extraction capabilities and its ability to effectively capture spatial
correlations by employing convolutional operations and weight sharing to
recognize feature distribution patterns in space. Yue et al. [27] proposed
DCNNs-LR classifier, which generated spectral feature maps by performing outer
products between different segments of the pixel’s spectra and stacked it with
spatial feature maps consisting of dominant PCA bands of image patch around
the pixel to form spatial-spectral input. Hyperspectral image can be
conceptualized as a three-dimensional voxel data, thus presenting an
opportunity to leverage the combined spectral and spatial knowledge through
3DCNN modeling. Furthermore, compared to 2DCNN, 3DCNN possess the
characteristic of scale invariance, enabling it to effectively handle HSI of
varying resolutions and sizes. Chen et al. [29] proposed a 3-D CNN-based FE
model to semantically extract spatial-spectral features by considering local
contextual information of HS images, achieving higher classification
accuracies. Li et al [30] proposed 3DCNN-Based HSI Classification which views
the HSI cube data altogether without relying on any preprocessing or post-
processing, extracting the deep spectral-spatial combined features
effectively. He et al [31] proposed M3D-DCNN model for HSI classification,
which the joint modeling of multiple spectral bands can be achieved, allowing
for the integration and fusion of diverse features across different scales.

As the mainstream backbone architecture, they have shown their powerful
ability in extracting spatially structural information and locally contextual
information from HSI. Nevertheless, they can hardly be capable of capturing
the sequence attributes well, particularly middle- and longterm dependencies.
This unavoidably meets the performance bottleneck in the HSI classification
task, especially when the classified categories are of a great variety and
extremely similar in spectral signatures. Very recently, the vision
transformer(ViT)[32] achieved the state-of-the-art effects on various vision
domain tasks, providing new insight, inspiration, and creative space on vision
related tasks. The introduction of Transformer effectively integrates spatial
and spectral information through the self-attention mechanism and enhances the
classification model’s understanding of comprehensive features. Moreover, the
self-attention mechanism enables each position (or spectral band) to interact
with all other positions, facilitating adaptive learning of feature
relationships. Hong et al. [33] proposed SpectralFormer, which is capable of
learning spectrally local sequence information from neighboring bands of HSI,
yielding groupwise spectral embeddings. However, this partitioning approach
may disrupt the continuity of spectral information, potentially impacting the
classification accuracy of the model. Furthermore, the augmented computational
costs and escalated memory demands are inevitable consequences arising from
the inherent high-dimensional characteristics and expanding data size of HSI.

Hence, taking into account the inherent long sequential nature of HSI, we put
forward a spectral-spatial axial aggregation transformer model named
SaaFormer, which aims to perform feature extraction and fusion across multiple
scales along the spectral dimension. This approach allows the model to
effectively attend to both the global information and the fine-grained
spectral details, thereby facilitating the extraction of more comprehensive
and richer feature representations from HSI. And we used axial aggregation
attention which not only allows efficient computation, but recovers large
receptive field in stand-alone attention models, that makes our model having
better generalization ability.

Figure 2:  Overview illustration of SaaFormer for HSI classification task.

The main contributions of this paper can be summarized as follows:

1)We propose a spectral-spatial axial aggregation transformer framework for
HSI classification problem, which performs multi-scale feature extraction and
fusion on the input data while utilizing spectral shift operations to ensure
information aggregation and feature extraction across different spectral
components.

2)We propose a new block-wise dataset partition approach to avoid data leakage
problem while maintaining nonlocal training samples. We re-evaluate the
classification performance of state-of-the-art backbone networks on
representative HS datasets;

3) We qualitatively and quantitatively evaluate the classification performance
of the proposed model on the six public datasets of Indian Pines, Pavia
University(PaviaU), Pavia Centre(Pavia), Salinas scene, Kennedy Space
Center(KSC) and Botswana with extensive ablation studies. The experimental
results demonstrate a significant superiority over classic transformer and
other state-of-the-art backbone networks.

The remaining of the paper is organized as follows. Section II proposes a new
evaluation method to alleviate the problem of data leakage, then we introduces
the overall framework of the proposed SaaFormer and the detailed structure of
each component. Otherwise, extensive experiments are conducted with ablation
studies and discussions. Section IV draws comprehensive conclusions.

##  II Methodology

###  II-A Block-wise Sampling

Due to difficulties in collection and costs associated with the analysis of
material’s chemical and physical properties, ground truth data is very scarce
in HSI datasets. For exploiting the spatial correlation, current paradigm
commonly employs pixel-wise random sampling to divide the datasets, where each
sample consists of a central pixel and its surrounding region. However, this
sampling approach results in a high overlap between adjacent samples, leading
to a significant similarity between the training and test sets and posing a
risk of data leakage. An experiment was conducted wherein the test set was
partitioned according to the extent of overlap with the samples in the
training set to validate this hypothesis. As can be seen from Figure 1, we
divide the test set into some parts according to the overlap rate, that is,
the overlap rate is higher than 50%, the overlap rate is between 0%-50% and no
overlap at all, and observe the change of accuracy on several different
models. The experimental results indicated that test regions with higher
overlap often exhibited higher accuracy. Therefore, this random sampling
approach is prone to data leakage issues.

Therefore, we proposed an innovative block-wise sampling approach, as
illustrated in Fig. 3, the hyperspectral dataset is divided into non-
overlapping blocks of equal size, followed by the insertion of gaps to create
distinct training and test sets. This partitioning method effectively
addresses the issue of data leakage and enables more accurate performance
evaluation of the model. By avoiding overlapping regions, we ensure the
independence between the training and test sets, thereby mitigating the risk
of information leakage. This novel block-wise sampling technique provides a
more reliable and precise approach for analyzing HSI.

(a)

(b)

(c)

(d)

Figure 3: Comparison diagram of PaviaU dataset by block-wise sampling and
random sampling

Taking into account each pixel of HSI represents the spectral curve of the
substance here, which contains the sequence properties of spectral
characteristics. This inspired us to adopt Transformer-based framework from a
sequential perspective. The proposed SaaFormer exhibits enhanced
generalization capability.

###  II-B Overview of SaaFormer

We put forward a multi-level axial transformer model, named SaaFormer, which
aims to perform feature extraction and fusion across multiple scales along the
spectral dimension. Fig. 2 illustrates an overview of SaaFormer, which
comprises two primary components: axial transformer block and multi-level
spectral extraction. Given a HS image (i.e. a hyperspectral dataset)
$X\in\mathbb{R}^{H\times W\times C}$, $x_{0}\in\mathbb{R}^{h\times w\times C}$
represents one of the samples, the input sample $x_{0}$ undergo dimensionality
transformation through linear embedding layer, followed by a multi-level
spectral feature partitioning process. And the resulting feature maps of
different scales are individually fed into axial transformer modules for
feature extraction, enabling the extraction of relevant features at various
spatial and spectral levels. These features are then fused to incorporate
information from different scales. Finally, the fused features are input to
the classification head to generate the output results. A detailed
introduction of each structure will be carried out in the next subsection.

(a)

(b)

Figure 4: Diagram of the Axial Transformer Block in the proposed SaaFormer.

###  II-C Axial Transformer Block

While the self-attention mechanism has showcased impressive performance across
a range of tasks, its limitation of increased parameter and computational
complexity cannot be disregarded, particularly when dealing with lengthy
sequential data. Presently, there exist numerous lightweight modifications
being explored to improve the self-attention mechanism[34, 35, 36, 37, 38,
39], with the axial attention mechanism standing out as a promising candidate
in various tasks[40, 41, 42, 43]. Hence, in addressing the challenge of HSI
classification, we incorporate the squeeze-enhanced axial attention
mechanism[43]. However, due to its imperfect alignment with the specific
characteristics of high-dimensional spectral data, we have undertaken further
improvements by designing the spatial-spectral axial aggregation attention
mechanism.

Our proposed SaaFormer uses an improved spectral-spatial axial attention
mechanism for feature extraction which can adaptively attend to different
spectral channels. And by considering the interrelationships among spectral
bands at each pixel location and aggregating features in the spatial domain,
it enables the extraction of meaningful context and relevance in HSI. In
contrast to conventional attention mechanisms, its efficiency enables us to
build models to learn long range continuous spectral curve characteristics.
Moveover, through the filtration of irrelevant information in input samples,
it enhances the utilization efficiency of crucial features. Consequently, it
reduces the model’s sensitivity to noise and interference, ultimately
enhancing its generalization capability.

Firstly, $\mathbf{q}=W_{q}\mathbf{x}$, $\mathbf{k}=W_{k}\mathbf{x}$, and
$\mathbf{v}=W_{v}\mathbf{x}$ are all linear projections of the input feature
map $\mathbf{x}\in\mathbb{R}^{h\times w\times c}$ where
$W_{q},W_{k},W_{v}\in\mathbb{R}^{C_{qkv}\times c}$ are all learnable matrices.
Then, we achieve horizontal extrusion by taking the maximum value of the query
feature map in the width direction. In the same way, we get the vertical
squeeze on the vertical direction. i.e. $\mathbb{q}_{row}={\rm
max}_{w}(\mathbb{q}),\mathbb{q}_{col}={\rm max}_{h}(\mathbb{q})$. The
aggregation operation on $\mathbf{q}$ also repeats on $\mathbf{k}$ and
$\mathbf{v}$, so we finally get
$\mathbf{q}_{row},\mathbf{k}_{row},\mathbf{v}_{row}\in\mathbb{R}^{h\times
1\times C_{qkv}}$,
$\mathbf{q}_{col},\mathbf{k}_{col},\mathbf{v}_{col}\in\mathbb{R}^{1\times
w\times C_{qkv}}$.

The aggregation operation in the proposed approach consolidates global
information by compressing it into a single axis, as demonstrated by Equation
1. By condensing the global information into a single axis, the aggregation
operation enables more efficient and effective extraction of global semantic
features.

$\begin{split}\mathbf{h}_{(i,j)}=\sum_{p=1}^{h}\operatorname{softmax}_{p}\left(\mathbf{q}_{(h)i}^{\top}\mathbf{k}_{(h)p}\right)\mathbf{v}_{(h)p}\\\
+\sum_{p=1}^{w}\operatorname{softmax}_{p}\left(\mathbf{q}_{(v)j}^{\top}\mathbf{k}_{(v)p}\right)\mathbf{v}_{(v)p}\end{split}$
(1)

It is worth mentioning that the spatial-spectral axial aggretation attention
mechanism lacks explicit position information, which limits its ability to
effectively utilize sequential context. To address this limitation, we follow
[43, 44, 35, 41] by including a relative position bias to each head. And apply
the learnable relative position coding directly to query, key, value. By
incorporating learnable relative position embeddings, denoted as
$\mathbb{r}_{row}^{q},\mathbb{r}_{row}^{k},\mathbb{r}_{row}^{v}\in\mathbb{R}^{h\times
1\times C_{qkv}}$, we enable the attention mechanism to be aware of the
positional relationships in the input feature map. This augmentation allows
the aggregation attention module to leverage the positional cues and enhance
its capability to capture sequential information effectively. Perform the same
operation in the vertical direction to get
$\mathbb{r}_{col}^{q},\mathbb{r}_{col}^{k},r_{row}^{v}\in\mathbb{R}^{1\times
w\times C_{qkv}}$.Thus, the positional-aware aggretation attention can be
expressed as Equation 2.

$\begin{split}\begin{aligned}
\mathbf{h}_{(i,j)}&=\sum_{p=1}^{h}\operatorname{softmax}_{p}(\left(\mathbf{q}_{(row)i}+\mathbf{r}_{(row)i}^{q}\right)^{\top}(\mathbf{k}_{(row)p}\\\
&+\mathbf{r}_{(row)p}^{k}))(\mathbf{v}_{(row)p}+\mathbf{r}_{(row)}^{v})\\\
&+\sum_{p=1}^{w}\operatorname{softmax}_{p}(\left(\mathbf{q}_{(col)j}+\mathbf{r}_{(col)j}^{q}\right)^{\top}(\mathbf{k}_{(col)p}\\\
&+\mathbf{r}_{(col)p}^{k}))(\mathbf{v}_{(col)p}+\mathbf{r}_{(col)}^{v})\end{aligned}\end{split}$
(2)

As is shown in the upper path of Fig. LABEL:aatten, the attention compresses
the feature map to a vertical or horizontal axis and compute self-attention
over a single axis at a time, then, add a horizontal and a vertical axial
attention module to obtain the global receptive field. Its efficiency enables
us to attend over large regions and build models to learn long range or even
global interactions. Otherwise, neighboring pixels are highly interdependent
because most of the land covers are much bigger than the size of the pixel. It
has been prevalent to utilize spatial information along with spectral data
during analysis to combat the problem of high dimensionality. Therefore, our
model integrates spatial constraint information while considering the
characteristics of spectral sequences making it well-applicable to the highly
accurate and fine classification.

The axial attention mechanism will pay more attention to the important
spectral characteristics of the sample space and extract global spectral curve
efficiently which may sacrifice the spatial local details. Hence, we apply an
auxiliary spatial feature extraction module which is made up of $3\times 3$
convolution and batch normalization to supplementary the spatial details and
maintain the spatial structure of the sample. As is shown in the down path of
Fig. LABEL:aatten, by using convolution layers, auxiliary spatical local
details can be aggregated. Finally, the enhancement feature will be fused with
the feature given by aggretation attention.

Figure 5:  Diagram of the Multi-level Spectral-spatial Extraction Structure.

###  II-D Multi-level Spectral-spatial Extraction

Multi-scale feature extraction has demonstrated significant benefits in image
processing, as it enables models to capture information at various levels and
granularities by extracting features, thereby enhancing their robustness.
[46]. Meanwhile, in view of the diverse absorption or reflection
characteristics of different materials for different wavelengths of light, as
well as the potential presence of noise interference in obtained spectral
curves, we propose a hierarchical feature extraction structure. It divides the
input samples into non-overlapping continuous spectral segments of varying
lengths, which are considered as spectral partition, as illustrated in Fig 5,
the feature map $x\in\mathbb{R}^{h\times w\times C_{qkv}}$ is divided into
spectral partitions of different lengths(i.e. $c_{1},c_{2},c_{3}$) in spectral
dimensions, ans axial transformer blocks is performed within each spectral
partition, followed by feature fusion. This approach reduces redundancy in HSI
and enables better utilization of spectral information, thereby enhancing the
model’s perception of different material components and its capability to
capture spectral details.

However, this approach has certain limitations as the connections across
different partitions is restricted. To introduce cross spectral partitions
connections while maintaining the efficient computation of non-overlapping
partitions, we propose a shifted spectral partition operation, as illustrated
in Fig 6. Specifically, after an initial feature extraction step, we perform a
subsequent re-aggregation of information between different segments in the
spectral dimension, followed by another round of feature extraction. The first
block uses a regular spectral partitioning strategy to get
$x_{1},x_{2}\in\mathbb{R}^{h\times w\times c_{1}}$. Then, $x_{1}$ is divided
into part A and part B, and $x_{2}$ is divided into part C and part D where
$A,B,C,D\in\mathbb{R}^{h\times w\times\frac{c_{1}}{2}}$. The part B and part C
are formed into new $x^{\prime}_{1}$. This operation enhances the capacity to
capture fine-grained spectral details and exploit the interdependencies among
different spectral partitions, thereby effectively leveraging the abundant
spectral information inherent.

Figure 6:  Diagram of the shifted spectral partitioning structure.

Due to the continuity of spectral curves in HS images, the spectral shift
operation aggregates the first and last spectral segments. However, from a
physical perspective, the first and last segments are completely unrelated and
represent distinct spectral features. To prevent erroneous feature extraction
and misclassification, we apply a masking process to the irrelevant aggregated
spectral segments, considering them as invalid information. This prevents the
model from erroneously utilizing these segments. Through this masking
operation, we preserve meaningful spectral features and ensure accurate
classification and recognition of relevant features.

With the shifted spectral partitioning approach, the specific flowchart is
shown in Fig. LABEL:aatten and consecutive Axial attention blocks are computed
as

$\begin{array}[]{l}\hat{f}^{l}=\operatorname{AxialAtten}\left(\mathrm{LN}\left(f^{l-1}\right)\right)+f^{l-1}\\\
f^{l}=\operatorname{FFN}\left(\mathrm{LN}\left(\hat{f}^{l}\right)\right)+\hat{f}^{l}\\\
\hat{f}^{l+1}=\operatorname{SAxialAtten}\left(\mathrm{LN}\left(f^{l}\right)\right)+f^{l}\\\
f^{l+1}=\operatorname{FFN}\left(\mathrm{LN}\left(\hat{f}^{l+1}\right)\right)+\hat{f}^{l+1}\end{array}$
(3)

##  III Experiments

In this part, we first introduce several well-known HSI datasets, and give the
datasets division method we proposed. Then compare SaaFormer with other state-
of-the-art methods in the new division method. Finally, extensive experiments
are conducted with ablation analysis to assess the HSI classification
performance of the proposed model.

###  III-A Data Description

Experiments are conducted on six hyperspectral datasets111All these datasets
can be obtained at
https://www.ehu.eus/ccwintco/index.php/Hyperspectral_Remote_Sensing_Scenes.
The detailed descriptions of these datasets are listed as follows.

  * •

The Indian Pines dataset was collected by the airborne visible/infrared
imaging spectrometer sensor (AVIRIS) over the Indian Pines test site in June
1992. It consists of 145 × 145 pixels and has 220 spectral bands. The dataset
has 20 m per pixel spatial resolutions and 10 nm spectral resolutions covering
a spectrum range of 200–2400 nm. In the experiment, after removing 20 noisy
and water absorption bands, 200 spectral bands are retained, i.e., 1–103,
109–149, and 164–219. The ground truth is composed of 16 vegetation classes.

  * •

The Pavia Centre and Pavia University dataset are acquired by the ROSIS sensor
during a flight campaign over Pavia, nothern Italy. The number of spectral
bands is 102 for Pavia Centre and 103 for Pavia University. Pavia Centre is a
1096*1096 pixels image, and Pavia University is 610*610 pixels, but some of
the samples in both images contain no information and have to be discarded
before the analysis. The geometric resolution is 1.3 meters. Both image
groundtruths differenciate 9 classes each.

  * •

The Salinas scene dataset was collected by the 224-band AVIRIS sensor over
Salinas Valley, California, and is characterized by high spatial resolution
(3.7-meter pixels). The area covered comprises 512 lines by 217 samples. As
with Indian Pines scene, this dataset discarded the 20 water absorption bands,
in this case bands: [108-112], [154-167], 224. This image was available only
as at-sensor radiance data.

  * •

The Kennedy Space Center (KSC) dataset was acquired by the NASA AVIRIS
(Airborne Visible/Infrared Imaging Spectrometer) instrument over the Kennedy
Space Center, Florida, on March 23, 1996. AVIRIS acquires data in 224 bands of
10 nm width with center wavelengths from 400-2500 nm. The KSC data, acquired
from an altitude of approximately 20 km, have a spatial resolution of 18 m.
After removing water absorption and low SNR bands, 176 bands were used for the
analysis. For classification purposes, 13 classes representing the various
land cover types that occur in this environment were defined for the site.

  * •

The Botswana dataset was acquired by the NASA EO-1 satellite over the Okavango
Delta, Botswana in 2001-2004. The Hyperion sensor on EO-1 acquires data at 30
m pixel resolution over a 7.7 km strip in 242 bands covering the 400-2500 nm
portion of the spectrum in 10 nm windows. Uncalibrated and noisy bands that
cover water absorption features were removed, and the remaining 145 bands were
included as candidate features: [10-55, 82-97, 102-119, 134-164, 187-220]. The
data analyzed in this study, acquired May 31, 2001, consist of observations
from 14 identified classes representing the land cover types in seasonal
swamps, occasional swamps, and drier woodlands located in the distal portion
of the Delta.

###  III-B Data Division

####  III-B1 Random Sampling

The method of random sampling in hyperspectral datasets refers to treating
each valid pixel along with its surrounding rectangular window of the same
size (e.g. 5x5) as a sample within a complete hyperspectral image.
Subsequently, all samples are randomly divided into train and test sets based
on a predetermined ratio. This identical training and testing set division
ratio is applied to all six hyperspectral datasets. 5% samples from each class
are randomly selected as the training set, and the remaining samples are used
for test.

####  III-B2 Block-wise Sampling

Block sampling method refers to dividing a complete hyperspectral dataset into
non-oversectangle areas according to a certain size, and each area is divided
into a training set or test set as a collation. In our experiment, We employ a
gapped-block method to partition the training set and the test set, select a
smaller sample area as the training set, and use the rest as the test set. The
selection of segmented size is to make each class of sample exist in the
minimum value of the training set and test set.

###  III-C Experimental Setting

1) Evaluation Metrics: We evaluate the classification performance of each
model quantitatively in terms of three commonly used indices, i.e., OA, AA,
and $\kappa$. Here, OA represents the number of correctly classified samples
out of the total test samples; AA represents the average of class-wise
classification accuracies; and Kappa is a metric of statistical measurement
which provides mutual information regarding a strong agreement between the
ground truth map and classification map.Expressed by formula as

$\displaystyle OA=\frac{1}{N}\sum_{i=1}^{row}x_{ii}$ $\displaystyle
AA=\frac{1}{N}\sum_{i=1}^{N}\frac{x_{ii}}{x_{i+}}$ $\displaystyle
Kappa=\frac{N\sum_{i=1}^{row}x_{ii}-\sum_{i=1}^{row}\left(x_{i+}\times
x_{+i}\right)}{N^{2}-\sum_{i=1}^{row}\left(x_{i+}\times x_{+i}\right)}$

where N is the total number of categories, and $x_{ij}$ indicates that the
real category is $i$, which is predicted to be the number of samples in
category $j$. And $r$ represents the number of samples correctly classified in
each category, $x_{i+}$ is the total number of samples in that category in the
predicted sample, and $x_{+i}$ is the total number of samples in that category
in the label.

2) Comparison With State-of-the-Art Backbone Networks: In order to verify the
classification performance of the proposed SaaFormer, five representative
methods based on convolutional neural network and recurrent neural network for
HSI classification, 1DCNN Classifier[22], 2DCNN-LR[29], 3-DCNN-based FE
model[30], M3D-DCNN[31], RNN[23] and two models based on Transformer, the
classical ViT[32] and SpectralFormer[33] are used as the comparison methods.
The parameter configurations of these compared methods are detailed as
follows.

  * •

For the 1DCNN Classifier, as suggested by the literature[22], one
convolutional block is defined as the basic network unit, including a set of
1D convolutional filters with the output size of 20, a maxpooling layer, and a
hyperbolic tangent function tanh as the activation function. Two fully
connected layers are finally added on the top of the 1DCNN.

  * •

For 2DCNN-LR, which has three 2D convolutional blocks and a softmax layer.
Each convolutional block of 2D CNN consists of a 2D conventional layer, a BN
layer, a max-pooling layer,and an ReLU activation function.

  * •

For 3-DCNN-based FE model, the spatial window size of 3D input and the
parameters of the 3D convolution kernel is set to the default value in [30].

  * •

For M3D-DCNN, as suggested by literature[31], uses 3D convolution kernel of
different sizes ($1\times 1\times 1,3\times 1\times 1,5\times 1\times
1,11\times 1\times 1$) for multi-scale feature extraction at the same time,
and then adds the extracted features.

  * •

For RNN framework, followed by [23], there are two recurrent layers with the
gated recurrent unit (GRU). Each of them has 128 neuron units.

  * •

For the Transformer model, we follow the ViT network architecture[32], i.e.,
The original self-attention mechanism is used for feature extraction for
input. In detail, four encoder blocks with eight attention heads are used in
the ViT-based network for HSI classification.

  * •

For SpectralFormer, the input size and the hyperparameters are set to the
default value in [33]

  * •

For SaaFormer, the embedded spectrum with 128 units is fed into two cascaded
transformer encoder blocks for HS image classification. Each encoder block
consists of a four-head axial aggregation attention layer, an MLP, and a GELU
nonlinear activation layer. The dropout layer is employed after encoding
positional embeddings and in MLPs for inhibiting 10% neurons. the patch size
is empirically set as $5\times 5$.

3) Implementation Details: Our proposed model was implemented on the PyTorch
platform using a workstation with i7-6850K CPU, 128-GB RAM, and an NVIDIA GTX
3090Ti 11-GB GPU. We adopt the Adam optimizer [45] with a minibatch size of
64. The learning rate is initialized with 5e-4 and decayed by multiplying a
factor of 0.9 after each one-tenth of total epochs. We roughly set the epochs
on the two datasets to 200.

###  III-D Quantitative Results and Analysis of HSI Classification Models

To verify the classification performance of the proposed model, experiments
are executed on aforementioned datasets which are set to be compared by block
division and random division respectively. The OA, AA, and kappa of
classification results are used as evaluating indicators under different
training set and test set.

(a)

(b)

(c)

Figure 7: Comparison of the results of different models on the three
indicators of OA, AA and Kappa on the six datasets.The lighter-colored bars
represent classification metrics of random sampling method. The darker-colored
bars represent metrics of block-wise sampling method.

As shown in Fig.7, three statistical graphs illustrate the comparative results
of different datasets using block-wise sampling and random sampling methods.
The dashed lines represent the accuracy of our proposed model under block-wise
sampling and random sampling. From the graph, it can be observed that the
light-colored bars of each model are higher than the dark-colored bars,
indicating the potential risk of data leakage when using random sampling.
Particularly, the 2DCNN model exhibits significant differences, indicating
overfitting issues and poor generalization capability. In contrast, SaaFormer
shows smaller variances and achieves the highest accuracy under block-wise
sampling among all the six tested datasets, demonstrating its excellent
generalization capability. Moreover, SaaFormer achieves the best performance
under random sampling method on Botswana, KSC, IndianPines, and Salinas
datasets. On PaviaU and PaviaC, the 2DCNN model has higher accuracy than
SaaFormer which is caused of serious overfitting.

Table I and Table II provide detailed data on the three evaluation metrics
obtained using different partitioning methods on the PaviaU and Botswana
datasets. The bold numbers indicate the best accuracy achieved. It is evident
that our proposed model achieves the best accuracy in almost all cases.
Furthermore, it is encouraging to note that the variations in the three
evaluation metrics across different partitioning methods are minimal, even
within a range of 1%. This indicates that our model not only maintains high
accuracy across different datasets but also exhibits consistent performance
under different partitioning methods. This demonstrates the excellent
generalization capability and robustness of the proposed SaaFormer.

TABLE I: Comparison of classification accuracy of random sampling and block division on PaviaU dataset | OA | AA | Kappa
---|---|---|---
Random | Block | Random | Block | Random | Block
1DCNN | 0.8820 | 0.7868 | 0.8619 | 0.8405 | 0.8416 | 0.6934
2DCNN | $\mathbb{0.9935}$ | 0.7307 | $\mathbb{0.9847}$ | 0.7774 | 0.8914 | 0.6216
3DCNN | 0.9121 | 0.9391 | 0.9130 | 0.9086 | 0.8825 | 0.9026
MS3DCNN | 0.9736 | 0.9673 | 0.9580 | 0.9261 | 0.9651 | 0.9477
RNN | 0.9057 | 0.9223 | 0.8833 | 0.8774 | 0.8743 | 0.8767
ViT | 0.9480 | 0.9462 | 0.9331 | 0.8993 | 0.9312 | 0.9160
SpectralFormer | 0.9625 | 0.8992 | 0.9499 | 0.8188 | 0.9503 | 0.8405
SaaFormer | 0.9875 | $\mathbb{0.9828}$ | 0.9795 | $\mathbb{0.9649}$ | $\mathbb{0.9833}$ | $\mathbb{0.9726}$
TABLE II: Comparison of classification accuracy of random sampling and block-wise sampling on Botswana dataset | OA | AA | Kappa
---|---|---|---
Random | Block | Random | Block | Random | Block
1DCNN | 0.8869 | 0.8151 | 0.8977 | 0.8611 | 0.8775 | 0.7973
2DCNN | 0.9818 | 0.5968 | 0.9757 | 0.6169 | 0.9803 | 0.5645
3DCNN | 0.9465 | 0.7810 | 0.9475 | 0.7585 | 0.9420 | 0.7609
MS3DCNN | 0.9345 | 0.6529 | 0.9306 | 0.6786 | 0.9290 | 0.6233
RNN | 0.8700 | 0.7406 | 0.8779 | 0.8094 | 0.8592 | 0.7193
ViT | 0.9102 | 0.8296 | 0.9136 | 0.7870 | 0.9028 | 0.8141
SpectralFormer | 0.9782 | 0.6700 | 0.9808 | 0.6770 | 0.9764 | 0.6431
SaaFormer | $\mathbb{0.9837}$ | $\mathbb{0.9223}$ | $\mathbb{0.9821}$ | $\mathbf{0.9232}$ | $\mathbb{0.9824}$ | $\mathbb{0.9144}$

Figure 8: Spatial distribution of testing sets using block-wise sampling, and
the classification maps obtained by different models on the Salinas dataset

Figure 9: Spatial distribution of testing sets using random sampling, and the
classification maps obtained by different models on the Salinas dataset

###  III-E Ablation Study

We investigate the performance gain of the proposed SaaFormer in terms of
classification accuracies by applying different structures (i.e. Multi-level
Spectral Extraction). For that, we conduct extensive ablation experiments on
the six datasets to verify the effectiveness of the structure in SaaFormer for
HSI classification applications, as listed in Table III.

TABLE III: ABLATION ANALYS IS OF THE PROPOSED SaaFormer WITH A COMBINATION OF DIFFERENT STRUCTURE ON THE SiX DATASETS. |  | Random Sampling | Block-wise Sampling
---|---|---|---
|  Multi-level Spectral Extraction | OA | AA | Kappa | OA | AA | Kappa
| $\checkmark$ | 0.9837 | 0.9821 | 0.9824 | 0.9223 | 0.9232 | 0.9144
Botswana | $\times$ | 0.9779 | 0.9804 | 0.9774 | 0.9041 | 0.9184 | 0.8947
| $\checkmark$ | 0.9676 | 0.9424 | 0.9640 | 0.8339 | 0.7740 | 0.8069
KSC | $\times$ | 0.9521 | 0.9339 | 0.9530 | 0.8276 | 0.7103 | 0.8031
| $\checkmark$ | 0.9727 | 0.9626 | 0.9683 | 0.8947 | 0.8713 | 0.8797
IndianPines | $\times$ | 0.9705 | 0.9612 | 0.9663 | 0.8807 | 0.8619 | 0.8698
| $\checkmark$ | 0.9875 | 0.9795 | 0.9833 | 0.9828 | 0.9649 | 0.9726
PaviaU | $\times$ | 0.9757 | 0.9772 | 0.9503 | 0.9824 | 0.9573 | 0.9720
| $\checkmark$ | 0.9944 | 0.9810 | 0.9921 | 0.9828 | 0.9649 | 0.9820
PaviaC | $\times$ | 0.9872 | 0.9776 | 0.9890 | 0.9765 | 0.9275 | 0.9649
| $\checkmark$ | $\mathbb{0.9928}$ | 0.9967 | $\mathbb{0.9920}$ | 0.9253 | 0.9653 | 0.9163
Salinas | $\times$ | $\mathbb{0.9559}$ | 0.9804 | $\mathbb{0.9510}$ | 0.9246 | 0.9637 | 0.9154

The experimental results from Table III demonstrate that the introduction of
the multi-level spectral extraction structure consistently leads to higher
accuracy in our model across different datasets (beyond around 3% OAs on the
Salinas dataset), independent of the data partitioning schemes. This finding
highlights the effectiveness of the multi-level spectral extraction structure
in capturing underlying patterns and correlations within the spectral curves,
enabling the model to exhibit strong generalization capabilities and
stability.

Figure 10: Spatial distribution of testing sets using block-wise sampling, and
the classification maps obtained by different models on the PaviaU dataset

Figure 11: Spatial distribution of testing sets using random sampling, and the
classification maps obtained by different models on the PaviaU dataset

###  III-F Visual Evaluation

We make a qualitative evaluation by visualizing the classification maps
obtained by different methods. Figs. 8 and 9 provide the obtained results for
Salinas dataset and Figs. 10 and 11 are the results for PaviaU dataset,
respectively.

From the visualization of the datasets, it can be observed that our proposed
SaaFormer obtains desirable classification maps In the case of block-wise
sampling and random sampling, especially in terms of texture and edge details.
Furthermore, we selected a region of interest (ROI) zoomed in two times to
highlight the differences in classification maps between different models,
further evaluating classification performance more intuitively. As can be seen
from these ROIs, a remarkable phenomenon is that our model shows more
realistic and finer details. Therefore, the visualization of the
classification results clearly demonstrates that our proposed SaaFormer,
incorporating both axial Transformer block and multi-level spectral extraction
structure, achieves remarkably favorable classification maps in terms of
texture and edge details.

##  IV Conclusion

HSI are typically represented as a data cube with spatial–spectral
information, which can be generally regarded as a sequence of data along the
spectral dimension. Due to the difficulty of obtaining hyperspectral data and
the limited data set, we find that there is a serious data leakage problem in
the evaluation method of the existing HSI classification model. This leads to
an inaccurate evaluation of the generalization of the model, thus, we propose
a block-wise sampling method to avoid data leakage at sample points while
maintaining non-local training samples. Otherwise, different from CNNs that
focus mainly on contextual information modeling, transformers have been proven
to be a powerful architecture in characterizing the sequential properties
globally. However, the classic ViT networks inevitably suffer from performance
degradation when processing HS-like data. This might be explained well by the
fact that ViT fails to model locally detailed spectral discrepancies. To this
end, in this paper, we proposed a novel transformer-based model, named
SaaFormer, which demonstrates robust generalization ability. And it comprises
two primary components: axial aggregation attention and multi-level spectral
extraction. The results on six publicly available datasets demonstrate that
our model exhibits comparable performance when using random sampling, while
significantly outperforming other methods when employing block-wise sampling
partition.

## References

  * [1] D. Hong et al., “Interpretable hyperspectral artificial intelligence: When nonconvex modeling meets hyperspectral remote sensing,” IEEE Geosci.Remote Sens. Mag., vol. 9, no. 2, pp. 52–87, Jun. 2021.
  * [2] M. Wang, Q. Wang, J. Chanussot, and D. Hong, “$L_{0}-L_{1}$ hybrid total variation regularization and its applications on hyperspectral image mixed noise removal and compressed sensing,” IEEE Trans. Geosci. Remote Sens., vol. 59, no. 9, pp. 7695–7710, Sep. 2021, doi:10.1109/TGRS.2021.3055516.
  * [3] J. Peng et al., “Low-rank and sparse representation for hyperspectral image processing: A review,” IEEE Geosci. Remote Sens. Mag., early access, Jun. 10, 2021, doi: 10.1109/MGRS.2021.3075491.
  * [4] F. Luo, T. Guo, Z. Lin, J. Ren, and X. Zhou, “Semisupervised hypergraph discriminant learning for dimensionality reduction of hyperspectral image,” IEEE J. Sel. Topics Appl. Earth Observ. Remote Sens., vol. 13, pp. 4242–4256, 2020.
  * [5] J. Yao, D. Meng, Q. Zhao, W. Cao, and Z. Xu, “Nonconvex-sparsity and nonlocal-smoothness-based blind hyperspectral unmixing,” IEEE Trans. Image Process., vol. 28, no. 6, pp. 2991–3006, Jun. 2019.
  * [6] Y. Yuan, Z. Zhang, and Q. Wang, “Improved collaborative non-negative matrix factorization and total variation for hyperspectral unmixing,” IEEE J. Sel. Topics Appl. Earth Observ. Remote Sens., vol. 13, pp. 998–1010, 2020.
  * [7] L. Gao, Z. Han, D. Hong, B. Zhang, and J. Chanussot, “CyCUnet: Cycle-consistency unmixing network by learning cascaded autoencoders,” IEEE Trans. Geosci. Remote Sens., early access, Mar. 23, 2021, doi: 10.1109/TGRS.2021.3064958.
  * [8] D. Hong et al., “Endmember-guided unmixing network (EGU-Net): A general deep learning framework for self-supervised hyperspectral unmixing,” IEEE Trans. Neural Netw. Learn. Syst., early access, May 28, 2021, doi: 10.1109/TNNLS.2021.3082289.
  * [9] J. Peng, W. Sun, and Q. Du, “Self-paced joint sparse representation for the classification of hyperspectral images,” IEEE Trans. Geosci. Remote Sens., vol. 57, no. 2, pp. 1183–1194, Feb. 2019.
  * [10] Q. Li, B. Zheng, B. Tu, J. Wang, and C. Zhou, “Ensemble EMDbased spectral-spatial feature extraction for hyperspectral image classification,” IEEE J. Sel. Topics Appl. Earth Observ. Remote Sens., vol. 13, pp. 5134–5148, 2020.
  * [11] Richards J A, Jia X, ”Using suitable neighbors to augment the training set in hyperspectral maximum likelihood classification,” IEEE Geoscience and Remote Sensing Letters, vol. 5, no. 4, pp. 774-777, 2008.
  * [12] Peng J, Li L, Tang Y Y, ”Maximum likelihood estimation-based joint sparse representation for the classification of hyperspectral remote sensing images,” IEEE transactions on neural networks and learning systems, vol. 30, no. 6, pp. 1790-1802, 2018.
  * [13] Kuching S, ”The performance of maximum likelihood, spectral angle mapper, neural network and decision tree classifiers in hyperspectral image analysis,” Journal of Computer Science, vol. 3, no. 6, pp. 419-423, 2007.
  * [14] Chakravarty S, Paikaray B K, Mishra R, et al., ”Hyperspectral image classification using spectral angle mapper,” 2021 IEEE International Women in Engineering (WIE) Conference on Electrical and Computer Engineering (WIECON-ECE). IEEE, pp. 87-90, 2021.
  * [15] F. Melgani and L. Bruzzone, ”Classification of hyperspectral remote sensing images with support vector machines,” IEEE Transactions on Geoscience and Remote Sensing, vol. 42, no. 8, pp. 1778–1790, Aug 2004\.
  * [16] Qingwang Wang, Yanfeng Gu, and Devis Tuia, ”Discriminative multiple kernel learning for hyperspectral image classification,” IEEE Transactions on Geoscience and Remote Sensing, vol. 54, no. 7, pp. 3912–3927, 2016.
  * [17] Qi Cheng, P.K. Varshney, and M.K. Arora, ”Logistic regression for feature selection and soft classification of remote sensing data,” IEEE Geoscience and Remote Sensing Letters, vol. 3, no. 4, pp. 491–494, Oct 2006.
  * [18] M. Khodadadzadeh, Jun Li, A. Plaza, and J.M. Bioucas-Dias, ”A subspace-based multinomial logistic regression for hyperspectral image classification,” IEEE Geoscience and Remote Sensing Letters, vol. 11, no. 12, pp. 2105–2109, Dec 2014.
  * [19] Yoshua Bengio, Aaron Courville, and Pierre Vincent, ”Representation learning: A review and new perspectives,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 35, no. 8, pp. 1798–1828, 2013\.
  * [20] Y. LeCun, Y. Bengio, and G. Hinton, “Deep learning,” Nature, vol. 521, pp. 436–444, May 2015.
  * [21] Yushi Chen, Zhouhan Lin, Xing Zhao, Gang Wang, and Yanfeng Gu, ”Deep learning-based classifi- cation of hyperspectral data,” IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, vol. 7, no .6, pp. 2094–2107, June 2014.
  * [22] Wei Hu, Yangyu Huang, Li Wei, Fan Zhang, and Hengchao Li, ”Deep convolutional neural networks for hyperspectral image classification,” Journal of Sensors, 501:258619, 2015.
  * [23] L. Mou, P. Ghamisi, and X. X. Zhu, ”Deep recurrent neural networks for hyperspectral image classification,” IEEE Transactions on Geoscience and Remote Sensing, vol. 55, no. 7, pp. 3639–3655, July 2017.
  * [24] Y. Yan, Y. Zhao, H.-f. Xue, X.-d. Kou, and Y. Liu, “Integration of spatial-spectral information for hyperspectral image classification,” in 2010 Second IITA International Conference on Geoscience and Remote Sensing, pp. 242–245, Qingdao, 2010.
  * [25] X. Xu, J. Li, Y. Zhang, and S. Li, “A subpixel spatial-spectral feature mining for hyperspectral image classification,” in IGARSS 2018 - 2018 IEEE International Geoscience and Remote Sensing Symposium, pp. 8476–8479, Valencia, Spain, 2018.
  * [26] Michael D Farrell and Russell M Mersereau, ”On the impact of PCA dimension reduction for hyperspectral detection of difficult targets,” IEEE Geoscience and Remote Sensing Letters, vol. 2, no. 2, pp. 192–195, 2005.
  * [27] Jun Yue, Wenzhi Zhao, Shanjun Mao, and Hui Liu, ”Spectral–spatial classification of hyperspectral images using deep convolutional neural networks,” Remote Sensing Letters, vol. 6, no. 6, pp. 468–477, 2015.
  * [28] Heming Liang and Qi Li, ”Hyperspectral imagery classification using sparse representations of convolutional neural network features,” Remote Sensing, vol. 8, no. 2, 2016.
  * [29] Y. Chen, H. Jiang, C. Li, X. Jia, and P. Ghamisi, “Deep feature extraction and classification of hyperspectral images based on convolutional neural networks,” IEEE Trans. Geosci. Remote Sens., vol. 54, no. 10, pp. 6232–6251, Oct. 2016.
  * [30] Li Y, Zhang H, Shen Q, ”Spectral–spatial classification of hyperspectral imagery with 3D convolutional neural network,” Remote Sensing, vol. 9, no. 1, pp. 67, 2017.
  * [31] M. He, B. Li and H. Chen, ”Multi-scale 3D deep convolutional neural network for hyperspectral image classification,” 2017 IEEE International Conference on Image Processing (ICIP), Beijing, China, 2017, pp. 3904-3908, doi: 10.1109/ICIP.2017.8297014.
  * [32] Dosovitskiy A, Beyer L, Kolesnikov A, et al., ”An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929, 2020.
  * [33] D. Hong et al., ”SpectralFormer: Rethinking Hyperspectral Image Classification With Transformers,” in IEEE Transactions on Geoscience and Remote Sensing, vol. 60, pp. 1-15, 2022, Art no. 5518615, doi: 10.1109/TGRS.2021.3130716.
  * [34] Shen, Zhuoran, et al., ”Efficient attention: Attention with linear complexities.” Proceedings of the IEEE/CVF winter conference on applications of computer vision. 2021.
  * [35] Liu, Ze, et al., ”Swin transformer: Hierarchical vision transformer using shifted windows.” Proceedings of the IEEE/CVF international conference on computer vision. 2021.
  * [36] Dong, Xiaoyi, et al., ”Cswin transformer: A general vision transformer backbone with cross-shaped windows.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.
  * [37] Han, Kai, et al. ”Transformer in transformer.” Advances in Neural Information Processing Systems 34 (2021): 15908-15919.
  * [38] Zhu, Chen, et al. ”Long-short transformer: Efficient transformers for language and vision.” Advances in neural information processing systems 34 (2021): 17723-17736.
  * [39] Huang, Zilong, et al., ”Ccnet: Criss-cross attention for semantic segmentation.” Proceedings of the IEEE/CVF international conference on computer vision. 2019.
  * [40] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans, ”Axial attention in multidimensional transformers,” arXiv preprint arXiv:1912.12180, 2019.
  * [41] Wang H, Zhu Y, Green B, et al., ”Axial-deeplab: Stand-alone axial-attention for panoptic segmentation,” Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part IV. Cham: Springer International Publishing, 2020: 108-126.
  * [42] Valanarasu J M J, Oza P, Hacihaliloglu I, et al., ”Medical transformer: Gated axial-attention for medical image segmentation,” Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part I 24. Springer International Publishing, 2021: 36-46.
  * [43] Wan, Qiang, et al., ”Seaformer: Squeeze-enhanced axial transformer for mobile semantic segmentation.” arXiv preprint arXiv:2301.13156 (2023).
  * [44] Ramachandran, Prajit, et al., ”Stand-alone self-attention in vision models.” Advances in neural information processing systems 32 (2019).
  * [45] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” 2014, arXiv:1412.6980.
  * [46] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov, “Dropout: a simple way to prevent neural networks from overfiting,” Journal of Machine Learning Research, vol. 15, no. 1, pp. 1929–1958, 2014.