On the Stability of Multi-branch Network

1 INTRODUCTION . Multi-branch architecture is a building block in state-of-the-art neural network models for many tasks , e.g. , the ResNeXt ( Xie et al. , 2017 ) for computer vision and the Transformer ( Vaswani et al. , 2017 ) for machine translation . It has been pointed out that the benefit of multi-branch architecture is the parameter efficiency ( Xie et al. , 2017 ) . The number of parameters grows linearly with the number of branches but quadratically with the width ( the number of neurons in one layer ) . It has also been argued that the multiple branches can bring diversity if branches are composed of sub-networks with different filters and depths ( Huang et al. , 2017 ; Li et al. , 2019 ) . To train multi-branch networks successfully , it usually requires careful designs and some hyperparameter tuning such as adding normalization layers , scaling down the initialization , and adjusting learning rate . As a verifying example , for a trainable single-branch network , simply adding branches multiple times and aggregating their outputs together often do not work as expected , e.g. , the training instability of sum aggregation in Figure 4 . This demonstrates the difficulty of training multi-branch network and also motivates us to do this study . In this paper , we try to understand the behavior of training multi-branch network . Specifically , we study the forward and backward process of multi-branch networks , which is believed to govern whether the network is easy to optimize by gradient-based methods . We find out that the aggregation scheme , i.e. , “ the way of combining the multi-branch outputs ” plays a central role in determining the behavior of training multi-branch network . We show that the sum aggregation would become unstable as the number of branches grows , which explains the bad performance of simply adding branches . Moreover , we characterize the condition on the aggregation scheme under which the forward and backward stability is guaranteed . Inspired by the theoretical analysis , we propose a “ STAM ” aggregation , that can STAbilize Multibranch network , which scales the sum of the branch outputs by a branch-aware factor α ( see the later part of Section 3.1 for details ) . We argue the benefit of STAM aggregation over the sum and average aggregations by analyzing the Hessian of the multi-branch network . We show that STAM permits the same gradient-based optimizer works for different settings , which could reduce lots of tuning burden for training network with flexible number of branches . We further examine the usual design wisdom through the stability lens . As a result , we find that scaling down initialization may control the forward or backward stability but not necessarily the both , which is verified in experiment . We also unveil a new role of normalization layer that it can stabilize the forward and backward process of multi-branch network besides the many wanted and unwanted properties that have been argued before ( Ioffe & Szegedy , 2015 ; Yang et al. , 2018 ; Santurkar et al. , 2018 ; Xiong et al. , 2020 ) . Apart from the usual feedforward multi-branch architecture , we analyze the multi-head attention layer , a multi-branch architecture widely used in natural language processing . We give an upper bound on the multi-head representations when the softmax operator is replaced with max operation . The upper bound unveils the relation between the head dimension and the length of the sequence , which interprets empirical observation well . This relation can not be discovered if assuming softmax outputs equal probability as in Xiong et al . ( 2020 ) . Overall , our contribution can be summarized as follows . • We analyze the forward/backward stability of multi-branch network , under which we can clearly interpret the benefit and potential problem of the practical wisdom , i.e. , scaling down initialization and adding normalization layer . • We propose a theoretically inspired STAM aggregation design for multi-branch network , which can handle arbitrary number of branches with a same optimizer . • We also analyze the forward/backward process of multi-head attention layer and identify its special property that has not been characterized before . 1.1 RELATED WORK . Multi-branch architecture , also known as split-transform-merge architecture , has been widely used in computer vision task , namely Inceptions ( Szegedy et al. , 2017 ; Chollet , 2017 ) , ResNeXt ( Xie et al. , 2017 ) , and many others ( Abdi & Nahavandi , 2016 ; Ahmed & Torresani , 2017 ) . In fact , the models in natural language tasks have also leveraged the multi-branch architecture including the BiLSTM ( Wu et al. , 2016 ; Zhou et al. , 2016 ) and the multi-head attention layer in Transformer ( Vaswani et al. , 2017 ; Anonymous , 2020 ) . Apart from the sum or average aggregation , recent works ( Li et al. , 2019 ; Zhang et al. , 2020 ) integrate the attention mechanism with the aggregation scheme , i.e. , the attentive aggregation , although only a small number ( 2 ∼ 3 ) of parallel branches are considered . Theoretically , Zhang et al . ( 2018 ) interpret the benefit of multi-branch architecture from reducing the duality gap or the degree of non-convexity . The theory of training general deep neural network has been widely studied , via the stability analysis ( Arpit et al. , 2019 ; Zhang et al. , 2019a ; c ; Yang & Schoenholz , 2017 ; Zhang et al. , 2019b ; Yang , 2019 ; Lee et al. , 2019 ) , neural tangent kernel ( Jacot et al. , 2018 ; Allen-Zhu et al. , 2018 ; Du et al. , 2018 ; Chizat & Bach , 2018 ; Zou et al. , 2018 ; Zou & Gu , 2019 ; Arora et al. , 2019 ; Oymak & Soltanolkotabi , 2019 ; Chen et al. , 2019 ; Ji & Telgarsky , 2019 ) . In contrast , we focus on the multi-branch network , which has not been studied theoretically before . 2 MODEL DESCRIPTION AND NOTATIONS . In practice , the multi-branch architecture is often used as a building block in a whole network . In this paper , we describe a multi-branch architecture/network N ( · ) as follows ( see Figure 1 ) . • N ( · ) has C branches { Bk } Ck=1 , input hin ∈ Rp and output hout ∈ Rd ; • The aggregation is parameterized with a vector α = ( α1 , . . . , αC ) T : hout : = N ( hin ) : = C∑ k=1 αk · Bk ( hin ) . ( 1 ) Each branch Bk often consists of multiple layers with various structures and flexible configuration : depth , width , kernel size for convolution layer , activation functions and normalization layer . The aggregation weight is given by α . Such description covers popular multi-branch architectures in state-of-the-art models , e.g. , Inception , ResNeXt and Transformer , if specifying Bk and α properly . Throughout the paper , we use ‖ · ‖ to denote the l2 norm of a vector . We further use ‖ · ‖ and ‖ · ‖F to denote the spectral norm and the Frobenius norm of a matrix , respectively . We denote a set of naturals with [ n ] : = { 1 , . . . , n } and [ c : d ] = { c , c+ 1 , . . . , d } . We use bold small case letters , e.g. , v , to denote vectors and bold capital case letters , e.g. , M , to denote matrices . Moreover , Id×d is the d× d identity matrix , 1d is d dimensional vector with all 1 ’ s , and Vec ( M ) stacks the row vectors of matrixM as a long vector . 3 STABILITY OF MULTI-BRANCH FEEDFORWARD NETWORK . To theoretically study the multi-branch network , we introduce a simplified multi-branch feedforward network . Specifically , we assume that each branch is a b-layer fully connected network with ReLU activation , and each layer has the same width m. Branches share the same structure and they differ from each other by the random initialization . One branch is given by Bk ( hin ) = W kbφ ( W k b−1 · · ·φ ( W k 1hin ) ) , ( 2 ) where W k1 ∈ Rm×p , W k b ∈ Rd×m and φ ( · ) is the ReLU activation φ ( · ) : = max { 0 , · } . We further introduce −→ W k : = ( W k1 , . . . , W k b ) that collects parameters of Bk , and −→ W : = ( −→ W 1 , . . . , −→ WC ) . Next we analyze the forward and backward propagation of the multi-branch network given by ( Equation 1 ) and ( Equation 2 ) , and characterize a condition that guarantees forward/backward stability . Based on the theoretical analysis , we propose the STAM aggregation that can stabilize the forward/backward process of multi-branch networks . We further argue why the practical wisdom of scaling down initialization and adding normalization layer works and when it could fail . 3.1 FORWARD AND BACKWARD PROCESS . We assume that the feedforward multi-branch network given by ( Equation 1 ) and ( Equation 2 ) adopts the Kaiming ’ s initialization ( He et al. , 2016 ) : entries of W a for a ∈ [ 1 : b− 1 ] are independently sampled from N ( 0 , 2m ) , and entries of W b are independently sampled from N ( 0 , 1 d ) . Then the forward norm is well concentrated around its mean value as follows . Theorem 1 . Suppose the multi-branch network N ( · ) is given by Equation 1 and Equation 2 with Kaiming ’ s initialization . For an input hin , the following holds with probability at least 1−O ( bC ) · e−Ω ( m 2/b ) over the initialization randomness of −→ W ‖N ( hin ) ‖ ∈ ( 1± ) √∑ k∈ [ C ] α2k‖hin‖ , ( 3 ) where C is the number of branches , and m , b are the width and depth of each branch , respectively . Proof . The proof is based on the Gaussianness of Bk ( hin ) and a concentration property , whose full version is deferred to Appendix A.2 . Theorem 1 is presented for one input sample . If we want such a result holds for all the training samples , the probability loses an n factor by the union bound and becomes 1−O ( nbC ) · e−Ω ( m 2/b ) . Remark 1 . With the same assumption as Theorem 1 , we have E‖N ( hin ) ‖ = √∑ k∈ [ C ] α2k‖hin‖ , ( 4 ) where C is the number of branches and the expectation is over the random initialization . These results are based on the bound of the forward propagation of one feed-forward branch given by ( Equation 2 ) , which is studied in Allen-Zhu et al . ( 2018 ) and restated in Appendix A . Furthermore , if the weight matrices follows the Gaussian distribution , then Bk ( hin ) is also roughly Gaussian and jointly Gaussian for different input samples , as the width m goes to infinity . Then the aggregation ( Equation 1 ) can be viewed as a sum of weighted Gaussian vectors . Hence at the initialization , we can characterize the aggregation of multiple branches as above . We next analyze the backward propagation of multi-branch feedforward network . We abuse “ gradient ” to refer to the values computed through back-propagation even for non-smooth function . We assume the loss function ` ( · , · ) is quadratic , i.e. , ` ( hout , y∗ ) = 12‖hout−y ∗‖22 . Hence , the objective function is L ( −→ W ) : = 1n ∑n i=1 Li ( −→ W ) , where Li ( −→ W ) : = ` ( N ( xi ) , y∗i ) . We next show the backward process is bounded for each individual sample for the multi-branch network N ( · ) . Theorem 2 . With probability at least 1− ( nb ) · exp ( −Ω ( m ) ) over the randomness of −→ W , it satisfies for every a ∈ [ b ] and k ∈ [ C ] , every i ∈ [ n ] , ∥∥∥∇WkaLi ( −→W ) ∥∥∥2F ≤ O ( Li ( −→ W ) α2k ×m d ) , ∥∥∥∇−→WLi ( −→W ) ∥∥∥2F ≤ O ( Li ( −→W ) × mbd ∑ k∈ [ C ] α2k ) . ( 5 ) Proof . We can compute the gradient with respect to intermediate layer outputs via the backward propagation procedure . The gradient upper bound is guaranteed if the intermediate layer outputs and their gradients are bounded across layers . The full proof is relegated to Appendix A.3 . If further assuming the gradient independence condition : weights in backward process can be assumed to be independent from weights in the forward pass ( Yang , 2019 ) , we can estimate the expectation of the gradient norm as follows . Remark 2 . Assuming the gradient independence , we have for every a ∈ [ b ] and k ∈ [ C ] , every i ∈ [ n ] , E ∥∥∥∇WkaLi ( −→W ) ∥∥∥2F = Li ( −→W ) α2k ×md , E∥∥∥∇−→WLi ( −→W ) ∥∥∥2F = Li ( −→W ) × mbd ∑ k∈ [ C ] α2k . ( 6 ) With Theorem 1 and 2 and two remarks , we can discuss the property of the forward and backward process of the multi-branch network . We can see that both the output of multi-branch network and the gradient are under control if ∑ α2k ≤ O ( 1 ) . Specifically , for the sum aggregation , we have∑ α2k = C which grows unbounded with the number of branches C. For the average aggregation , we have ∑ α2k = 1/C which diminishes with the number of branches C. There exists a better choice of αk : αk = 1/ √ C for k ∈ [ C ] that keeps ∑ α2k = 1 constant as the number of branches varies . We call it “ STAM ” aggregation , abbreviating STAble Multibranch aggregation . We plot the output norm of the first residual block in multi-branch ResNets at initialization in Figure 2 . Multi-branch ResNets are generated by varying the number of branches in the residual block with batch normalization removed . We can see that the forward norm of STAM aggregation roughly remains the same , while that of the sum aggregation explodes and that of the average aggregation diminishes , as the number of branches grows . We also analyze the Hessian of different aggregation schemes . We find that the spectral norm of the Hessian , which determines the smoothness of the objective , proportionally scales with the square root of the number of branches for the sum aggregation , while reciprocally scales with the square root of the number of branches for the average aggregation . In contrast , the Hessian for the STAM aggregation keeps unchanged as the number of branches varies . Hence with STAM , the same learning rate works for network with different number of branches . We present the details in Appendix B .

Summary: This study focuses on the stability of multi-branch networks. It analyzes the forward and backward stability of multi-branch network, and builds the relations with some widely-adopted initialization and normalization schemes. A simple new aggregation method is proposed that enjoys better stability than the sum and average aggregations. The method is extended to the multi-head attention layer in Transformer. Experiments on image classification and machine translation tasks using ResNeXt and Transformer are conduced to show the efficacy.