Test-time Training for Data-efficient UCDR

The Semantic Neighbourhood and Mixture Prediction Network (SnMpNet) has a deep neural architecture, with SE-ResNet50 as its backbone feature extractor. The network implies a semantic feature (through a linear projection layer) of 300-d, on top of the SE-ResNet (2048-d). SnMpNet aims to learn this semantic feature to be domain-invariant, as well as meaningful, so that domain-generalization and category-wise generalization can be achieved simultaneously. Following CuMix [16], SnMpNet also treats the input data in a mixed format, where the mixing may be performed inter or intra-domain. Thus, for given samples $\mathbf{x}_{i}$, $\mathbf{x_{j}}$ and $\mathbf{x}_{k}$ from training set $\mathcal{D}_{tr}=\{(\mathbf{x}_{m},y_{m})\}_{m=1}^{N}$ of $N$-samples, the input to the network is computed as, $\mathbf{x}=\alpha\mathbf{x}_{i}+(1-\alpha)[\lambda\mathbf{x}_{j}+(1-\lambda)\mathbf{x}_{k}]$. Here, $\alpha\sim Beta(\beta,\beta)$ and $\lambda\sim Bernoulli(\gamma,\gamma)$, with $\beta$ and $\gamma$ being hyper-parameters. To obtain a domain-invariant representation of such mixed samples, the mixture-prediction loss  $\mathcal{L}_{Mp}$ is introduced, which only predicts the correct ratio of the component categories present in $\mathbf{x}$, and forgets about their domain. Additionally, SnMpNet also minimizes a semantic neighbourhood loss $\mathcal{L}_{Sn}$, which is essentially a cross-entropy loss computed between the latent-space representation of $\mathbf{x}$ and its components’ semantic ground-truth data (eg., word2vec representations of their category-name). Combining both these losses, the network learns meaningful semantic-representation of data, which are domain-agnostic. Final retrieval is performed in the learned semantic-space on the basis of the Euclidean distance between the query sample and search-set instances.

In our work, we retain the architecture and training methodology of SnMpNet unmodified, which provisions for this base model to be replaced by any other cross-domain retrieval algorithm with a shared-space representation learning technique. Next, we discuss the test-time training proposed on top of such a retrieval algorithm to enhance its performance for UCDR. We begin this discussion by providing a brief insight into the self-supervision techniques explored for the same.

Self-supervised learning (SSL) has become a popular choice when learning from unlabeled or unstructured data. Here we specifically explore the following three SSL-loss components in this regard. We choose RotNet [10] and Jigsaw[17] because of their simplicity, as well as Barlow twins [28] for its effectiveness. We first discuss the details on the loss functions and then describe the adaptation process followed using each of these losses in this work.

RotNet Loss: RotNet [10] has been a very popular choice for introducing self-supervision in a feature learning network, due to its simplicity and effectiveness. It uses four rotations of an input image, in angles of $\{0^{\circ},90^{\circ},180^{\circ},270^{\circ}\}$, and learns to predict the rotation-index of any such sample from its corresponding feature representations. Thus, the RotNet-loss is computed as,

$\displaystyle\mathcal{L}_{SSL}^{rotnet}$

$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}{\mathcal{L}_{CE}(r_{i},h_{i})}$

(1)

where $\mathcal{L}_{CE}$ is the cross-entropy loss computed between the true rotation-index $r_{i}$ and the predicted index $h_{i}$. The loss is averaged over the total $n$-number of samples present in the test set. This loss is minimized without the direct class supervision of the samples, but indirectly learns the semantic content of the data through the rotation prediction.

Jigsaw Puzzles Loss: Similar to RotNet, jigsaw puzzle [17] is another very effective self-supervision component, which has been used in transfer learning [17], domain-generalization [26] etc. Here, any input image is broken down into a number of patches, based on the fixed grid size. For example, following the authors’ approach in [17], we resize and then break down an image into $9$-patches using a $3\times 3$ grid. These $9$-patches are then shuffled and a total of $31$ different combinations (out of a possible $9!$) of jigsaw images are created through permutations. Now, the network is trained to predict the permutation index of such jigsaw images, forming the following jigsaw-puzzle loss,

$\displaystyle\mathcal{L}_{SSL}^{jigsaw}$

$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}{\mathcal{L}_{CE}(p_{i},h_{i})}$

(2)

where this cross-entropy loss is computed between the true permutation-index $p_{i}$ and the predicted index $h_{i}$, and averaged over the total number of samples present in the test set.

Barlow Twins Loss: We follow a similar formulation of this loss, as proposed in [28]. For each image in the test set, $X_{i}$, we create two differently augmented version of the same as $X^{(1)}_{i}$ and $X^{(2)}_{i}$. Augmented versions are created through various image operations, such as Gaussian blur, grayscale transformation, solarization, etc. A cross-correlation matrix $\mathcal{C}\in\mathbb{R}^{d\times d}$ is computed between the feature representations of these augmented versions as $\mathbf{x}^{(1)}_{i}\in\mathbb{R}^{d}$ and $\mathbf{x}^{(2)}_{i}\in\mathbb{R}^{d}$ and the following loss function is minimized,

