Towards Improved Proxy-based Deep Metric Learning via Data-Augmented Domain Adaptation

Deep Metric Learning (DML) Consider a set of data samples $\mathcal{S}=\{I_{i},y_{i}\}_{i=1}^{N}$ with raw images $I_{i}$ and its corresponding class label $y_{i}\in\{1,\dots,C\}$; we learn a projection function $f_{G}:\mathcal{S}\stackrel{{\scriptstyle f}}{{\rightarrow}}\mathcal{X}$, which project the input data samples to a hidden embedding space (or metric space) $\mathcal{X}$. We define the projected features set as $X=\{x_{i}\in\mathbb{R}^{d}\}_{i=1}^{N}$. The primal goal of Deep Metric Learning (DML) is to refine the projector function $f_{G}(\cdot)$, which is usually constructed with convolutional deep neural networks (CNN) as the backbone, to generate the projected features that can be easily measured with defined distance metric $d(x_{i},x_{j})$ based on the semantic similarity between sample $I_{i}$ and $I_{j}$. Here we adopt the distance metric $d(\cdot)$ as the cosine similarity. Before delivering features to any loss, we use L2 normalization to eliminate the effect of differing magnitudes.

Proxy-based DML To boost the learning efficiency, a group of DML methods pre-define a set of learnable representations $P=\{p_{i}\in\mathbb{R}^{d}\}_{i=1}^{C}$, named proxy, to represent subsets or categories of data samples. Typically there is one proxy for each class so that the number of proxies is the same as the number of classes $C$. The proxies are also optimized with other network parameters. The first proxy-based method, Proxy-NCA (Movshovitz-Attias et al. 2017), or its improved version Proxy-NCA++ (Teh, DeVries, and Taylor 2020), utilizes the Neighborhood Component Analysis (NCA) (Goldberger et al. 2004) loss to conduct this optimization. The later loss Proxy-Anchor (PA) (Kim et al. 2020) inversely sets the proxy as the anchor and measures all proxies for each minibatch of samples. The PA loss $\mathcal{L}_{proxy}(X,P)$ can be presented as

where $X_{p}^{+}$ denotes the set of positive samples for a proxy $p$; $X_{p}^{-}$ is its complement set; $\tau$ is the scale factor; and $\delta$ is the margin. Since the PA updates all proxies for each mini-batch, the model has higher learning efficiency in capturing the structure of samples beyond the mini-batches. We propose these two fundamental proxy-based losses (PNCA++ and PA) that achieve competitive results as our baselines.

To reduce the distribution gap between the data samples and the proxies, we transform the proxy-based DML into a domain adaptation problem. We regard the data samples as data points in the source domain, while the initialed proxies are data points in the target domain. We noticed that the number of proxies in the target domain is especially limited compared to the data samples because the basic proxy method only assigns a single proxy for each class. The unbalanced samples and proxies would cause learning biases in modeling the distributions. Also, the proxies are initialized from a normal distribution and do not contain any related semantic information, which also causes difficulty in aligning their distribution to the data sample domain.

To overcome these difficulties, we propose a novel data augmentation strategy to create an intermediate domain to balance the amount of data points for domain adaptation. Specifically, we interpolate the space with mixed features from data set $X$ and proxy set $P$. For each data sample $x_{i}$ and its corresponding proxy $p_{i}$, we create a data feature $\hat{d}$:

where $\lambda\sim Beta(\alpha,\beta)$ is the linear interpolation coefficient that sampled from beta distribution with $\alpha>0$ and $\beta>0$ that decide its probability density function. The new data contains semantic information between the data sample and proxies and shares their distribution statistics. Therefore pushing $\hat{d}$ is equal to pushing both samples $x$ and proxies $p$, and their distribution is closer to the data sample domain than the original proxies.

In addition, we further propose extending the number of training instances by augmenting the data-proxy pairs within the same class. For each pair of samples $(x_{i},x_{j})$ and their corresponding augmented data $(\hat{d}_{i},\hat{d}_{j})$ inside the mini-batch, we propose the following mixing:

