Few-Shot Adaptation of Pre-Trained Networks for Domain Shift

1 Introduction

While deep neural networks have demonstrated remarkable ability in representation learning, their performance relies heavily on the assumption that training (source domain) and test (target domain) data distributions are the same. However, as real-world data collection can be difficult, time-consuming or expensive, it may not be feasible to adequately capture all potential variation in the training set, such that test samples may be subject to domain shift (also known as covariate shift) due to factors such as illumination, pose and style [Gulrajani and Lopez-Paz, 2021; Koh et al., 2020]. To prevent severe performance degradation when models are deployed, timely adaptation to the target test distribution is needed.

∗Work done while at Institute for Infocomm Research.

To carry out this adaptation, a range of methods with varying requirements on the availability of source and target domain data have been developed. In the classic domain adaptation (DA) setting, methods assume source and target data are jointly available for training [Wilson and Cook, 2020], which compromises the privacy of source domain data. To address this, methods for the source-free DA setting [Qiu et al., 2021; Yang et al., 2021; Liang et al., 2020] instead adapt a pretrained source model using only unlabeled target data, but they still require access to the entire unlabeled target dataset like traditional DA methods. This can delay adaptation, or even make it impractical, when collecting the unlabeled target data is costly in terms of time or other resources.

Very recently, test-time adaptation methods [Nado et al., 2020; Schneider et al., 2020; Wang et al., 2021] have been proposed to adapt pre-trained source models on-the-fly without first collecting a large target domain dataset, by updating batch normalization (BN) layers in pre-trained source models with every mini-batch of target domain data. While these methods reduce the delay to initiate adaptation, they face three main challenges: 1) there is no guarantee that the selfsupervised or unsupervised objectives used in these methods can correct domain shift without using target domain labels, 2) performance is dependent on having large mini-batches to obtain good estimates of BN parameters and statistics, 3) test samples need to be class-balanced, which is not always the case in real-world deployments. We show in Figure 1 on the VisDA benchmark that while in ideal conditions (large batch sizes, balanced classes) test-time adaptation methods can improve performance of source models, their performance is severely affected outside of these ideal conditions; in fact, for one of the methods, we observe catastrophic failure where all outputs "collapse" to one or a few classes.

In this work, we propose a different solution to the higherlevel objective of source-free domain adaptation with limited target domain data. We suggest that adaptation with a few labeled target domain samples instead of mini-batches of unlabeled data can address the identified challenges with testtime adaptation methods. We propose a new method under a new source-free k-shot adaptation setting, that adapts a pretrained source model using k-shot labeled (or support) samples from the target domain; a comparison with existing DA settings is shown in Table 1. Our proposed method adapts batch normalization (BN) layers of deep source models with the support samples using a reparameterization of BN-layers and a low-dimensional approximation of the optimal target domain BN statistics. Although BN layer modulation has been explored for DA [Chang et al., 2019; Nado et al., 2020; Schneider et al., 2020; Wang et al., 2021], reliably optimizing BN layers with extremely few support samples (as few as one per class) is a new and challenging problem. Naively optimizing high-dimensional network parameters directly risks a severely ill-posed problem caused by data scarcity. Our method approximates optimal target domain BN statistics with a small set of spanning vectors representing both source and support datasets and finetunes only the combination coefficients, significantly reducing the number of parameters that need to be learned during adaptation. Our method is inspired by the success of controlling sample stylization through BN layers [Huang and Belongie, 2017; Zhou et al., 2021; Nam and Kim, 2018], and we aim to approximate the optimal style to stylize the target samples to address domain shift.

More broadly, our proposed source-free k-shot adaptation setting addresses salient challenges associated with data availability and inference efficiency in real-world deployments. Specifically, the setting helps to protect the privacy of source domains, and has low requirements for target data availability of only k labeled samples per class during adaptation. During inference, test batches can be of any size and class composition with no restrictions. As model parameters are not further updated at test-time after k-shot adaptation, our proposed setting enables more efficient inference in comparison to the test-time adaptation setting.

