Title: P -BN: TOWARDS EFFECTIVE BATCH NORMALIZA-TION IN THE PATH SPACE

Abstract: Neural networks with ReLU activation functions have demonstrated their success in many applications. Recently, researchers noticed a potential issue with the optimization of ReLU networks: the ReLU activation functions are positively scaleinvariant (PSI), while the weights are not. This mismatch may lead to undesirable behaviors in the optimization process. Hence, some new algorithms that conduct optimizations directly in the path space (the path space is proven to be PSI) were developed, such as Stochastic Gradient Descent (SGD) in the path space, and it was shown that, SGD in the path space is superior to that in the weight space. However, it is still unknown that whether other deep learning techniques beyond SGD, such as batch normalization (BN), could also have their counterparts in the path space. In this paper, we conduct a formal study on the design of BN in the path space. According to our study, the key challenge is how to ensure the forward propagation in the path space, because BN is utilized during the forward process. To tackle such challenge, we propose a novel re-parameterization of ReLU networks, with which we replace each weight in the original neural network, with a new value calculated from one or several paths, while keeping the outputs of the network unchanged for any input. Then we show that BN in the path space, namely P-BN, is just a slightly modified conventional BN on the re-parameterized ReLU networks. Finally, our experiments on two benchmark datasets, CIFAR and ImageNet, show that the proposed P-BN can significantly outperform the conventional BN in the weight space.

1. INTRODUCTION:

In recent years, neural networks with rectified linear unit (ReLU) activation functions (abbrev. ReLU networks) , have been successfully applied to many domains, such as image classification (He et al., 2016) , game playing (Mnih et al., 2015) , and text processing (Kim, 2014) . Recently, it has been noticed that the feedforward neural networks with ReLU, namely ReLU networks, are positively scale-invariant (PSI) (Neyshabur et al., 2015a; 2016) , which means when the incoming and outgoing weights of a hidden node, are respectively multiplied and divided by a positive constant, the outputs of ReLU networks will keep unchanged for any input. However, the vector space, composed of weights, in which the conventional optimization algorithms, e.g., Stochastic Gradient Descent (S-GD), optimize the neural networks, is not PSI, and such mismatch may lead to undesirable behaviors in the optimization process (Neyshabur et al., 2015a) .

To tackle such challenge, some recent studies have shown that, one ReLU network can be optimized in a completely new PSI parameter space, i.e., the path space, instead of its original weight space (Meng et al., 2019) . Detailedly, while regarding ReLU networks as directed acyclic graphs (DAGs), we can calculate the outputs of ReLU networks by using path-values, i.e., the multiplication of weights along each input-output path in the DAG. Here, path-values are invariant to the positive rescaling of weights, which exactly matches the PSI property of ReLU networks, and the optimization algorithms (e.g., SGD) in the path space (i.e., the vector space composed of path-values), are shown to be superior to those in the weight space (Neyshabur et al., 2015a; 2016; Meng et al., 2019) .

There are already some optimization algorithms developed in the path space, however, to the best of our knowledge, it is still unknown whether other deep learning techniques beyond SGD, can also have their counterparts in the path space. For example, BN is one of the most widely used normal-ization approaches (Ioffe & Szegedy, 2015; Santurkar et al., 2018) , and is crucial for facilitating the training process of neural networks. With the help of BN, the training process of neural networks can be accelerated, by both decoupling the scale and direction of weight vectors (Arora et al., 2018) and smoothing the optimization landscape (Santurkar et al., 2018) . In this paper, we conduct a formal study on the design of BN in the path space. First, it is known that, BN is successfully utilized during the forward process in the weight space, but how to ensure the forward propagation still remains unclear in the path space, i.e., the way to calculate the outputs of hidden nodes via path-values layer by layer is unknown. For example, in Fig. 1 , we show a network and its outputs calculated by both weights and path-values.

𝑣

Here, the outputs of this network can be simply calculated by using path-values, because a path starts from an input node and ends up with an output node, and the path-value is just the multiplication of weights along this path. Nevertheless, the outputs of hidden nodes cannot be calculated by path-values in this obvious way. Thus, to solve this problem, we propose a re-parameterization for the weights in Re-LU networks, which can replace each weight in the original network, with a new value calculated by one or several path-values, and the outputs of the re-parameterized network will keep unchanged for any input. Therefore, the forward propagation in the path space can be easily conducted after re-parameterizing the network.

Then, we show that BN in the path space is just a slightly modified conventional BN on the reparameterized network. Considering that BN in the weight space, can make the gradient propagation more effective, we study BN in the path space from the perspective of effective propagation. In details, we apply the conventional BN to the re-parameterized network, and find that the radients will explode during the backward propagation in the path space. Thus, we propose a novel BN in the path space, namely P-BN, which slightly modifies the conventional BN, via excluding the term whose coefficient is fixed by re-parameterization from normalization. Particularly, we theoretically prove and experimentally show that compared with the conventional BN, P-BN can propagate gradients in the path space more effectively.

Finally, we conduct extensive experiments on multiple datasets, including CIFAR, ImageNet, with multiple network structures, including ResNet and PlainNet. The results demonstrate that the P-BN based optimization algorithms can significantly outperform the ones with the conventional BN.

2. BACKGROUND:



2.1. RELATED WORK:

In recent years, some researchers have conducted many theoretical studies on the path of ReLU networks. On the one hand, in terms of optimization, for example, Neyshabur et al. (2015a) proposed a new regularization term, i.e., path-norm, and the Path-SGD algorithm for minimizing the regularized loss. Meng et al. (2019) proposed G-SGD algorithm to directly optimize the model in the path space, by utilizing the gradients w.r. However, the aforementioned studies have not involved normalization approaches (Ba et al., 2016; Wu & He, 2018) , such as BN (Ioffe & Szegedy, 2015) , which are proved to be much crucial for training neural networks. In particular, as for ReLU networks, BN is especially widely used in computer vision domain (He et al., 2016; Li et al., 2016; Huang et al., 2017) . Furthermore, there are also many theoretical studies on BN (Santurkar et al., 2018; Arora et al., 2018; Kohler et al., 2019) .

Accordingly, based on the studies on optimizing ReLU networks in the path space, and the crucial normalization technique, i.e., BN, in this paper, we will give a formal study on designing BN in the path space. We will next briefly introduce the optimization of ReLU networks without BN in the path space.

