Regularizing activations in neural networks via distribution matching with the Wasserstein metric

1 INTRODUCTION . Training of deep neural networks is very challenging due to the vanishing and exploding gradient problem ( Hochreiter , 1998 ; Glorot & Bengio , 2010 ) , the presence of many flat regions and saddle points ( Shalev-Shwartz et al. , 2017 ) , and the shattered gradient problem ( Balduzzi et al. , 2017 ) . To remedy these issues , various methods for controlling hidden activations have been proposed such as normalization ( Ioffe & Szegedy , 2015 ; Huang et al. , 2018 ) , regularization ( Littwin & Wolf , 2018 ) , initialization ( Mishkin & Matas , 2016 ; Zhang et al. , 2019 ) , and architecture design ( He et al. , 2016 ) . Among various techniques of controlling activations , one well-known and successful path is controlling their first and second moments . Back in the 1990s , it has been known that the neural network training can be benefited from normalizing input statistics so that samples have zero mean and identity covariance matrix ( LeCun et al. , 1998 ; Schraudolph , 1998 ) . This idea motivated batch normalization ( BN ) that considers hidden activations as the input to the next layer and normalizes scale and shift of the activations ( Ioffe & Szegedy , 2015 ) . Recent works show the effectiveness of different sample statistics of activations for normalization and regularization . Deecke et al . ( 2019 ) and Kalayeh & Shah ( 2019 ) normalize activations to several modes with different scales and translations . Variance constancy loss ( VCL ) implicitly normalizes the fourth moment by minimizing the variance of sample variances , which enables adaptive mode separation or collapse based on their prior probabilities ( Littwin & Wolf , 2018 ) . BN is also extended to whiten activations ( Huang et al. , 2018 ; 2019 ) , and to normalize general order of central moment in the sense of Lp norm including L0 and L∞ ( Liao et al. , 2016 ; Hoffer et al. , 2018 ) . In this paper , we propose a projected error function regularization ( PER ) that regularizes activations in the Wasserstein probability distribution space . Specifically , PER pushes the distribution of activations to be close to the standard normal distribution . PER shares a similar strategy with previous approaches that dictates the ideal distribution of activations . Previous approaches , however , deal with single or few sample statistics of activations . On the contrary , PER regularizes the activations by matching the probability distributions , which considers different statistics simultaneously , e.g. , all orders of moments and correlation between hidden units . The extensive experiments on multiple challenging tasks show the effectiveness of PER . 2 RELATED WORKS . Many modern deep learning architectures employ BN as an essential building block for better performance and stable training even though its theoretical aspects of regularization and optimization are still actively investigated ( Santurkar et al. , 2018 ; Kohler et al. , 2018 ; Bjorck et al. , 2018 ; Yang et al. , 2019 ) . Several studies have applied the idea of BN that normalizes activations via the sample mean and the sample variance to a wide range of domains such as recurrent neural network ( Lei Ba et al. , 2016 ) and small batch size training ( Wu & He , 2018 ) . Huang et al . ( 2018 ; 2019 ) propose normalization techniques whitening the activation of each layer . This additional constraint on the statistical relationship between activations improves the generalization performance of residual networks compared to BN . Although the correlation between activations are not explicitly considered , dropout prevents activations from being activated at the same time , called co-adaptation , by randomly dropping the activations ( Srivastava et al. , 2014 ) , the weights ( Wan et al. , 2013 ) , and the spatially connected activations ( Ghiasi et al. , 2018 ) . Considering BN as the normalization in the L2 space , several works extend BN to other spaces , i.e. , other norms . Streaming normalization ( Liao et al. , 2016 ) explores the normalization of a different order of central moment with Lp norm for general p. Similarly , Hoffer et al . ( 2018 ) explores L1 and L∞ normalization , which enable low precision computation . Littwin & Wolf ( 2018 ) proposes a regularization loss that reduces the variance of sample variances of activation that is closely related to the fourth moment . The idea of controlling activations via statistical characteristics of activations also has motivated initialization methods . An example includes balancing variances of each layer ( Glorot & Bengio , 2010 ; He et al. , 2015 ) , bounding scale of activation and gradient ( Mishkin & Matas , 2016 ; Balduzzi et al. , 2017 ; Gehring et al. , 2017 ; Zhang et al. , 2019 ) , and norm preserving ( Saxe et al. , 2013 ) . Although the desired initial state may not be maintained during training , experimental results show that they can stabilize the learning process as well . Recently , the Wasserstein metric has gained much popularity in a wide range of applications in deep learning with some nice properties such as being a metric in a probability distribution space without requiring common supports of two distributions . For instance , it is successfully applied to a multilabeled classification ( Frogner et al. , 2015 ) , gradient flow of policy update in reinforcement learning ( Zhang et al. , 2018 ) , training of generative models ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ; Kolouri et al. , 2019 ) , and capturing long term semantic structure in sequence-to-sequence language model ( Chen et al. , 2019 ) . While the statistics such as mean and ( co ) variance are useful summaries of a probability distribution , they can not fully represent the underlying structure of the distribution ( Fig . 1 ) . Therefore , regular- izing or normalizing activation to follow the target distribution via statistics can be ineffective in some cases . For instance , normalizing activations via single mean and variance such as BN and decorrelated BN ( Huang et al. , 2018 ) can be inadequate in learning multimodal distribution ( Bilen & Vedaldi , 2017 ; Deecke et al. , 2019 ) . This limitation motivates us to investigate a more general way of regularizing the distribution of activations . Instead of controlling activations via statistics , we define the target distribution and then minimize the Wasserstein distance between the activation distribution and the target distribution . 3 PROJECTED ERROR FUNCTION REGULARIZATION . We consider a neural network with L layers each of which has dl hidden units in layer l. Let D = { ( xi , yi ) } ni=1 be n training samples which are assumed to be i.i.d . samples drawn from a probability distribution Px , y . In this paper , we consider the optimization by stochastic gradient descent with mini-batch of b samples randomly drawn from D at each training iteration . For i-th element of the samples , the neural network recursively computes : hli = φ ( W lhl−1i + b l ) ( 1 ) where h0i = xi ∈ Rd0 , hli ∈ Rdl is an activation in layer l , and φ is an activation function . In the case of recurrent neural networks ( RNNs ) , the recursive relationship takes the form of : hlti = φ ( W lrech l t−1i +W l inh l−1 ti + b l ) ( 2 ) where hlti is an activation in layer l at time t and h l 0i is an initial state . Without loss of generality , we focus on activations in layer l of feed-forward networks and the mini-batch of samples { ( xi , yi ) } bi=1 . Throughout this paper , we let f l be a function made by compositions of recurrent relation in equation 1 up to layer l , i.e. , hli = f l ( xi ) , and f lj be a j-th output of f l. This paper proposes a new regularization loss , called projected error function regularization ( PER ) , that encourages activations to follow the standard normal distribution . Specifically , PER directly matches the distribution of activations to the target distribution via the Wasserstein metric . Let µ ∈ P ( Rdl ) be the Gaussian measure defined as µ ( A ) = 1 2dl/2 ∫ A exp ( − 12 ‖ x ‖ 2 ) dx and νhl = 1 b ∑ i δhli ∈ P ( R dl ) be the empirical measure of hidden activations where δhli is the Dirac unit mass on hli . Then , the Wasserstein metric of order p between µ and νhl is defined by : Wp ( µ , νhl ) = ( inf π∈ ∏ ( µ , ν hl ) ∫ Rdl×Rdl dp ( x , y ) π ( dx , dy ) ) 1/p ( 3 ) where ∏ ( µ , νhl ) is the set of all joint probability measures on Rdl × Rdl having the first and the second marginals µ and νhl , respectively . Because direct computation of equation 3 is intractable , we consider the sliced Wasserstein distance ( Rabin et al. , 2011 ) approximating the Wasserstein distance by projecting the high dimensional distributions onto R ( Fig . 2 ) . It is proved by that the sliced Wasserstein and the Wasserstein are equivalent metrics ( Santambrogio , 2015 ; Bonnotte , 2013 ) . The sliced Wasserstein of order one between µ and νhl can be formulated as : SW1 ( µ , νhl ) = ∫ Sd−1 W1 ( µθ , νhlθ ) dλ ( θ ) = ∫ Sd−1 ∫ ∞ −∞ ∣∣∣∣∣Fµθ ( x ) − 1b b∑ i=1 1〈hli , θ〉≤x ∣∣∣∣∣ dxdλ ( θ ) ( 4 ) where Sdl−1 is a unit sphere in Rdl , µθ and νhlθ represent the measures projected to the angle θ , λ is a uniform measure on Sd−1 , and Fµθ ( x ) is a cumulative distribution function of µθ . Herein , equation 4 can be evaluated through sorting { 〈hli , θ〉 } i for each angle θ . While we can directly use the sliced Wasserstein in equation 4 as a regularization loss , it has a computational dependency on the batch dimension due to the sorting . The computational dependency between samples may not be desirable in distributed and large-batch training that is becoming more and more prevalent in recent years . For this reason , we remove the dependency by applying the Algorithm 1 Backward pass under PER Input The number of Monte Carlo evaluations s , an activation for i-th sample hi , the gradient of the loss ∇hiL , a regularization coefficient λ 1 : g ← 0 2 : for k ← 1 to s do 3 : Sample v ∼ N ( 0 , I ) 4 : θ ← v/ ‖ v ‖2 5 : Project h′i ← 〈hi , θ〉 6 : gk ← erf ( h′i/ √ 2 ) 7 : g ← g + gkθ/s 8 : end for 9 : return ∇hiL+ λg Minkowski inequality to equation 4 , and obtain the regularization loss Lper ( νhl ) : SW1 ( µ , νhl ) ≤ ∫ Sd−1 ∫ ∞ −∞ 1 b b∑ i=1 ∣∣∣Fµθ ( x ) − 1〈hli , θ〉≤x∣∣∣ dxdλ ( θ ) = 1 b b∑ i=1 ∫ Sd−1 ( 〈hli , θ〉erf ( 〈hli , θ〉√ 2 ) + √ 2 π exp ( −〈h l i , θ〉2 2 ) ) dλ ( θ ) = Lper ( νhl ) ( 5 ) whose gradient with respect to hli is : ∇hliLper ( νhl ) = 1 b Eθ∼U ( Sdl−1 ) [ erf ( 〈θ , hli/ √ 2〉 ) θ ] ( 6 ) where U ( Sdl−1 ) is the uniform distribution on Sdl−1 . In this paper , expectation over U ( Sdl−1 ) is approximated by the Monte Carlo method with s number of samples . Therefore , PER results in simple modification of the backward pass as in Alg . 1 . Encouraging activations to follow the standard normal distribution can be motivated by the natural gradient ( Amari , 1998 ) . The natural gradient is the steepest descent direction in a Riemannian manifold , and it is also the direction that maximizes the probability of not increasing generalization error ( Roux et al. , 2008 ) . The natural gradient is obtained by multiplying the inverse Fisher information matrix to the gradient . In Raiko et al . ( 2012 ) and Desjardins et al . ( 2015 ) , under the independence assumption between forward and backward passes and activations between different layers , the Fisher information matrix is a block diagonal matrix each of which block is given by : Fl = E ( x , y ) ∼ ( x , y ) [ ∂L ∂vec ( W l ) ∂L ∂vec ( W l ) T ] = Ex [ hl−1hl−1 T ] E ( x , y ) [ ∂L ∂al ∂L ∂al T ] ( 7 ) where vec ( W l ) is vectorizedW l , hl−1 = f l−1 ( x ) , and al = W lf l−1 ( x ) + bl for x ∼ x . Since computing the inverse Fisher information matrix is too expensive to perform every iterations , previous studies put efforts into developing reparametrization techniques , activation functions , and regularization losses to makeF l close to I , thereby making the gradient close to the natural gradient . For instance , making zero mean and unit variance activations ( LeCun et al. , 1998 ; Schraudolph , 1998 ; Glorot & Bengio , 2010 ; Raiko et al. , 2012 ; Wiesler et al. , 2014 ) and decorrelated activations ( Cogswell et al. , 2016 ; Xiong et al. , 2016 ; Huang et al. , 2018 ) make E [ hl−1hl−1 T ] ≈ I , and these techniques result in faster training and improved generalization performance . In this perspective , it is expected that PER will enjoy the same advantages by matching νhl to N ( 0 , I ) . 3.1 COMPARISON TO CONTROLLING ACTIVATIONS IN Lp SPACE In this subsection , we theoretically compare PER with existing methods that control activations in Lp space . Lp ( Rd0 ) is the space of measurable functions whose p-th power of absolute value is Lebesgue integrable , and norm of f ∈ Lp ( Rd0 ) is given by : ‖ f ‖p= ( ∫ Rd0 |f ( x ) |pdPx ( x ) ) 1/p < ∞ ( 8 ) where Px is the unknown probability distribution generating training samples { xi } ni=1 . Since we have no access to Px , it is approximated by the empirical measure of mini-batch samples . The Lp norm is widely used in the literature for regularization and normalization of neural networks . For instance , activation norm regularization ( Merity et al. , 2017a ) penalizes L2 norm of activations . As another example , BN and its p-th order generalization use Lp norm such that the norm of the centralized activation , or pre-activation , is bounded : ψ ( hlij ) = γ l jξ ( h l ij ) + β l j , ξ ( h l ij ) = hlij − µ̄j ( ∑ k 1 b |h l kj − µ̄j |p ) 1/p ( 9 ) where hlij is j-th unit of h l i , µ̄j = 1 b ∑ k h l kj is the sample mean , β l j is a learnable shift parameter , and γlj is a learnable scale parameters . Herein , we have ‖ ξ ◦ f lj ‖p= 1 for any unit j and any empirical measure , thus ‖ ψ ‖p≤‖ γljξ ◦ f lj ‖p + ‖ βlj ‖p= |γlj |+ |βlj | . PER differs from Lp norm-based approaches in two aspects . First , PER can be considered as Lp norm with adaptive order in the projected space because it is very similar to the Pseudo-Huber loss in one-dimensional space ( Fig . 3 ) . Herein , the Pseudo-Huber loss is a smooth approximation of the Huber loss ( Huber , 1964 ) . Therefore , PER smoothly changes its behavior between L1 and L2 norms , making the regularization loss sensitive to small values and insensitive to outliers with large values . However , the previous approaches use predetermined order p , which makes the norm to change insensitively in the near-zero region when p ≤ 1 or to explode in large value region when p > 1 . Second , PER captures the interaction between hidden units by projection vectors , unlike Lp norm . To see this , let ‖ f l ‖pp= 1b ∑ i , j |hlij |p = 1 b ∑ i , j |〈hli , ej〉|p where { ej } dl j=1 is the natural basis of Rdl . That is , the norm computes the regularization loss , or the normalizer , of activations with the natural basis as a projection vector . However , PER uses general projection vectors θ ∼ U ( Sdl−1 ) , capturing the interaction between hidden units when computing the regularization loss . These two differences make PER more delicate criterion for regularizing activations in deep neural networks than Lp norm , as we will show in the next section .

Summary: This submission belongs to the general field of neural networks and sub-field of activation regularisation. In particular, this submission proposes a novel approach for activation regularisation whereby a distribution of activations within minibatch are regularised to have standard normal distribution. The approach, projected error function regularisation (PER), accomplishes that by minimising an upper-bound on 1-Wasserstein distance between empirical and standard normal distributions.