$\displaystyle\mathcal{L}_{SSL}^{BT}$

$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}{[\sum_{a=1}^{d}{(1-\mathcal{C}_{aa})^{2}+\lambda\sum_{a=1}^{d}{\sum_{b\neq a}{\mathcal{C}_{ab}^{2}}}}]}$

(3)

where,

$\displaystyle\mathcal{C}_{ab}$

$\displaystyle=\frac{\sum_{i=1}^{n}{\mathbf{x}^{(1)}_{i,a}\mathbf{x}^{(2)}_{i,b}}}{\sqrt{\sum_{i=1}^{n}{(\mathbf{x}^{(1)}_{i,a}})^{2}}\sqrt{\sum_{i=1}^{n}{(\mathbf{x}^{(2)}_{i,b}})^{2}}}$

(4)

Here, the main idea is to make the diagonal term of the cross-correlation matrix 1, so that the embedding becomes invariant to any distortion. Additionally, the off-diagonal terms are pushed towards zero to de-correlate the different vector components of the embedding [28]. It has been reported to be particularly successful in image classification problems in the low-data regime.

In our implementation, these self-supervised losses only become effective to adapt the pre-trained base model on the test sets during retrieval. Thus, we introduce an additional auxiliary classifier with existing pre-trained base-model (on top of the semantic embedding (300-d) of SnMpNet) to compute the preferred loss-variant. However, this additional classifier is required only for RotNet and Jigsaw loss computations; for Barlow Twins, we leverage the 300-d embeddings directly from SnMpNet to compute the $\mathcal{C}$-matrix. Next, we discuss our adaptation or test-time training in detail.

Here, we discuss the proposed test-time training or adaptation framework to address the UCDR in data-efficient manner. In our setup, since the training data does not have samples from many numbers of domains or categories, the generalization ability of the base model SnMpNet would be low. In other words, the network may be biased towards the job of sketch-based image retrieval, in case sketch and image are the domains present during training. But its performance degrades (details in the experiments section, Table 1) when a cartoon or painting is presented as query. We hypothesize that under such a condition, any information or clue extracted from test query may help the network to adapt and retrieve better for that query domain.

Towards that goal, we propose to perform a single-step parameter update of the network using the gradient computed through any of the SSL-loss components discussed before. Thus, for any test sample $X_{te}$, we perform a forward-propagation to compute the $\mathcal{L}_{SSL}^{v},v\in\{rotnet,jigsaw,BT\}$, and then compute the corresponding gradients and back-propagate through the network components (SE-ResNet-50 based feature extractor, linear projection layer for learning the semantic-embedding and the auxiliary classifier) just once to tune the already trained SnMpNet on the basis of the test sample itself. We make the final inference on the retrieval list for the query sample, using this updated model instead of the previously-trained SnMpNet. The advantage of this proposed test-time adaptation is that any pre-trained retrieval network can be used as the base model in this case, and no modification in the training strategy or additional data / computation during training is required.

We experimented with two large-scale datasets:
Sketchy extended [21] contains $75,471$ sketches and $73,002$ images across $125$ categories. Out of these total $125$-categories, $93$, $11$, and $21$-classes are used for training, validation, and testing, respectively (following the ZS-SBIR setup [6]). The presence of annotated samples from only two domains makes this dataset an excellent choice for data-efficient generalized retrieval scenario. Thus, we largely use this dataset for pre-training the base model, and unknwon samples from other domains (except sketch and image from DomainNet) to validate the effectiveness of proposed test-time adaptation in data-efficient manner.

DomainNet [19] has approximately $6,00,000$ samples from $345$ categories, collected in six domains, namely, Clip-art, Sketch, Real, Quickdraw, Infograph, and Painting. Following [16], $245$ and $55$ classes are used for training and validation [16], respectively. To simulate the unknown domain, we leave one domain out and rests are used for training. During testing, the samples from this left-out domain is used as query. We train the base model using this setting for evaluating the traditional UCDR only.

It is not be noted that, for fair comparison, we have always reported results of SnMpNet and proposed variants on the same training and test data-splits.

We use PyTorch 1.1.0 (following [18]) and a single Nvidia Tesla v100 GPU for all experiments. Maintaining that any pre-trained baseline model can be used for the purpose, we use the same training parameters for SnMpNet with either just two training domains (for data-efficient UCDR), or with five domains (standard UCDR). For our test-time adaptation, we update the pre-trained model for single-iteration per query sample. To prevent huge parameter divergence from the learned distribution, we use a lower learning rate (1e-6 to 1e-5) and use SGD optimizer with weight decay and Nesterov momentum of 0.9, with a batch-size of $32$. This batch is created by applying standard data augmentation techniques, like random resized-crop, horizontal-flip, and color-jitter on the query sample, during the single iteration of training at test-time.