We extensively evaluate our proposed method on 5 image classification and 4 semantic segmentation benchmark datasets and show that it achieves competitive adaptation performance through comprehensive experimental validations and comparisons with state-of-the-art methods.

2 RelatedWork

Domain adaptation (DA). In the DA setting, a network is trained jointly on labeled source dataset and unlabeled target dataset to optimize task performance on the target domain, as surveyed in [Wilson and Cook, 2020]. Existing sourcefree DA approaches adapt a pre-trained source model with unlabeled target dataset by generating source-like representations [Qiu et al., 2021], making use of clustering structures in target features for classification [Yang et al., 2021; Liang et al., 2020], or aligning target samples to source hypothesis through entropy minimization and information maximization [Liang et al., 2020; Kundu et al., 2020; Li et al., 2020]. While these methods are similar to our proposed method in not accessing source data, they use a large number of unlabeled target samples for adaptation.

Domain generalization (DG). DG methods aim to learn a model robust to unseen domain shifts by training only on labeled source datasets [Gulrajani and Lopez-Paz, 2021; Koh et al., 2020]. Our proposed method can work directly on source models learned by DG strategies, and hence complements existing DG methods to further close the generalization gap on specific domains of interest.

Test-time adaptation. Test-time adaptation methods update a pre-trained source model continuously during test time with unlabeled target samples to improve performance on the target distribution. [Sun et al., 2020] update network parameters guided by image rotation estimation. However, the method requires training with the self-supervised task on the source dataset, and hence does not work directly with any pre-trained source model. Recent works propose adapting BN layers of the source model instead [Nado et al., 2020; Schneider et al., 2020; Li et al., 2016]. These test-time BN methods reestimate BN statistics on each mini-batch with mini-batch evaluation to correct domain shifts due to shifts in first and second moments of the data distribution [Nado et al., 2020; Schneider et al., 2020]. Tent [Wang et al., 2021] further updates batch-norm weights and biases with entropy minimization to obtain more confident predictions. However, the reliability of test-time adaptation depends on the number and class-distribution of samples in each mini-batch.

Few-shot transfer learning. There is a wide range of works that learn a metric space specifically for k-shot tasks, and use meta-learning to learn adaptation strategies [Pan and Yang, 2010]. These approaches typically need specific network architectures, loss functions, training strategies requiring multiple source domains or joint training with source and support data together. Since they do not directly work with a pretrained source model, we do not focus on them here.

Another popular strategy is model finetuning or weight transfer from a pre-trained source model, and our proposed method also falls in this category. Directly finetuning all source model weights on a limited support set with small k is known to severely overfit, so a support set of at least k = 100

is often required [Yosinski et al., 2014; Scott et al., 2018]. Recent works constrain the dimensionality of learnable parameter space. FLUTE [Triantafillou et al., 2021] finetunes BN parameters with a nearest-centroid classifier. [Yoo et al., 2018] clusters network parameters and constrains all parameters in a cluster to share the same update, but this method requires activations on the source dataset for clustering.

3 ProposedMethod

We first revisit BN layer optimization for DA in Section 3.1 and establish that optimizing the BN layer can be achieved by optimizing BN statistics {µ, σ}. Then, we introduce our proposed method that optimizes {µ,σ} on a k-shot support set Lspt with a supervised loss in Section 3.2 in the context of classification. For classification tasks, the support set contains k labeled target samples (x, y) per class, with image x and one-hot class vector y. For semantic segmentation, the support set contains k labeled target samples in total, and y is the segmentation map corresponding to image x. Implementation details are described in Section 3.3.

3.1 Revisiting BN Layer Optimization for DA

The batch-norm layer operation f, shown below in Equation 1, is a function acting on input features Z parameterized by four variables: BN statistics {µ,σ} and BN parameters {γ, β}. Here µ = E[Z] ∈ RC and σ = √V[Z] + ϵ ∈ RC are channel-wise feature statistics for C channels, and γ, β ∈ RC are layer weight and bias parameters.