where $\mu_{1},\mu_{2}$ are also sampled from $Beta$ distribution. Then we mix the new samples $\tilde{x}_{i}$ and $\tilde{d}_{i}$ inside the mini-batch to ensure the number of data samples with the same label $n\geq 2$. Combined with the original mini-batch, the augmented data sample set and the augmented proxy set are noted as $\tilde{X}=X\cup\{\tilde{x}_{1},\tilde{x}_{2}\cdots\}$ and $\tilde{D}=\{\hat{d}_{1},\hat{d}_{2}\cdots\}\cup\{\tilde{d}_{1},\tilde{d}_{2}\cdots\}$, and the size of mini-batch is also extended accordingly. We then normalize the composed features in $\tilde{X}$ and $\tilde{D}$ with L2 normalization to constrain them on a unit hypersphere embedding space where the magnitude is fixed to 1.

Based on the augmented data, our goal is to refine the set $\tilde{X}$, $\tilde{D}$ and the original proxy set $P$ to domain invariant representations that share the same distribution to help the proxy-based losses. We follow the principle idea of adversarial domain adaptation (Ganin et al. 2016) to estimate the domain divergence by learning a domain-level discriminator. Specifically, we learn a classifier $f_{D}(\cdot)$ that minimizes the risk of domain prediction (to predict if the data comes from a unique domain) between the set $\tilde{X}$, $\tilde{D}$ and $P$.

Generally, we would label the data from a specific domain with the one-hot label as the prediction target. Since we have three different domains including the augmented data domain and our labeling space is symmetric, we would simply assume the features $\tilde{x}\in\tilde{X}$ are labeled as $y_{0}=\overline{001}$ while $\tilde{d}\in\tilde{D}$ are labeled as $y_{1}=\overline{010}$, and the initial proxies $P$ are labeled as $y_{2}=\overline{100}$ for convenience. Specifically, we estimate the domain classifier $f_{D}(\cdot)$ as an MLP with a single hidden layer and a ReLU function. The hidden layer is then projected to a 3-dimensional head as the logits prediction of the domains. To optimize the $f_{D}(\cdot)$ with a low prediction risk on the labeling space, we conduct the cross-entropy objective $\mathcal{L}_{adv}$ as follows

where $\mathcal{L}_{ce}$ is the cross entropy loss and $\tilde{N}$ is the total number of samples after the data augmentation. The parameters of the classifier $f_{D}(\cdot)$ are optimized to minimize the adversarial loss $\mathcal{L}_{adv}$ in training. Recall that the feature $x$ is generated from the projection function $f_{G}(\cdot)$. Thus, the parameters of generator $f_{G}(\cdot)$ are optimized to fool the discriminator $f_{D}(\cdot)$ in the opposite direction. Since $\tilde{D}$ in the target domain contains features that mixed from $X$ and proxies $P$, optimizing the $\tilde{D}$ equals optimizing the generator $f_{G}(\cdot)$ in the source domain while updating the original proxies $P$. Thus, the adversarial learning signal of $\mathcal{L}_{adv}$ would help both generator $f_{G}(\cdot)$ and the original proxies $P$ to maintain the domain invariant representations to fool the classifier.

One drawback of the domain-level discriminator described above is that the discriminate information, especially the inter-class correlation, is ignored in the optimization process. Losing the discriminative information will cause all data points to be concentrated on a local area or a surface, which would cause inter-class ambiguity and confuse the metric learning losses. To solve this problem, we further propose a category-level discriminator that learns to predict the class of data samples and compare the discrepancy of predictions between the data samples and mixture proxies.

Specifically, we optimize a classifier $f_{C}(\cdot)$ with the feature generator $f_{G}(\cdot)$ to predict the category label $Y=\{y_{0},y_{1},\dots\}$ from mixture data samples $\tilde{X}$ with the classification loss $\mathcal{L}_{cls}(\tilde{X},Y)$ as

The cross-entropy loss $\mathcal{L}_{ce}$ would provide a supervised learning signal to $f_{G}(\cdot)$ to maintain the discriminative information during the DML training process.

We note that the data samples that share the distributions would also share the labeling space with the target proxy domain. To further align the distributions, we propose to constrain the samples from the source domain and augmented data from the target domain to have a low discrepancy of predictions from our category classifier $f_{C}(\cdot)$. Thus, one additional goal of $f_{C}(\cdot)$ is to learn the maximized discrepancy of the category prediction between the data samples $\tilde{X}$ and mixture proxies $\tilde{D}$ while the $\tilde{D}$ are later optimized to minimize this discrepancy.

