Test-time Fourier Style Calibration for Domain Generalization

1 Introduction

Over the past few years, deep convolutional neural networks (CNNs) have been shown to perform excellently in a variety of computer vision applications. However, CNNs have demonstrated a drastic performance drop when both the training (source domain) and testing (target domain) data come from different distributions. This incapability causes numerous problems in real-world fields such as autonomous driving and medical imaging analysis. As a consequence, domain adaptation (DA) has been extensively researched as a mechanism of transferring knowledge from source domains to labeled/unlabeled target domains. Yet, DA approaches rely on known target domains at training time, which fails to generalize well to unknown testing target domains. Domain generalization (DG) has lately attracted much attention as a more challenging setting for enhancing CNNs' generalizability to unknown testing data. Specifically, DG is to train models that are resistant to distribution shifts between domains (i.e., domain shifts) so that models can generalize well when evaluated on any previously unseen target domain, provided that several source domains are available for training. To address the problem, various methods were proposed and achieved promising results, including adversarial learning [Zhou et al., 2020; Sicilia et al., 2021], meta learning [Li et al., 2018a; Balaji et al., 2018], and data augmentation [Xu et al., 2021; Zhou et al., 2021]. Among these, approaches based on data augmentation have gained popularity due to their simplicity and effectiveness, which can expose models to a wide variety of source domains to learn generalizable representation. However, most existing data augmentation approaches cannot perform well on test data that are out of the augmented source data distribution. This begs the question of how we can reduce domain gaps when the augmented source domains are insufficient to cover the target domain.

In this paper, we propose a novel yet simple method called Test-time Fourier Style Calibration (TF-Cal) to calibrate the target style features at inference-time. Our method is driven by the assumption that the lack of generalizability is due to the difference in styles of images (e.g., sketches, cartoon, photo) between the source and target [Zhou et al., 2021]. Based on this idea, we employ Fourier transformation to decompose latent features from early CNN layers into amplitude features and phase features. The amplitude encodes the intensities (style information), while the phase encodes the spatial positions (semantic information) [Piotrowski and Campbell, 1982]. To achieve test-time style calibration, we propose to convexly calibrate the target amplitude features using the source amplitude prototype estimated at train-time. Hence we narrow the style gap between the target and source domains on the fly when evaluating, while retaining the key semantic feature for classification. We additionally show that increasing the diversity of features can complement TF-Cal to further boost models' generalizability based on the convex hull theory proposed by [Albuquerque et al., 2019]. Inspired by [Xu et al., 2021] which augments amplitude in image space, we propose augmenting amplitude in feature space (AAF) that diversifies higher-level styles, while preserving semantics. Then we have the Test-time Augmentated Fourier Style Calibration (TAF-Cal) by combining TF-Cal and AAF.

Moreover, our algorithm does not require domain labels, allowing for more data collection flexibility.

The proposed work is evaluated on multiple DG benchmarks with different backbones on image classification and medical image segmentation. Through extensive experiments, we demonstrate that our method significantly improves the generalizability of CNNs and outperforms multiple state-of-the-art methods. We also show that TF-Cal is useful in conjunction with augmentation techniques. Finally, we present qualitative analysis and quantitative ablation studies to verify the effectiveness of our method.

2 RelatedWork

2.1 Domain Generalization

DG works in recent years can be roughly categorized into: domain-invariant learning, meta-learning and data augmentation approaches.

Domain-invariant learning aims to map input data to features that are invariant to source domain shift in order to be resilient to unseen target domain shift. The domain adversarial neural network (DANN) [Ganin et al., 2016], which was originally developed for DA, consists of a domain classifier trained to discriminate between the source and target distributions, and a feature extractor trained to learn features that are not only beneficial for the task (e.g., classifcation) but also capable of confusing the domain classifier. [Albuquerque et al., 2019; Sicilia et al., 2021] leverage the DANN to align the representation of source domains to learn the domain invariance in DG. [Li et al., 2018b] propose to use Maximum Mean Discrepancy with adversarial autoencoder to align domain distributions, mapping the aligned distribution to a prior distribution adversarially. Another line of works focus on minimizing contrastive loss to learn domain-invariant representation. MASF [Dou et al., 2019] minimizes contrastive losses with positive pair of same class and negative pair of different class.

Meta-learning based approach MLDG [Li et al., 2018a] simulates domain shift by dividing data from multiple sources into meta-train and meta-test sets. [Balaji et al., 2018] propose a meta-regularizer learned to be robust to domain shift by meta-learning. [Dou et al., 2019] regularize semantic structure of features while optimizing models by metalearning. While meta-learning methods can simulate the domain shift during training by dividing source domains to meta-train and meta-test sets, it is still likely that the simulated domain shift cannot cover the target domain shift.