2.2. OPTIMIZING RELU NETWORKS IN THE PATH SPACE:

𝑂 1 0 𝑂 2 0 𝑂 3 0 𝑂 1 1 𝑂 2 1 𝑂 3 1 𝑂 1 3 𝑂 2 3 𝑂 1 2 𝑂 2 2 𝑂 3 2 Input Output 𝑂 3 3 𝑤 22 1 𝑤 12 2 𝑤 31 3 𝑤 1 𝑤 2 𝑤 3

Hidden Figure 2 : An example of a feedforward ReLU network.

In this subsection, we introduce the feedforward ReLU network and the optimization algorithm in the path space.

Firstly, as shown in Fig. 2 , consider an L-layer feedforward ReLU network without BN1 whose weight matrices are

{w l ; l = 1, • • • , L}, we use O l j , l = 0, • • • , L

to denote the j-th node in the l-th layer. Based on these above, we further use w l jk to represent the weight on the edge which connects O l-1 k and O l j . As for the outputs of hidden nodes, suppose that given a

d-dimensional input x = (x1, • • • , x d ),

then the output of O l j and the hidden vector in the l-th layer, can be denoted as o l j (x) and o l (x), respectively. Thus, the outputs can be propagated as o l (x) = g(w l o l-1 (x)), where g(•) = max(•, 0) is a ReLU activation function.

Secondly, we will show the concept of a path in ReLU networks, and briefly introduce the optimization algorithm in the path space. In details, regarding the network structure as a DAG consisting of nodes and edges, a path can be defined as: a list of nodes or edges, starting from an input node, successively passing by several hidden nodes along the edges, and finally ending up with an output node. For example, in Fig. 2 As for the optimization algorithm, we introduce two typical ones, i.e., Path-SGD (Neyshabur et al., 2015a) and G-SGD (Meng et al., 2019) . Path-SGD is to add a new regularizer, i.e., path-norm, into the loss function, and then heuristically utilizes a coordinate-wised solution for minimizing the pathnorm regularized loss. It is shown to be superior to SGD for MLP, but with huge computational costs when applied to Convolutional Neural Networks (CNNs), due to the complicate correlations among path-values. Besides, Path-SGD only utilizes the PSI property to propose the regularizer, but cannot optimize the network in the path space. Comparatively, G-SGD can solve the computation problem, by decoupling correlations between path-values. Besides, G-SGD can optimize the network via updating path-values by using gradients w.r.t. them instead of weights. Thus, G-SGD can serve as the SGD in the path space, and the detailed on G-SGD are listed in Appendix C.

3. TO DESIGN BN IN THE PATH SPACE:

In this section, we first introduce a re-parameterization process, to ensure the forward propagation in the path space, and calculate the outputs of hidden nodes by using path-values. Then, we will show that the conventional BN cannot effectively propagate the error signals in the path space, which calls us to design an effective BN in the path space.

3.1. RE-PARAMETERIZATION:

In this subsection, we introduce a re-parameterization for the parameters of ReLU networks, which can replace each weight with a path-value or a ratio of path-values, and the outputs of the re- parameterized network will keep unchanged for any input. First of all, we present the process of re-parameterization in Theorem 3.1.

Theorem 3.1 Consider an L-layer ReLU network with weight matrices {w l ; l = 1, • • • , L}, if the diagonal elements in matrix w l , l > 2 are positive, the hidden outputs of the ReLU network can be calculated by using path-values as follows:

EQUATION

EQUATION

Here, g(•) denotes the ReLU activation function, and v 1 kk is the path-value of the path which only contains the k-th diagonal elements in the weight matrices, i.e., v 1 kk = L l=1 w l kk , and v l jk , k = j is the path-value of the path which contains one non-diagonal weight, i.e., v l jk = ( s<l w s kk ) • w l jk • ( s>l w l jj ).

Proof: Based on the PSI property, we will prove Theorem 3.1 by designing a re-parameterization process which consists several weight rescaling steps. After these steps, the ReLU network can be equivalently re-parameterized, which means that the outputs will keep unchanged for any input after re-parameterization. Here, we will next demonstrate the re-parameterization process by showing that it can be conducted for a 3-layer MLP, as well as generalized to a multi-layer setting.

As for a 3-layer MLP, the process is shown in Fig. 3 , and described in the following two steps:

(1) The weights connected with nodes O 2 j , j = 1, 2 are rescaled, i.e., the incoming and outgoing weights of node O 2 j are multiplied and divided by w 3 jj , respectively. To be specific, w 3 11 , w 3 21 and w 2 11 , w 2 12 are divided and multiplied by w 3 11 , and similarly, w 3 22 , w 3 12 and w 2 22 , w 2 21 are also rescaled by using w 3 22 . According to the definition of a path value, we have w 3 21

w 3 11 = w 3 21 w 2 11 w 1 11 w 3 11 w 2 11 w 1 11 = v 3 21 v 1 11 , w 3 12 w 3 22 = w 3 12 w 2 22 w 1 22 w 3 22 w 2 22 w 1 22 = v 3 12 v 1 22

, and hence, each weight in the 3rd layer can be replaced by a constant or a ratio of path-values. (2) The weights connected with nodes O 1 j , j = 1, 2 are rescaled, i.e., the incoming and outgoing weights of node O 1 j are multiplied and divided by the corresponding outgoing diagonal weight (i.e., w 2 jj w 3 jj ). Afterwards, each weight in the middle layer can be replaced by a constant or a ratio of path-values, and each weight in the 1st layer can be replaced by a path-value. Therefore, each weight in this network are replaced by a path-value or a ratio of path-values.

As for a multi-layer setting, the re-parameterization process is similar to that of a 3-layer MLP. In particular, this process can be proceeded by rescaling weights orderly, from which is connected with the nodes in the last hidden node layer to the first hidden node layer, and the scalar of the rescaling operation at each node, is the corresponding outgoing diagonal weight. After such rescaling steps, the weights can be re-parameterized as follows: (1)

w 1 jk → v 1 jk ; (2) w l kk → 1; (3) w l jk → v l jk v 1 kk , l = 1, j = k.