To measure the discrepancy of the category probabilities, we empirically adopt the discrepancy introduced in (Chen et al. 2022) that utilizes the Nuclear-norm Wasserstein Distance (NWD). The NWD is demonstrated to be the upper bound of the Frobenius-norm, which estimates the correlations of the predictions (Cui et al. 2020). Thus, we compare the NWD between the logistic predictions of $f_{C}(\cdot)$ from the augmented samples $\tilde{X}$ and data $\tilde{D}$. The loss $\mathcal{L}_{d}(\tilde{X},\tilde{D})$, which measures the NWD can be described as,

where $||x||_{\star}=\sum\sigma(x)$ denotes the nuclear-norm of $x$, which is defined as the sum of its singular values.

We adopt the paradigm of adversarial learning to alternatively update the gradient of our feature generator $f_{G}(\cdot)$ and the discriminators $f_{D}(\cdot)$ and $f_{C}(\cdot)$ discussed above. To this end, we train our combined loss by playing the min-max game as follows,

where $\eta$ is the pre-defined hyperparameter that balances the contribution between the domain-level and category-level discriminators. Empirically, we do not set another weight between classification loss $\mathcal{L}_{cls}$ and discrepancy loss $\mathcal{L}_{d}$. We also need the original proxy-based loss $\mathcal{L}_{proxy}$ in Eq. 1 to do the basic DML of the sample-proxy pair in training. Note that the augmented data set $\tilde{D}$ is only for domain adaptation progress; the original $\mathcal{L}_{proxy}$ only operates $\tilde{X}$ and original proxies $P$. Thus, our combined training progress can be described as the following two sub-processes:

where parameters $\theta_{f_{D}}$ and $\theta_{f_{C}}$ are updated in first phase and $\theta_{f_{G}}$ and the proxies $P$ are updated with $\tilde{D}$ in the second phase. Even if gradient reversal layers are accepted for achieving adversarial training in earlier domain adaptation works, we empirically conclude that a separate training phase would be more feasible for us in our search for stable training parameters. The full training progress can be referred to in Algorithm 1.

We use the standard benchmarks CUB-200-2011 (CUB200) (Wah et al. 2011) with 11,788 bird images and 200 classes, and CARS196 (Krause et al. 2013) that contains 16,185 car images and 196 classes. We also evaluate our method on larger Stanford Online Products (SOP) (Oh Song et al. 2016) benchmark that includes 120,053 images with 22,634 product classes, and In-shop Clothes Retrieval (In-Shop) (Liu et al. 2016) dataset with 25,882 images and 7982 classes. We follow the data split that is consistent with the standard settings of existing DML works (Teh, DeVries, and Taylor 2020; Kim et al. 2020; Venkataramanan et al. 2022; Zheng et al. 2021b; Roth, Vinyals, and Akata 2022; Lim et al. 2022; Zhang et al. 2022). We adopt the Recall@K (K=1,2,4 in CUB200 and CARS196, K=1,10,100 in SOP, and K=1,10,20,30 in In-Shop) proposed in existing works to evaluate the accuracy of ranking. We also evaluate it with Mean Average Precision at R (MAP@R) which is based on the idea of MAP and R-precision, which is a more informative DML metric (Musgrave, Belongie, and Lim 2020).

We train our model in a machine that contains a single RTX3090 GPU with 24GB memory. The Implementation is based on the existing RDML (Roth et al. 2020)

Backbones and Preprocessing. In this paper, we propose two basic backbones to evaluate our learning algorithm: the ResNet50(He et al. 2016) and the InceptionBN (Ioffe and Szegedy 2015). They are pre-trained on ImageNet1K(Deng et al. 2009) and are widely used in DML works for performance evaluation, where we resize the image to $224\times 224$, do random resized cropping, and random horizontal flipping. In the test phase, the images are first resized to $256\times 256$, then cropped back to $224\times 224$. A linear head embeds the feature from the second last layer of the backbones to a 512-dimension hidden space. We follow the standard pre-processing introduced in other deep metric learning works (Venkataramanan et al. 2022; Zheng et al. 2021b; Roth, Vinyals, and Akata 2022; Lim et al. 2022; Zhang et al. 2022). We also adopt global max and average pooling with layer normalization on CNN backbones suggested by Teh et al. (Teh, DeVries, and Taylor 2020) to further improve the generalization of features.

