REGULARIZING ACTIVATIONS IN NEURAL NETWORKS VIA DISTRIBUTION MATCHING WITH THE WASSER-
Year: 2020

Abstract:
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 bothL andL 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 first work to regularize activations via distribution matching in the probability distribution space. We evaluate the proposed method on the image classification task and the word-level language modeling task.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORKS
text: 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 cannot 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.

heading: 3 PROJECTED ERROR FUNCTION REGULARIZATION
text: 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.

heading: 4 EXPERIMENTS
text: 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 first and second moments and VCL regularizing the fourth moments. PER is also compared with L1 and L2 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 coefficient for the regularization losses computed in each layer.

heading: 4.1 IMAGE CLASSIFICATION IN CIFAR-10, CIFAR-100, AND TINY IMAGENET
text: We evaluate PER in image classification task in CIFAR (Krizhevsky et al., 2009) and a subset of ImageNet (Russakovsky et al., 2015), called tiny ImageNet. We first evaluate PER with ResNet (He et al., 2016) in CIFAR-10 and compare it with BN and a vanilla network initialized by fixup 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. (2019), and we obtain similar performances presented in the papers. Herein, we search the regularization coefficient 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. (2016). 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. (2019). By doing so, we obtain similar results presented in Littwin & Wolf (2018). Again, a search space of the regularization coefficient 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 filters in each layer is doubled. Due to the limited time and resources, we
conduct experiments only with ELU that gives good performances for PER, BN, and VCL in CIFAR. As shown in Table 2, PER is also effective in the larger model in the larger image classification dataset.

heading: 4.2 LANGUAGE MODELING IN PTB AND WIKITEXT2
text: We evaluate PER in word-level language modeling task in PTB (Mikolov et al., 2010) and WikiText2 (Merity et al., 2017b). We apply PER to LSTM with two layers having 650 hidden units with and without reuse embedding (RE) proposed in Inan et al. (2017) and Press & Wolf (2016), and variational dropout (VD) proposed in Gal & Ghahramani (2016). We used the same configurations with Merity et al. (2017a) and failed to reproduce the results in Merity et al. (2017a). 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.25 instead of 10 based on the default hyperparameter of the PyTorch word-level language model1 that is also mentioned in Merity et al. (2017a). We also train the networks for 60 epochs instead of 80 epochs since validation perplexity is not improved after 60 epochs in most cases. In this task, PER is compared with recurrent BN (RBN; Cooijmans et al., 2017) because BN is not directly applicable to LSTM. We also compare PER with L1 and L2 activation norm regularizations. Herein, the search space of regularization coefficients of PER, L1 regularization, and L2 regularization is {3e-4, 1e-4, 3e-5 }. For L1 and L2 penalties in PTB, we search additional coefficients over { 1e-5, 3e-6, 1e-6, 3e-6, 1e-6, 3e-7, 1e-7 } because the searched coefficients 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.

heading: 4.3 ANALYSIS
text: 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.
1Available in https://github.com/pytorch/examples/tree/master/word_language_ model

heading: 4.3.1 COMPUTATIONAL COMPLEXITY
text: 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(bdls) for projection operation in each layer l. On the other hand, BN and VCL have O(bdl) 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.

heading: 4.3.2 CLOSENESS TO THE STANDARD NORMAL DISTRIBUTION
text: To examine the effect of PER on the closeness toN (0, I), we analyze the distribution of activations in 11-layer CNN in different perspectives. We first analyze the distribution of a single activation hlj 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 L2 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 νh6j 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 νh3j and νh9j at the early stage of learning.
To precisely evaluate closeness to the standard normal distribution, we also analyze SW1(N (0, I), νhl) 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. (2011). As similar to the previous result, while BN βlj = 0 and γ l j = 1 at initial state gives small SW1(N (0, I), νhl) in early stage of training, PER also can effectively control
the distribution without such normalization. This confirms that PER prevents the distribution of activation to be drifted away from the target distribution.

heading: 5 CONCLUSION
text: 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 classification and language modeling experiments, PER gives marginal but consistent improvements over methods based on sample statistics (BN and VCL) as well as L1 and L2 activation regularization methods. The analysis of changes in activations’ distribution during training verifies 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-specific 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 the function L2 space of logits (Benjamin et al., 2019).

heading: ACKNOWLEDGMENTS
text: We would like to thank Min-Gwan Seo, Dong-Hyun Lee, Dongmin Shin, and anonymous reviewers for the discussions and suggestions.

Review summaries:
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.