Data augmentation method is most related to our work, attempting to augment the training images or features to enhance the generalization performance on unseen data. [Zhou et al., 2021] probabilistically mix the instance statistics of training features. [Xu et al., 2021] propose to augment amplitude of images with a regularization to form consistency between original and augmented images. To generate diverse novel training images, [Zhou et al., 2020] design an image generation network with domain-adversarial training plus optimal transport based metrics. A concurrent work [Zhang et al., 2021] has a similar motivation as ours to augment/modify target data during evaluation, developing a multi-view regualizeried meta-learning that encourages test-time augmentation with multi-view. Although these methods achieve favorable results on unseen targets, they might fail when the augmented source cannot cover the unseen target. Our method mitigates this issue by calibrating the target with a source prototype, allowing the calibrated target to be covered by the source.

3 Methodology

3.1 ProblemDefinition

We define X and Y as the input and label space. Given K source domains Ds = {D1 s , D2 s , . . . ,DK s } (1 ≤ i ≤ K) that are available at train-time, we indicate that each Dis contains N input and label pairs (xj, yj)Ni j=1 where x ∈ X , y ∈ Y and Ni is the number of inputs in the source domain Dis . The goal of DG is to train a function learned from data of Ds to perform well on an unseen target domain Dt which share the same label space Y as Ds. The function can be defined as Rθ,σ = Tσ ◦ Fθ, where Tσ is a task-specific network parameterized by σ and Fθ = {F 1 θ , F 2 θ , . . . , F L θ } (1 ≤ i ≤ L) which is a feature encoder with L layers parameterized by θ.

3.2 Background

To reduce the domain shift, we need to first understand the difference between two distributions which are measured by H-divergence proposed by [Ben-David et al., 2009]:

where classifier h : X → {0, 1}. Following that, [Albuquerque et al., 2019] defines the convex hull Λs of Ds that is a set of mixture of source domain distributions:

where the π represents non-negative coefficient in the (K - 1)-dimensional simplex ∆K-1. Then, an ideal case ¯Dt ∈ Λs is assumed that the ideal target ¯Dt lies in the source domain convex hull Λs. Under this assumption, the risk ϵt(h) on the target domain is upper-bounded [Albuquerque et al., 2019] by:

where, on the right hand side, the first term corresponds to the risks over all source domains, and the second term γ = dH∆H( ¯Dt,Dt) is the H-divergence between the ideal target ¯Dt and the real target Dt, and the third term ζ = supD′s ,D′′ s ∈Λs dH∆H (D′s , D′′ s ) which is the largest Hdivergence between any pair of source domains. H∆H corresponds to {h(x) ⊕h′ (x) | h, h′ ∈ H}. The first term can be minimized by empirical risk minimization (ERM), and the second term is hard to minimized due to having no access to the target domain at training, and the third term can be minimized by removing the source domain-discrimintive information which is the style information in the context of this paper. According to [Cha et al., 2021], ERM is a strong baseline to learn source domain invariance compared to other methods such as DANN [Ganin et al., 2016].

3.3 Motivation

As the Eq.(3) states, a model can achieve acceptable performance on the target domain by using ERM to train on the source domains when the ideal target domain ¯Dt is assumed to be covered by source convex hull Λs. However, in reality, this assumption typically does not hold. This is the reason why domain-adversarial based methods [Zhou et al., 2020; Albuquerque et al., 2019] may fail in DG as they only reduce source divergence. Therefore, augmenting images or features of source domains can possibly extend the Λs (i.e., γ → 0), which motivating a variety of works [Zhou et al., 2021; Sicilia et al., 2021] to diversify training data. Still, as illustraed in Fig.1b, it is likely that the augmented source distributions (dashed circles) may not accurately reflect the previously unseen target domain (red rhombi). This difficulty arises mostly due to the absence of target data during training. Instead of extending Λs arbitrarily, we aim to pull Dt closer to ¯ Dt with the help from source data.

Therefore, it motivates us to calibrate the target features at test-time based on what a model learned from source features. By doing so, we pull the target closer to the source convex hull Λs, and reduce the difference between Dt and ¯ Dt. To further encourage the reduction in source-target divergence, we augment the source features to expand the convex hull.

3.4 Method Overview

