Augmenting Prototype Network with TransMix for Few-shot Hyperspectral Image Classification

For the problem of few-shot HSI classification, there are two datasets given: source dataset and target dataset. In practical application, all the samples in the source dataset are labeled, which are used for training the classification model. While, there are only few labeled samples in the target dataset. The samples are the pixels of the HSIs. In order to fully preserve spatial information, the fixed-size patches centering around each pixel instead of the pixeles themselves are often treated as the samples. The few-shot learning task is to predict the classes of these unlabeled samples in the target dataset by using all these labeled samples in source dataset and target dataset. Nevertheless, another dataset in which all the samples are also labeled is often selected as the target dataset in the research environment, to facilitate the performance evaluation of the few-shot learning methods. Such dataset can be also called testing dataset.

To elaborate the problem formally, it is often assumed that there are $C_{s}$ classses in the source dataset $D_{s}$ and $C_{t}$ classses in the target dataset $D_{t}$, where $C_{t}$ is smaller than $C_{s}$. To imitate the setting of target dataset, lots of the tasks will be constructed by randomly selecting $C\times(K+M)$ samples each time, and used to train the classification model in an episode manner. The $C$ refers to the number of classes selected, which is often set to the number of classes in the target dataset. $K$ is the number of few labeled samples per class, and $M$ is the number of query samples per class to be classified, where $K$ is smaller than $M$. The set consisiting of the $C\times K$ labeled samples is often called support set $\{({x_{i}},{y_{i}})\}_{i=1}^{{C\times K}}$, while these $C\times M$ query samples will form the query set $\{({x_{i}},{y_{i}})\}_{i=1}^{{C\times M}}$. Such few-shot tasks are often called the $C$-way $K$-shot tasks, whose purpose are to predict the classes of these query samples by using these support samples.

The architecture of the proposed APNT method and the flowchart of the learning process is shown in Fig. 2. APNT model mainly consists of three components: query sample mixing, transformer based feature extractor, and TransMix based few-shot learning loss. Given an input task consisting of the support and query set, the query sample mixing module will mix up each query sample with another randomly selected query sample in the way of CutMix[35], and output the synthetic query samples. Subsequently, the transformer based feature extractor extracts the embedding features of the support samples and the synthetic query samples, and also returns the attention maps reflecting the importance of each pixel of the original query samples. Following that, the mean of the embedding features of the support samples in the same classes will be caculated as the prototpyes of each class. In the meanwhile, the TransMix based few-shot learning loss module computes the cross entropy loss of the synthetic query samples according to their synthetic labels which are derived by using the attention maps returned by transformer.

Similar with the work of RPCL[31], We adopt the fine-tuning strategy to train the network. That is, we first pre-train the network by using the tasks constructed from $D_{s}$, and then fine-tune it by using the tasks from $D_{t}$. This also means there are some labeled samples selected from $D_{t}$ to participate in the training. Consistent with previous works, there are five samples per class selected, which are augmented to 200 samples per class in our experiments. The tasks from $D_{t}$ will be constructed from these augmented smaples. Finally, the remaining samples in the target dataset $D_{t}$ are used for testing. And after extracting sample features with the trained transformer, the classes of these query samples are pedicted by using KNN algorithm during testing. It is worth noting that because the spectral dimensions of the samples from source and target datasets are different, there are also the mapping modules to unify the dimensionality.

When given an input task consisting of a support set and a query set, either the support or query samples will be first processed by the mapping modules to transform their channel dimensionality. After that, each sample in the query set is mixed up with another randomly selected sample to generate the synthetic query samples. We choose to only mix up the query samples without support samples for the purpose of ensuring the correctness of the class prototypes caculated from these support samples. For two query samples $x_{i}$ and $x_{k}$ with the labels of $y_{i}$ and $y_{k}$, the mixing up process can be formally stated as Eq. 1.

where $\bm{M}\in\left\{0,1\right\}^{W\times H}$ denotes a binary mask indicating where to drop out the pixels from the sample with the spatial size of $W\times H$. $\bm{1}$ is a binary mask filled with ones, and $\odot$ is the element-wise multiplication. $\tilde{x}$ is the synthetic query sample, and $\tilde{y}$ is the corresponding synthetic label, where $\lambda_{1}$ and $\lambda_{2}$ are the proportion of $y_{i}$ and $y_{k}$ in the synthetic label.

The proposed method follows CutMix [35] to mix up the query samples. That is, a randomly sampled patch in $x_{i}$ is removed and filled with the patch cropped from the same region of $x_{k}$. CutMix uses the proportion of the sampled patch to the entire sample as the value of $\lambda_{1}$ (i.e., $\lambda_{2}=1-\lambda_{1}$) to mix up the labels. Differently, we will caculate the values of $\lambda_{1}$ and $\lambda_{2}$ by using the attention maps returned from transformer, to take the different importance of the pixels into consideration. The details will be described in the following subsection of $E$.