We denote f(Z;µt,σt, γt, βt) as the BN operation with optimal target domain BN statistics and parameters {µt,σt, γt, βt}. By rewriting this operation in terms of source model BN parameters βs and γs

we see that it is equivalent to changing BN statistics in the source model BN operation f(Z; ˜µ, ˜σ, γs, βs) from ˜σ = σs and ˜µ = µs to ˜σ = σtγs γt and ˜µ = µt - (βt-βs)σt γt , where {µs,σs, γs, βs} are BN statistics and parameters in the pre-trained source model. This observation implies that we can obtain the optimal target domain BN layer f(Z;µt,σt, γt, βt) by optimizing only the BN adaptation statistics {˜µ, ˜σ}. Edge cases of zero-valued elements in γs can be avoided by adding a small number to such elements. Setting elements as zero in γt implies that the corresponding features are completely ignored, which we assume unlikely given a well-trained source feature extractor. A detailed explanation is provided in Appendix B.

3.2 LowDimensional Approximation for BN Layer Adaptation

Overall formulation. In our setting with extremely few support samples, optimizing BN layers by estimating the highdimensional BN parameters risks solving an ill-posed problem caused by data scarcity. We hypothesize that source and support set statistics are related to those of the target domain due to shared classes and the relatedness of domains, and propose to approximate target BN statistics in the linear span of these available statistics. We assume a linear relationship as it can be interpreted as style mixing and for computational efficiency. Specifically, for a BN layer with C channels, we approximate the optimal target BN statistics by

where M, Σ ∈ RC×(n+1) contain n+1 spanning vectors for first and second moment statistics, and η, ρ ∈ Rn+1 are the learnable Linear Combination Coefficients for batch normalization Statistics (LCCS). Our proposed formulation generalizes the original BN operation where ηspt = ρspt = 0 and ηs = ρs = 1. It also generalizes channel-wise finetuning of BN statistics and parameters: by setting support set spanning vectors as basis vectors and n = C, our method is equivalent to unconstrained optimization of BN layers. The key advantage of our approach is that when n << C, we can optimize BN layers using significantly fewer learnable parameters.

Spanning vectors. We first describe how source and support set spanning vectors are obtained.

Source domain spanning vectors (µs,σs): We directly utilize BN statistics {µs,σs} from the pre-trained source model. We incorporate source representations to benefit from knowledge shared between source and target domains.

Support set spanning vectors (Mspt,Σspt): The support set is represented by n spanning vectors extracted by aggregation functions or dimensionality reduction methods such as singular value decomposition (SVD). We choose vectors from the orthogonal basis vectors of per-sample feature statistics at the BN layer. Specifically, forMspt (and similarly for Σspt), let ˜Zµ ∈ RC×|Lspt | denote the matrix of channelwise feature mean for each sample. We set the first spanning vector as the overall channel-wise feature mean µspt computed across all support samples. We then subtract this first vector as ˜Z(1) µ = ˜Zµ - ˜ZµµsptµTspt /∥µspt∥2 , factorize ˜Z(1) µ = USV T using SVD, and extract the top n - 1 vectors in US as the remaining spanning vectors. The resulting features post-BN are then ZBN = ( Z-µLCCS σLCCS ) γs + βs.

Learning LCCS parameters. We finetune LCCS parameters {η, ρ} from all BN layers by minimizing cross-entropy on the support set. After finetuning the LCCS parameters, we use the estimated BN adaptation statistics µLCCS and σLCCS on the target domain for inference. Note that we do not adjust the BN statistics further during test time. Hence, unlike test-time BN, our adapted model is not affected by testtime mini-batch size and class distribution.