Our method TF-Cal is to calibrate style features. To this end, we use Fourier transformation to decompose features into phase features that contain semantics and amplitude features that contain styles. Then we let a model to learn with a style calibrating function C (TF-Cal) for reducing the gap between the target domain and source domains prototype. We further introduce an augmenting function A (AAF) for expanding the source convex hull by diversifying style features. The model then is trained in a two-stage manner: in the early stage, we train the model with AAF to learn a great diversity of style features, and in the late stage, we train the model with both AAF and TF-Cal (i.e., TAF-Cal) that uses the estimated source prototype from diverse styles. During testing, we leverage the learned C to calibrate target styles.

Our method is easy to implement as a plug-and-play module without adding any additional network. The overall pipeline is illustrated in Fig.1a and TF-Cal is illustrated in Fig.1b, and the detail of the method will be explained as follows.

3.5 Fourier Transformation of Features

Given a latent feature tensor zi = F i (zi-1 ) with zi ∈ RN×C×H×W computed from i-th layer of the feature extractor, we want to extract style information from the feature. In our implementation, we choose the early layers since they can extract style information well [Zhou et al., 2021]. Fourier transformation F(zi ), which can be implemented using Fast Fourier transformation (FFT) [Nussbaumer, 1982], is able to decompose the feature zi into the phase spectrum PH(zi ) encoding the spatial information and amplitude spectrum AM(zi ) encoding the intensity information. As illustrated in Fig.1, the amplitude component of the horse features depicts the intensity of texture, while phase component preserves the shape of object. Thus, it is natural to associate the amplitude with styles and the phase with semantics. Since our method only modifies the amplitude, the semantic information in the phase will not be affected.

The Fourier transformation F(zi ) is done on each channel C of zi :

where F(zi )(m, n) represents the value of the features in the frequency space corresponding to the coordinates m and n. With F, we can retrieve the amplitude and phase features with real Re and imaginary Im parts:

To reconstruct phase and amplitude spectrum features back into the spatial feature space, we can use the inverse Fourier Transformation F -1 (AM(zi ), PH(zi )) = zi .

3.6 Test-time Fourier Style Calibration

Estimating Source Prototype

The first step in our algorithm TF-Cal is to estimate the source domain prototype Pts ∈ R1×C×H×W that will be used for the calibrating function C. The Pts can be estimated by taking the average over the amplitude features of all data points ND from source domains along the batch dimension:

Since the Pts is estimated via all training samples, it will be a proper representation of the centroid of source styles. We leverage Pts as the anchor to pull a style closer to the source convex hull Λs. In practice, we store the AM(zis ) from each mini-batch in a memory bank, and compute the mean of all amplitude features from the bank at the end of one epoch. We will verify calibrating style with source prototype instead of a random prototype in the ablation studies section.

Training with the Style Calibrating Function

The TF-Cal is proposed to reduce the divergence between source and target styles during testing. We need to first train a model with a calibrating function C with the input of the source prototype Pts and source amplitude features AM(zis ). Therefore, we let the model's parameters be associated with the source prototype that is the common information shared by the target and source styles. The C is used to train the model by linearly interpolating between AM(zis ) and the Pts:

where η is the hyperparameter to control the strength of calibration. The calibrating function C is activated with a probability pcal during training. Before the training is finished, we store the source prototype Pts from the final epoch for the usage of calibrating target style at test-time.

Calibrating the Target Style on the Fly

During the evaluation, we leverage the C function to calibrate the target style AM(zit ) with the source prototype Pts:

where the τ is the hyperparamter to control the strength of calibration during testing. The reason for controlling τ is that, due to CNNs' inductive biases toward styles [Geirhos et al., 2019], target amplitude may contain useful style information for classification. We set the calibration probability pcal to 1 during testing. Our method requires only one test sample zit ∈ R1×C×H×W for calibration which is realistic in practical applications. Through calibration, the target style is pulled closer to the source style, helping the model to achieve a similar performance on target data as it does on source data. The advantage of our method is that we let the target and source styles share the Pts that is controllable and at hand. Thus, we can offset the source-target divergence with the guidance from the source prototype at test-time. Instead, other DG methods only rely on the source data, hoping to cover the target distribution or learn domain-invariance that overfits to source domains during training.

3.7 Augmenting Amplitude Features

