Compressed Self-Attention for Deep Metric Learning with Low-Rank Approximation

1 Introduction

Metric learning aims to construct well-structured distance metrics, which can be used to perform various tasks, such as k-NN classification, clustering, and information retrieval. Recently, deep metric learning with CNNs has shown a large improvement in learning embedding features that have small intra-class and large inter-class distance. Deep metric learning has a wide range of application in computer vision, such as person re-identification [Ye et al., 2018; Wang et al., 2018c], face recognition [Deng et al., 2019; Liu et al., 2018], and keypoint descriptor learning [Mishchuk et al., 2017; Xu et al., 2019].

∗Corresponding Author. This work was supported in part by the National Natural Science Foundation of China under Grants 61822113, 62041105, the Science and Technology Major Project of Hubei Province (Next-Generation AI Technologies) under Grant 2019AEA170, the Natural Science Foundation of Hubei Province under Grants 2018CFA050, the Fundamental Research Funds for the Central Universities under Grant 413000092 and 413000082. The numerical calculations in this paper have been done on the supercomputing system in the Supercomputing Center of Wuhan University.

However, due to various geometric and photometric changes such as scale change, viewpoint change, and illumination change, and the limited receptive field of convolutional kernels, the learned embedding features are not discriminative enough to ensure the intra-class compactness and the inter-class discrepancy, which would affect the performance of deep metric learning.

To tackle this problem, we enhance the discriminative power of CNNs with self-attention (SA) mechanism [Vaswani et al., 2017], which can capture long-range contextual dependencies adaptively. SA calculates the response at each position as a weighted sum of all the feature vectors, where the weights (i.e., attention maps) are determined by the pairwise similarities among all the feature vectors. However, the pairwise similarity measurement in SA leads to high computation and memory costs, making it infeasible for large inputs. Although [Chen et al., 2020] proposed a compressed self-attention (CSA) module, it lacks the theoretical guarantee for the accurate reconstruction of the original attention maps.

Therefore, we propose a compressed self-attention with low-rank approximation (CSALR) module, which reduces the computation and memory costs greatly without sacrificing the accuracy compared with the original SA. Taking advantage of the property that the feature vectors in a feature map are redundant, especially for the spatially adjacent ones, we decompose the complete attention map into a landmark attention map and a combination coefficient map with a small number of landmark feature vectors, which are sampled from the input feature map by average pooling. Then we apply the landmark attention map to the input feature map to produce the landmark output, and then apply the combination coefficient map to the landmark output to obtain the output feature map.

The high efficiency of CSALR brings additional benefits. First, we can apply CSALR to high-resolution shallow convolutional layers which lack long-range dependencies. Second, we can implement a multi-head form of CSALR where we partition the feature maps into several groups along the channel dimension and apply CSALR in each group independently, which makes the features more discriminative. We evaluate the proposed CSALR on person re-identification, which is a typical metric learning task. Extensive experiments validate the significance of CSALR in deep metric learning.

2 RelatedWork

2.1 DeepMetric Learning

Deep metric learning employs deep convolutional neural networks (CNNs) to learn embedding features which have small intra-class and large inter-class distance. However, due to the limited receptive field of CNNs and the challenging geometric and photometric changes, the learned embedding features are not discriminative enough. Recently, some researchers focus on designing a more effective loss function to force the networks to learn more representative embedding features. For example, [Fan et al., 2019] proposed a modified softmax function to learn a discriminative hypersphere manifold embedding for person re-identification. [Deng et al., 2019] proposed an additive angular margin loss to enhance the discriminative power of the network for face recognition.

Other researchers focus on designing a more robust network architecture. For instance, [Sun et al., 2018] proposed a body partition strategy and a partition refinement method for person re-identification. [Kalayeh et al., 2018] adopted semantic parsing strategy to extract the features of human body parts for person re-identification. There are also some works applying SA in deep metric learning, such as [Si et al., 2018] and [Han et al., 2018]. However, due to the high computation and memory costs of SA, SA is only applied to the part level or global level features, which is difficult to make full use of SA mechanism. [Chen et al., 2020] proposed a compressed form of self-attention, however, it cannot guarantee that the reconstructed attention maps can actually approximate the original attention maps, which may degrade the performance.