Comparison of parameter counts. In Table 2, we summarize the number of network parameters versus LCCS parameters when n = 1 for common network architectures. Model finetuning involves updating network parameters according to a new objective. Without sufficient labeled samples, updating a large number of parameters is an ill-posed problem. While adapting only the BN parameters allows a smaller number of parameters to be tuned, this is still a large number of learnable parameters in deep networks. Our proposed method dramatically reduces this number to only 4 LCCS parameters per BN layer when n = 1, a greater than 100-fold reduction compared to the number of BN parameters. As a result, even when provided with only an extremely small support set, our proposed method is less prone to overfitting.

3.3 Implementation Details

The overall adaptation objective is to finetune LCCS parameters to minimize cross-entropy loss L(η, ρ) = - ∑(x,y)∈Lspt y log h(x; η, ρ) on the support set, where x and y are input and one-hot class encoding of support samples Lspt , and h is the source model with learnable LCCS parameters {η, ρ}. The proposed method comprises an initialization stage and a gradient update stage.

Initialization stage. We search for initialization values for LCCS to warm start the optimization process. We first compute the support BN statistics µspt and σspt by exponential moving average (EMA) for m epochs to allow µspt and σspt to update smoothly. Then, we conduct a onedimensional grid search on the LCCS parameters by setting ηspt = ρspt = [v 0 · · · 0]T where v ∈ {0, 0.1, . . . , 1.0} and ηs = ρs = 1 - v with values tied across all BN layers. The initialization value that minimizes cross-entropy loss on the support samples is selected.

Gradient update stage. We compute support set spanning vectors Mspt and Σspt with initialized LCCS parameters, and further optimize the parameters with gradient descent for m epochs. In this stage, parameter values are not tied across BN layers and we do not impose sum-to-one constraints on the coefficients to allow more diverse combinations.

We set m = 10 epochs in all our experiments. Support samples are augmented with the same data augmentations for source model training. We retain the pre-trained source classifier with extremely small support sets, and update it when sufficient support samples are available to represent the target domain and learn a new target classifier. In our experiments, we use the nearest-centroid (NCC) classifier following [Triantafillou et al., 2021] as the default classifier for k ≥ 5.

4 Experiments

We evaluate on image classification and segmentation tasks using publicly available benchmark datasets and compare our proposed method to existing source-free methods. For each dataset, our source models are the base models trained on source domain(s) using empirical risk minimization (ERM) or DG methods that have state-of-the-art performance on that dataset. Support samples are randomly selected and experiment results are averaged over at least 3 seeds. We use a mini-batch size of 32 for classification and 1 for segmentation, and use the Adam optimizer with 0.001 learning rate for finetuning LCCS. We provide brief descriptions of taskspecific experimental setups in the respective sections; detailed descriptions of the datasets, their corresponding source models and implementation are provided in Appendix C.1.

4.1 Image Classification

We evaluate on 5 classification datasets covering various types of shifts: PACS (7 classes in 4 domains) for style shift, VisDA (12 classes in 2 domains) for synthetic-to-real shift, Camelyon17 (2 classes in 2 domains) and iWildCam (182 classes in 2 domains) for naturally-occurring shifts, and Office (31 classes in 3 domains) for objects photographed in different environments. We adopt evaluation metrics used by recent works: accuracy for PACS and Camelyon17, macroF1 for iWildCam, average precision for VisDA and average per-class accuracy for Office. For LCCS we set n = 1 on iWildCam due to memory constraints, and n = k × # classes otherwise.

Comparison to test-time adaptation. We compare our method with the baseline source models and test-time adaptation methods that also augment only batch normalization layers. These are Test-time BN [Nado et al., 2020] and the state-of-the-art Tent [Wang et al., 2021] that adapts BN parameters by entropy minimization with default hyperparameters: Adam or SGD with momentum with mini-batch size 128 (larger than support set size with k = 10 on datasets evaluated) and learning rate 0.001.