As we can see from Eq.(7) and Eq.(8), the common part of the input is the Pts and the different part is the amplitude features AM(zit ) and AM(zis ). Therefore, to further boost the generalizability with TF-Cal, we want AM(zit ) and AM(zis ) to be as similar as possible. An effective way is to augment amplitude in feature space (AAF) hoping to expand the source convex hull Λs to decrease the source-target divergence. Therefore, based on the strategy of mixing and swapping amplitude of images [Xu et al., 2021], we propose to mix and swap amplitude feature pairs (AM(zis ), AM(zi′ s )) randomly from any source domains:

where δ ∼ Beta(α, α). For swapping amplitude between two samples, we set δ to 0. Compared to [Xu et al., 2021] that augment styles in the image space, AAF utilizes higher level style information from features with much more channels, which provides additional statistics for diversifying styles. In our implementation, we choose amplitude features from early layer of feature extractors since they encode style information about the inputs. Our results in the ablation studies demonstrate that augmenting amplitude on feature is better than on image, which empirically verify that AAF helps TF-Cal.

4 Experiments

4.1 Datasets and Settings

Our method is evaluated on three DG benchmarks: (1) PACS [Li et al., 2017] is an object recognition dataset that contains 4 different domains (i.e., Photo, Art-Painting, Cartoon, Sketch), including 9,991 images and 7 classes. We follow the train-val split provided by [Li et al., 2017]. The domain shift is primarily contributed by the style difference. (2) Office-Home [Venkateswara et al., 2017] is an object recognition dataset that contains 4 different domains (i.e., Art, Clipart, Product, Real-World), including 15,550 images and 65 classes. The dataset is split randomly into 90% training set and 10% validation set. The domain shift is primarily contributed by the style and viewpoint differences, with less diversity than PACS. (3) MICCAI WMH Challenge [Kuijf et al., 2019] is a medical image segmentation dataset for White Matter Hyperintensity that contains 3 different domains (i.e., Utrecht, Amsterdam, Singapore), including 20 subjects in each domain. For each subject, there are 47 slices in Utrecht and Singapore, and 82 slices in Amsterdam. The domain shift is mainly due to scanner and protocol differences.

For evaluation, we use the leave-one-domain-out protocol. In other words, the model is trained on source domains, and is evaluated on a held-out target domain. For PACS and OfficeHome, we report the classification accuracy. For WMH Challenge, we report the Dice Similarity Coefficient (DSC) (higher is better). All results are averaged over 3 runs with different random seeds. We use DeepAll as the ERM baseline with no DG mechanisms.

4.2 Implementation Details

(1) For PACS and OfficeHome, we employ the ResNet [He et al., 2016] pretrained on the ImageNet as the backbone. The networks are trained with the learning rate of 1e-3 for feature extractors, the learning rate of 1e-4 for classifiers, batch size of 64, and epochs of 50. We use the features from the second residual block as the input for AAF and TF-Cal. We adopt the same data augmentation as [Matsuura and Harada, 2020]. (2) For WMH Challenge, we use the U-Net [Ronneberger et al., 2015] trained from scratch as the backbone with the the learning rate of 2e-4, batch size of 30, and epochs of 300. We use the features from the second downsamling layer as the input for our method. We randomly augment each training slices by rotation, scaling, and shearing. All networks are optimized by SGD with weight decay of 5e-4, and all learning rates are decayed by 0.1 after 80% of the epochs. For all experiments: For TF-Cal, we set the calibration strength η and τ to 0.5, and set the pcal to 0.5. For AAF activated with a probability of 0.5, we set δ ∼ Beta(0.2, 0.2). The AAF is applied for all epochs, and TF-Cal is trained after 70% of the epochs.

4.3 Evaluation on PACS

From the results in Table 1, our method outperforms the SOTA methods on both ResNet-18 and ResNet-50, improving the DeepAll by a large margin, which demonstrates the effectiveness of our method on unseen target domains. In comparison to the meta-learning methods MLDG [Li et al., 2018a] and MASF [Dou et al., 2019], as well as the domainadversarial method MMLD [Matsuura and Harada, 2020], our method clearly outperforms them significantly with simple feature transformations and no additional network. With the advantage of alleviating the overfit to source domains, our method also surpasses data augmentation methods DDAIG [Zhou et al., 2020] and MixStyle [Zhou et al., 2021], as well as a more recent method, RSC [Huang et al., 2020]. FACT [Xu et al., 2021] is most related to our work in terms of Fourier transformation and augmentation. Still, our method outperforms FACT: (1) TF-Cal can calibrate target style at test time, which solves the problem of training with no knowledge of target domains. (2) AAF can encourage more diverse styles of source domains by providing additional signals in higher-level style information via features from early CNN layers. MVRML [Zhang et al., 2021] is relevant to our work in terms of test-time augmentation because it relies on augmenting target images with multiple views. Our method beats MVRML since we can calibrate target images with guidance from the source prototype.