In summary, a network with weight at each edge (e.g., Step 1 in Fig. 3 ), can be equivalently reparameterized into another network with a path-value or a ratio of path-values at each edge (e.g.,

Step 4 in Fig. 3 ), by rescaling its weights layer by layer. After the re-parameterization, the outputs of hidden nodes can be calculated by using path-values, in the same way as they are calculated by weights, and hence, Theorem 3.1 is established.

Remark::

The positive constraint is not essential for proving Theorem 3.1, because when the diagonal weights are negative, the re-parameterization can still be conducted correctly by using the absolute value of the corresponding weight as the rescaling scalar. Considering the robustness and computational efficiency, this constrain exists in the optimization algorithm (Meng et al., 2019) , and hence, we follow it. Besides, this constrain will not bring much influence on the model expressiveness, because the number of the constrained weights is tiny, compared with the total weights. Moreover, we discuss some other choices on the re-parameterization and the influence in Appendix D.

3.2. ANALYZING THE BN IN THE PATH SPACE:

With the re-parameterization in Section 3.1, the forward propagation can be ensured in the path space, and the outputs of hidden nodes can be calculated by using path-values. Hence, in this subsection, we will design BN in the path space, via applying and analyzing the conventional BN to the re-parameterized network.

Firstly, we briefly introduce the conventional BN. Given a mini-batch of inputs

{x i ; i = 1, • • • , m}, BN can normalize each output of the hidden node as BN (o l j (x i )) = γ l j • o l j (x i )-µ l j σ l j + β l j , where µ l j = 1 m m i=1 o l j (x i ) and σ l j = 1 m m i=1 (o l j (x i ) -µ l j ) 2 +

represent the mean and standard deviation of the minibatch outputs, and γ l j , β l j are the scale and shift term2 which are independent with w, respectively. Then, we can denote the output after BN and ReLU as z

l j (x i ) = g BN (o l j (x i )) .

Secondly, we analyze the gradient propagation in the re-parameterized network with the the conventional BN. Here, we use L and

∇ z l (x i ) L = ∂L ∂z l 1 (x i ) , • • • , ∂L ∂z l d (x i )

, to denote the cross-entropy loss and the gradient vector w.r.t hidden output in layer l, respectively. Then, we can give an estimation of the norm of ∇ z l (x i ) L in Theorem 3.2. Theorem 3.2 We use W l to denote the parameter matrix in layer l after the re-parameterization. Suppose z l j (x i ) < O(1), ∀i, j, l, and m > 64. The norm of the gradient w.r.t z l (x i ) can be upper bounded as

∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • γ s • (σ s ) -1 • W s ,

where

σ l = diag(σ l 1 , • • • , σ l d ), γ l = diag(γ l 1 , • • • , γ l d ), and D l (x i ) = diag(D l 1 (x i ), • • • , D l d (x i )). Here, diag(a1, • • • , a d ) represents a diagonal matrix whose diagonal elements are a1, • • • , a d , and D l j (x i ) = 1 if z l j (x i ) > 0, otherwise D l j (x i ) = 0.

The proof of Theorem 3.2 can be found in Appendix B.

Discussion: According to Theorem 3.2, the gradient norm exponentially depends on depth L. If the spectral norm of matrix

D s (x i ) • γ s • (σs) -1

• W s roughly equals to 1, the error signals can be effectively propagated in the backward process (i.e., not vanished or exploded). On the one hand, if there is non-zero elements in D s (x i ), its spectral norm will equal to 1. On the other hand, in practice, γ s j is initialized to be 1, and will keep its value around 1 during the training process, so the spectral norm of γ s will be around 1. Hence, the magnitudes of W s and σ -1 s is crucial for the gradient norm. For matrix W s , according to Theorem 3.1, the diagonal elements of W s (s ≥ 2) are degenerated to be constant 1, and these constant terms will not be trained. Besides, the trained parameters are initialized to be much smaller than 1 (cf. Appendix D.3), so matrix W s will approach to an identity matrix. For matrix σ -1 s , if each σ s j is smaller than 1, it can cause gradient exploding during the back propagation process. We will conduct observational experiments in Section 5.2, and it shows that σ s is less than 1 and gradients becomes larger during the back propagation process.

Therefore, such discussion result motivates us to design BN in the path space by modifying the conventional BN, in order to propagate the error signals effectively. We will next detailedly introduce the proposed BN in the path space, i.e., P-BN.

4. P -BN: EFFECTIVE BN IN THE PATH SPACE:

In this section, we design the BN in the path space, namely P-BN, via slightly modifying the forward process of the conventional BN. Motivated by the calculation for the outputs of hidden node in Theorem 3.1, as well as the discussion of Theorem 3.2, we propose to normalize the terms related to the trained coefficients (i.e., Eq. 1 and the 2nd term in Eq. 2), and exclude the term related to the constant coefficient (i.e., the 1st term in Eq. 2). Specifically, we use z l j (x i ) to denote the output after P-BN and ReLU, and detailedly describe the forward process of P-BN in the following two steps.

(1) For the 1st layer, the operation of P-BN is the same as the conventional BN, i.e., z

l j (x i ) = z l j (x i ).

(2) For the l-th layer (l ≥ 2), as shown in Fig. 4 , there are three sub-steps for P-BN as follows. First, calculate the partial weighted summation. Here, the input related to the diagonal constant element in the parameter matrix is excluded, i.e.,

BN:

EQUATION

Second, normalize the partial weighted summation above, i.e.,

EQUATION

Third, add the excluded term z l-1 j (x i ) into the normalized partial weighted summation in Eq. 4, i.e.,

z l j (x i ) = g BN (o l,/ j (x i )) + z l-1 j (x i ) (5)

Consequently, P-BN can be easily combined with the optimization algorithm in the path space for ReLU networks, and we will show it in Algorithm 1 in Appendix E.

Then, we will show that the gradients can be effectively propagated by using P-BN in Theorem 4.1.

Theorem 4.1 Suppose z l j (x i ) < O(1), ∀i, j, l, and m > 64, the norm of gradient w.r.t z l (x i ) can be upper bounded as

∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • (I + γ s,/ • (σ s,/ ) -1 • Ŵ s ) ,

where

Ŵ s = W s -I, σ l,/ = diag(σ l,/ 1 , • • • , σ l,/ d ), γ l,/ = diag(γ l,/ 1 , • • • , γ l,/ d ) and D l (x i ) = diag(D l 1 (x i ), • • • , D l d (x i )) where D l j (x i ) = 1 if z l j (x i ) > 0, otherwise D l j (x i ) = 0.