We see from Table 3 that our proposed method improves over source models even with a single example per class for all datasets evaluated, and outperforms Test-time BN and Tent on most cases when k ≥ 5. We observe that Tent performance is dependent on optimizer choice, with Adam outperforming SGD by approximately 2% on PACS and SGD outperforming Adam by as much as 39.2% on VisDA. Since the better performing optimizer is dataset-dependent, this makes the practical usage of Tent challenging. On PACS, our proposed method outperforms the best case of Tent when the support set has 5 or 10 samples per class. On Camelyon17, all BN adaptation methods improve over source model accuracy by at least 18% in the best case, implying that replacing source BN statistics with target BN statistics is effective in addressing domain shift in this dataset. On the 182-class iWildCam dataset, the test dataset is imbalanced: 102 classes have at least 1 sample, and only 87 classes have at least 5 samples. Hence, to prevent introducing more class imbalance in the adaptation and evaluation processes, we only evaluate the setting where k = 1. Our proposed method improves over the source model, while all test-time adaptation methods degrade performance. Similarly, on VisDA, our proposed method obtains the best performance on the target domain.

Effect of test batches: Although test-time adaptation methods can improve the source model without the supervision of labeled samples, their performance relies on streaming conditions and severely degrades with smaller mini-batch size and class-imbalanced mini-batches as shown in Table 4, and in detail in Appendix C.2. We constructed long-tailed imbalanced PACS and VisDA following the procedure in [Cao et al., 2019], where sample sizes decay exponentially from the first to last class with α being the ratio of samples in the largest to smallest class. The good performance of test-time methods is dependent on having large, class-balanced minibatches whereas our method is independent of streaming conditions and can more reliably adapt to the target domain.

Comparison to few-shot transfer learning. As far as we know, there are no existing methods designed and evaluated specifically for our source-free few-shot DA setting. We thus apply existing methods in DA and few-shot transfer learning to provide benchmark performance. We compare with AdaBN [Li et al., 2016] by replacing source BN statistics with those calculated on target support set, finetuning the source model on BN parameters, classifier, or feature extractor, finetuning the entire model with L2 or L2 -SP [Li et al., 2018] or DELTA [Li et al., 2019] regularization, Late Fusion [Hoffman et al., 2013] which averages scores from source and target classifier, and FLUTE [Triantafillou et al., 2021] which optimizes BN parameters with nearest-centroid classifier. FLUTE assumes the availability of multiple source datasets to train multiple sets of BN parameters for further blending to initialize the finetuning process. Since we only have access to the pre-trained source model in our setting, we reduce FLUTE to the single source dataset case and initialize FLUTE with single source BN parameters. We follow learning settings in [Li et al., 2019] for regularized finetuning and use the SGD optimizer with momentum and weight decay 0.0004, and L2 regularization is also added to finetuning on classifier or feature extractor to prevent overfitting. For all other methods that finetune on BN layers, we use the Adam optimizer following [Wang et al., 2021]. We set learning rate 0.001, mini-batch size 32 and epochs 10 for all methods and datasets evaluated.

From Table 5, we observe that regularized finetuning tends to adapt well at k = 1, but performance can lag behind with larger support sets. AdaBN does not consistently improve adaptation; we show in Appendix D that completely replacing source BN statistics degrades performance for VisDA. Overall, our proposed method has the best performance in most cases. We also compare our method to L2 and FLUTE on Office in Table 6, and our method consistently outperforms on all 6 source-target combinations. Though the transfer learning methods produce comparable performance on the relatively easy dataset Camelyon17 (two classes, two domains), LCCS outperforms on more difficult datasets, which demonstrates the effectiveness of the proposed low-dimensional finetuning strategy.

Comparison to source-free unsupervised domain adaptation. We additionally compare our proposed method with source-free unsupervised DA methods which adapt with the entire unlabeled target dataset, including AdaBN [Li et al., 2016], SHOT [Liang et al., 2020], SFDA [Kim et al., 2020] and SDDA [Kurmi et al., 2021] in Table 6 on all 6 sourcetarget pairs in Office. We observe that self-supervision over the entire unlabeled target can produce good adaptation performance. However, despite using only 5 samples per class, our proposed method with finetuned linear classifier (source classifier finetuned for 200 epochs) has best adaptation performance in 5 out of 6 domain pairs, with an average accuracy of 88.9% compared to 88.5% obtained by state-of-the-art source-free unsupervised DA method SHOT.