Once the synthetic query samples are generated, the original query samples, the synthetic query samples, and the support samples are input into the transformer module. The embedding features of the support samples and the synthetic query samples will be extracted. At the same time, the attention maps of the original query samples are also required. The whole process can be seen in Fig. 3.

When inputting the samples, i.e., the patches with the three-dimensional size of $W\times H\times D$, into transformer, each sample will be first pull into a sequence $x\in\mathbb{R}^{WH\times D}$ of which the composing items are the spectral vectors of the pixels. And then, the token vectors with $D$ dimension, which are used to learn global features of each sample, are randomly generated and appended into the sequences of each sample. After that, the features and the attention maps of each sample are obtained by directly inputting the updated sequences $x\in\mathbb{R}^{(WH+1)\times D}$ into transformer module.

A shown in Fig. 3, the transformer module of APNT is composed of a set of stacked transformer encoders. Each encoder consists of a self-attention layer and a feedforward layer. For the self-attention layer, there exist three learnable projection matrixs $\bm{W}_{q}\in\mathbb{R}^{D\times D^{{}^{\prime}}}$, $\bm{W}_{k}\in\mathbb{R}^{D\times D^{{}^{\prime}}}$ and $\bm{W}_{v}\in\mathbb{R}^{D\times D^{{}^{\prime}}}$. They will project the input sequences into different embedding spaces, as shown in Eq. 2.

where $\bm{Q}$, $\bm{K}$ and $\bm{V}$ are the projected matrices of the sequence $x$. With these projected matrices, the attention maps of each input sequence $a_{x}\in\mathbb{R}^{(WH+1)\times(WH+1)}$, which capture the attention of each item in the sequence to the other items, can be caculated as Eq. 3. And the features of each item in the sequence are updated by Eq. 4.

There are multiple heads of self-attention compution in transformer. The features returned from each head will be concatenated together and restored to the original dimension through a linear operation. The updated features will be further processed through the feedforward layer without changing the dimensions. That is, a new sequence of $x^{{}^{\prime\prime}}\in\mathbb{R}^{(WH+1)\times D}$ will be produced by a transformer encoder for each sample. At the same time, the attention maps $a_{x}$ from different heads will be averaged. Through two transformer encoders, the updated taken vectors in the sequence $x^{{}^{\prime\prime}}$ generated by the last transformer encoder are selected as the feature vectors of the samples. And the attention vectors capturing attentions of the samples to their constituent pixels will be also selected from the attention maps $a_{x}$ returned by the last transformer encoder. When reshaping the obtained attention vectors into the shape of $W\times H$, the required attention maps reflecting the importance of each pixel are obtained.

After obtaining the attention maps $\bm{A}\in\mathbb{R}^{W\times H}$ of these query samples, the sum of the attention values belonging to the mixed region is directly used to derive the values of $\lambda_{1}$ and $\lambda_{2}$ for label mixing shown in Eq. 1. This is because that the sum of all values in a attention map $\bm{A}$ is 1. This derivation process can be experssed as Eq. 5.

where $\bm{A_{i}}$ and $\bm{A_{k}}$ represent the attention maps of the query smaple of $x_{i}$ and $x_{k}$. $\bm{M}$ represents the binary mask mentioned in Eq. 1. From Eq. 5, it can be seen that the values of $\lambda_{1}$ and $\lambda_{2}$ represent the probability that the synthetic sample $\tilde{x}$ shown in Eq. 1 belongs to the classes of $y_{i}$ and $y_{k}$.

In addition, assuming the feature of each support smaple is $z_{i}$, the $c^{th}$ class prototype is computed as Eq.6.

where $S^{c}$ represents the set of features of $c^{th}$ class support samples.

Then, through predicting the synthetic query sample $\tilde{x}$ as the classes denoted by the labels of $y_{i}$ and $y_{k}$, the few-shot learning losses are caculated according to Eq. 7 and Eq. 8. The $\tilde{z}_{h}$ in Eq. 7 and 8 represents the embedding feature of synthetic query sample $\tilde{x}_{h}$. And $p(\tilde{y}_{h}=y_{i}|\tilde{x}_{h})$ means the probability that the model accurately predicts the synthetic query sample $\tilde{x}$ to the class of $y_{i}$. Meanwhile, $d(\tilde{z}_{i},p_{y_{i}})$ denotes the Euclidean distance between the feature of synthetic query sample $\tilde{x}_{i}$ and the class prototype of $y_{i}$.