The proof of Theorem 4.1 can be found in Appendix B. Then, we give Corollary 4.2, via summarizing from Theorem 3.2 and Theorem 4.1.

Corollary 4.2 Using the notations in Theorem 3.2 and Theorem 4.1 and initializing γ l and γ l,/ to be identity matrix, we have the following upper bound for gradient norm at initialization point Suppose λ = Ŵ s . A sufficient condition for that ∇ z l (x i ) L and ∇ z l (x i ) L can be upper bounded by constant (i.e., will not diverge as

∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • (σ s ) -1 • (I + Ŵ s ) , ∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • I + (σ s,/ ) -1 • Ŵ s . Dataset Method

L → ∞) is σ s j ∈ [ L(1+λ) L+1 , L(1+λ) L-1 ] and σ s,/ j ∈ [λL, ∞) ∀j, respectively. 3

Discussion: We can conclude from Corollary 4.2 that the range for σ s,/ j is much larger than that for σ s j . Furthermore, if λ ≤ 1 L , we have

EQUATION

L-1 ] ⊂ [λL, ∞) (i.e., the range for σ s j is fully contained in the range σ s,/ j ), which indicates the backward stability of P-BN. We will observe the variance in each layer in Section 5.2, and the results demonstrate that the variance becomes larger with P-BN. Thus, to some extent, the gradient propagated more stable for NN with P-BN.

5. EXPERIMENTS:

In this section, we conduct experiments by first comparing the performance of P-BN with the baselines on various datasets with some network structures, and then showing some experimental observations to support theoretical analyses.

5.1. PERFORMANCE EXPERIMENTS:

In this subsection, we evaluate the performance of P-BN on training deep neural networks by conducting experiments on three datasets, CIFAR-10, CIFAR-100, and ImageNet. We will next introduce the network structures and compared methods first.

5.1.1. NETWORK STRUCTURES AND COMPARED METHODS:

In this subsection, we describe the network structures and compared methods. First, as for the network structures, we apply P-BN to train ResNet and PlainNet (He et al., 2016) . Second, we show the setting details of four compared methods as follows.

(1) SGD+BN: We use SGD to train the network with the conventional BN.

(2) G-SGD+BN: We use G-SGD to train the network with the conventional BN.

(3) G-SGD+BN(wnorm): We use a method which was intuitively proposed to handle the conventional BN without any theoretical analysis (Meng et al., 2019) . Here, the gradient of the path-value, is normalized by the L2-norm of the weights pointing to the same hidden unit.

(4) G-SGD+P-BN: Our proposed P-BN targets on BN in the path space, and the network with P-BN should be optimized by the algorithms in the path space. Thus, we use G-SGD to train the network with P-BN. (3) G-SGD+P-BN outperforms G-SGD+BN(wnorm), which shows the benefit and significance of our proposed formal analyses.

5.1.2. CIFAR:



5.1.3. IMAGENET:

We conduct experiments on ImageNet dataset (Russakovsky et al., 2015) . Here, we train ResNet of 50, 101 and 152 layers, to demonstrate that the proposed P-BN can help to train very deep networks on a large dataset, and we compare the performance results between SGD+BN and G-SGD+P-BN.

The method for tuning the hyper-parameters is similar to that in Section 5.1.2, and more details can refer to Appendix F.1. In Table 2 , we list top-1 and top-5 test error rates, and in Fig. 5 , plot the training curve and test accuracy during training ResNet of 50, 101, and 152 layers. These results well demonstrate the superiority of G-SGD+P-BN, and show that for the large dataset and very deep networks, P-BN can still outperform the conventional BN.

5.2. OBSERVATIONS ON STANDARD DEVIATIONS AND GRADIENTS:

After analyzing the advantage of P-BN over the conventional BN theoretically in Section 3.2 and 4, we now provide some experimental observations to support our analyses. Here, we will show the magnitude of standard deviations in BN, and the gradients of ReLU networks with BN and P-BN.

In this experiments, we use a stacked CNN of 50 layers with 128 channels in each layer, and use CIFAR-10 as training data, with the batch size of 128. Then, we add the conventional BN or P-BN at the end of each hidden layer in the stacked CNN, which is denoted as BN and P-BN, respectively. For the model with the conventional BN, it is randomly initialized. For the model with P-BN, the diagonal elements in parameter matrices (expect for the 1st layer) are initialized to 1 (cf. Section 3.1), and other elements are initialized with small random values (cf. Appendix D.3). Because the theoretical analyses are not related to the training process, we will just observe some quantities at initialization as follows.

First, we show some observations on standard deviation σ l j in Theorem 3.2. Here, we log σ l j for each hidden node during the forward process of the 1st mini-batch, and then calculate the average standard deviation among the nodes in each layer. In Fig. 6 , we plot the curves regarding standard deviations, and we can observe that: (1) For the conventional BN, the standard deviations are smaller than 1, which matches the claim in Theorem 3.2, and is one of the reasons of the gradient exploding.

(2) For P-BN, the standard deviation grows when the layer index becomes larger, and most of them are larger than 1. Thus, in Theorem 4.1, σ -1 s • W s are smaller than I , and ∇ z l L can be upper bounded, which is one of the reasons for that the gradient exploding can be weaken by P-BN.

Second, we show some observations on gradients, to investigate that whether the gradient exploding can be weaken by P-BN in practice. Here, we log the L2-norm of gradients of loss w.r.t. the outputs in each layer, and plot the log of the gradient norm in Fig. 6 . We can observe that, the gradients indeed explode in the ReLU network with both the conventional BN and P-BN, and the exploding speed is much slower with P-BN, which demonstrates that P-BN can help to weaken the gradient exploding.

6. CONCLUSION:

In this paper, we conduct a formal study on the design of BN in the path space. First, considering that BN is adopted during the forward process, we ensure the forward propagation in the path space, via proposing a re-parameterization for weights in ReLU networks. Here, the re-parameterization can equivalently replace each weight in the original network with a new value calculated by one or several path-values, and hence, the outputs of hidden nodes can be calculated by using path-values. Then, we design BN in the path space, by applying and analyzing the conventional BN to the reparameterized network. Here, we analyze the gradient propagation after BN is applied, and then notice that the gradient will explode in such networks. Thus, we propose a novel BN in the path space, i.e., P-BN, which excludes the term whose coefficient is fixed by re-parameterization from normalization. Finally, we conduct experiments to verify the proposed method, and results on both CIFAR and ImageNet dataset, demonstrate that P-BN can enhance the performance significantly.

