Test-time Fourier Style Calibration for 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 meta-learning. 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.

We define $\mathcal{X}$ and $\mathcal{Y}$ as the input and label space. Given $K$ source domains $\mathcal{D}_{s}=\left\{D^{1}_{s},D^{2}_{s},\ldots,D^{K}_{s}\right\}(1\leq i\leq K)$ that are available at train-time, we indicate that each $D^{i}_{s}$ contains $N$ input and label pairs $(x_{j},y_{j})_{j=1}^{N_{i}}$ where ${x}\in\mathcal{X},y\in\mathcal{Y}$ and $N_{i}$ is the number of inputs in the source domain $D^{i}_{s}$. The goal of DG is to train a function learned from data of $\mathcal{D}_{s}$ to perform well on an unseen target domain $\mathcal{D}_{t}$ which share the same label space $\mathcal{Y}$ as $\mathcal{D}_{s}$. The function can be defined as $R_{\theta,\sigma}=T_{\sigma}\circ F_{\theta}$, where $T_{\sigma}$ is a task-specific network parameterized by $\sigma$ and $F_{\theta}=\left\{F^{1}_{\theta},F^{2}_{\theta},\ldots,F^{L}_{\theta}\right\}(1\leq i\leq L)$ which is a feature encoder with $L$ layers parameterized by $\theta$.

To reduce the domain shift, we need to first understand the difference between two distributions which are measured by $\mathcal{H}$-divergence proposed by Ben-David et al. (2009):

$$d_{\mathcal{H}}(\mathcal{D}_{s},\mathcal{D}_{t})=2\sup_{h\in\mathcal{H}}\left|\operatorname{Pr}_{x\sim{\mathcal{D}_{s}}}[h(x)=1]-\operatorname{Pr}_{x\sim{\mathcal{D}_{t}}}[h(x)=1]\right|$$

(1)

where classifier $h:\mathcal{X}\rightarrow\{0,1\}$. Following that, Albuquerque et al. (2019) defines the convex hull $\Lambda_{s}$ of $\mathcal{D}_{s}$ that is a set of mixture of source domain distributions:

$$\Lambda_{s}=\left\{\sum_{i=1}^{K}\pi_{i}\mathcal{D}_{s}^{i}\mid\pi\in\Delta_{K-1}\right\}$$

(2)

where the $\pi$ represents non-negative coefficient in the $(K-1)$-dimensional simplex $\Delta_{K-1}$. Then, an ideal case $\bar{\mathcal{D}}_{t}\in\Lambda_{s}$ is assumed that the ideal target $\bar{\mathcal{D}}_{t}$ lies in the source domain convex hull $\Lambda_{s}$. Under this assumption, the risk $\epsilon_{t}(h)$ on the target domain is upper-bounded Albuquerque et al. (2019) by:

$$\epsilon_{t}(h)\leq\sum_{i=1}^{K}\pi_{i}\epsilon^{i}_{s}(h)+\gamma+\zeta$$

(3)