Because the synthetic query samples only partially belong to the classes of the samples used for mixing, the total loss used for updating the model is caculated as Eq. 9, which uses the values of $\lambda_{1}$ and $\lambda_{2}$ to balance the losses of $L^{i}_{fsl}$ and $L^{k}_{fsl}$.

To verify the performance of the proposed method, we selected several typical classification methods for comparison, including DFSL[22], DCFSL[27], CMFSL[29], Gia-CFSL[30], RPCL[31], MRLF[32], and HFSL[33]. In the experiment, according to the specific training process of each model, DFSL only used these augmented samples from target dataset for training. DCFSL, CMFSL, Gia-CFSL, MRLF, RPCL and our proposed method used both the samples from source datasets and target datasets, while HFSL used the natural image dataset of Mini-ImageNet for pre-training and was fine-tuned on these augmented samples from target datasets. In addition, in order to prove that the proposed method still had good performance without pre-training on source datasets, the performance of APNT* method which is only trained on these augmented samples from target datasets is also evaluated.

Table II to Table IV report the specific classification results of each class, as well as the AA, OA, and kappa values of each method on the three target datasets. From the comparison results, it can be seen that APNT* and APNT have superior classification results compared to other methods. Among them, APNT outperforms the other results by $1.5\%$, $1.7\%$, and $0.5\%$ on the IP, PU, and SA datasets, respectively, reaching the optimal results on class 1, 4, 6, 7, 9, 10, and 11 of IP dataset, class 1, 3, and 7 of PU dataset, and class 2, 9, 11, and 15 of SA dataset. This result demonstrates the effectiveness of the proposed method. In the meanwhile, APNT* outperforms the other models by $1.1\%$, $0.9\%$, and $0.8\%$ on the IP, PU, and SA datasets respectively, reaching the optimal results on class 1, 2, 7, 14, and 15 of IP dataset, class 4 and 6 of PU dataset, and class 3 and 16 of SA dataset. Because APNT* does not use the sampes from source dataset for pre-training, this result shows that the proposed method can eliminate the dependence on existing large-scale datasets and can achieve comparative performance by only using these few-shot labeled samples in target tasks. This will facilitate the use of the proposed method in a range of applications.

The classification maps corresponding to each method are shown in Figures 8 to 10. In these maps, colored pixels are labeled, and black pixels are unlabeled and displayed as the background. By observing the results displayed in the classification maps, it can be observed that the classification maps generated by the proposed method are most similar to the real ground map. In particular, we can notice the reduction of misclassified boundary points. These results also confirm the effective improvement of the classification results by the proposed method.

To validate the improvement of the predicting on boundary pixels, we processed the datasets and extracted all the boundary patches for testing. In the hyperspectral datasets, when all the pixels in the patches don’t belong to the same class, we define such patches as boundary patches. Generally speaking, these boundary patches are the hard samples, and it is more difficult to predict their classes.

We calculated the overall accuracy of various methods on all the boundary patches, which are shown in Fig. 11. As can be seen from the results, consistent with the assumption, APNT has the highest overall prediction performance on the boundary patches over three datasets.

While applying transformer to pay different attention to different pixels and extract better features, the proposed method also mixs up the query samples by following TransMix and use the synthetic samples to train the model. We divide our contribution into two parts: the transformer based feature extractor and the TransMix based sample mixing. To verify the contribution of each part, the ablation experiments were also conducted. By comparing the proposed method with the method with 3DCNN feature extractor which is used in the work of [31], we validate the contribution of transformer based feature extractor. At the same time, by comparing the proposed method with that of using CutMix [35] instead of TransMix [34] to mix up query samples, the contribution of TransMix based query sample mixing is also validated.

Table V shows the results of the ablation experiments, where Trans denotes the transformer. From the first and second column, it can be seen that when adopting the same mixing technique of CutMix, better performance are achieved in most cases by adopting transformer as feature extractor. This indicates that compared with the 3DCNN feature extractor, transformer model can learn more point-to-point relationships in the patches. Furthermore, when using transformer as the feature extractor but adopting TransMix instead of CutMix for query sample mixing, there are also better performance which can be seen from the second and third column of Table V. This proves that compared with CutMix which mixs up the labels of two samples just according to the proportion of synthetic areas, TransMix can generate better mixing labels by using the attention on each pixel which reflect their relative importance.

There are several parameters which affect the performance of the proposed method. We have done the experiments to study how sensitive the proposed method is to these parameters. In this section, we report the analysis results of two kinds of parameters, namely the number of transformer encoders and the size of the HSI patch.

To further illustrate the performance of the proposed method on time and space consumption, Table VI lists the size of model parameters (M) and the inference time including the training and testing time (s) of these comparsion methods. It can be seen that although our APNT method is not the lightest compared to other few-shot methods, it has a faster training speed and competitive testing time. This may be attributed to the parallel execution characteristics possessed by the Transformer model.