On the Stability of Multi-branch Network
Authors: Huishuai Zhang
Da Yu
Wei Chen
Tie-Yan Liu
Venue:  (2021)
Keywords: stability
multi-branch network
backward propagation

Abstract: Multi-branch architectures are widely used  in state-of-the-art neural networks. Their empirical success relies on some design wisdom, like adding normalization layers or/and scaling down the initialization.  In this paper, we investigate the multi-branch architecture from the stability perspective. Specifically, we establish the forward/backward stability of multi-branch network, which leads to several new findings. Our analysis shows that only scaling down the initialization may not be enough for training multi-branch network successfully because of the uncontrollable backward process. We also unveil a new role of the normalization layer in terms of stabilizing the multi-branch architectures. More importantly, we propose a new design ``STAM aggregation" that can guarantee to STAbilize the forward/backward process of Multi-branch networks irrespective of the number of branches. We demonstrate that with STAM aggregation, the same training strategy is applicable to models with different numbers of branches, which can reduce the hyper-parameter tuning burden. Our experiments verify our theoretical findings and also demonstrate that the STAM aggregation can improve the performance of multi-branch networks considerably.

Under review as a conference paper at ICLR 2021
ON THE STABILITY OF MULTI-BRANCH NETWORK
Anonymous authors Paper under double-blind review
ABSTRACT
Multi-branch architectures are widely used in state-of-the-art neural networks. Their empirical success relies some design wisdom, like adding normalization layers or/and scaling down the initialization. In this paper, we investigate the multi-branch architecture from the stability perspective. Specifically, we establish the forward/back ward stability of multi-branch network, which leads to several new findings. Our analysis shows that only scaling down the initialization may not be enough for training multi-branch network successfully because of the un- controllable backward process. We also unveil a new role of the normalization layer in terms of stabilizing the multi-branch architectures. More importantly, we propose a new design "STAM aggregation" that can guarantee STAbilize the forward/backward process of Multi-branch networks irrespective of the number of branches. We demonstrate that with STAM aggregation, the same training strategy is applicable models with different numbers of branches, which reduce the hyper-paramen parameter tuning burden. Our experiments verify our theoretical findings and also demonstrate that the STAM aggregation can improve the performance of multi-branch networks considerably.
INTRODUCTION
Multi-branche architecture a building block in state-of-the-an art neural network models for many tasks, e.g., the ResNeXt (Xie et al.) (017) for computer vision and the Transformer ( (Vaswani et al. (017)) for machine translation. It has been pointed out that the benefit of multi-branch architecture 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). 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 hyper- parameter tuning such as adding normalization layers, scaling down the initialization, and adjusting learning rate. As verifying example, for 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 This demonstrates the difficulty of training multi-branch network and also motivates us to do this study.
In this paper, we try 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 that the aggregation scheme, i.e., "the way of combining the multi-branch outputs" plays central role in determining the behavior of training baalti-bannchhhc 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 Multi- branch 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 branch network. We show that STAM permits the same gradient-bart based optimizer works for different settings, which could reduce lots of tuning burden for training network with flexible number of branches.

---