where, on the right hand side, the first term corresponds to the risks over all source domains, and the second term $\gamma=d_{\mathcal{H}\Delta\mathcal{H}}(\bar{\mathcal{D}}_{t},D_{t})$ is the $\mathcal{H}$-divergence between the ideal target $\bar{\mathcal{D}}_{t}$ and the real target $D_{t}$, and the third term $\zeta=\sup_{D_{s}^{\prime},D_{s}^{\prime\prime}\in\Lambda_{s}}d_{\mathcal{H}\Delta\mathcal{H}}\left(D_{s}^{\prime},D_{s}^{\prime\prime}\right)$ which is the largest $\mathcal{H}$-divergence between any pair of source domains. $\mathcal{H}\Delta\mathcal{H}$ corresponds to $\left\{h(x)\oplus h^{\prime}(x)\mid h,h^{\prime}\in\mathcal{H}\right\}$. 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).

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 $\bar{\mathcal{D}}_{t}$ is assumed to be covered by source convex hull $\Lambda_{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 $\Lambda_{s}$ (i.e., $\gamma\rightarrow 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 $\Lambda_{s}$ arbitrarily, we aim to pull $D_{t}$ closer to $\bar{D_{t}}$ 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 $\Lambda_{s}$, and reduce the difference between $D_{t}$ and $\bar{D_{t}}$. To further encourage the reduction in source-target divergence, we augment the source features to expand the convex hull.

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 $\mathcal{C}$ (TF-Cal) for reducing the gap between the target domain and source domains prototype. We further introduce an augmenting function $\mathcal{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 $\mathcal{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.

Given a latent feature tensor $\mathbf{z}^{i}=F^{i}(\mathbf{z}^{i-1})$ with $\mathbf{z}^{i}\in\mathbb{R}^{N\times C\times H\times 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 $\mathcal{F}(\mathbf{z}^{i})$, which can be implemented using Fast Fourier transformation (FFT) Nussbaumer (1982), is able to decompose the feature $\mathbf{z}^{i}$ into the phase spectrum $\mathcal{PH}(\mathbf{z}^{i})$ encoding the spatial information and amplitude spectrum $\mathcal{AM}(\mathbf{z}^{i})$ 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 $\mathcal{F}(\mathbf{z}^{i})$ is done on each channel $C$ of $\mathbf{z}^{i}$:

$$\mathcal{F}(\mathbf{z}^{i})(m,n)=\sum_{h=0}^{H-1}\sum_{w=0}^{W-1}\mathbf{z}^{i}(h,w)e^{-j2\pi}\left(m\frac{h}{H}+n\frac{w}{W}\right)$$

(4)

where $\mathcal{F}(\mathbf{z}^{i})(m,n)$ represents the value of the features in the frequency space corresponding to the coordinates $m$ and $n$. With $\mathcal{F}$, we can retrieve the amplitude and phase features with real $\operatorname{Re}$ and imaginary $\operatorname{Im}$ parts:

		$\displaystyle\mathcal{PH}(\mathbf{z}^{i})=\arctan\frac{\operatorname{Im}(\mathcal{F}(\mathbf{z}^{i})(m,n))}{\operatorname{Re}(\mathcal{F}(\mathbf{z}^{i})(m,n))}$		(5)
		$\displaystyle\mathcal{AM}(\mathbf{z}^{i})=\sqrt{\operatorname{Im}(\mathcal{F}(\mathbf{z}^{i})(m,n))^{2}+\operatorname{Re}(\mathcal{F}(\mathbf{z}^{i})(m,n))^{2}}$

To reconstruct phase and amplitude spectrum features back into the spatial feature space, we can use the inverse Fourier Transformation $\mathcal{F}^{-1}(\mathcal{AM}(\mathbf{z}^{i}),\mathcal{PH}(\mathbf{z}^{i}))=\mathbf{z}^{i}$.

As we can see from Eq.(7) and Eq.(8), the common part of the input is the $\mathbf{Pt}_{s}$ and the different part is the amplitude features $\mathcal{AM}(\mathbf{z}^{i}_{t})$ and $\mathcal{AM}(\mathbf{z}^{i}_{s})$. Therefore, to further boost the generalizability with TF-Cal, we want $\mathcal{AM}(\mathbf{z}^{i}_{t})$ and $\mathcal{AM}(\mathbf{z}^{i}_{s})$ to be as similar as possible. An effective way is to augment amplitude in feature space (AAF) hoping to expand the source convex hull $\Lambda_{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 ($\mathcal{AM}(\mathbf{z}^{i}_{s}),\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s})$) randomly from any source domains:

$$\mathcal{A}(\mathcal{AM}(\mathbf{z}^{i}_{s}),\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s}))=\delta\mathcal{AM}(\mathbf{z}^{i}_{s})+(1-\delta)\mathcal{AM}(\mathbf{z}^{i^{\prime}}_{s})$$

(9)

where $\delta\sim\operatorname{Beta}(\alpha,\alpha)$. For swapping amplitude between two samples, we set $\delta$ 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.

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.

(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 $\eta$ and $\tau$ to 0.5, and set the $p_{cal}$ to 0.5. For AAF activated with a probability of 0.5, we set $\delta\sim\operatorname{Beta}(0.2,0.2)$. The AAF is applied for all epochs, and TF-Cal is trained after 70% of the epochs.

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 domain-adversarial 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.

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 source-target shifts, which improves the performance on Clipart by a large margin over the DeepAll.

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.