2.2 Self-AttentionMechanism

Self-attention (SA) mechanism was originally proposed for machine translation [Vaswani et al., 2017], then it has been successfully applied in computer vision, such as video analysis [Wang et al., 2018b], image segmentation [Fu et al., 2019], and image generation [Zhang et al., 2018]. SA also has a close relationship to community search[Fang et al., 2020b; Fang et al., 2020a]. The implementation of SA mechanism for CNNs is as follows.

As shown in Figure 1, given an input feature map I ∈ RC×H×W , where C, H and W represent the channels, height and width of I, respectively, we first map I to two embedding feature maps Q,K ∈ RC′×H×W , and a new feature map V ∈ RC×H×W by 1 × 1 convolutions. Then, the attention map A ∈ RHW×HW between Q and K is computed as follows:

where Qi,Kj ∈ RC denotes the feature vectors at the i-th and j-th position of Q and K, respectively, and Aij is a pairwise similarity measuring the j-th position's impact on the i-th position. Then, we apply the attention map A to the feature map V, and perform a gated residual connection with the input I to obtain the output O ∈ RC×H×W as follows:

where α is a learnable scale parameter which is initialized to 0 and gradually learns its value during training.

In contrast to convolutions, SA can capture long-range contextual dependencies by aggregating all the feature vectors based on the pairwise similarities among them. However, the pairwise similarity measurement leads to massive consumption of computation and memory, especially for a large input, making it challenging to utilize SA mechanism fully.

3 Compressed Self-Attention with Low-Rank Approximation

In this section, we first propose a low-rank approximation to any square matrix. Then we present a compressed selfattention with low-rank approximation (CSALR) module. Finally, we introduce a general framework for deep metric learning with the CSALR modules.

3.1 Low-Rank Approximation via Sampling

Given a square matrix An,n ∈ Rn×n , we propose to approximate it by randomly sampling m rows and m columns of An,n without replacement, and then setting An,n ≈ An,mA-1 m,mAm,n, where An,m is the n × m block of An,n, and with similar definitions for other blocks. The proposed method is similar to the Nystr¨ommethod [Baker, 1977], however, it does not require An,n to be a positive (semi-)definite matrix. We give the proof as follows.

Given the eigensystem of a full matrix An,n ∈ Rn×n :

where Un,n, Λn,n ∈ Rn×n are the eigenfunction matrix and the diagonal eigenvalue matrix of An,n, respectively.

Following the Nystr¨om method [Baker, 1977], we approximate (3) by randomly sampling m columns of An,n as follows:

Then we match (4) against (3), and arrive at

randomly choose m rows of An,m in (5), and arrive at

multiply both sides of (6) by A-1 m,m, An,m (from left) and U -1 n,n (from right) successively, and arrive at

match (4) against (7), and arrive at

match (8) against (3), and arrive at

Equation (9) is the proposed low-rank approximation to the square matrix An,n. It doesn't require An,n to be positive (semi-)definite or symmetric. We can sample different rows and columns of An,n to form the block Am,m, as long as Am,m is full rank.

3.2 CSALRModule

In this section, we propose a compressed self-attention with low-rank approximation (CSALR) module, where we compress the original attention maps with the proposed low-rank approximation to reduce both the computation and memory costs. We observe that the feature vectors in a feature map are inherently redundant, especially those spatially adjacent to each other. Thus we can decompose the complete attention map into a landmark attention map and a combination coefficient map with a small number of landmark feature vectors sampled from the input feature map by average pooling.

As shown in Figure 2, given an input feature map I ∈ RC×H×W , we first obtain the feature maps Q, K and V as the original SA mechanism in Figure 1 (Section 2.2). Same with (1), the complete attention map An,n is computed as follows:

where n = H ×W is the total number of the feature vectors in Q or K. We consider the normalization with softmax in the following low-rank approximation for numerical stability.

Then we approximate the complete attention map An,n via sampling based on Eq.(9). Leveraging the spatial redundancy of the feature vectors, we sample the landmark feature map Q', K' ∈ RC′×H′×W ′ by applying average pooling to Q and K, respectively, where the kernel size and the stride of the pooling layer both equal to ( H H′ , W W ′ ).

Similar to (10), the partial attention maps An,m between Q and K', Am,m between Q' and K', and Am,n between Q' and K are computed as follows:

where Q′ i and K′ j are the landmark feature vectors at the i-th and j-th position of the landmark feature map Q' and K', respectively, m = H′ ×W ′ is the total number of the feature vectors in Q' or K'. It is worth noting that since the normalization is considered, An,m is just a rescaled approximations of the corresponding n ×m block of An,n (same for Am,m), which is reasonable considering the low-rank property of An,n.

Then, according to (9), the complete attention map An,n can be approximated with the partial attention maps An,m, Am,m and Am,n as follows:

where B = Am,n ∈ Rm×n represents the landmark attention map between the landmark feature vectors and all the feature vectors, and W = An,mA-1 m,m ∈ Rn×m represents the combination coefficient map based on the relationship between all the feature vectors and the landmark feature vectors. The complete attention map An,n is reconstructed by applying the combination coefficient map W to the landmark attention map B.

Since the adjacent feature vectors in space are more similar, the sub attention map of each feature vector is more similar to that of the landmark feature vector corresponding to the same pooling patch, whose combination coefficient should be larger. Thus we transform An,m and Am,m as follows:

where 1(Qi,K′ j) denotes that if the feature vector Qi and the landmark feature vector K′ j correspond to the same pooling patch, its value is 1, otherwise, its value is 0 (same for 1(Q′ i,K′ j)). It is worth noting that (16) guarantees that A-1 m,m always exists.

Then we apply the landmark attention map B to the feature map V to produce the landmark output E ∈ RC×H′W ′ , and then apply the combination coefficient map W to the landmark output E, and perform a gated residual connection with the input I to obtain the output feature map O ∈ RC×H×W as follows:

where α is a learnable scale parameter which is initialized to 0 and gradually learns its value during training.

The low-rank decomposition of the original attention maps in (14) significantly reduces the computation and memory costs, and the inherent redundancy of the feature vectors avoids performance degradation caused by compression. In experiments, we find that setting m ≪ n does not decrease the accuracy.

The high efficiency of CSALR allows us to implement a multi-head form of CSALR. We partition the feature maps into several groups along the channel dimension and apply CSALR in each group independently. Then we fuse the different groups with 1 × 1 convolution. In this way, each group can handle a specific relationship between the features alone, which helps to diversify the long-range interactions and makes the features more discriminative.

3.3 CSALR in DeepMetric Learning

The proposed CSALR module can be easily applied to the existing backbones to capture long-range contextual dependencies efficiently, providing a general framework for deep metric learning. The computational and memory efficiency of CSALR allows us to apply CSALR to high-resolution shallow convolutional layers and implement a multi-head form of CSALR.

As shown in Figure 3, the framework includes a backbone network, CSALR modules, and deep supervision branches. We apply the CSALR module to each stage of the backbone network. In this way, the high-resolution shallow convolutional layers, whose receptive filed is limited, can benefit more from the long-range interactions. To diversify the long-range interactions, we implement the multi-head form of CSALR. To assist the learning of attention maps, we apply deep supervision to each CSALR module. It's worth noting that the residual connection of CSALR is drawn out of the CSALR module in Figure 3.

In each deep supervision branch, we first perform an element-wise multiplication between the input feature map and the output of CSALR module and obtain a new feature map, then we apply a global average pooling layer to the feature map and obtain a feature vector, and then we project the feature vector linearly with a fully connected layer and apply a loss function to the result. In experiments, we find that the element-wise multiplication is better than the elementwise summation, so we implement deep supervision before the residual connection. The final loss function of the framework is as follows:

where LF is the loss of the main branch, and LDSi is the loss of the i-th deep supervision branch. N is the total number of CSALR modules. λ is a balance factor between LF and LDS .

When applying the proposed framework to boost the performance of the specific metric learning algorithms, we just modify the backbone networks of the original models without changing the remaining settings. Thus, the loss LF is the same as the original algorithm. The loss LDS is cross entropy in our experiments. When comparing CSALR with SA or CSA, we just replace the CSALR module with the SA or CSA module in the framework, without changing the deep supervision branches or the corresponding loss functions.

4 Experiments

We validate the proposed CSALR on person re-identification, which is a typical metric learning task. We choose two representative methods in person re-identification, i.e., PCB [Sun et al., 2018] and PCB-RPP [Sun et al., 2018], as our baselines. We modify the backbone networks of the baseline models without changing other settings. We provide both qualitative and qualitative comparisons to demonstrate the effectiveness and efficiency of CSALR in deep metric learning.

4.1 Datasets

We use three datasets for evaluation, i.e., Market-1501 [Zheng et al., 2015], DukeMTMC-reID [Ristani et al., 2016; Zheng et al., 2017], and CUHK03-NP [Zhong et al., 2017]. Market-1501 contains 751 training IDs with 12,936 images and 750 query IDs with 3,368 query images and 19,732 gallery images, which are captured by 6 cameras. DukeMTMC-ReID contains 702 training IDs with 16,522 images and 702 query IDs with 2,228 query images and 17,661 gallery images, which are captured by 8 cameras. CUHK03NP contains 767 training IDs with 7,365 images and 700

query IDs with 1,400 query images and 5,332 gallery images, where each identity is captured by 2 disjoint cameras. It offers two kinds of bounding boxes, one is hand-labeled, the other is DMP-detected, and we use the latter.

For evaluation, we use the cumulative matching characteristic (CMC) for rank-1, rank-5, rank-10 and the mean average precision (mAP) as our metrics. We evaluate all the methods under the single-query mode. For simplicity, we do not use re-ranking [Zhong et al., 2017] for post-processing, which can improve mAP by a large margin.

4.2 Implementation Details

We use a similar setting to [Sun et al., 2018]. The backbone is ResNet-50 with pre-trained weights from ImageNet. We use the proposed framework with CSALR to enhance the backbone. We apply the CSALR module to each residual block of the backbone except the last one, leading to 3 CSALR modules in total. In Eq. (9), the loss LF is the same as PCB and PCB-RPP, and the loss LDS is cross entropy. The balance factor λ is set to 1.0. We implement the multi-head form of CSALR. For each CSALR module, the number of sampled points in each group is set to 96, and the number of groups is set to 2. When comparing CSALR with SA or CSA, we replace the CSALR module with the SA or CSA module in the framework without changing other settings. It's worth noting that we do not implement the multi-head form of SA, since it costs too much computation and memory resources.

For data processing, the input images are all resized to 384 × 128. We apply random horizontal flipping with 0.5 and normalization as data augmentation. For training, we set the batch size to 64 and the total number of training epochs to 100. The optimizer is set to Stochastic Gradient Descent (SGD) with momentum of 0.9 and weight decay of 10-4 . The base learning rate is initialized to 0.1 and decayed by 0.1 after every 40 epochs, and the learning rate for the backbone network is set to 0.1× the base learning rate. For PCB-RPP, we first train PCB for 40 epochs, and then train RPP for another 60 epochs with weights initialized from PCB.

4.3 Performance Comparison

