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 paper studies the training of multi-branch networks, i.e. networks formed by linearly combining multiple disjoint branches of the same architecture. The core contribution in this paper is the “STAM” aggregation rule which is to set the combination coefficient to $1/\sqrt{C}$ for a network with $C$ branches. This aggregation rule is justified by (1) theoretical analysis on the function values and gradient norms at initialization, and (2) experiments on residual networks and transformers showing that this rule performs better than the baseline rules (such as sum or average).