MetaNorm: Learning to Normalize Few-Shot Batches Across Domains

1 INTRODUCTION . Batch normalization ( Ioffe & Szegedy , 2015 ) is crucial for training neural networks , and with its variants , e.g. , layer normalization ( Ba et al. , 2016 ) , group normalization ( Wu & He , 2018 ) and instance normalization ( Ulyanov et al. , 2016 ) , has thus become an essential part of the deep learning toolkit ( Bjorck et al. , 2018 ; Luo et al. , 2018a ; Yang et al. , 2019 ; Jia et al. , 2019 ; Luo et al. , 2018b ; Summers & Dinneen , 2020 ) . Batch normalization helps stabilize the distribution of internal activations when a model is being trained . Given a mini-batch B , the normalization is conducted along each individual feature channel for 2D convolutional neural networks . During training , the batch normalization moments are calculated as follows : µB = 1 M M∑ i=1 ai , σ 2 B = 1 M M∑ i=1 ( ai − µB ) 2 , ( 1 ) where ai indicates the i-th element of the M activations in the batch , M = |B| ×H ×W , in which H and W are the height and width of the feature map in each channel . We can now apply the normalization statistics to each activation : a′i ← BN ( ai ) ≡ γâi + β , where , âi = ai − µB√ σ2B + , ( 2 ) where γ and β are parameters learned during training , is a small scalar to prevent division by 0 , and operations between vectors are element-wise . At test time , the standard practice is to normalize activations using the moving average over mini-batch means µB and variance σ2B . Batch normalization is based on an implicit assumption that the samples in the dataset are independent and identically distributed . However , this assumption does not hold in challenging settings like few-shot learning and domain generalization . In this paper , we strive for batch normalization when batches are of small size and suffer from distributions shifts between source and target domains . Batch normalization for few-shot learning and domain generalization problems have so far been considered separately , predominantly in a meta-learning setting . For few-shot meta-learning ( Finn et al. , 2017 ; Gordon et al. , 2019 ) , most existing methods rely critically on transductive batch normalization , except those based on prototypes ( Snell et al. , 2017 ; Allen et al. , 2019 ; Zhen et al. , 2020a ) . However , the nature of transductive learning restricts its application due to the requirement to sample from the test set . To address this issue , Bronskill et al . ( 2020 ) proposes TaskNorm , which leverages other statistics from both layer and instance normalization . As a non-transductive normalization approach , it achieves impressive performance and outperforms conventional batch normalization ( Ioffe & Szegedy , 2015 ) . However , its performance is not always performing better than transductive batch normalization . Meanwhile , domain generalization ( Muandet et al. , 2013 ; Balaji et al. , 2018 ; Li et al. , 2017a ; b ) suffers from distribution shifts from training to test , which makes it problematic to directly apply statistics calculated from a seen domain to test data from unseen domains ( Wang et al. , 2019 ; Seo et al. , 2019 ) . Recent works deal with this problem by learning a domain specific normalization ( Chang et al. , 2019 ; Seo et al. , 2019 ) or a transferable normalization in place of existing normalization techniques ( Wang et al. , 2019 ) . We address the batch normalization challenges for few-shot classification and domain generalization in a unified way by learning a new batch normalization under the meta-learning setting . We propose MetaNorm , a simple but effective meta-learning normalization . We leverage the metalearning setting and learn to infer normalization statistics from data , instead of applying direct calculations or blending various normalization statistics . MetaNorm is a general batch normalization approach , which is model-agnostic and serves as a plug-and-play module that can be seamlessly embedded into existing meta-learning approaches . We demonstrate its effectiveness for few-shot classification and domain generalization , where it learns task-specific statistics from limited data samples in the support set for each few-shot task ; and it can also learn to generate domain-specific statistics from the seen source domains for unseen target domains . We verify the effectiveness of MetaNorm by extensive evaluation on few-shot classification and domain generalization tasks . For few-shot classification , we experiment with representative gradient , metric and model-based meta-learning approaches on fourteen benchmark datasets . For domain generalization , we evaluate the model on three widely-used benchmarks for cross-domain visual object classification . Last but not least , we introduce the challenging new task of few-shot domain generalization , which combines the challenges of both few-shot learning and domain generalization . The experimental results demonstrate the benefit of MetaNorm compared to existing batch normalizations . 2 RELATED WORKS . Transductive Batch Normalization For conventional batch normalization under supervised settings , i.i.d . assumptions about the data distribution imply that estimating moments from the training set will provide appropriate normalization statistics for test data . However , in the meta-learning scenario data points are only assumed to be i.i.d . within a specific task . Therefore , it is critical to select the moments when batch normalization is applied to support and query set data points during meta training and meta testing . Hence , in the recent meta-learning literature the running moments are no longer used for normalization at meta-test time , but instead replaced with support/query set statistics . These statistics are used for normalization , both at meta-train and meta-test time . This approach is referred to as transductive batch normalization ( TBN ) ( Bronskill et al. , 2020 ) . Competitive meta-learning methods ( e.g. , Gordon et al. , 2019 ; Finn et al. , 2017 ; Zhen et al. , 2020b ) rely on TBN to achieve state-of-the-art performance . However , there are two critical problems with TBN . First , TBN is sensitive to the distribution over the query set used during meta-training , and as such is less generally applicable than non-transductive learning . Second , TBN uses extra information for multiple test samples , compared to non-transductive batch normalization at prediction time , which could be problematic as we are not guaranteed to have a set of test samples available during training in practical applications . In contrast , MetaNorm is a non-transductive normalization . It generates statistics from the support set only , without relying on query samples , making it more practical . Meta Batch Normalization To address the problem of transductive batch normalization and improve conventional batch normalization , meta-batch normalization ( MetaBN ) was introduced ( Triantafillou et al. , 2020 ; Bronskill et al. , 2020 ) . In MetaBN , the support set alone is used to compute the normalization statistics for both the support and query sets at both meta-training and meta-test time . MetaBN is non-transductive since the normalization of a test input does not depend on other test inputs in the query set . However , Bronskill et al . ( 2020 ) observe that MetaBN performs less well for small-sized support sets . This leads to high variance in moment estimates , which is similar to the difficulty of using batch normalization with small-batch training ( Wu & He , 2018 ) . To address this issue , Bronskill et al . ( 2020 ) proposed TaskNorm , which learns to combine statistics from both layer normalization and instance normalization , with a lending parameter to be learned at meta-train time . As a non-transductive normalization , TaskNorm achieves impressive performance , outperforming conventional batch normalization . However , it can not always perform better than transductive batch normalization . TaskNorm indicates non-transductive batch normalization estimates proper normalization statistics by involving learning in the normalization process . We also propose to learn batch normalization within the meta-learning framework , but instead of employing a learnable combination of existing normalization statistics , we directly learn to infer statistics from data . At meta-train time , the model learns to acquire the ability to generate statistics only from the support set and at meta-test time we directly apply the model to infer statistics for new tasks . Batch Normalization for Domain Adaptation and Domain Generalization Domain adaption suffers from a distribution shift between source and target domains , which makes it sub-optimal to directly apply batch normalization ( Bilen & Vedaldi , 2017 ) . Li et al . ( 2016 ) proposed adaptive batch normalization to increase the generalization ability of a deep neural network . By modulating the statistical information of all batch normalization layers in the neural network , it achieves deep adaptation effects for domain-adaptive tasks . Nado et al . ( 2020 ) noted the possibility of accessing small unlabeled batches of the shifted data just before prediction time . To improve model accuracy and calibration under covariate shift , they proposed prediction-time batch normalization . Since the activation statistics obtained during training do not reflect statistics of the test distribution , when testing in an out-of-distribution environment , Schneider et al . ( 2020 ) proposed estimating the batch statistics on the corrupted images . Kaku et al . ( 2020 ) demonstrated that standard non-adaptive feature normalization fails to correctly normalize the features of convolutional neural networks on held-out data where extraneous variables take values not seen during training . Learning domain-specific batch normalization has been explored ( Chang et al. , 2019 ; Wang et al. , 2019 ) . Wang et al . ( 2019 ) introduced transferable normalization , TransNorm , which normalizes the feature representations from source and target domain separately using domain-specific statistics . Along a similar vein , Chang et al . ( 2019 ) proposed a domain-specific batch normalization layer , which consists of two branches , each in charge of a single domain exclusively . The hope is that , through the normalization , the feature representation will become domain invariant . Nevertheless , these normalization methods are specifically designed for domain adaptation tasks , where data from target domains are available , though often unlabelled . This makes them inapplicable to domain generalization tasks where data from target domains are inaccessible at training time . Seo et al . ( 2019 ) proposed learning to optimize domain specific normalization for domain generalization tasks . Under the meta-learning settings , a mixture of different normalization techniques is optimized for each domain , where the mixture weights are learned specifically for different domains . Instead of combining different normalization statistics , MetaNorm learns from data to generate adaptive statistics specific to each domain . Moreover , we introduce an even more challenging setting , i.e. , few-shot domain generalization , which combines the challenges of few-shot classification and domain generalization . Conditional Batch Normalization de Vries et al . ( 2017 ) proposed conditional batch normalization to modulate visual processing by predicting the scalars γ and β of the batch normalization conditioned on the language from an early processing stage . Conditional batch normalization has also been applied to align different data distributions for domain adaptation ( Li et al. , 2016 ) . Oreshkin et al . ( 2018 ) applies conditional batch normalization to metric-based models for the few-shot classification task . Tseng et al . ( 2020 ) proposed a learning-to-learn method to optimize the hyper-parameters of the feature-wise transformation layers by conditional batch normalization for cross-domain classification . Unlike conditional batch normalization , we use extra data ( the query set ) to generate normalization statistics under the meta-learning setting , rather than the scalars .

Summary: This paper proposes to replace batch normalization statistics, which are typically computed as the batch moments during training or a fixed training average during testing, with the outputs of learned neural networks. These networks are trained to minimize the KL divergence between their output and the expected or desired batch statistics. In this way, the statistics computation is amortized and can hopefully generalize in the face of small batches and distribution shift.