A NOTATIONS:

Notation Object 

= ( s<l w s kk ) • w l jk • ( s>l w l jj ) z l j (x)

Output after BN and ReLU of the j-th hidden node in the l-th layer

z l j (x)

Output after P-BN and ReLU of the j-th hidden node in the l-th layer

W l

Parameter matrix in layer l after the re-parameterization. 

B SUPPLEMENTARY PROOFS:

Theorem B.1 We use W l to denote the parameter matrix in layer l after the re-parameterization. Suppose z l j (x i ) < O(1), ∀i, j, l, and the batch size m > 64. The norm of gradient w.r.t z l (x i ) can be upper bounded as

∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • γ s • (σ s ) -1 • W s ,

where

σ l = diag(σ l 1 , • • • , σ l d ), γ l = diag(γ l 1 , • • • , γ l d ) and D l (x i ) = diag(D l 1 (x i ), • • • , D l d (x i )). Here, diag(a1, • • • , a d ) represents a diagonal matrix whose diagonal elements are a1, • • • , a d , and D l j (x i ) = 1 if z l j (x i ) > 0, otherwise D l j (x i ) = 0.

Proof: We concat the gradient vector of different instances in a minibatch into one (d × m)dimensional vector as

EQUATION

and hence, according to chain rule, we have

EQUATION

Here, on the one hand, as for matrix ∂z s ∂o s , it donates

EQUATION

So we have

∂z s j (x i ) ∂o s j (x i ) = D s j (x i ) • γ s j σ s j (1 - 1 m - 1 m • (z s j (x i )) 2 ), if j 1 = j 2 = j and i 1 = i 2 = i; (9) ∂z s j (x i1 ) ∂o s j (x i2 ) = D s j (x i1 ) • γ s j σ s j (- 1 m - 1 m • z s j (x i1 ) • z s j (x i2

)), if j 1 = j 2 = j and i 1 = i 2 ; (10)

EQUATION

Obviously, most of the elements in the matrix equals to zero. Under the assumption that z s j (x i ) < O(1) for all i, j, l and m > 64, the above matrix approaches to the diagonal matrix

D s • γ s • (σ s ) -1 .

On the other hand, as for matrix ∂o s ∂z s-1 , it donates

EQUATION

Thus, using the fact that AB ≤ A B , we can get the result in the theorem. 1), ∀i, j, l, and m > 64. The norm of gradient w.r.t z l (x i ) can be upper bounded as

Theorem B.2 Suppose z l j (x i ) < O(

∇ z l (x i ) L ≤ O L s=l+1 D s (x i ) • (I + γ s,/ • (σ s,/ ) -1 • Ŵ s ) , where Ŵ s = W s -I, σ l,/ = diag(σ l,/ 1 , • • • , σ l,/ d ), γ l,/ = diag(γ l,/ 1 , • • • , γ l,/ d ) and D l (x i ) = diag(D l 1 (x i ), • • • , D l d (x i )) where D l j (x i ) = 1 if z l j (x i ) > 0, otherwise D l j (x i ) = 0.

Proof: According to the chain rule, we have

EQUATION

where ∂z s ∂o s,/ denotes

EQUATION

)

So we have

EQUATION

)

Obviously, most of the elements in the matrix equals to zero. Under the assumption that z l j (x i ) is bounded for all i, j, l, the above matrix tend to be the diagonal matrix

D l •γ l •(σ l ) -1 if m approaches ∞.

On the one hand, as for matrix ∂o s,/ ∂z s-1 , we have

EQUATION

On the other hand, as for matrix ∂o s,/ ∂z s-1 , we have

EQUATION

Thus, using the fact that AB ≤ A B , we can get the result in the theorem.

C DETAILS ON OPTIMIZATION ALGORITHMS IN THE PATH SPACE:

We provide the detailed update rules on G-SGD (Meng et al., 2019) in this section. Here, we use w j and w i to denote the diagonal weight and non-diagonal weight at layer 1, respectively, and use w j and w i to denote the diagonal weight and non-diagonal weight at layer l, l ≥ 2, respectively. The gradient of w w.r.t loss l, can be represented by δ w . Then we have

EQUATION

EQUATION

where event w j → w i , means w j and w i are contained in one of the paths, containing one nondiagonal weights at most. In this section, we first show the method for finding weights used for rescaling in Theorem 3.1, which is denoted as "red weights" here, for the network with unequal width of the input and output layers.

As shown in Fig. 7 , for weight matrix {w l ; l = 2, • • • , L -1}, the diagonal elements are selected as red weights. For weight matrix w 1 and w L with size (D, h1) and (hL-1, K), where D and K are number of input and output nodes, respectively, and hi is the number of nodes in the i-th hidden layer, then, elements {w 1 (i mod D, i); i = 1, • • • , h1} and {w L (i, i mod K); i = 1, • • • , hL-1} are selected as red weights, respectively.

Besides, all methods can be easily applied to CNN. As we all know, a filter in CNN transforms a input feature map to the output feature map. Thus, each feature map can be regarded as a hidden node in MLP, and then the filter connects two feature maps in CNN corresponds to the weight connects two hidden nodes in MLP. Besides, the method to identify red weights is similar with that for MLP unless the diagonal elements of the 4-triple of the weight matrices are selected. Suppose the 4-triple of weight matrices is [64, 64, 3, 3] where the first two elements denote the number of input channels and output channels, the last two elements denote the size of filter. The element [63, 63, 1, 1] are selected to be red weights. In this way, the re-parameterization method in Theorem 3.1 is easy to extend to CNN by rescaling the weights using the selected red weights.

[0, 0, 1, 1], [1, 1, 1, 1], [2, 2, 1, 1], • • • ,

D.2 OTHER CHOICES ON RE-PARAMETERIZATION:

In this section, we discuss other choices on re-parameterization and the influence. The reparameterization method using path-values showed in the proof of Theorem 3.1 is not unique. For example, we can obtain another kind of re-parameterization, via multiplying and dividing the incoming and outgoing weights of O 1 1 by v 1 11 in step 4 of Figure .3. This will not influence analyses much in this paper for the following two reasons. First, whatever re-parameterization method used, each reparameterized network can serve as a sufficient representation for ReLU networks in the path space, because the outputs of the network will keep unchanged for any input after re-parameterization. Second, our studies are based on Theorem 3.1, which can be generalized to other re-parameterization methods.

D.3 RELATIONSHIP BETWEEN RE-PARAMETERIZATION AND OPTIMIZATION PROCESS:

In this section, we show how the network should be optimized after it has been re-parameterized. Specifically, after re-parameterization, we can optimize the loss function according to gradients with respect to the path-values using the optimization algorithm such as SGD as follows. First, we denote the parameter matrix of the re-parameterized network at the l-th layer as W l , whose elements are constant, path-value, and ration of path-values. Next, we can use back propagation to obtain the gradient of W l . According to the chain rule, for the path-values that will not appear in the denominator, its gradient can be calculated as:

EQUATION

and for the path-values that will appear in the denominator, its gradient can be calculated as:

∇ v 1 kk L = ∇ w 1 kk L - d j=1,k=1 ∇ w l jk L v 1 kk (23)

Then, we can derive the update rule of SGD in path space mentioned in Section C. Here, please note that the update rules derived from re-parameterization is consistent with that provided in (Meng et al., 2019) , and re-parameterization provides a much simpler and straightforward way to explain the complex update rules.

Based on the process above, we now provide some analyses on the magnitude of parameter matrices. Specifically, for the 2nd term in Eq. 23, the numerator contains many terms, which tends to be larger than the gradient of other paths, and thus, the denominator, i.e., v 1 kk , (k = 1, • • • , d) is initialized to be large in practice to match the magnitude of the numerator. Therefore, for the parameter matrix W l , v 1 kk , (k = 1, • • • , d) appears in the denominator for most of its elements, and hence, these elements are initialized to be much smaller than 1, which supports analyses in Theorem 3.2.

E FULL ALGORITHM FOR OPTIMIZING RELU NETWORKS WITH P -BN:

In Algorithm 1, we show the full algorithm for optimizing ReLU networks with P-BN in the path space. Besides, please check more details of P athOptimizer in Section 2.2 and Appendix C.

F DETAILS ON EXPERIMENTS F.1 EXPERIMENTAL SETTING DETAILS:

In particular, following the settings in Meng et al. (2019) , we apply G-SGD or P-BN with residual blocks, because there is no positive scaling invariance across residual blocks. In addition, we introduce a coefficient λ(λ < 1) for z l-1 j (x i ) in Eq. 5 for experiments of ResNet, to prevent the output exploding during the forward process, and λ for all layers in ResNet is set to be 0.1.

Algorithm 1 Optimization Algorithm in the Path Space of ReLU networks with P-BN.

Require: A mini-batch of inputs: B = {x 1 , • • • , x B }, initialization w (0) , and P athOptimizer.

for t = 0, 1, • • • , T do Forward Process 1. Normalize the output in the 1st layer in the same way as the conventional BN. for l = 2, • • • , L do 2. Normalize the partial weighted summation by using the conventional BN according to Eq. 3.

3. Calculate the mean and standard deviation of the partial weighted summation according to Eq. 4. 4. Add the excluded term into the normalized partial weighted summation according to Eq. 5. end for 5. Calculate the value of loss function l(w (t) ) by using the final output. Backward Process 6. Calculate the gradient of w (t) : ∇ w l(w (t) ) ← Back Propagation (l(w (t) ); w (t) ) Update Rule of Path Optimizer 7. w (t+1) = P athOptimizer(w (t) , ∇ w l(w (t) )). end for Ensure: w (T ) .

Here we describe pre-processing steps, which is same as He et al. (2016) . For CIFAR-10 and CIFAR-100, we randomly crop the input image by size of 32 with padding size of 4, and normalize every pixel value to [0, 1]. Then the random horizontal flipping to the image is applied. For ImageNet, we randomly crop the input image by size of 224, and normalize every pixel value to [0, 1]. Then the random horizontal flipping to the image is also applied.

We use 1 NVIDIA Tesla P100 or P40 GPU to run the experiments on CIFAR, and use 4 NVIDIA Tesla P100 or P40 GPUs in one machine to run the experiments on ImageNet. All experiments are averaged over 5 runs with different random seeds, and PyTorch (Paszke et al.) is used for implementation.

Here we list hyper-parameters used for CIFAR-10 and CIFAR-100. As for the hyper-parameters of methods mentioned in their corresponding papers, i.e., SGD+BN (He et al., 2016) and (G-SGD)+BN(wnorm) (Meng et al., 2019) , we use the same settings as their original ones. Specifically, for SGD+BN, the initial learning rate is set to 0.1, and for G-SGD+BN (wnorm), the initial learning rate is set to 1.0. Besides, we tune hyper-parameters for (G-SGD)+BN and the proposed (G-SGD)+(P-BN). Specifically, the initial learning rate is searched from {0.1, 0.2, 0.5, 1.0}. Then, we use the method proposed in (Zheng et al., 2018) as the weight decay in the path space, and set the coefficient of the weight decay to 1 × 10 -4 , for these two methods. Besides, we use the SGD without momentum in our experiments, because the way to utilize momentum in the path space remains unclear now. Moreover, for all models and algorithms, the mini-batch size is set to be 128 and the training process is conducted for 64k iterations. The learning rates are multiplied by 0.1 after 32k and 48k iterations in all experiments, and the coefficient of weight decay for methods is set to be 0.0001.

Here we list hyper-parameters used for ImageNet. For all experiments here, the mini-batch size is set to be 256 and the hyper-parameter for weight decay is set to be 0.0001 He et al. (2016) . For SGD+BN, the initial learning rate is set to 0.1 He et al. (2016) . For G-SGD+P-BN, we only tune the initial learning rate from {0.1, 0.2, 0.5}. The training process is conducted for 90 epochs, and the learning rate is multiplied by 0.1 after 30 and 60 epochs.

As for the stacked CNN used for observations in Section 5.2, we use the following settings: kernel size 3, stride 1, padding 1, and no bias. For CNN, the network width is the the number of channels, which is set to 128 for all hidden layers.