4.4 Evaluation on Office-Home

As Table 2 shows, our method has the best performance. For this dataset, we have similar observations and conclusions as the PACS. Our method shows improvement consistently over all target domains. Note that, Clipart has the most style difference compared to other domains. Our method demonstrates its ability to calibrate the target style to offset the sourcetarget shifts, which improves the performance on Clipart by a large margin over the DeepAll.

4.5 Evaluation onWMHChallenge

We evaluate our method on the medical image segmentation task, which is important for real-world applications and is often ignored by other DG works. Our method outperforms other related methods: (1) BigAug [Zhang et al., 2020] applies heavy data augmentation. (2) UDA [Li et al., 2020] is an unsupervised domain adaptation method. (3) MixDANN [Zhao et al., 2021] is a method that combines mixup and DANN. We notice that the DeepAll fails to generalize to the Utrecht due to the significant scanner-difference. Our method significantly improves the overall performance of Utrecht by a considerable margin compared to previous methods.

4.6 Further Analysis

Ablation Studies of Different Components

We investigate a thorough ablation studies of different components in our method. With only AAF, our method already outperforms most of the methods listed in Table 1. The AAF also achieves similar performance as FACT. As showed in [Xu et al., 2021], FACT's augmentation component obtains accuracy of 83.44 which is lower than our AAF, which empirically verifies the merit of AAF in generating more diverse styles in features. To prove that using TF-Cal does not serve as a way for augmenting source data, we show that only applying TF-Cal at training does not have much difference in performance with the DeepAll. The similar observation is also found when both AAF and TF-Cal (Train) are activated. When we activate TF-Cal at both train-time and test-time, we see the great improvement, which also beats most of the methods in Table 1. Moreover, we see that AAF and TF-Cal can complement each other and further boost the generalizability. Finally, we empirically conduct the experiment on the importance of the source prototype. We randomly sample a prototype from Gaussian distribution for both TF-Cal (Train) and TF-Cal (Test), and we see a significant performance drop.

Visualization of Features

To qualitatively analyze our contribution in calibrating style features, we applied t-SNE to compare the features learned by the DeepAll and our approach on PACS dataset with sketch as the target domain. From the Fig. 2, DeepAll has a satisfactory capacity for categorizing features, but it fails to align the source domain Photo and the target domain Sketch with other source domains. Specifically, AAF can well align the Sketch with the source domains, but with unclear cluster boundaries due to the mixing of source styles. TF-Cal precisely calibrates the target domain to the source domains, with distinct cluster boundaries. With the TF-Cal's advantage of aligning styles and AAF's advantage of diversifying styles, we observe a robust representation toward the target domain.

Hyperparameter Sensitivity

We conduct the hyperparameter experiment (Fig.3a) on the calibration stength hyperparameter η and τ for TF-Cal with ResNet-18 and ResNet-50 on PACS. As we define in the method section, η and τ share the same value. We choose η and τ from {0.1, 0.3, 0.5, 0.75, 1}. The larger η and τ are, the more source prototype and the less target style in the target features. Both ResNet-18 and ResNet-50 have the similar trend. We find that TF-Cal performs the best when the η and τ are equal to 0.5. When the η and τ are smaller, the target style is less calibrated. When the η and τ are larger, classdiscriminative information in target styles might be reduced.

Where to Apply the Method

We explore applying TF-Cal+AAF after different residule blocks of ResNet-18 on PACS. From the result in Fig.3b, we see that applying our method after the second residual block performs the best, and the reason is that early layers extract more domain information than the later layers. As we can see that applying our method to the third and the fourth residual blocks, the performance drops. Because the these two blocks are closer to the classification layer, augmenting or calibrating the features will unnecessarily affect the classdiscriminative features.

5 Conclusion

In this paper, we propose a test-time calibration function for target styles extracted by Fourier transformation, and a feature style augmentation technique that complements style calibration. With TF-Cal, we alleviate the stress of having no knowledge about target domains during training by calibrating target styles with the source prototype at test-time. Additionally, AAF efficiently diversifies source amplitudes using higher-level style information extracted from features to help TF-Cal, outperforming image-based amplitudes augmentation. While our method TAF-Cal can be easily integrated into any CNN network, we achieve state-of-the-art performance on two primary tasks (i.e., classification and segmentation), and we conduct extensive experiments and qualitative analysis that validate our method's effectiveness and robustness.