R EGULARIZING ACTIVATIONS IN NEURAL NETWORKS VIA DISTRIBUTION MATCHING WITH THE W ASSERSTEIN METRIC
Anonymous authors
Paper under double-blind review
A BSTRACT
Regularization and normalization have become indispensable components in training deep neural networks, resulting in faster training and improved generalization performance.
We propose the projected error function regularization loss (PER) that encourages activations to follow the standard normal distribution.
PER randomly projects activations onto one-dimensional space and computes the regularization loss in the projected space.
PER is similar to the Pseudo-Huber loss in the projected space, thus taking advantage of both L 1 and L 2 regularization losses.
Besides, PER can capture the interaction between hidden units by projection vector drawn from a unit sphere.
By doing so, PER minimizes the upper bound of the
Wasserstein distance of order one between an empirical distribution of activations and the standard normal distribution.
To the best of the authors’ knowledge, this is the ﬁrst work to regularize activations via distribution matching in the probability distribution space.
We evaluate the proposed method on the image classiﬁcation task and the word-level language modeling task.
I NTRODUCTION
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 ﬂat 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, 2015;
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 ﬁrst and second moments.
Back in the 1990s, it has been known that the neural network training can be beneﬁted 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.
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 L p norm including L 0 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.
Speciﬁcally, 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.,
[Figure]
all orders of moments and correlation between hidden units.
The extensive experiments on multiple challenging tasks show the effectiveness of PER.
R ELATED 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 L 2 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 L p norm for general p .
Similarly, Hoffer et al.
(2018) explores Land 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, 2015; 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 classiﬁcation (Frogner et al., 2015), gradient ﬂow 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 cannot fully represent the underlying structure of the distribution (Fig. 1).
Therefore, regularizing 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 deﬁne the target distribution and then minimize the Wasserstein distance between the activation distribution and the target distribution.
P ROJECTED ERROR FUNCTION REGULARIZATION
We consider a neural network with L layers each of which has d l hidden units in layer l .
Let
T = { ( x i , y i ) } ni =1 be n training samples which are assumed to be i.i.d.
samples drawn from a probability distribution P x , y .
In this paper, we consider the optimization by stochastic gradient descent with mini-batch of b samples randomly drawn from T at each training iteration.
For i -th element of the samples, the neural network recursively computes:
[Equation]
where h
i
= x i ∈ R d 0 , h li ∈ R d l 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:
[Equation]
(2) where h lt i is an activation in layer l at time t and h l 0 i 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 { ( x i , y i ) } 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., h li =
f l ( x i ) , and f lj be a j -th output of f l .
This paper proposes a new regularization loss, called PER, that encourages activations to follow the standard normal distribution.
Speciﬁcally, PER directly matches the distribution of activations to the target distribution via the Wasserstein metric.
Let µ ∈ P
( R d l ) be the Gaussian measure deﬁned as µ
[Equation]
A
[Equation]
=
[Equation]
and ν
[Equation]
P ( R d l ) be the empirical measure of hidden activations where δ h li is the Dirac unit mass on h li .
Then, the Wasserstein metric of order p between µ and ν h l is deﬁned by:
[Equation]
where (cid:81) ( µ, ν h l ) is the set of all joint probability measures on R d l × R d l having the ﬁrst and the second marginals µ and ν h l , 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 ν h l can be formulated as:
[Equation]
where S d l − 1 is a unit sphere in R d l , µ θ and ν h l θ represent the measures projected to the angle θ , λ is a uniform measure on S d − 1 , and F µ θ ( x ) is a cumulative distribution function of µ θ .
Herein, equation 4 can be evaluated through sorting (cid:8) (cid:104)
h li , θ (cid:105) (cid:9) 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
Input The number of Monte Carlo evaluations s , an activation for i -th sample h
i , the gradient of the loss ∇ h
i L , a regularization coefﬁcient λ 1: g ←2: for k ← 1 to s do 3:
Sample v ∼ N ( 0 , I ) 4: θ ← v / (cid:107) v (cid:107)5:
Project h
[Equation]
6: g
[Equation]
(cid:1) 7: g ← g + g k θ /s 8: end for 9: return ∇ h
i L
+ λ g
Minkowski inequality to equation 4, and obtain the regularization loss L per ( ν h l ) :
[Equation]
(cid:19)(cid:33)
[Equation]
whose gradient with respect to h li is:
[Equation]
where U ( S d l − 1 ) is the uniform distribution on S d l − 1 .
In this paper, expectation over U ( S d l − 1 ) is approximated by the Monte Carlo method with s number of samples.
Therefore, PER results in simple modiﬁcation 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.
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:
[Equation]
where vec ( W l ) is vectorized W l , h l − 1 = f l − 1 ( x ) , and a l = W l f l − 1 ( x ) + b l 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 make F 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; Wiesler et al., 2014; Glorot & Bengio, 2010; Raiko et al., 2012) and decorrelated activations (Huang et al., 2018; Cogswell et al., 2015; Xiong et al., 2016) make E (cid:104)
h l − 1 h l − 1 T (cid:105)
≈ 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 ν h l to N (0 , I ) .
3. C OMPARISON TO CONTROLLING ACTIVATIONS IN L p SPACE
In this subsection, we theoretically compare PER with existing methods that control activations in
L p space.
L p ( R d 0 ) is the space of measurable functions whose p -th power of absolute value is
Lebesgue integrable, and norm of f ∈ L p ( R d 0 ) is given by:
[Equation]
where P x is the unknown probability distribution generating training samples { x i } ni =1 .
Since we have no access to P x , it is approximated by the empirical measure of mini-batch samples.
The L p norm is widely used in the literature for regularization and normalization of neural networks.
For instance, activation norm regularization (Merity et al., 2017) penalizes L 2 norm of activations.
As another example, BN and its p -th order generalization use L p norm such that the norm of the centralized activation, or pre-activation, is bounded:
[Equation]
where h lij is j -th unit of h li , ¯ µ j = 1 b (cid:80)
k
h lkj is the sample mean, β lj is a learnable shift parameter, and γ lj is a learnable scale parameters.
Herein, we have (cid:107) ξ ◦ f lj (cid:107) p = 1 for any unit j and any empirical measure, thus (cid:107) ψ (cid:107) p ≤(cid:107)
γ lj ξ ◦ f lj (cid:107) p + (cid:107) β lj (cid:107) p = | γ lj | + | β lj | .
PER differs from L p norm-based approaches in two aspects.
First, PER can be considered as L p 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 L 1 and Lnorms, 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 L p norm.
To see this, let (cid:107)
f l (cid:107) pp
= 1 b (cid:80) i,j | h l ij | p = 1 b (cid:80) [Table] , e
[Table]
| p where { e j } d l j =1 is the natural basis of
[Table] [Table]
R d l .
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 ( S d l − 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 L p norm, as we will show in the next section.
E XPERIMENTS
This section illustrates the effectiveness of PER through experiments on different benchmark tasks with various datasets and architectures.
We compare PER with BN normalizing the ﬁrst and second moments and VCL regularizing the fourth moments.
PER is also compared with L 1 and L 2 activation norm regularizations that behave similarly in some regions of the projected space.
We then analyze the computational complexity PER and the impact of PER on the distribution of activations.
Throughout all experiments, we use 256 number of slices and the same regularization coefﬁcient for the regularization losses computed in each layer.
4. I MAGE CLASSIFICATION IN CIFAR-10, CIFAR-100, AND TINY I MAGE N ET
We evaluate PER in image classiﬁcation task in CIFAR (Krizhevsky et al., 2009) and a subset of
ImageNet (Russakovsky et al., 2015), called tiny ImageNet.
We ﬁrst evaluate PER with ResNet (He et al., 2016) in CIFAR-10 and compare it with BN and a vanilla network initialized by ﬁxup initialization (Zhang et al., 2019).
We match the experimental details in training under BN with He et al.
(2016) and under PER and vanilla with Zhang et al.
and we obtain similar performances presented in the papers.
Herein, we search the regularization coefﬁcient over { 3e-4, 1e-4, 3e-5, 1e5 } .
Table 1 presents results of CIFAR-10 experiments with ResNet-56 and ResNet-110.
PER outperforms BN as well as vanilla networks in both architectures.
Especially, PER improves the test errors by 0.49 % and 0.71 % in ResNet-56 and ResNet-110 without BN, respectively.
We also performed experiments on an 11-layer convolutional neural network (11-layer CNN) examined in VCL (Littwin & Wolf, 2018).
This architecture is originally proposed in Clevert et al.
Following Littwin & Wolf (2018), we perform experiments on 11-layer CNNs with ELU,
ReLU, and Leaky ReLU activations, and match experimental details in Littwin & Wolf (2018) except that we used 10x less learning rate for bias parameters and additional scalar bias after ReLU and Leaky ReLU based on Zhang et al.
By doing so, we obtain similar results presented in
Littwin & Wolf (2018).
Again, a search space of the regularization coefﬁcient is { 3e-4, 1e-4, 3e-5, 1e-5 } .
For ReLU and Leaky ReLU in CIFAR-100, however, we additionally search { 3e-6, 1e-6, 3e-7, 1e-7 } because of divergence of training with PER in these setting.
As shown in Table 3, PER shows the best performances on four out of six experiments.
In other cases, PER gives compatible performances to BN or VCL, giving 0.16 % less than the best performances.
Following Littwin & Wolf (2018), PER is also evaluated on tiny ImageNet.
In this experiment, the number of convolutional ﬁlters in each layer is doubled.
Due to the limited time and resources, we
[Table]
conduct experiments only with ELU that gives good performances for PER, BN, and VCL in CIFAR.
As shown in Table
[Table]
PER is also effective in the larger model in the larger image classiﬁcation dataset.
4. L ANGUAGE MODELING IN PTB AND W IKI T EXT
We evaluate PER in word-level language modeling task in PTB (Mikolov et al., 2010) and WikiText2 (Merity et al., 2016).
We apply PER to LSTM with two layers having 650 hidden units with and without reuse embedding (RE) proposed in Inan et al.
and Press & Wolf (2016)
[Table]
and variational dropout (VD) proposed in Gal & Ghahramani (2016).
We used the same conﬁgurations with Merity et al.
and failed to reproduce the results in Merity et al.
Especially, when we rescale gradient when its norm exceeds 10, we observed divergence or bad performance (almost 2x perplexity compared to the published result).
Therefore, we rescale gradient with norm over 0.
instead of 10 based on the default hyperparameter of the PyTorch word-level language model 1 that is also mentioned in Merity et al.
We also train the networks for 60 epochs instead ofepochs since validation perplexity is not improved after 60 epochs in most cases.
In this task, PER is compared with recurrent BN (RBN) because BN is not directly applicable to LSTM (Cooijmans et al., 2016).
We also compare PER with L 1 and L 2 activation norm regularizations.
Herein, the search space of regularization coefﬁcients of PER, L 1 regularization, and L 2 regularization is { 3e-4, 1e-4, 3e-5 } .
For L 1 and L 2 penalties in PTB, we search additional coefﬁcients over { 1e-5, 3e-6, 1e-6, 3e-6, 1e-6, 3e-7, 1e-7 } because the searched coefﬁcients seem to constrain the capacity.
We list in Table 4 the perplexities of methods on PTB and WikiText2.
While all regularization techniques show regularization effects by giving improved test perplexity, PER gives the best test perplexity except LSTM and RE-VD-LSTM in the PTB dataset wherein PER is the second-best method.
We also note that naively applying RBN often reduces performance.
For instance, RBN increases test perplexity of VD-LSTM by about 5 in PTB and WikiText2.
4. A NALYSIS
In this subsection, we analyze the computational complexity of PER and its impact on closeness to the standard normal distribution in the 11-layer CNN.
[Table]
4.3. C OMPUTATIONAL COMPLEXITY
PER has no additional parameters.
However, BN and VCL require additional parameters for each channel and each location and channel in every layer, respectively; that is, 2.5K and 350K number of parameters are introduced in BN and VCL in the 11-layer CNN, respectively.
In terms of time complexity, PER has the complexity of O ( bd l s ) for projection operation in each layer l .
On the other hand, BN and VCL have O ( bd l ) complexities.
In our benchmarking, each training iteration takes 0.071 seconds for a vanilla network, 0.083 seconds for BN, 0.087 for VCL, and 0.093 seconds for PER on a single NVIDIA TITAN X. Even though PER requires slightly more training time than
BN and VCL, this disadvantage can be mitigated by computation of PER is only required in training and PER does not have additional parameters.
4.3. C LOSENESS TO THE STANDARD NORMAL DISTRIBUTION
To examine the effect of PER on the closeness to N ( 0 , I ) , we analyze the distribution of activations in 11-layer CNN in different perspectives.
We ﬁrst analyze the distribution of a single activation h lj for some unit j and layer l (Fig. 4).
We observe that changes in probability distributions between two consecutive epochs are small under BN because BN bound the L 2 norm of activations into learned parameters.
On the contrary, activation distributions under vanilla and PER are jiggled between two consecutive epochs.
However, PER prevents the variance explosion and pushes the mean to zero.
As shown in Fig.
4, while variances of ν h 6 j under both PER and Vanilla are very high at the beginning of training, the variance keeps moving towards one under PER during training.
Similarly,
PER recovers biased means of ν h 3 j and ν h 9 j at the early stage of learning.
To precisely evaluate closeness to the standard normal distribution, we also analyze
SW 1 ( N ( 0 , I ) , ν h l ) at each epoch (Fig. 5).
Herein, the sliced Wasserstein distance is computed by approximating the Gaussian measure using the empirical measure of samples drawn from N ( 0 , I ) as in Rabin et al.
As similar to the previous result, while BN β lj = 0 and γ lj = 1 at initial state gives small SW 1 ( N ( 0 , I ) , ν h l ) in early stage of training, PER also can effectively control
[Figure]
the distribution without such normalization.
This conﬁrms that PER prevents the distribution of activation to be drifted away from the target distribution.
C ONCLUSION
We proposed the regularization loss that minimizes the upper bound of the 1-Wasserstein distance between the standard normal distribution and the distribution of activations.
In image classiﬁcation and language modeling experiments, PER gives marginal but consistent improvements over methods based on sample statistics (BN and VCL) as well as L 1 and L 2 activation regularization methods.
The analysis of changes in activations’ distribution during training veriﬁes that PER can stabilize the probability distribution of activations without normalization.
Considering that the regularization loss can be easily applied to a wide range of tasks without changing architectures or training strategies unlike BN, we believe that the results indicate the valuable potential of regularizing networks in the probability distribution space as a future direction of research.
The idea of regularizing activations with the metric in probability distribution space can be extended to many useful applications.
For instance, one can utilize task-speciﬁc prior when determining a target distribution, e g., the Laplace distribution for making sparse activation.
The empirical distribution of activations computed by a pretrained network can also be used as a target distribution to prevent catastrophic forgetting.
In this case, the activation distribution can be regularized so that it does not drift away from the activation distribution learned in the previous task as different from previous approaches constrains the changes in the parameter l 2 space (Kirkpatrick et al., 2017) or the function L 2 space (Benjamin et al., 2018)