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 work proposes a technique for domain generalization by mixing style of images from different domains. This work adopts a mix up style approach [A] for domain generalization. Different from [A], the paper proposes to conduct mix-up in the intermediate layers, in particular, instance normalization layers. The proposed approach diversifies the data implicitly and the experimental results show that the mix-style can improve domain generalization.