Domain Generalization with MixStyle

1 INTRODUCTION . Key to automated understanding of digital images is to compute a compact and informative feature representation . Deep convolutional neural networks ( CNNs ) have demonstrated remarkable ability in representation learning , proven to be effective in many visual recognition tasks , such as classifying photo images into 1,000 categories from ImageNet ( Krizhevsky et al. , 2012 ) and playing Atari games with reinforcement learning ( Mnih et al. , 2013 ) . However , it has long been discovered that the success of CNNs heavily relies on the i.i.d . assumption , i.e . training and test data should be drawn from the same distribution ; when such an assumption is violated even just slightly , as in most realworld application scenarios , severe performance degradation is expected ( Hendrycks & Dietterich , 2019 ; Recht et al. , 2019 ) . Domain generalization ( DG ) aims to address such a problem ( Zhou et al. , 2021 ; Blanchard et al. , 2011 ; Muandet et al. , 2013 ; Li et al. , 2018a ; Zhou et al. , 2020b ; Balaji et al. , 2018 ; Dou et al. , 2019 ; Carlucci et al. , 2019 ) . In particular , assuming that multiple source domains containing the same visual classes are available for model training , the goal of DG is to learn models that are robust against data distribution changes across domains , known as domain shift , so that the trained model can generalize well to any unseen domains . Compared to the closely related and more widely studied domain adaptation ( DA ) problem , DG is much harder in that no target domain data is available for the model to analyze the distribution shift in order to overcome the negative effects . Instead , a DG model must rely on the source domains and focus on learning domain-invariant feature representation in the hope that it would remain discriminative given target domain data . A straightforward solution to DG is to expose a model with a large variety of source domains . Specifically , the task of learning domain-invariant and thus generalizable feature representation becomes easier when data from more diverse source domains are available for the model . This would reduce the burden on designing special models or learning algorithms for DG . Indeed , model training with large-scale data of diverse domains is behind the success of existing commercial face recognition or vision-based autonomous driving systems . A recent work by Xu et al . ( 2021 ) also emphasizes the importance of diverse training distributions for out-of-distribution generalization . However , collecting data of a large variety of domains is often costly or even impossible . It thus can not be a general solution to DG . In this paper , a novel approach is proposed based on probabilistically mixing instance-level feature statistics of training samples across source domains . Our model , termed MixStyle , is motivated by the observation that visual domain is closely related to image style . An example is shown in Fig . 1 : the four images from four different domains depict the same semantic concept , i.e . dog , but with distinctive styles ( e.g. , characteristics in color and texture ) . When these images are fed into a deep CNN , which maps the raw pixel values into category labels , such style information is removed at the output . However , recent style transfer studies ( Huang & Belongie , 2017 ; Dumoulin et al. , 2017 ) suggest that such style information is preserved at the bottom layers of the CNN through the instance-level feature statistics , as shown clearly in Fig . 4 . Importantly , since replacing such statistics would lead to replaced style while preserving the semantic content of the image , it is reasonable to assume that mixing styles from images of different domains would result in images of ( mixed ) new styles . That is , more diverse domains/styles can be made available for training a more domain-generalizable model . Concretely , our MixStyle randomly selects two instances of different domains and adopts a probabilistic convex combination between instance-level feature statistics of bottom CNN layers . In contrast to style transfer work ( Huang & Belongie , 2017 ; Dumoulin et al. , 2017 ) , no explicit image synthesis is necessary meaning much simpler model design . Moreover , MixStyle perfectly fits into modern mini-batch training . Overall , it is very easy to implement with only few lines of code . To evaluate the effectiveness as well as the general applicability of MixStyle , we conduct extensive experiments on a wide spectrum of datasets covering category classification ( Sec . 3.1 ) , instance retrieval ( Sec . 3.2 ) , and reinforcement learning ( Sec . 3.3 ) . The results demonstrate that MixStyle can significantly improve CNNs ’ cross-domain generalization performance.1 2 METHODOLOGY . 2.1 BACKGROUND . Normalizing feature tensors with instance-specific mean and standard deviation has been found effective for removing image style in style transfer models ( Ulyanov et al. , 2016 ; Huang & Belongie , 2017 ; Dumoulin et al. , 2017 ) . Such an operation is widely known as instance normalization ( IN , Ulyanov et al . ( 2016 ) ) . Let x ∈ RB×C×H×W be a batch of tensors , with B , C , H and W 1Source code can be found at https : //github.com/KaiyangZhou/mixstyle-release . denoting the dimension of batch , channel , height and width , respectively , IN is formulated as IN ( x ) = γ x− µ ( x ) σ ( x ) + β , ( 1 ) where γ , β ∈ RC are learnable affine transformation parameters , and µ ( x ) , σ ( x ) ∈ RB×C are mean and standard deviation computed across the spatial dimension within each channel of each instance ( tensor ) , i.e . µ ( x ) b , c = 1 HW H∑ h=1 W∑ w=1 xb , c , h , w , ( 2 ) and σ ( x ) b , c = √√√√ 1 HW H∑ h=1 W∑ w=1 ( xb , c , h , w − µ ( x ) b , c ) 2 . ( 3 ) Huang & Belongie ( 2017 ) introduced adaptive instance normalization ( AdaIN ) , which simply replaces the scale and shift parameters in Eq . ( 1 ) with the feature statistics of style input y to achieve arbitrary style transfer : AdaIN ( x ) = σ ( y ) x− µ ( x ) σ ( x ) + µ ( y ) . ( 4 ) 2.2 MIXSTYLE . Our method , MixStyle , draws inspiration from AdaIN . However , rather than attaching a decoder for image generation , MixStyle is designed for the purpose of regularizing CNN training by perturbing the style information of source domain training instances . It can be implemented as a plug-andplay module inserted between CNN layers of , e.g. , a supervised CNN classifier , without the need to explicitly generate an image of new style . More specifically , MixStyle mixes the feature statistics of two instances with a random convex weight to simulate new styles . In terms of implementation , MixStyle can be easily integrated into mini-batch training . Given an input batch x , MixStyle first generates a reference batch x̃ from x . When domain labels are given , x is sampled from two different domains i and j , e.g. , x = [ xi , xj ] ( xi and xj have the same batch size ) . Then , x̃ is obtained by swapping the position of xi and xj , followed by a shuffling operation along the batch dimension applied to each batch , i.e . x̃ = [ Shuffle ( xj ) , Shuffle ( xi ) ] . See Fig . 2 ( a ) for an illustration . In cases where domain labels are unknown , x is randomly sampled from the training data , and x̃ is simply obtained by x̃ = Shuffle ( x ) ( see Fig . 2 ( b ) ) . Fig . 4 shows that sub-domains exist within each domain , so even if two instances of the same domain are sampled , new domain could be synthesized . After shuffling , MixStyle computes the mixed feature statistics by γmix = λσ ( x ) + ( 1− λ ) σ ( x̃ ) , ( 5 ) βmix = λµ ( x ) + ( 1− λ ) µ ( x̃ ) , ( 6 ) where λ ∈ RB are instance-wise weights sampled from the Beta distribution , λ ∼ Beta ( α , α ) with α ∈ ( 0 , ∞ ) being a hyper-parameter . Unless specified otherwise , we set α to 0.1 throughout this paper . Finally , the mixed feature statistics are applied to the style-normalized x , MixStyle ( x ) = γmix x− µ ( x ) σ ( x ) + βmix . ( 7 ) In practice , we use a probability of 0.5 to decide if MixStyle is activated or not in the forward pass . At test time , no MixStyle is applied . Note that gradients are blocked in the computational graph of µ ( · ) and σ ( · ) . MixStyle can be implemented with only few lines of code . See Algorithm 1 in Appendix A.1 for the PyTorch-like pseudo-code . 3 EXPERIMENTS . 3.1 GENERALIZATION IN CATEGORY CLASSIFICATION . Dataset and implementation details . We choose the PACS dataset ( Li et al. , 2017 ) , a commonly used domain generalization ( DG ) benchmark concerned with domain shift in image classification . PACS consists of four domains , i.e . Art Painting , Cartoon , Photo and Sketch , with totally 9,991 images of 7 classes . As shown in Fig . 1 , the domain shift mainly corresponds to image style changes . For evaluation , a model is trained on three domains and tested on the remaining one . Following prior work ( Li et al. , 2019 ; Zhou et al. , 2020a ) , we use ResNet-18 ( He et al. , 2016 ) as the classifier where MixStyle is inserted after the 1st , 2nd and 3rd residual blocks . Our code is based on Dassl.pytorch ( Zhou et al. , 2020c ) .2 Baselines . Our main baselines are general-purpose regularization methods including Mixup ( Zhang et al. , 2018b ) , Manifold Mixup ( Verma et al. , 2019 ) , DropBlock ( Ghiasi et al. , 2018 ) , CutMix ( Yun et al. , 2019 ) and Cutout ( DeVries & Taylor , 2017 ) , which are trained using the same training parameters as MixStyle and the optimal hyper-parameter setup as reported in their papers . We also compare with the existing DG methods which reported state-of-the-art performance on PACS . These include domain alignment-based CCSA ( Motiian et al. , 2017 ) and MMD-AAE ( Li et al. , 2018b ) , Jigsaw puzzle-based JiGen ( Carlucci et al. , 2019 ) , adversarial gradient-based CrossGrad ( Shankar et al. , 2018 ) , meta-learning-based Metareg ( Balaji et al. , 2018 ) and Epi-FCR ( Li et al. , 2019 ) , and data augmentation-based L2A-OT ( Zhou et al. , 2020a ) . Comparison with general-purpose regularization methods . The results are shown in Table 1 . Overall , we observe that the general-purpose regularization methods do not offer any clear advantage over the vanilla ResNet-18 in this DG task , while MixStyle improves upon the vanilla ResNet18 with a significant margin . Compared with Mixup , MixStyle is 5.2 % better on average . Recall that Mixup also interpolates the output space , we further compare with a variant of Mixup in order to demonstrate the advantage of mixing style statistics at the feature level over mixing images at the pixel level for DG—following Sohn et al . ( 2020 ) , we remove the label interpolation in Mixup and sample the mixing weights from a uniform distribution of [ 0 , 1 ] . Still , MixStyle outperforms this new baseline with a large margin , which justifies our claim . MixStyle and DropBlock share some commonalities in that they are both applied to feature maps at multiple layers , but MixStyle significantly outperforms DropBlock in all test domains . The reason why DropBlock is ineffective here is because dropping out activations mainly encourages a network to mine discriminative patterns , but does not reinforce the ability to cope with unseen styles , which is exactly what MixStyle aims to achieve : by synthesizing “ new ” styles ( domains ) MixStyle regularizes the network to become more 2https : //github.com/KaiyangZhou/Dassl.pytorch . robust to domain shift . In addition , it is interesting to see that on Cartoon and Photo , MixStyle w/ random shuffle obtains slightly better results . The reason might be because there exist sub-domains in a source domain ( see Fig . 4 ( a-c ) ) , which allow random shuffling to produce more diverse “ new ” domains that lead to a more domain-generalizable model . Comparison with state-of-the-art DG methods . Overall , MixStyle outperforms most DG methods by a clear margin , despite being a much simpler method . The performance of MixStyle w/ domain label is nearly 1 % better on average than the recently introduced L2A-OT . From a data augmentation perspective , MixStyle and L2A-OT share a similar goal—to synthesize data from pseudo-novel domains . MixStyle accomplishes this goal through mixing style statistics at the feature level . Whereas L2A-OT works at the pixel level : it trains an image generator by maximizing the domain difference ( measured by optimal transport ) between the original and the generated images , which introduces much heavier computational overhead than MixStyle in terms of GPU memory and training time . It is worth noting that MixStyle ’ s domain label-free version is highly competitive : its 82.8 % accuracy is on par with L2A-OT ’ s .

Summary: This paper proposed a simple regularization technique for domain generalization tasks, termed MixStyle, based on the observation that domains are determined by image styles. By mixing styles of different instances, which generates synthesized domain samples while preserving the content features, the proposed method achieves the generalizability of the trained model. The MixStyle was applied to numerous applications, such as category classification, instance retrieval, and reinforcement learning, and attained the state-of-the arts. The MixStyle is relatively simple to implement, but effective.