As shown in Table 1, the proposed framework with CSALR modules can effectively boost the performance of all the original algorithms on Market-1501, DukeMTMC-ReID and CUHK03-NP with respect to both rank-1 accuracy and mAP. On the three datasets, for PCB with CSALR, the increase in rank-1 accuracy is +1.6%, +4.4%, and +7.9%, respectively; the increase in mAP is +5.0%, +7.9%, and +10.8%, respectively; for PCB-RPP with CSALR, the increase in rank-1 accuracy is +0.2%, +2.1%, and +6.0%, respectively; the increase in mAP is +2.2%, +4.6%, and +8.6%, respectively. This suggests that CSALR can effectively capture long-range contextual dependencies and make the extracted features more discriminative. It's worth noting that the improvement for mAP is larger than that for rank-1 accuracy. In reality, rank-1 accuracy represents the ability to retrieve the easiest match in the gallery, while mAP represents the ability to find all the matches, which indicates that CSALR helps to find more challenging matches.

We also compare CSALR with SA and CSA. As shown in Table 1, it's interesting that although the attention maps are compressed in CSALR, the performance of CSALR is better than that of SA. There are two reasons for this phenomenon. On the one hand, the inherent redundancy of the feature vectors guarantees an accurate approximation of the original attention maps. On the other hand, the multi-head form of CSALR helps to diversify the long-range interactions, which compensates for the performance degradation caused by compression. We can also see that the performance of CSALR is better than that of CSA, which suggests that the approximation to the attention maps in CSALR is more reasonable and more accurate than that in CSA.

4.4 Different Number of Sampled Points

We fix the number of groups to 2 and vary the number of sampled points in each group to see its influence. As shown in Table 2, the performance rises as the number of sampled points increases. The reason behind this is that the more the sampled points, the more accurate the approximation to the original attention maps and the less the performance degradation. However, the performance improvement is slight when the number of sampled points increases from 96 to 192. This further demonstrates the significant redundancy of the feature vectors in a feature map, so that we can approximate the complete attention maps accurately with only a small number of landmark feature vectors.

4.5 Different Number of Groups

We fix the number of sampled points to 96 and vary the number of groups to see its influence. As shown in Table 3, the performance is best when the number of groups is 2. The performance drops when the number of groups is smaller or larger than 2, indicating that the multi-head form of CSALR helps to diversify the long-range interactions and makes the extracted features more discriminative. However, too many groups and too few channels in each group will lead to insufficient expression of the specific semantics, which harms the performance of CSALR.

4.6 CSALR on Different Layers

As shown in Table 4, applying the CSALR modules to both the shallow and deep layers of the backbone leads to the best performance. Applying the CSALR modules to shallow layers is better than deep layers, which demonstrates that the shallow layers lack long-range dependencies extremely. So it's more beneficial to apply SA mechanism to shallow layers, and our CSALR module is very efficient considering the large size of the shallow layers.

4.7 Resource Costs Comparison

As shown in Table 5, for SA and CSALR with 1, 2, 4, 8 groups, compared with the original method, the relative increase in memory cost is +118.45%, +19.88%, +39.67%, +61.27%, and +85.08%, respectively; the decrease in FPS is -54.17%, -12.70%, -20.86%, -28.10%, and -34.52%, respectively. We can see that the methods with CSALR consume much less computation and memory resources than the methods with SA, even if the number of groups is set to 8. Thus, we can apply CSALR with more flexibility.

4.8 Analysis of AttentionMaps

As shown in Figure 4, for each selected point (marked in a specific color) in the image, the corresponding attention map focus on the most relevant parts in the whole image, which indicates that the SA mechanism can capture long-range contextual dependencies adaptively. We can also see that the reconstructed attention maps are very similar to the original attention maps, which indicates that the feature vectors in a feature map have significant redundancy. We can approximate the complete attention maps accurately with only a small number of landmark feature vectors.

5 Conclusion

In this paper, we aim to boost the performance of deep metric learning with self-attention (SA) mechanism, which can capture long-range contextual dependencies adaptively. Accordingly, we propose a compressed self-attention with lowrank approximation (CSALR), which significantly reduces the computation and memory costs without sacrificing the accuracy. Qualitative and quantitative experiments demonstrate the significance of CSALR in deep metric learning.