Training Details. Our optimization is done using Adam ($\beta_{1}=0.5,\beta_{2}=0.999$) (Kingma and Ba 2015) with a decay of $1\cdot 10^{-3}$. We set the learning rate at $1.2\cdot 10^{-4}$ for the feature generator $f_{G}(\cdot)$ and $5\cdot 10^{-4}$ for our discriminators. We adopt the learning rate $4\cdot 10^{-2}$ for the proxies as suggested in (Roth, Vinyals, and Akata 2022). For most of the experiments, we fixed the batch size to 90 as a default setting, which is consistent with (Kim et al. 2020). Empirically we apply batch normalization on the domain-level discriminator to reduce its correlation variance within the batch. For all experiments, the first layer of $f_{C}(\cdot)$ is set to 512. For the second layer, we assigned 128 dimensions to the CUB200 and CARS196 datasets, 8192 dimensions to the SOP datasets, and 4096 dimensions to the In-Shop datasets. We set $\{\eta=0.005,\gamma=0.0075\}$ for CUB200, and $\{\eta=0.01,\gamma=0.0075\}$ for CARS196. We select $\{\eta=0.01,\gamma=0.005\}$ for both SOP and In-Shop datasets.

Comparing with Proxy Baselines. We compare the performance of our approach with the existing proxy-based metric learning methods and the recent state-of-the-art metric learning methods on the popular benchmarks introduced above (refer to Table 1). We observe that our DADA frameworks can significantly improve the performance of the original proxy-based DML methods (marked with $\vartriangleright$) by a large margin. Specifically, comparing with the original PA method on ResNet50, our proposed PA+DADA outperforms $3.2pp$ ($4.6\%$) on the recall@1 of CUB200 and $4.4pp$ ($5.0\%$) on the recall@1 of CARS196. On the larger datasets (SOP and In-Shop), our method is also better than the original PA and PNCA++.

Comparing with state-of-the-art. We further compare the performance of our method with the state-of-the-art methods based on the CNN backbones as listed in Table 1 and 2. For the CARS196 dataset, our method reaches $92.1$ on Recall@1, which has a $0.9pp$ improvement over the previous state-of-the-art AVSL (Zhang et al. 2022) on the ResNet50 backbone. For CUB200, our method outperforms the previous state-of-the-art AVSL $0.6pp$ on Recall@1, $0.2pp$ on Recall@2, and $0.1pp$ on Recall@4. We observe that our performance on SOP and In-Shop is limited but very close to the previous state-of-the-art IBC (Seidenschwarz, Elezi, and Leal-Taixé 2021), CRT (Kan et al. 2022), and HIST (Lim et al. 2022) on a few metrics. The lesser improvement in the high-value recall of these two datasets is mainly due to the large number of classes (11318 and 3997) and the limited number of samples in each class (less than 10). This causes some difficulty for our category-level discriminator to learn the discriminative information. Nevertheless, our method still achieves good performance comparable to those of the state-of-the-art methods in all metrics and outperforms other proxy-related methods on these two datasets. We will investigate techniques to overcome this limitation in our future works.

Contributions of the Objective Components. We analyze the ablation study to evaluate the contribution of each objective component of our proposed framework based on both ProxyAnchor (Kim et al. 2020) and ProxyNCA++ (Teh, DeVries, and Taylor 2020) on the CUB200 as listed in Table 3. We first notice that the data augmentation strategy (Aug) does not improve our baseline significantly in the absence of $\mathcal{L}_{adv}$ and $\mathcal{L}_{cls}$. This is because, without those regularization losses, Aug simply boosts some redundant positive samples and the mixed features do not take part in training. We conclude that the domain-level discriminator with $\mathcal{L}_{adv}$ has higher efficiency when the category-level discriminator with $\mathcal{L}_{cls}$ helps regularize the space and avoid the inter-class ambiguity. It increases the improvement to $+2.3pp$ on R@1 and $+1.7pp$ on MAP@R from $+1.1pp$ on R@1 and $+0.8pp$ on MAP@R in comparison with the single $\mathcal{L}_{adv}$ setting. We also demonstrate that the efficiency of the category-level classifier ($+\mathcal{L}_{cls}$) can be further improved by comparing the discrepancy of class prediction ($+\mathcal{L}_{d}$) between the source data and target proxies in adversarial learning. Comparing the general discrepancy L1 distance ($+\mathcal{L}_{1}$), the proposed NWD also shows increasing performances on both R@1 and MAP@R. A similar conclusion can also be driven by the results based on ProxyNCA. Therefore, we conclude that the combination of the domain and the category-level discriminator is more suitable for proxy-based DML than the settings with any single discriminator. We also study the impact of our hyperparameters and the combination of data groups that apply domain adaptation in the Appendix.