SHOT outperforms on a more challenging OfficeHome dataset (71.8% vs 67.8%), reflecting the difficulty of adaptation with limited data. Nonetheless, our proposed method outperforms other few-shot methods evaluated by at least 4.5% on average, which demonstrates its effectiveness in the few-shot setting. Full results are provided in Appendix C.4.

4.2 Semantic Segmentation

For semantic segmentation, we transfer from source dataset GTAV (synthetic) to 4 target datasets: SYNTHIA (photorealistic), Cityscapes (real, urban), BDD-100K (real, urbandriving), and Mapillary (real, street-view). We use the meanIntersection-over-Union (mIoU) metric to evaluate performance, and set n = 1 for LCCS due to memory constraints. We use a total of 5 labeled images as the support set. We compare our method with SW [Pan et al., 2019], IBN-Net [Pan et al., 2018], IterNorm [Huang et al., 2019] and ISW [Choi et al., 2021] which adapt on each sample by instance whitening or standardization, as well as L2 -regularized finetuning on ISW pre-trained models. From Table 7, LCCS finetuning further improves over the state-of-the-art ISW and outperforms ISW + L2 on all 4 target datasets. The improvement on Cityscapes is as much as 7%, and is visualized in Figure 2.

4.3 Further Analysis

We provide further analysis of our method in the classification setting with base configuration n = 1 and source classifier. Additional analyses are provided in Appendix D.

Initialization and gradient update stages. We conduct ablation studies on the initialization and gradient update stages of the proposed method on the VisDA dataset. From Table 8, we see that each stage independently improves the base model's performance, showing that both stages help adaptation.

Design choices for optimizing LCCS. We further experiment with different algorithm design choices in optimizing LCCS parameters. Initializing LCCS with values tied across all BN layers (67.0% avg prec) is better than greedily initializing each BN layer sequentially starting from the shallowest layer (66.1% avg prec). In the gradient update stage, linear combination of statistics (67.0% avg prec) in Equations 3 and 4 performs better than restricting to a convex combination (66.1% avg prec).

Choice of n. From Table 9, we see that with larger support sets, more vectors can be used to represent the target domain and performance improves with a larger number of adaptable LCCS parameters. We observe in general that setting n = k × # classes for k ≥ 5 obtains the best performance, while the choice at k = 1 is dataset-dependent.

Computational cost. Our method requires little extra training time compared to that of the original network, especially with small k. For simplicity, consider one epoch of gradientbased training. Computation time for our proposed BN layer is O(N(A+Cn)) for constant A denoting the time for backpropagating from the loss to the BN layer's output, versus O(N(A+C)) for the original BN layer. We expect A >> C for large neural networks. Empirically, on VisDA using a Tesla V100-SXM2 GPU, the average training time per epoch on a vanilla ResNet-101 is 3.18s, 3.56s and 4.32s for k = 1, 5 and 10 respectively. With our default LCCS parameters of n = k ×# classes, average times are 3.48s, 4.43s and 6.40s.

5 Conclusion

In this work, we proposed the source-free k-shot domain adaptation setting and LCCS finetuning method. LCCS uses a low-dimensional approximation of BN statistics to significantly reduce the number of adaptation parameters, enabling robust adaptation using only a limited support set. Our method adapts source models with as few as one sample per class, is not affected by streaming conditions, and performs well across classification and semantic segmentation benchmarks. These characteristics make our proposed solution a favourable option for data-efficient adaptation, and provide a useful foundation for future work on this challenging topic.

Acknowledgments

This research is supported by the Agency for Science, Technology and Research (A*STAR) under its AME Programmatic Funds (Grant No. A20H6b0151).