Mode Normalization

1 INTRODUCTION . A challenge in optimizing deep learning models is the change in input distributions at each layer , complicating the training process . Normalization methods , such as batch normalization ( BN , Ioffe & Szegedy , 2015 ) aim to overcome this issue — often referred to as internal covariate shift ( Shimodaira , 2000 ) .1 When applied successfully in practice , BN enables the training of very deep networks , shortens training times by supporting larger learning rates , and reduces sensitivity to parameter initializations . As a result , BN has become an integral element of many state-of-the-art machine learning techniques ( He et al. , 2016 ; Silver et al. , 2017 ) . It can be difficult to standardize the activations in a neural network exposed to heterogeneous or multi-modal data . When training a deep neural network on images that come from a diverse set of visual domains , each with significantly different statistics , BN is not effective at normalizing the activations with a single mean and variance ( Bilen & Vedaldi , 2017 ) . In this paper we relax the assumption that the entire mini-batch should be normalized with the same mean and variance . Our new normalization method , mode normalization ( MN ) , first assigns samples in a mini-batch to different modes via a gating network , and then normalizes each sample with estimators for its corresponding mode ( Figure 1 ) . We further show that MN can be incorporated into other normalization techniques such as group normalization ( GN , Wu & He , 2018 ) by learning which filters should be grouped together . The proposed methods can easily be implemented as layers in standard deep learning libraries , and their parameters are learned jointly with the other parameters of the network in an end-to-end manner . We evaluate MN on multiple classification tasks where it achieved a consistent improvement over currently available normalization approaches . 2 RELATED WORK . Normalization . Normalizing input data ( LeCun et al. , 1998 ) or initial weights of neural networks ( Glorot & Bengio , 2010 ) are known techniques to support faster model convergence , and have been studied extensively . More recently , normalization has become part of functional layers that adjust the internal activations of neural networks . Local response normalization ( LRN ) ( Lyu & Simoncelli , 2008 ; Jarrett et al. , 2009 ) is used in various models ( Krizhevsky et al. , 2012 ; Sermanet et al. , 2014 ) to perform normalization in a local neighborhood , and thereby enforce competition 1The underlying mechanisms are still being explored from a theoretical perspective ( Bjorck et al. , 2018 ; Kohler et al. , 2018 ; Santurkar et al. , 2018 ) . between adjacent pixels in a feature map . BN ( Ioffe & Szegedy , 2015 ) implements a more global normalization along the batch dimension . Unlike LRN , BN requires two distinct training and inference modes of operation . At training time , samples in each batch are normalized with the batch statistics , while during inference samples are normalized using precomputed statistics from the training set . Heterogeneity and small batch sizes can both lead to inconsistencies between training and test data . Our proposed method alleviates the former issue by better dealing with different modes in the data , simultaneously discovering these and normalizing the data accordingly . Several recent normalization methods ( Ba et al. , 2016 ; Ulyanov et al. , 2017 ; Ioffe , 2017 ) have emerged that perform normalization along the channel dimension ( Ba et al. , 2016 ) , or over a single sample ( Ulyanov et al. , 2017 ) to overcome the limitations of BN . Ioffe ( 2017 ) proposes a batch renormalization strategy that clips gradients for estimators by using a predefined range to prevent degenerate cases . While these methods are effective for training sequential and generative models respectively , they have not been able to reach the same level of performance as BN in supervised classification . Alongside these developments , BN has started to attract attention from theoretical viewpoints ( Kohler et al. , 2018 ; Santurkar et al. , 2018 ) . More recently , Wu and He ( Wu & He , 2018 ) have proposed a simple yet effective alternative to BN by first dividing the channels into groups and then performing normalization within each group . The authors show that group normalization ( GN ) can be coupled with small batch sizes without any significant performance loss , and delivers comparable results to BN when the batch size is large . We build on this method in Section 3.2 , and show that it is possible to automatically infer filter groupings . An alternative normalization strategy is to design a data-independent reparametrization of the weights in a neural network by implicitly whitening the representation obtained at each layer ( Desjardins et al. , 2015 ; Arpit et al. , 2016 ) . Lastly , Kalayeh & Shah ( 2019 ) propose fitting a Gaussian mixture model to account for the modality in intermediate feature distributions of some layers . While these methods show promising results , they do not generalize to arbitrary non-linearities and layers , or suffer from runtime restrictions . Mixtures of experts . Mixtures of experts ( MoE ) ( Jacobs et al. , 1991 ; Jordan & Jacobs , 1994 ) are a family of models that involve combining a collection of simple learners to split up the learning problem . Samples are thereby allocated to differing subregions of the model that are best suited to deal with a given example . There is a vast body of literature describing how to incorporate MoE with different types of expert architectures such as SVMs ( Collobert et al. , 2002 ) , Gaussian processes ( Tresp , 2001 ) , or deep neural networks ( Eigen et al. , 2013 ; Shazeer et al. , 2017 ) . Most similar to ours , Eigen et al . ( 2013 ) propose using a different gating network at each layer in a multilayer network to enable an exponential number of combinations of expert opinions . While our method also uses a gating function at every layer to assign the samples in a mini-batch to separate modes , it differs from the above MoE approaches in two key aspects : i ) we use the assignments from the gating functions to normalize the data within a corresponding mode ; ii ) the normalized data is forwarded to a common module ( i.e . a convolutional layer ) rather than to multiple separate experts . Our method is also loosely related to Squeeze-and-Excitation Networks ( Hu et al. , 2018 ) , that adaptively recalibrate channel-wise feature responses with a gating function . Unlike their approach , we use the outputs of the gating function to normalize the responses within each mode . Multi-domain learning . Our approach also relates to methods that parametrize neural networks with domain-agnostic and specific layers , and transfer the agnostic parameters to the analysis of very different types of images ( Bilen & Vedaldi , 2017 ; Rebuffi et al. , 2017 ; 2018 ) . In contrast to these methods , which require the supervision of domain knowledge to train domain-agnostic parameters , our method can automatically learn to discover modes both in single and multi-domain settings , without any supervision . Recent work studies the role of adaptive intermediate features based on task-conditioned inputs ( Perez et al. , 2017 ) . From this viewpoint MN can be understood as a conditional layer , however it has the different focus of accounting for modality in intermediate feature distributions . 3 METHOD . We first review the formulations of BN and GN in Section 3.1 , and introduce our method in Section 3.2 . 3.1 BATCH AND GROUP NORMALIZATION . Our goal is to learn a prediction rule f : X → Y that infers a class label y ∈ Y for a previously unseen sample x ∈ X . For this purpose , we optimize the parameters of f on a training set { xi } Ndi=1 for which the corresponding label information { yi } Ndi=1 is available , where Nd denotes the number of samples in the data . Without loss of generality , in this paper we consider image data as the input , and deep convolutional neural networks as our model . In a slight abuse of notation , we also use the symbol x to represent the features computed by layers within the deep network , producing a three-dimensional tensor X = C ×H ×W where the dimensions indicate the number of feature channels , height and width respectively . Batch normalization ( BN ) computes estimators for the mini-batch { xn } Nn=1 ( usually N Nd ) by average pooling over all but the channel dimensions.2 Then BN normalizes the samples in the batch as BN ( xn ) = α ( xn − µ σ ) + β , ( 1 ) where µ and σ are the mean and standard deviation of the mini-batch , respectively . The parameters α and β are |C|-dimensional vectors representing a learned affine transformation along the channel dimensions , purposed to retain each layer ’ s representative capacity ( Ioffe & Szegedy , 2015 ) . This normalizing transformation ensures that the mini-batch has zero mean and unit variance when viewed along the channel dimensions . Group normalization ( GN ) performs a similar transformation to that in ( 1 ) , but normalizes along different dimensions . GN first separates channels c = 1 , . . . , |C| into fixed groups Gj , over which it then jointly computes estimators , e.g . for the mean µj = |Gj |−1 ∑ xc∈Gj xc . Because GN does not average the statistics along the mini-batch dimension , it is appealing when it is not practical to use large batch sizes . A potential problem when using GN is that channels that are being grouped together might get prevented from developing distinct characteristics in feature space . In addition , computing estimators from manually engineered rules as those found in BN and GN might be too restrictive under some circumstances , for example when jointly learning on multiple domains . 3.2 MODE NORMALIZATION . The heterogeneous nature of complex datasets motivates us to propose a more flexible treatment of normalization . Before the actual normalization is carried out , the data is first organized into modes to which it likely belongs . To achieve this , we reformulate the normalization in the framework of 2How estimators are computed is what differentiates many of the normalization techniques currently available . Wu & He ( 2018 ) provide a detailed introduction . Algorithm 1 Mode normalization , training phase . Input : parameters λ , K , batch of samples { xn } , small ε , learnable α , β , Ψ : C → RK . Compute expert assignments : gnk ← [ softmax ◦Ψ ( xn ) ] k for k = 1 to K do Determine new component-wise statistics : Nk ← ∑ n gnk 〈x〉k ← 1Nk ∑ n gnkxn 〈x2〉k ← 1Nk ∑ n gnkx 2 n Update running means : 〈x〉k ← λ〈x〉k + ( 1− λ ) 〈x〉k 〈x2〉k ← λ〈x2〉k + ( 1− λ ) 〈x2〉k end for for n = 1 to N do Normalize samples with component-wise estimators : µk ← 〈x〉k σ2k ← 〈x2〉k − 〈x〉2k ynk ← gnk xn−µk√ σ2k+ε end for Return : { α ∑ k ynk + β } n=1 , ... , N Algorithm 2 Mode normalization , test phase . Input : refer to Alg . 1 . Compute expert assignments : gnk ← [ softmax ◦Ψ ( xn ) ] k for n = 1 to N do Normalize samples with running average of component-wise estimators : µk ← 〈x〉k σ2k ← 〈x2〉k − 〈x〉 2 k ynk ← gnk xn−µk√ σ2k+ε end for Return : { α ∑ k ynk + β } n=1 , ... , N Algorithm 3 Mode group normalization . Input : parameter K , sample x ∈ C , small ε , learnable α , β , Ψ : R→ RK . Compute channel-wise gates : gck ← [ softmax ◦Ψ ( xc ) ] k for k = 1 to K do Update estimators and normalize : µk ← 〈x〉k σ2k ← 〈x2〉k − 〈x〉2k yk ← x−µk√ σ2k+ε end for Return : αK ∑ k yk + β mixtures of experts ( MoE ) . In particular , we introduce a set of simple gating functions { gk } Kk=1 where gk : X → [ 0 , 1 ] and ∑ k gk ( x ) = 1 . In mode normalization ( MN , Alg . 1 ) , each sample in the mini-batch is then normalized under voting from its gate assignment : MN ( xn ) , α ( K∑ k=1 gk ( xn ) xn − µk σk ) + β , ( 2 ) where α and β are a learned affine transformation , just as in standard BN.3 The mean µk and variance σk estimates are weighted averages under the gating network . For example , the k ’ th mean is estimated from the batch as µk = 〈x〉k = 1 Nk ∑ n gk ( xn ) · xn , ( 3 ) where Nk= ∑ n gk ( xn ) . In our experiments , we average overW andH , and parametrize the gating functions via an affine transformation Ψ : C → RK of the channels , which is jointly learned alongside the other parameters of the network . This is followed by a softmax activation , reminiscent of attention mechanisms ( Denil et al. , 2012 ; Vinyals et al. , 2015 ) . Mode normalization generalizes BN , which can be recovered as a special case by setting K=1 , or if the gates collapse : gk ( xn ) =const . ∀ k , n. To demonstrate how the extra flexibility helps , Fig . 2 3Learning individual { ( αk , βk ) } Kk=1 for each mode did not improve performance in preliminary experiments .

Summary: The paper proposes a generalisation of Batch Normalisation (BN) under the assumption that the statistics of the unit activations over the batches and over the spatial dimensions (in case of convolutional networks) is not unimodal. The main idea is to represent the unit activation statistics as a mixture of modes and to re-parametrise by using mode specific means and variances. The "posterior" mixture weights for a specific unit are estimated by gating functions with additional affine parameters (followed by softmax). A second, similar variant applies to Group Normalisation, where the statistics is taken over channel groups and spatial dimensions (but not over batches).