First, we analyze the effect of limited training data for UCDR. Towards that goal, we pre-train the SnMpNet model with the training-split of Sketchy-extended dataset [6], and test on the additional domain query data from DomainNet [18]. This model is treated as the base-model in data-efficient UCDR scenario for the rest of the paper. We compare this performance with the results reported in [18], where the same model is trained with multi-domain data from DomainNet. Following [18], these results are reported in the form of mAP@200 and Prec@200 in Table 1. Also, we perform experiments for two different configurations of the search set: (1) when the search set contains samples from unseen categories only; and (2) when both seen and unseen categories samples are present in the search set. We exclude the Sketch domain from DomainNet as query in this experiment, as it is already present in the training set of Sketchy-extended and hence isn’t effectively unseen-domain to the model.

As we can observe from Table 1, the performance of SnMpNet drops significantly when the number of training domains becomes restricted. This result is on par with our expectation, since traditional SnMpNet is not originally designed to learn generalization in a data-efficient manner. This regression in performance has been observed for all unseen query domains, except for Quickdraw. This may be because the contents of Quickdraw is basically Sketch (but very abstract in nature); thus, SnMpNet can retrieve for this case easily, compared to other drastically different (eg. Painting or Infograph) query domains. Thus, this observation justifies the need for further effort in this direction.

Next, we explore the effectiveness of test-time training in data-efficient learning context. We observe the performance of all 3 SnMpNet variants under limited training domains in Table 1. We can observe that for all 4 query domains (sketch excluded), test-time training has improved the performance of SnMpNet [18]. Particularly, rotation-SnMpnet significantly outperforms the base model on all four query domains and almost matches the original SnMpNet [18] performance (trained on all five-domains). These results are encouraging and support our previous hypothesis that any information/hint extracted from the unknown test sample could help improve the model’s generalization performance under such a challenging condition. Also, the proposed adaptation does not require any modification on the pre-trained base network, which makes it easy to integrate with any data-efficient base model.

Next, we study the effectiveness of such test-time adaptation for traditional UCDR. Following [18], we assume that sufficient multi-domain data is available to the model for training, and aim to explore if proposed test-time adaptation can further improve the performance of the model for UCDR. Thus, the SnMpNet is pre-trained with training samples from five domains (leaving one out) from DomainNet, and used as base model. While testing the query from unknown domain, we apply the test-time adaptation discussed in Section 3. We summarize our observations in Table 2.

We see that at least one of the proposed variants outperforms SnMpNet on Painting, and Clipart domains for both configurations of the search set. However, performance drops for Sketch, Infograph and Quickdraw by small amounts. This also impacts the overall Average performance. Only BT-SnMpNet outperforms SnMpNet on Prec@200 for unseen-class search set. We note this to be a limitation of test-time adaptation in traditional UCDR. In case sufficient training data is available during training, the proposed self-supervision-based test-time adaptation cannot improveme on the pre-trained base model’s retrieval performance. It can be noted that such a failure case is on the similar line of analysis for TTT reported by TTT++ [15] in the context of traditional image-classification.

With these results, we now move on to the in-depth analysis of the proposed approach in the next section.

We first analyze the impact of different design parameters during test-time training of SnMpNet. We primarily identify the learning rate, embedding-dimension, optimizer configuration, no. of iterations during test-time update as some of the crucial parameters effecting the final performance of model. We summarize the effect of each of these parameters in Table 3, using Infograph as the query domain. We observe that the test-time adaptation strategy is almost robust to the number of update-iterations, optimizer configuration or embedding-dimension etc., since the reported mAP@200 and Prec@200 remain similar against these conditions. However, learning rate of this adaptation seems to affect the retrieval results noticably. We see that a learning rate of 1e-6 for the base model and 1e-5 for the auxiliary classifier performs much better than 1e-4 and 1e-3, respectively. For context, pre-trained SnMpNet final learning rate is 1e-6 for our experiments. This finding is on par with recommendations in [23] of TTT learning rate being at the same scale of the pre-trained model’s final epoch of training.

Here, we explore variants of the test-time update strategy based on the literature [23]. We refer to the adaptation followed in this work as Standard adaptation, where the base-model parameters are updated based on each query sample, and then discarded for the next query. Essentially, for each query, the original pre-trained SnMpNet is being adapted and accordingly retrieval is performed. We also explore the retrieval results using the online-version [23] of adaptation, where the updates from each sample get accumulated and the pre-trained base model keeps updating itself sequentially for each incoming query. This strategy has been reported to be effective in the classification setting for gradually changing distribution shifts at test time. However, the evaluation summarized in Table 4 clearly depict that it is not as effective as the standard adaptation process for data-efficient UCDR. This may be because that UCDR is much more complex compared to traditional classification problem. In UCDR, each query can potentially belong to a different domain and category, and thus a cumulative update from each of such samples can lead to complete parameter divergence for the model. Thus the assumption of gradually changing distribution shifts that hold for traditional TTT benchmarks [23], won’t be valid for UCDR.