F.2 SUPPLEMENTARY RESULTS:

In Fig. 8 , plot the training curve and test accuracy of 50-layer ResNet. Here, we add other baseline, i.e., SGD+P-BN, in which we use SGD to train the network with P-BN, and in the test error for SGD+P-BN. We can observe that G-SGD+P-BN outperforms SGD+P-BN, which demonstrate the benefits of optimizing in the path space.

F.3 EXPERIMENTS FOR PATH-SGD:

As mentioned in Section 2.2, Path-SGD (Neyshabur et al., 2015a) is another PSI-related optimization algorithm, but not the optimizer in the path space.

Method:

Validation Error First, we introduce the network structure and some basic settings. Here, because Path-SGD is computationally cost for CNN, we follow the experimental settings in (Neyshabur et al., 2015a) , and use an MLP of 5 layers, whose number hidden nodes in each layer is 4000, and the conventional BN or P-BN is added at the end of every hidden layer. Besides, we use CIFAR-10 as the training data, and the batch size is set to 64. Moreover, all experiments are run for 150 epochs, and results are obtained by averaging over 5 runs. The only hyper-parameter needs to be tuned is learning rate, which is searched from {1, 5} -{1,2,3} .

Then, we show the setting details of 4 compared methods as follow:

• SGD+BN: We use SGD to train the network with the conventional BN.

• Path-SGD: We use Path-SGD to train the network without BN.4 

• Path-SGD+BN: We use Path-SGD to train the network with the conventional BN.

• Path-SGD+P-BN: We use Path-SGD to train the network with P-BN.

In Table 5 , we list the validation error of these 4 methods. Results show that: (1) Path-SGD+P-BN outperforms other baselines, which well demonstrate the superiority of P-BN. (2) Path-SGD+P-BN and Path-SGD+BN perform better than SGD+BN, which show the benefits of optimizing in the path space.

:

For ease of presentation, we only discuss the MLP with equal width, and in this paper, the analyses and algorithms can be applied to the MLP and CNN with unequal width (cf. Appendix D.1).

In the following analysis, referring toSanturkar et al. (2018), we assume that both of the scale and shift term are constant.

Please note that the conditions are tight in the sense that if we change L in the ranges to be L 1-, the upper bound of both the gradient norm for BN and P-BN will diverge.

We do not apply this setting for ResNet and PlainNet, because the performance will be badly hurt, and training will fail when BN is removed.

Figures and Tables:

Figure 1: An example of an MLP network. a 1 and a 2 are 1 when corresponding ReLU is activated, otherwise they are 0.

Figure 3: An example to demonstrate the re-parameterization process for an MLP.

Figure 4: An example of the 2nd step in P-BN.

Figure 5: Curves for training ImageNet.

Figure 6: Experimental observations.

input, i.e., x = (x1, • • • , xd) Output of the j-th hidden node in the l-th layer o l (x) Hidden vector in the l-h layer w l jk Weight on the edge which connects the k-th node in the (l -1)-th layer and the j-th node in the l-th layer v 1 kk Path-value of the path which only contains the k-th diagonal elements in weight matrices, i.e., v 1 kk = L l=1 w l kk v l jk , k = j Path-value of the path which contains one non-diagonal weight, i.e., v l jk

Mean of the minibatch outputs of the j-th hidden node in the l-th layer σ l j Standard deviation of the minibatch outputs of the j-th hidden node in the l-th layer

Figure 7: An example of feedforward ReLU network.

Figure 8: Training curves and test accuracy for training ResNet of 50 layers on CIFAR.

-BN 34.45% (±0.29) Table 5: Validation error for Path-SGD experiments.

𝑓 1 = 𝑎 1 𝑣 1 + 𝑎 2 𝑣 2 𝒙 𝟏 + 𝑎 1 𝑣 3 + 𝑎 2 𝑣 4 𝒙 𝟐 𝑓 1 = 𝑅𝑒𝐿𝑈 𝒙 𝟏 𝑤 1 + 𝒙 𝟐 𝑤 3 𝑤 5 + 𝑅𝑒𝐿𝑈 𝒙 𝟏 𝑤 2 + 𝒙 𝟐 𝑤 4 𝑤 7 = 𝑎 1 𝑤 1 𝑤 5 + 𝑎 2 𝑤 2 𝑤 7 𝒙 𝟏 + 𝑎 1 𝑤 3 𝑤 5 + 𝑎 2 𝑤 4 𝑤 7 𝒙 𝟐

Table content: <table><tr><td>Input</td><td/><td colspan="2">Hidden</td><td colspan="2">Output</td></tr><tr><td>𝒙 𝟏</td><td>𝑤 2</td><td>𝑤 1</td><td>𝑤 6</td><td>𝑤 5</td><td>𝒇 1</td></tr><tr><td>𝒙 𝟐</td><td>𝑤 3</td><td>𝑤 4</td><td>𝑤 7</td><td>𝑤 8</td><td>𝒇 2</td></tr><tr><td colspan="3">Calculate with weights:</td><td colspan="2">The Outputs of Hidden Nodes</td><td/></tr><tr><td>1</td><td/><td>𝑣 2</td><td/><td>𝑣 3</td><td>𝑣 4</td></tr><tr><td colspan="3">Calculate with path-values:</td><td/><td/><td/></tr></table>

t. path-values. On the other hand, in terms of generalization, for example, Neyshabur et al. (2015b) and Zheng et al. (2018) gave the relationship between generalization error and path-norm or basis path-norm, respectively. Besides, E et al. (2018) leveraged the path representation of ReLU networks, to analyze the generalization error of them.

Table content: <table/>

, path p starts from O 0 2 , passes O 1 2 and O 2 1 , and ends up with O 3 3 . And then, the path-value of path p can be defined as a multiplication of the weights along p, i.e., vp = w 2 vp • ap • x k . Here, P is a set which contains the paths staring from O 0 k and ending up with O L j , and ap, the ReLU activation status of path p, can be calculated as ap =

Table content: <table><tr><td>22 w 2 12 w 3 31 . jp ∈ p represents p∈P k,j O j (x) can be calculated by using path-values as o L Therefore, output o L j (x) = d k=1 lp jp ∈p I(o lp jp (x) &gt; 0), where O lp</td></tr><tr><td>that O jp is passed by p . lp</td></tr></table>