Under review as a conference paper at ICLR 2021
We further examine the usual design wisdom through the stability lens. As result, 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 nulti-branch network besides the many wanted and unwanted properties that have been argued before (loffe & Szegedy 2013 Yang al. 2018 Sannturar et al.) 2018 Xiong et al. 2020).
Apart from the usual feedforward multi-branch architecture, we analyze the multi-head attention layer, multi-branch architecture widely used in natural language processing. We give an upper bound on the multi-head representations when the softmax operator 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 cannot be discovered 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-brance 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 theoretically inspired STAM aggregation design for maati- 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- architecture, has been widely used in computer vision task, namely Inceptions (Szegedy al. 2017; Chollei. 2017). ResNeXt (Xie et al.
2017). and many others (Abdi & Nahavandi] 2016 Ahmed & Torresani (011]) 7). In fact, the models natural language tasks have also leveraged the bralic architecture including the BiLSTM (Wu et al. 2016) Zhou et al. 2016) and the multi-head attention layer in Transformer (Vaswani et al. 3017]|nooyya
2020). Apart from the sum or average aggregation, recent works (Li al.j 2019 Zhang al., (020) integrate the attention mechanism with the aggregation scheme, i.c., 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 training general deep neural network has been widely studied, via the stability analysis (Arpit et al. 2019 Zhang et al.) 2019a Yang &
Schoenholz 2017 Zhang et al. 2019b Yang 20199 Lee et al. 2211)) ), neural tangent kernel ( Jacot
al. 2018 Allen Zhu et al. 2018:
2018/Chizz & Bach | Zou al. 2018 Zou & Gu
2019 Arora et al. 2019) 9) Oymak& Soltanolkotabi 2019) Chen et al. 2019 Ji & Jelgarsky,211 (2019). In
contrast, we focus on the multi-branch network, which has not been studied theoretically before.
MODEL DESCRIPTION AND NOTATIONS
In practice, the multi-branch architecture is often used as a building block whole network. In this paper, we describe a multi-branch architecture/networkk (-)) as follows (see Figure
(- has C' branches {B input hin E R and output hout E IR
The aggregation is parameterized with a vector = "cT
hout N(hin)
ak B²t (hin).
k=1
Each branch 3 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 a. Such description covers popular multi-branch architectures in state-of-the-art models, e.g., Inception, ResNeXt and Transformer, if specifying B and a properly.
2

---

Under review as a conference paper at ICLR 2021
Forward norm
10a
x
average
sum
C branches
norm rabo
100
B² (x)
BC (x)
Aggregation
10
4
N (x)
number of branches
C-branch block
Figure 2: Output norm after first aggregation in multi-branch ResNets at initialization. The norm ratio with respect the case of C 4.
Figure 1: A multi-branch network.
Throughout the paper, we use || - || to denote the le norm of a vector. We further use || and || Il denote the spectral norm and the Frobenius norm of matrix, respectively. We denote set of naturals with [n] := {1..... ....,
!1. n and [c: d] fc, + 1. 1....,d? dd. We use bold small case letters, e.g.. v, to denote vectors and bold capital case letters, e-g.. M, denote matrices. Moreover, Iaxa the
d x d identity matrix. 1a is d dimensional vector with all I's, and Vec(M) stacks the row vectors of matrix M as a long vector.
STABILITY OF MULTI-BRANCH FEEDFORWARD NETWORK
To theoretically study the multi-branch network, introduce 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 width m. Branches share the same structure and they differ from each other by the random initialization. One branch is given by
B (h... WW 6-1 o W hii)))
where W E Rmxp, W E Rdxx and <(-) is the ReLU activation o( max J. We further introduce W = (W..... w that collects parameters B* and W = (W WO).
Next we analyze the forward and backward propagation of the multi-branch network given by (Equation MT and (Equation 2). and characterize a condition that guarantees forward/backwwrrd stability. Based on the theoretical analysis, propose the STAM aggregation that can stabilize the forward/backwwar process brant- 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 bulth-branch network given by (Equation I and (Equation 2 adopts the Kaiming's initialization (He al. 2016): entries of W. for E [1 : 1] are independently sampled from (0, 2), and entries of W are independently sampled from No 1). Then the forward norm is well concentrated around its mean value as follows.
Theorem 1. Suppose the branch network (-) is given by Equation 7 and Equation 2 with Kaiming's 's initialization. For an input hin the following holds with probability least 1 OTC))- (mmme /b) over the initialization randomness of W
[[ )ll (1 (1
Σ kEIC]
(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 B (ha and a concentration property, whose full version is deferred Appendix A.2

---

Under review as a conference paper at ICLR 2021
Theorem I is presented for one input sample. If we want such result holds for all the training samples, the probability loses an n factor by the union bound and becomes 1 O(nCC)) -Q(mm /b).
Remark 1. With the same assumption as Theorem 7 we have
Σ RE(C)
of
(4)
where C' the number of branches and the expectation 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 studied in AlleennhZ al. (2018) and restated in Appendix [A] Furthermore, the weight matrices follows the Gaussian distribution, then B* (han) is also roughly Gaussian and jointly Gaussian for different input samples, as the width m goes 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 bulti- feedforward network. We abuse "gradient" to refer to the values computed through back- proppaaopagttion even for non-smooth function. We assume the loss function ((.,)) -) is quadratic, i.c., [houu,, 3²) = ouu y 112. Hence, the objective function is C(W) := 1 Σn= (W),, where (.WW) == ((V(z;), ).). We next show the backward process is bounded for each individual sample for the multi-branch network (.).
Theorem 2. With probability at least 1- 1-(bb)- exbd))exp(-S(mm)) over the randomness of W, it satisfies for every E [b] and & E [C], every? i E [rz].
Vwe ( W Ip so so
(2 (W) at x xm) d
|/ s so (W) X mb
(5)
kE(C
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 E [b] and k E [C], i E [n)].
every
E (wy L.
x m d
C. (
mbb Σ
o
(6)
CE
With Theorem aand 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 aaz < O(1). Specifically, for the sum aggregation, we have Σ aa C which grows unbounded with the number of branches C. For the average aggregation, we have Σ af = 1/C which diminishes with the number of branches C.
There exists a better choice of as: as = 1/vC for k E [C] that keeps Σの = constant as the number of branches varies. We call it "STAM" aggregation, abbreviating STAble Multi- branch aggregation. We plot the output norm of the first residual block in bratti-raanch ResNets at initialization in F 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
4

---

Under review as a conference paper at ICLR 2021
3.2 UNDERSTANDING THE PRACTICAL WISDOM
In practice people have design wisdom to stabilize the multi-branch network. We next analyze them through the stability lens.
One can scale down the initialization for multi-brance network. For example, one can initialize each W with a scaling-down factor C' 25 for all E [b] and k E [C] that the forward norm of (hi,)) has a bound around | ll. irrelevant with C. However, with this initialization, the norm of the output update induced by one gradient descent step, scales with BC Alternatiell, we suggest initializing each W with a scaling-down factor (bC) কहা for all / [6] and & E [C] to stabilize the backward process such that the expected update on the output is irrelevant with b and C. This indicates that constant learning rate can train multi-branch network with any b and C, although the forward process
has diminishing output norm scaling with b- à (bC) 215-1), More detailed discussion is presented in Appendix A.4. We empirically verify this in Section
Another widely used technique is adding normalization layers. It is easy layer can stabilize the forward process as it normalizes the output to be with mean 0 and variance 1. Moreover, the normalization layer can also stabilize the backward process the error signal is divided by the standard deviation which is proportional show a new role of the normalization layer that automatically stabilizes the forward and backward process of the multi-branch network or in general stabilizes the architecture, besides the previous understanding of increasing the smoothness (Santurkar et al. (2018) or handling the covariance shift (Iofe & Szegedy 2015) or implicit structure bias (De & Smith 2020).
that normalization
VC when propagating backward. Therefore,
We next examine concrete structures. One popular aggregation scheme for multi-branch outputs is concatenation followed by a linear transformation. This is equivalent
sum aggregation. The default Xavier initialization (Glorot & Bengio, 2010) would scale down the initialization of the linear transformation by roughly 1/VC' to stabilize the forward process but the training could still be unstable because of the unbounded backward process as discussed above and verified in Section
A well-known multi-branch structure is ResNeXt for image classification task. It uses batch normal- ization on the output of the residual branch to stabilize the forward and backward processes. However, if we use the pre-act form, i.e., remove the batch normalization on the output of the residual branch and add a batch normalization on the input of the residual branch, the training becomes unstable as the number of branches increases (see Section 5.2 With the pre-act residual block, one needs normalization layer on the final output of the network, which has been studied in (Xiong et al.) 2020). It could still be unstable for the deep half- saff-precision ResNet, the output may overflow before the final normalization layer. This is interesting but beyond the scope of this paper.
Transformer also uses multi-branch structure in the attention layer and fully-connected layer. It does not use normalization layer on the output of the residual block because of the practical performance concerns. Hence it is observed unstable training deteriorating performance when adding the number of heads in attention layer or adding fully-connected branches (see Section 5.3).
FORWARD AND BACKWARD PROCESS OF MULTI-HEAD ATTENTION LAYER
In this section, we analyze the multi-head attention layer, a multi-branch structure used in Transformer (Vaswani et al. 2017) with each head viewed as one branch. It is worthy to note that the multi-head attention layer behaves differently from the feed-forward network and at the same time it is in general very hard to analyze, because of the softmax operator and the inter-symbol dependence. Previous work assumes that the softmax outputs equal weights geights/Xiong et al. (2020) and ignores the inter-symbol dependence, which does not fully reflect the attention behavior.
Suppose that the input is a sequence of symbols X = (#I.... ) and each row of X is symbol representation 2; Rd.. The multi-head self-attention layer, given by
MultiHead QQ K, V) = Concat ( ho)O,
QKT Von
h softmax
V,
for e [C],
(8)
5

---

Under review as a conference paper at ICLR 2021
where Q= k V and Q Rd ERRdd V. R Oe Id xd are projection (trainable) matrices and C' is the number of heads, du dy are dimensions.
Raxdx
The multi-head attention computes each head representation (with dimension d.))) independently, and then concatenates these heads (C' in total) together to form a large representation (Cd, dimensional vector), finally uses a fully-connected layer O E IR (Cdx) projecting back to IR (d is the model embedding dimension). This indicates that the original multi-heat attention uses the sum aggregation. is convenient to first consider the case of one head with dimension de for studying the multi-head attention layer. We give an upper bound the forward process when the softmax behaves more like a "max" operator. The softmax output could be rather spreadout even after training which is usually the case after initial training.
Proposition 1. Suppose a self-attention layer as described above and assume || Va i E [m]] If the parameters Q, K, V are randomly initialized with N N(,, 2), then for symbol æ; and its head representation h E Rul,, we have
El S (log 1logn-10 1+ v 2(dd, 1) log n).
where "<" means that the bound holds asymptotically and the expectation is over the initialization.
Proof. The proof is based on estimating the extreme Chi-square values, deferred to Appendix C
We can see that besides du,, the norm also scales with log n where is the number of symbols. In the Transformer literature, one usually fixes C' d, = d, and then the upper bound of head norm grows with C' because of the second term in the right hand side of Equation This term is ignored in previous analysis, which can dominate when d, < 22log n). We plot the norms of self attention layer (in the first decoder block) in Transformers with varying number of heads on IWSLT de-en task in Figure to illustrate this point. We can see at initialization the norm does not change with number of heads because of the expected behavior initialization (Xiong et al., 2020). However, this is not true after training (see the curve at Epoch 10). The assumption that softmax outputs uniform simplex does not hold. With large C/ and a small de. the second term in Equation 9 hurts the forward/baud hrwwardbbackwer bat//bbckwwrrrr stability and the training process, matching the observation in practice.
dominates, which may
For the backward process, we can estimate the gradient at initialization with the gradient independence assumption (Yang, 2019). Suppose the backward signal is (e1, E2, ..., en the gradient on Ô is Vo = Σis(n) e and E = Em (E]) //,l))))le 12, which scales with the head norm. Large forward process results in large gradients, which may lead to unstable training procedure.
5 EXPERIMENTS
In this section, we first verify the theoretical claims in Section 3 on the popular ResNeXt (Xie et 2017) model for image classification and on Transformer [Vaswani et al. 2017) to demonstrate its efficacy.
We then apply the STAM aggregation
5.1
THEORETICAL VERIFICATION
We conduct some experiments to verify the theoretical claims in Section Specifically, we use a Fixup ResNet20 (Zang et al., (219a) as backbone network, which uses Fixup initialization without all batch normalization (BN) layers and can be trained to decent accuracy. The multi-branch ResNets are generated by varying the number branches in each residual block (Figure We train these networks on the CIFAR 1//IFAR 100 classification task (Krizhevsky & Hinton] (0099) and record the top validation error in Figure 4
with a standard training procedure, i.c., 128 batch size, 200 epoch, initial learning rate 0.1 and step-wise learning rate decay by 1/ 10 at epoch 100 and 150 respectively.
We compare several approaches: the STAM aggregation with a = L/VC; the sum aggregation with a = 1; the average aggregation with a = 1/C; the backward scaling-down initialization y = I/VC, referred to "Back-scaled init" in Figure 4 the concatenation then aggregation by a linear layer with scaling-down initialization 7 1 CV referred to as "linear layer aggregation" in Figure 4 As

---

Under review as a conference paper at ICLR 2021
20.0
Epech Epech 10
rage
Sumh
15.0
Bacisscaled lnit
representation
12.5
linear layer aggregaton
vaildation
10.0
7.5
5.0
0
16 64
256 512
0.0
2
16
32
50
number attention heads
number parallel branches
Figure 3: Head representation norm of the self attention layer in the first decoder block. We fix hd, d and vary h. When d small, the representation norm grows with h.
Figure Top1 validation error of multi- branch ResNets on CIFARIO. Models trained with STAM aggregation consis- tently benefit from increasing branches.
Table I: Topl validation accuracy (in %) of ResNeXt models. The numbers with are reported in Xie et al. 2017) and other numbers are our implementation (averaged over 3 repeats).
Dataset
Setting
Param.
Basline
+STAM
8x64d
34.4M
96.35"
96.51(0.07)
CIFAR1O
16x64d
68.1M
96.42*
96.57(0.12)
32x64d
135.9M
96.37(0.11)
96.73(0.13)
8x64d
34.4M
82.23"
83.17(0.11)
CIFARI00
16x64d
68. MM
82.69*
83.63(0.09)
32x64d
135.9M
83.26(0.14)
83.86(0.12)
shown in Figure 4 the performance of model trained with STAM aggregation consistently improves as the number of branches C/ grows. In contrast, the models with sum aggregation fail to converge for C > 4, as predicted in Theorem [ - The performance of average aggregation degrades as the number of branches is beyond a certain value. In addition, the backward maling-down initialization works well terms of stability as we argued in Section 3.2 However, the scheme of concatenation then aggregation with linear transformation becomes unstable as the number of branches increases even with scaled-down initialization.
5.2 APPLY STAM AGGREGATION TO TRAIN RESNEXT
We next investigate a well-known multi-branch network, the ResNeXt (Xie al., (017) model, which divides the bottleneck residual block into multiple branches. It uses 1x1 convolution to aggregate the branch outputs, which is equivalent to sum aggregation if taking the Ixl convolution as part of each branch (see Figure 3 in Xie et al. (017)). ResNext uses BN layer to normalize the aggregated result. We have argued how normalization layer can stabilize the forward and backward process in Section It can be observed that ResNeXt fails to converge for C > 2 without such BN layer.
To verify the efficacy of STAM aggregation on ResNeXt, we use separate BN layer for each branch, equivalent viewing the BN layer a part of branch, and then apply the STAM aggregation these outputs. The hyperparameters and the data augmentation are the same as in Xie et al.] (2017). From Table we see that our STgM-aaggegattinn is applicable to the cases with many branches (more than the original paper) and improves over the strong baselines on both CIFAR IFFAR1O/CIFRRI 100. More details about the experiments can be found in Appendix D
5.3 APPLY STAM AGGREGATION TO TRAIN TRANSFORMER WITH MANY BRANCHES
In this section, investigate the aggregation scheme Transformer models Vaswani al. 2017 with machine translation tasks. A Transformer model consists several encoder and decoder layers, with each layer stacking or multi-head attention blocks and one fully connected (FC) block.
7

---

Under review as a conference paper at ICLR 2021
Table 2: BLEU scores IWSLT14 de-en test sets. The higher, the better. is baseline setting.
Table 3: BLEU scores newstest2014 for
WMT En-De. The higher, the better. The number
BLEU
with is reported in (Ott et al.) 2018).
h
Param.
BLEU
(+STAM)
BLEU
h
Param.
BLEU
4/512/1024
36.7M
35.21*
N/A
(+STAM)
8/256/1024
18.4M
34.96
35.30
16/1024/4096
209.8M
29.3"
N/A
12/256/1536
26.3M
35.0
35.79
16/512/4096
105. 1M
28.6
28.8
16/256/2048
34.2M
34.74
36.05
32/512/8192
193.2M
24.3
29.4
32/256/4096
65.7M
34.30
36.09
36/512/9216
213.2M
23.5
29.8
There are several hyperparameters about the model setting: model/embedding dimension d, head dimension d,, number of heads h, intermediate FC dimension dfc. More detailed introduction on Transformer can be found in Appendix C
We conduct two sets of experiments: IWSLT14 German- EermannEEgglih (de-en) task and the WMT16 (en-de) task. For baseline models, choose the default Transformer and the big Transformer for IWSLT14 de-en and WMT16 en-de, respectively. For multi-brance versions of Transformer, we first construct mini model, and then generate multi-branch Transformers by copying mini-models multiple times. This can explore many branches while controlling model size.
For 1WSLT14 de-en task, the baseline Transformer configuration is d 512, d, 128, dfc 1024, h = 4. We choose the mini model for IWSLT14 de-en with configuration d' = 256, da 512, and the generated multi-branch Transformers are with configurations (d, dr, h. de) = (256, 64, 8, 1024). (256, 64, 12, 1536), .... These models are trained with the Adam (Kingma & Ba 2014 optimizer with initial learning rate 1e-3, 31 = 0.9, B2 0.98 and the inverse sqrt learning rate scheduler, which are standard hyperparemeter choices following Ott et al.) (2018). For the STAM aggregation, we set a = 1/VC as C is the equivalent number of mini models. We evaluate the translation quality by BLEU scores by using multi-bleu.per We train every model on a single Tesla P40 GPU for 200 epochs. We test on a single model that achieves best BLEU score on the validation set. Other detailed configuration can be found in the code.
64, h 4, 4.
For the WMT16 English-German (en-de) task, the baseline big Transformer model has configuration 1024, dn = 64, fe = 4096, h = 16. We choose the mini model for WMTI6 en-de task with configuration d = 512, d, 64, ers are with configurations (d, dr, h, d = (512, 64, 16, 4096), (512, 64, 32, 8192), ... For the STAM aggregation, we set = 1/VC where C' is the equivalent number of mini models. All models are trained for 150 epochs on single node with Tesla V100 GPUs. We average the model parameters of last 5 checkpoints for evaluation.
8, dfc 2048, and the generated bulti-branch Transform-
We report the BLEU scores of all the models on test set in Table 2 and We can see that the performance of original architecture degrades as the number of branches increases, while with STAM aggregation the performance gradually improves as the number of branches increases.
CONCLUSION
In this paper, we study the training process of multi-branch network, especially the forward and backward stability. The theory tells that the aggregation is not stable as the number of branches increases. Motivated by the theoretical analysis of the multi-branch network, we propose the STAM aggregation that not only guarantee the forward/backward stability but also allows a same training strategy works for networks with different number of branches. We show that with STAM aggregation, the models can consistently benefit from increasing the number of branches. We believe that our analysis and the proposed STAM aggregation gives practitioners a new direction to design new multi-brane bilti-branch models.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Masoud Abdi and Saeid Nahavandi. Multi-residual networks: Improving the speed and accuracy of residual networks. arXiv preprint wrXiv:1609.06677 2016.
Karim Ahmed and Lorenzo Torresani. Connectivity learning in multi-branch networks, arXiv preprint arXiv:1709.09582, 2017.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. arXiv preprint arXiv:1811.03992y) 2018.
Anonymous Anonymous. Hybrid attentive transformer. 2020.
Sanjeev Arora, Simon $ Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, 2019.
Devansh Arpit, Victor Campos, and Yoshua Bengio. How to initialize your network? robust initialization for weightnorm & resnets. In Advances in Neural Information Processing Systems, 2019.
ZD BAT and YQ YIN. Limit of the smallest eigenvalue a large dimensional sample covariance matrix. Annals of probubility, 21(3): 1993.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much over- paramtterettoooooonnn is sufficient to learn deep ReLU networks? arXiv preprint arXiv:1911.12360, 2019.
Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31. 2018.
François Chollet. Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1251-1258, 2017.
S De and SL Smith. Batch normalization biases residual blocks towards the identity function in deep networks. arXiv preprint WiXi::2002.1044: 2020.
Simon Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv: I8I1 11811.03804, 2018.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, Proceedings of Machine Learning Research, pp. 249-256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Gao Huang, Zhuang Liu, Kilian Q Weinberger, and Laurens van der Maaten. Densely connected convolutional networks. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), volume 1, pp. 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning (ICML), pp- 448-456 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. Advances Neural Information Processing Systems, pp. 8571-8580, 2018.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test enror with shallow ReLU networks. arXiv preprint arXiv: /2292, 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 2014.

---

Under review as a conference paper at ICLR 2021
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. 2009.
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in neural information processing systems, pp. 8570-8581, 2019.
Xiang Li, Wenhai Wang, Xiaolin Hu, and Jian Yang. Selective kernel networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 510-519, 2019.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation ( (WMT), 2018.
Samet Oymak and Mahdi Soltanolkotabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning (ICMI., 2019.
Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normal- ization help optimization? In Advances in Neural Information Processing Systems, pp. 2483-2493,
2018.
Christian Szegedy, Sergey loffe, Vincent Vanhoucke, and Alexander A Alemi. Inception-v4, inception- resnet and the impact of residual connections on learning. In Thirty-first AAAI conference on artificial intelligence, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoe V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. Google's neural machine translation sys- tem: Bridging the gap between human and machine translation. arXiv preprint arXiv: 1tt9:086001144 2016.
Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5987-5995. IEEE, 2017.
Ruibin Xiong, Yunchang Yang, Di Hc, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie- Yan Liu. On layer normalization in the transformer architecture. arXiv preprint 2001v:2002.04745 2020.
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of chaos. In Advances in Neural Information Processing Systems, pp. 7103-7114, 2017.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv: 1902.04760, 2019.
Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representations, 2018.
Hang Zhang, Chongruo Wu, Zhongyue Zhang, Yi Zhu, Zhi Zhang, Haibin Lin, Yue Sun, Tong He, Jonas Mueller, R Manmatha, et al. Resnest: Split-attention networks. arXiv preprint 2rX4v:2004.08955. 2020.
Hongyang Zhang, Junru Shao, and Ruslan Salakhutdinov. Deep neural networks with architectures are less non-convex. arXiv preprint 1806.0848 2018.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representations (ICLR), 2019a.
Huishuai Zhang, Da Yu, Mingyang Yi, Wei Chen, and Tie- Yan Liu. Stability and convergence theory for learning resnet: A full characterization. arXiv preprint 1rXivv1903.077 2019b.
10

---

Under review as a conference paper at ICLR 2021
Jingfeng Zhang, Bo Han, Laura Wynter, Kian Hsiang Low, and Mohan Kankanhalli. Towards robust resnet: A small step but a giant leap. In International Joint Conferences on Artificial Intelligence (IJCAI), 2019c.
Jie Zhou, Ying Cao, Xuguang Wang, Peng Li, and Wei Xu. Deep recurrent models with fast-forward connections for neural machine translation. Transactions of the Association for Compulational Linguisties, 4:371-383, 2016.
Difan Zou and Quanquan Gu. An improved analysis of training over peeeraaramtteized deep neural networks. In Advances in Neural Information Processing Systems, 2019.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes 2018.
over-parameterized deep relu networks. arXiv preprint

---

Under review as a conference paper at ICLR 2021
A PROOFS IN SECTION 3
For fordard branches given by equation we adopt the Kaiming's initialization He et al. (2016): entries of Wa a for E [1 b 1] are independently sampled from NO,, 2) and entries of W 6 independently sampled from (0, 2)-
A.1 ONE BRANCH RESULT
Lemma 1. [Lemma 7.1 in Allen-Zhu al. [0188) For one branch B(han) given by equation with Kaiming's 's initialization, there exists some small constant e such that with probability at least
1 -(b) -Q(me (b) over the initialization randomness of W (RRmxmy the following holds
Va E [b ||all E (1 + e)llbanll:
(10)
where ha =- for 1,... - 1 with hg = hin.
We note that ||aal| E // means (1 - //m - |/all s (1 +)) The proof this lemma is based on the randomness of W and the concentration property.
A.2 PROOF OF THEOREM I
Theorem 1. Suppose the
network (-) is given by Equation 7 and Equation 2 with Kaiming's 's initialization. For an input hin. the following holds with probability least 1 TTC)). -0mme /b) over the initialization randomness of W
[[ E
Σ
a aill
(3)
KETC
where C' is the number of branches, and m, b are the width and depth of each branch, respectively.
Proof. In the following derivation, we fix a specific sample E [n]] and the input is hs,in. For notation simplicity, we ignore the index. We define an event & = {1 ll E (1 + e)]l/ vk e [C]]- (0me /b) On the event E, using Whi. h are independent Gaussian across k and
By Lemmal 1 we have that holds with probability 1 only the randomness of W E Rdxmm hg
B[Ag]] = 0, Var[hg] = || 1 Taxa for l: E [C].
Hence, we have (hhm)) Exejoy B² (hin) is Gaussian with mean 0 and covariance matrix Execcy až 1/ Laxd By law of large number, we have
[I E 6)² Σ a (hanl
(11)
The claim is proved.
Proof of Remark 1
The proof of Remark is straightforward In expectation we always have E|4 = O, [ = || 1| for & E [C] and E [b]. Furthermore as higs are independent across k. then ((h.)) Ex ag |h
A.3 PROOF OF THROREM 2
Given the branch equation 2 we further introduce notations to denote intermediate representation:
h 2-11: and h == )(g.)) for = 1,...,b b.
with hi hin and h2 = Φf for all [C]. Let D be diagonal matrix representing the activation
state with (D)),s = if (gt), 0,0, otherwise. Let gh 12. ..t)) E Rdxb
.
12

---

Under review as a conference paper at ICLR 2021
Theorem 2. With probability at least (nb) Exxp(-S((m)) over the randomness of W, it satisfies for every a [b] and k E [C], every i [n].
o x m
at
(5)
kE(C)
Deep neural networks are often trained with gradient based optimization algorithms. The gradient with respect to the parameter is computed through back-propagation, e.g., oW, 0h of 1- where H. represents the gradient of the objective with respect to :. Therefore, the gradient upper bound guaranteed if 91-1 and Ōh, are bounded across layers and iterations. We next show the backward process is bounded for each individual sample for the multi-branch network (-)... (-.)..
Proof. The gradient is computed as follows via back pack-propagation. We ignore the sample index for notational simplicity.
Olout
hout = Fhoc,
= Or Taxds
(12)
0gh
Obout Oyt
T
0gt =
Thout Ohout:
(13)
0h (W (WtTTagtt ow Og.- (h--2)²
(14)
0gh = D Sh.phh--...
(15)
0h (Wh)" ow? Oyti (1²)T.
(16)
For quadratic loss and sample E [ni]. we have that |ah 112 2C, (W).
We have
V wT C.wCC(WW) MIe |10g%.a Af.--1le
D' ... ( Ohs ||
< <O(Vm/d) on V L. (W) (()][zze],,
(17)
where the last inequality is due to Theorem I and the Alleen -hu et al. 2018 Lemma7.4b)
Thus, with high probability O(nbC) exp(-S((m)) for all i [ni]. E [6] and [C], we have
C. W 11 S 0 (WW
d at m) under the assumption || 1.
Proof Remark 2
To establish the expectation value on the backward process, it requires the gradient independence assumption (Yang 2019), which can be verified and argued for certain cases. With this assumption, we can take the expectation on the forward pass and the backward pass independently in (Equation 17). easy to obtain the expectation estimation Remark
A.4 MORE DISCUSSION ON PRACTICAL WISDOM
We argue the choices of scaling down initialization in detail.
Without loss of generality, we first initialize the parameters following the Kaiming's 's scheme, the same as in Section 3 Then we multiply the parameters multi-branch block by a same scaling down factor 7. We use the sum aggregation og = 1 for & E [C).
For forward stability, we need (h, < 0(1), which requires Ex < 1. Therefore
we obtain 7 < CC-2. 26.
For the backward process, we focus on the output update induced by one gradient descent step. That is (han: W V -hhan W). Suppose the backward error signal h with |ell 1.
13

---

Under review as a conference paper at ICLR 2021
Then following the proof of Theorem have
E V We || 70-1 Vm/d ||el| || 12(7 12(7611))
(18)
E B (han: W + v/V v;+ B² (hsn: Wx)l = 2/
(19)
where the second equality is due to Taylor expansion and the forward step has factor 6-1 Then the forward output update of branch is
Bh (bin: W vV w.) 1hth (Nin: Wh)
Bh (hens W W + 1 w/ - Ba (hini W(r),
ME(M
where the "~" ignores second order perturbation. Hence E] BI Wikh + 17V Wi) Bº (han: Wlll| = ((bby)((-1))) Summing over C branches, we have
(hi W + V w - (him W = ((1bC(20001)
(20)
Hence, to make the expected output update irrelevant with and 6, it requires y = (6C) 2,
With (6C)) 2, the forward process will not explode but obtain diminishing output norm
a16 (bC) क(
B HESSIAN ANALYSIS
Next to obtain a refined understanding on the property of different aggregation schemes, we analyze the Hessian of multi-branch network. Specifically for simplicity, assume that 8 has only one linear layer: B (h..) Whinn We can compute the Hessian of the objective with respect to the input hin and the learnable parameter W.
H hin
Σ oh W
Σ a Wh), H w = (aaa" @ Laxd @ (Eh
(21)
where W = Vee w and Vec W) = (Vec WT)T, Vec(W3)7 ...., Vec(WW is the long vector stacking W. is the Kronecker product, and E is average over the training samples.
Fact 1. For B² (him) w hin and a alc, we have the spectral norm of Hessian,
E H Q(a²C) and H = a Cl hT.l
(22)
where the first expectation over the randomness at initialization and the second expectation is over the training samples.
Proof. The Hessians are written as
Σ w T
Σ (Iww))
H Hbm. = CE or why
O Woswt aiza o²
k Hw- (aaa) & Laxa & (Ehinh
agza a² (1cxc) (Iuxb @ (Ehin
where & represents the Kronecker product and H w has dimension Cbp x Cbp.
Because entries of w are initialized with Gaussian distribution (0, 1/d), E. W follows Gaus- sian distribution (0, C/d). The singular value of Gaussian matrix is given by the following Bai-Yin's law (BAT YIN 1993).
Lemma 2. Let A E RNNN and entries of A are independent standard Gaussian random variables. Suppose that dimensions N and grow infinity while the aspect ratio n/N converges a constant in [0, 11. one has
Smax (A) = Vn v(v)) almost surely.
(23)
where Smax (A) the largest singular value A.
14

---

Under review as a conference paper at ICLR 2021
Hence E| H a2C(1 / m/dd² C) if viewing m/d some fixed hyyee--paameee
For the spectral norm of H w. we use the property of Kronecker product Amax (A @ B) = Amax (A)
Amax (B), where Amax (-) is the largest eigenvalue of the positively semi-definite matrix :, Hence
W 1mmammxxxxx 1- Amax (Bhin ((nn)) a² (ECin
(24)
where the second equality is due to the fact (max(11xx)) C and Amax (mxx)) =1
The spectral norm of the Hessian corresponds to the smoothness or the Lipschitz coefficient of the gradient. It determines the largest allowable learning rate for gradient descent algorithm. For the sum aggregation, the spectral norm | ll = Si(C) and | W both are C times larger than those in the case with only one branch. This indicates the learning rate for the gradient descent needs properly scaling down with the number of branches to guarantee convergence. For the average aggregation, initialization, the spectral norm H hen M(1/C)) and
C'IlEA h ll- Hence they
|| Φ T h I. They both are shrunk by 1/C compared to the case with only one branch. This indicates the learning rate or the gradient needs properly scaling up with the number of branches. Thus one cannot apply the same training strategy to train network with varying number of branches and the sum or average aggregation.
In contrast, for the STAM aggregation, the spectral norm H They both remain the same when the number of branches varies. This indicates that we can apply the same training strategy for varying number of branches.
h.n. || = 2(1) and H
| A ll-
C MORE DISCUSSION ON TRANSFORMER MODEL
C.I FORWARD PROCESS
Proof of Proposition 1
Without loss of generality, we assume that there are only one head and ignore the head subscript
for simplicity. For a specific symbol representation x, the query is = x Q, the keys are k xu K, kn = ? K (i.e., K X ), the values are t1 = x We note that q, k ;, are row vectors. Let (pI....,Ph)
On = x (i.e., V = XV.
softmax
g Vda"
qk 9k Vdl Va.
We note that for symbol h = (pi,.... (P.).....P).V. V and then we have an upper bound on the head representation | ( Pm) XVI I maxtej]] VI).
V
Suppose ||> dd for all i E [n). which is reasonable assumption because of the layer normaliza-
tion. Then ||x; 1I a Chi-square x². variable. Next we estimate the maximum of n independent Chi-square variables by using the following asymptotic relation. For random variable U ~ x2 and large v, we have 20 V2v I I ~ (0, 1). We use the estimation of extreme value of Gaussian variables. For Z maxieln] V; and Vi (O 1), we have B(Z) < v2 log n. Based the above estimation, we have asymptotically max;eln] || 12 S dy - log + 22(2d, 1) log We can also estimate the maximum of sequence of x2 variables by using the relation that the extreme value of x2 distribution (a case of Gamma distribution) asymptotically converges to Gumbel distribution and get similar result. Hence, we have ||A||2 S d + (log n /2(2d, 1) log n).
C.2 BACKWARD PROCESS
We focus the case with one head and omit the subscript We assume d, du dg =: dn. We rewrite the forward process as follows. Given the parameters Q E Rdxda K E Rddx V
15

---

Under review as a conference paper at ICLR 2021
Rdxd
0 R
X (0), .. kn)²
with each row being symbol embedding, E Raxd
(25)
V XV, E Raxdi
(26)
Q = XQ, E R Ruxds
(27)
K = XK, E Raxd,
(28)
P QK dv
(29)
P softmax(P), with each row being probability simplex, E R
(30)
H PV, with each row being the head representation, RR=x
(31)
= HŌ, with each row being the new representation of a symbol, E Raxd
(32)
For the backward process, it can be written as follows. Suppose the error signals on the output O is E (e1,..., e,,), with each being
signal symbol representation.
80 E, E Raad
(33)
0H = 80 o",
a0 = H 80,
(34)
0P 0H VT,
OV p OH,
(35)
8y X OV. v OV x"
(36)
OP 0PJ
(37)
AQ aP K/ dp OK OP" Q/ Q/
(38)
0q IQ Q iQ X aQ.
(39)
0gX = OK K OK X OK,
(40)
0X = 0y + 0g X + 0
(41)
We can further compute the expected norm of the gradient as in the proof of Proposition Jnder the gradient independence assumption (Yang 2019)
C.3 MORE ABOUT TRANSFORMER
Apart from the multi-head attention block, the fully-connected (FC) block is also multi-branch architecture with sum aggregation. We use dfc to denote the intermediate dimension of fully con- nected block. The fully connected block consists of two FC layers with weights w Bdxd and W E Xd practice, dfc is several times larger than For example, the "big" Transformer in Vaswani et al. (2017) uses dfc- 4d. These sum aggregation operations restrict practitioners to explore Transformer models with large and djc-
D MORE DETAILS ABOUT EXPERIMENTS
The training procedure of ResNeXt model in Section 5]is the same as Xie et al. (2017). Specifically, we train the models GPUs with batch size 128, total number of epochs 300 and weight decay 0.0005. We use Nesterov momentum with coefficient 0.9. The learning rate is multiplied by 0.1 at the 1500th and 225-th epoch. For data augmentation, we take a random crop with size 32x32 from zero-padded 40x40 image or its horizontal flipping.
Now we introduce more details of Transformers on translation tasks Section The IWSLT14 De-En dataset is collected from Fairseq official site?] which contains 156K/7K/7K sentence pairs for traiiingyalliaationntess respectively. For WMT
task, we use the same data setup as Ott et al. (2018). Specifically, train the model using the training data of WMT16. which contains 4.5M sentence pairs. The validation and test sets are newstest13 and newstest14, respectively.
16

---

Under review as a conference paper at ICLR 2021
We use head-level dropout which drop heads representation (before the output projection of attention block). The head-level dropout probability is p. We set p = 0.3 for IWSLT14 De-En task and 0.1 for WMT16 En En-De task, respectively. Our source code is available at anonymous GitHub pagq
/ttn://ggthhiioooomooooooossssss 17


Reviews:
[{"review_id": "r82N-eV6_8M", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "[Summary]\n\nThis 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).\n\n[Pros]\n\n--- The problem this paper focuses on (multi-branch architectures) is important and I think worth more studies since they are used more widely these days. From a theoretical perspective, the number of branches seems like a similar thing as the width of a network, yet the understanding is not yet as comprehensive as our understanding on the width (for which we know things like what are their effects, how should we initialize a wide network, etc).\n\n--- The experimental results are interesting and could be a good reference point for designing these multi-branch architectures. Specifically, it is interesting to see it is better to explicitly do a fixed $1/\\sqrt{C}$ scaling, rather than making it trainable or embedding it into the individual linear layers in each block. \n\n--- The paper is generally clearly written and easy to follow. The notation between the theory and empirical parts matched quite well.\n\n[Cons]\n\n--- I feel like the core contribution of the $1/\\sqrt{C}$ scaling rule is rather straightforward and not necessarily justified as a new method or theory, especially with the analogy to width in mind: It is standard to initialize weight matrices to have scale $1/\\sqrt{d}$ where $d$ is the either the input or output dimension as this gives the right normalization. Even with the other “baselines” (making the $1/\\sqrt{C}$ trainable or embedding it into the individual blocks) performing worse according to this paper, I still feel like a researcher or practitioner who encounters the issue of choosing the aggregation rule could still end up with the same solution proposed in this paper. \n\n--- The theoretical results also build on the exact same normalization idea ($1/\\sqrt{C}$ scaling gives the right normalization) as for width, and do not seem to provide us with much new theoretical understandings for what is unique about the multi-branch architecture. From my (maybe biased) perspective, I would have been more interested in some more in-depth understandings (either theoretical or empirical), such as what is the effect of adding more branches (and where does diminishing return kick in). Reading this paper leaves me more questions than answers about these multi-branch networks.\n\n------\nAfter rebuttal: Thank the authors for the response. I think my original concern largely stands and would like to keep my original evaluation of the paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "mxKzMCG3Lq_", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nPros:\nThe paper is well-written. Its motivation is clear and the analyses are technically correct. The proposed method is well interpretable and easy to follow. The connections with some initialization and normalization methods are well explained. \n\nCons & questions:\n1.\tThe paper should offer the definition of “stability” first. Does it refer to the numerical stability of features or the robustness to noise? What does $\\epsilon$ in Eq (3) mean? Is it the perturbation to the input feature $h_{in}$? \n2.\tHow to quantize stability? The experiments on image classification and machine translation show that the proposed method STAM is able to improve classification accuracy and BLEU score. I think these metrics are more related to representation ability. How do we know the stability is improved? \n3.\tThe extendibility is limited. The analyses are based on the assumption that all branches have the same structure without normalization layers. Does it work for network with variant branches? How to generalize this study to network with BN layers? What about the results of experiment 5.1 when BN layers are enabled? Since BN layers have been successful in solving the training difficulties of deep networks (not only multi-branch), could the authors explain the superiority of this study over BN in practical implementation? Besides, are the analyses and proposed methods also applicable to concatenation aggregation? \n4.\tThe improvements on CIFAR-10 and CIFAR-100 over baseline are insignificant. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "m1Ed1KfjdU2", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the multi-branch neural networks. It analyzes the feedforward/backward dynamics of multi-branch networks, proposes an aggregation design named STAM, and shows the feedforward/backward dynamics of multi-head attention layers, and discusses some of its properties.\n\nHere are some issues that might need the authors’ attention.\n\n1. In the introduction, the paper states that simply adding branches multiple times does not often work as expected. I am aware that there are experimental results in the later sections that support this argument. But putting it in the introduction without any concrete examples might be confusing.\n\n2. I think it would also be helpful if there are some brief descriptions of what the proposed STAM does and how it improves over other designs in the introduction.\n\n3. The paper seems to lack a strong main theme and all the contributions and results are loosely connected.\n\n4. The theoretical results are mostly based on the assumption that the weights are randomly initialized. However, the distributions and the properties of the weights might be different after the training begins. Does this imply that the significance of some of the results is limited?\n\n5. The stability analysis is mostly done on norm analysis. The proposed STAM is also inspired by the norm analysis where the goal is to keep the norm of activations and gradients within a reasonable range. However, it seems to me that the theoretical results are a little bit weak: it’s not hard to imagine that the relationship between norms of activations/gradients and the norm of $\\alpha$ is somewhat linear, and sum/mean becomes improper when the number of branches grows.\n\n6. Missing words in paragraph 5 of section 3.2: … as the number of branches (see Section 5.2).\n\n7. The aggregation form that the paper has studied is the weighted sum of multiple branches. The paper claims that concatenation is equivalent to sum aggregation for ResNeXt. I’m not fully convinced that this is true considering it’s impractical for a sum aggregation to perform matrix multiplication, batch normalization, and optional ReLU. And I don’t see an easy way to absorb the last 1x1 conv+bn into the branches while satisfying the assumptions of the theoretical results, plus bn will also cancel the effect of $\\alpha$ if my understanding is correct.\n\n8. ResNeXt uses the batch normalization layers, which automatically adjust $\\alpha$. How is this related to the proposed STAM? Do they serve similar purposes?\n\nSummary:\nThe paper presents some interesting findings. However, it seems lacking a strong main theme and the novelty/significance of the theoretical results is not strong, either. The proposed STAM only provides incremental design changes to the networks.\n\nUpdates:\nThanks for the authors' response and the revisions. However, the paper still lacks strong and significant results. Therefore, I keep my previous rating. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "UjExLoGCTQg", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##### Summary\nThis paper proposes to scale down the features of multiple branch networks by $1/\\sqrt{C}$ during aggregation where $C$ is the number of branches. The authors provide some theoretical insights to back up the method. However, the experiments are not very convincing. I had a good impression of this paper when I read the method sections but until the experiment sections. There are many problems in the experimental setup that makes me hesitate to accept this paper. The authors have a very unclear definition of the number of branches which gives me the impression that the authors just introduce a hyperparameter $\\alpha = 1/\\sqrt{C}$ and they can pick whatever $\\alpha$ or $C$ they want in the experiments tailored for the story.\n\n##### Strength\n- The arguments and theory in sections 3 & 4 are sound and interesting.\n- The proposed method is quite simple.\n- The proposed method leads to improvements on CIFAR 10 & 100.\n- By tuning the model architecture with the proposed method, the authors also show improvements on IWSLT 14 & nestest2014.\n\n\n##### Weaknesses\n\n- This paper is lack of comparisons with a series of related works that sets the outputs of residual branches to zeros such as FixUp (Zhang et al., 2019) and Re-Zero (Bachlechner et al., 2020).\n- It is very suspicious that the authors do not report the BLEU (+STAM) of the first row in Tables 2 & 3. It gives me the impression that the authors try to hide the fact that the proposed method does not do well in this setting and they have to tune the model structure and training settings to find some cases that the proposed method works.\n- In Table 3, the baseline (Ott et al., 2018) is trained for about 43 epochs. In my experience, it starts to overfit after that. However, the authors train all their models for 150 epochs. Otherwise, it is possible that the Transformers without STAM BLEU converge faster and overfit at epoch 150, which makes them look worse than the ones with STAM. Based on Appendix D, the authors also use a new technique head-level dropout which was not used by Ott et al., 2018. I would like to see the authors use the same setting as the baseline or report curve of BLEU scores at different epochs to convince me that this is not the case.\n- The author set $\\alpha = 1/\\sqrt{C/8}$ for Transformers and $\\alpha = 1/\\sqrt{C/4}$ for ResNeXt and just provide some very hand waiving explanation, which gives an impression that the theory doesn't work and all we need is to tune a hyperparameter $\\alpha$.\n- On page 6, \"We give an upper bound on the forward process when the softmax behaves more like a “max” operator, which is usually the case after initial training.\" This is a false claim. It is more often that the attentions assign large weights to multiple positions. Making the attention sparse usually restricts the models and leads to worse results.\n- The claim in Appendix C.3 makes the authors' definition of $C$ very arbitrary. The author claim that the 2-layer FFN with $d_{fc} = 4d$ in a Transformer-big has 4 branches. I don't see why there is a reason that this should view as having four branches. This should be considered as having only one branch only. Otherwise, since it has only 2 layers, we can view each hidden neuron in the middle layer as one branch. The authors just take their advantage to group ever $d$ neurons as one branch, but some other people can also group them in any other ways and say there are $C$ branches take doing them arbitrarily. For example, I can group ever 4 neurons as one branch and say this FFN has $d$ branches, or even group the first $2d$ dimensions as one branch and the rest of the neurons as $2d$ branches and say this FFN has $2d + 1$ branch. Overall, this makes me feel that the authors are playing a game of definitions and choosing whatever $C$ or $\\alpha$ as they want.\n\n\n##### Questions & Suggestions\n- How does STAM work when applied to depthwise separable convolution in MobileNets and EfficientNets which can be viewed as an extreme case of multi-branch convolution.\n- It would be more exciting to see if this method can be applied to the field of neural architecture search (NAS) where people explicitly train a gigantic multi-branch network and has problems with having some branches undertrained.\n\n##### References\n- Zhang, Hongyi, Yann N. Dauphin, and Tengyu Ma. \"Fixup initialization: Residual learning without normalization.\" arXiv preprint arXiv:1901.09321 (2019).\n- Bachlechner, Thomas, et al. \"Rezero is all you need: Fast convergence at large depth.\" arXiv preprint arXiv:2003.04887 (2020).", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the training of multi-branch networks, i.e., networks formed by linearly combining multiple disjoint branches of the same architecture.  The four reviewers seem to reach a consensus that the paper is not ready for publication for ICLR. 
Decision: Reject