Test error rate on CIFAR-10 and CIFAR-100.

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td colspan="2">PlainNet</td><td/><td>ResNet</td></tr><tr><td/><td/><td/><td/><td/><td/><td>18</td><td/><td/><td>34</td><td>18</td><td>34</td><td>50</td></tr><tr><td/><td/><td/><td colspan="2">SGD+BN</td><td/><td colspan="3">6.93% (±0.12)</td><td>7.76% (±0.22)</td><td>6.76% (±0.10)</td><td>6.40% (±0.09)</td><td>6.55% (±0.24)</td></tr><tr><td/><td/><td/><td colspan="2">G-SGD+BN</td><td/><td colspan="3">6.66% (±0.12)</td><td>6.74% (±0.13)</td><td>6.84% (±0.30)</td><td>6.54% (±0.06)</td><td>6.62% (±0.19)</td></tr><tr><td colspan="2">CIFAR-10</td><td/><td colspan="6">G-SGD+BN(wnorm) 6.26% (±0.17)</td><td>6.67% (±0.26)</td><td>6.31% (±0.14)</td><td>6.33% (±0.15)</td><td>6.31% (±0.14)</td></tr><tr><td/><td/><td/><td colspan="3">G-SGD+P-BN ours</td><td colspan="3">5.99% (±0.13)</td><td>6.04% (±0.16)</td><td>6.04% (±0.14)</td><td>5.66% (±0.10)</td><td>5.99% (±0.12)</td></tr><tr><td/><td/><td/><td colspan="2">SGD+BN</td><td/><td colspan="4">28.10% (±0.19) 33.37% (±0.41) 26.97% (±0.12) 26.42% (±0.23) 25.62% (±0.18)</td></tr><tr><td/><td/><td/><td colspan="2">G-SGD+BN</td><td/><td colspan="4">27.08% (±0.35) 28.19% (±0.35) 27.13% (±0.39) 26.61% (±0.10) 25.99% (±0.40)</td></tr><tr><td colspan="3">CIFAR-100</td><td colspan="7">G-SGD+BN(wnorm) 26.76% (±0.05) 27.61% (±0.24) 26.60% (±0.17) 26.72% (±0.27) 25.65% (±0.22)</td></tr><tr><td/><td/><td/><td colspan="3">G-SGD+P-BN ours</td><td colspan="4">26.70% (±0.10) 27.04% (±0.12) 26.39% (±0.10) 26.24% (±0.29) 25.29% (±0.19)</td></tr><tr><td>1.00 1.25 1.50 1.75 2.00 Training Loss</td><td/><td/><td colspan="2">50 SGD+BN 50 -SGD+ -BN ours 101 SGD+BN 101 -SGD+ -BN ours 152 SGD+BN 152 -SGD+ -BN ours</td><td colspan="2">0.76 0.68 0.70 0.72 0.74 Test accuracy</td><td/><td/></tr><tr><td>0.75</td><td>20</td><td>40</td><td>60 Epoch</td><td colspan="2">80 100</td><td>20 0.66</td><td>40</td><td>60 Epoch</td><td>80</td></tr></table>

Test error rate on ImageNet.

Table content: <table><tr><td>We conduct experiments on CIFAR-10 dataset and</td></tr><tr><td>CIFAR-100 dataset (Krizhevsky et al., 2009). Here,</td></tr><tr><td>we train ResNet of 18, 34, and 50 layers, and</td></tr><tr><td>PlainNet of 18 and 34 layers, respectively. As</td></tr><tr><td>for the hyper-parameters of methods mentioned</td></tr><tr><td>in their corresponding papers, i.e., SGD+BN (He</td></tr><tr><td>et al., 2016) and G-SGD+BN(wnorm) (Meng et al.,</td></tr><tr><td>2019), we use the same settings as their original</td></tr><tr><td>ones. Besides, we tune the hyper-parameters for</td></tr><tr><td>other methods, i.e., G-SGD+BN, and the proposed</td></tr></table>

Main notations in this paper.

Table content: <table/>

Table 4, we show CIFAR-100 SGD+P-BN 28.57% (±0.46) 32.98% (±0.81) 26.71% (±0.22) 26.80% (±0.51) 25.68% (±0.28)

Table content: <table><tr><td>Dataset</td><td>Method</td><td>PlainNet</td><td/><td>ResNet</td><td/></tr><tr><td/><td/><td>18</td><td>34</td><td>18</td><td>34</td><td>50</td></tr><tr><td>CIFAR-10</td><td colspan="2">SGD+P-BN 7.00% (±0.08)</td><td>7.68% (±0.21)</td><td>6.78% (±0.15)</td><td>6.57% (±0.06)</td><td>6.47% (±0.23)</td></tr></table>

Test error rate on CIFAR-10 and CIFAR-100 for SGD+P-BN.

Table content: <table/>

Peer review (iclr2020):
This paper analyzes a reparametrization of the network that migrates from the weight space to the path space. It enables an easier way to understand the batch normalization (BN). Then the authors propose a variant of BN on the path space and empirically show better performance with the new proposal. To study BN in the reparameterized space is well-intuited and a natural idea. Theorem 3.1 itself is interesting and has some value in understanding BN. However, the main contribution of the paper, i.e., the proposal of the P-BN, is not motivated enough. It is merely mentioned in the beginning of section 4 and it*s not clear why this modification is better compared to conventional BN. This is not verified by theory either. By comparing theorem 3.2 and theorem 4.1, it seems P-BN even gives even worse upper bound of gradient norm. Plus we don*t actually care that much about the issue of gradient exploding since one could always do gradient clipping. The notorious gradient vanishing problem on the other hand, is not address in the theorems. The formulation of the P-BN seems to be closely related to ResNet, since it sets aside the identity mapping and only normalizes on the other part. It would be better to have some discussions. Also, the reparameterization and P-BN seems only to apply to fully connected layer from Eqn. (3-5) where they are proposed, but the experiments applies to ResNet. It would be better to describe the method in a broader sense. How would you do this P-BN in more complicated networks? Finally, it*s very unclear to me the value of Theorem 3.1 and the proof that takes almost one page in the main context. The assumption of diagonal elements of matrix w to be all positive is very restrictive and simply removes the effect of ReLU activations. Therefore I think the paper has some room for improvement and is not very suitable for publication right now.