On the Convergence of Shallow Neural Network Training with Randomly Masked Neurons
Authors: Fangshuo Liao
Anastasios Kyrillidis
Venue: Accepted by TMLR ()

Abstract: With the motive of training all the parameters of a neural network, we study why and when one can achieve this by iteratively creating, training, and combining randomly selected subnetworks. Such scenarios have either implicitly or explicitly emerged in the recent literature: see e.g., the Dropout family of regularization techniques, or some distributed ML training protocols that reduce communication/computation complexities, such as the Independent Subnet Training protocol. While these methods are studied empirically and utilized in practice, they often enjoy partial or no theoretical support, especially when applied on neural network-based objectives.

In this manuscript, our focus is on overparameterized single hidden layer neural networks with ReLU activations in the lazy training regime. By carefully analyzing $i)$ the subnetworks' neural tangent kernel, $ii)$ the surrogate functions' gradient, and $iii)$ how we sample and combine the surrogate functions, we prove linear convergence rate of the training error --up to a neighborhood around the optimal point-- for an overparameterized single-hidden layer perceptron with a regression loss. Our analysis reveals a dependency of the size of the neighborhood around the optimal point on the number of surrogate models and the number of local training steps for each selected subnetwork. Moreover, the considered framework generalizes and provides new insights on dropout training, multi-sample dropout training, as well as Independent Subnet Training; for each case, we provide convergence results as corollaries of our main theorem.

Published in Transactions on Machine Learning Research (08/2022)
On the Convergence of Shallow Neural Network Training with Randomly Masked Neurons
Fangshuo Liao Deportment of Computer Science
Fangshuo-Liao Liao@riceoeeee
Rice University
Anastasios Kyrillidis
anastasios@ricee.ct
Deportment Computer Science
Rice University
Reviewed on OpenReview: openreveee. net/ forum? id= e Y7mYYMSy2H
Abstract
With the motive of training all the parameters of neural network, we study why and when one can achieve this by iteratively creating, training, and combining randomly selected sub- networks. Such scenarios have either implicitly or explicitly emerged the recent literature: sec c.g., the Dropout family of regularization tochniques, or some distributed ML training protocols that reduce Subnet Training protocol. While these methods are studied empirically and utilized in prac- ticc, they often enjoy partial or no theoretical support, especially when applied on neural network-basee objectives.
complexities, such the Independent
In this manuscript, our focus is overparameterrzeed single hidden layer neural networks with ReLU activations in the lazy training regime. By carefully analyzing i) the subnetworks' neural (angent kernel, ii) the surrogate functions' gradient, and vii how we sample and combine the surrggate functions, we prove linear convergence rate of the training error uup to a neighborhood around the optimal point- for an overparameterized single-hidden layer perceptron with regression loss. Our analysis reveals a dependency of the size of the neighborhood around the optimal point on the number of surrogate models and the number of local training steps for each selected subnetwork. Moreover, the considered framework generalizes and provides as well as Independent Subnet Training; for each case, provide convergence results as corollaries of our main theorem.
insights dropout training, multi-sample dropout training,
1 Introduction
Overparameterized neural networks have led to both unexpected empirical success in deep learning
Zhang
et al.
2021
Goodfellow et al.
2016: Arpit et al.
2017
Recht et al.
2019
Toneva et al,
2008).
and new
techniques in analyzing neural network training Kawagucht
al. 201
7 Bartlett al. 2017 7
Neyshabur Neyshabur
et al.
2017 7 Golowich et al.) 2018 Liang et al.)
2019
Arora al.) 2018
Dziugaite & Roy
2017 7
ct al.
2018
Zhou et al.] 2018 Soudry al. 2018
Shah et al.] 20200 Belkin et al. 2019
201
8
Feldman
2020
Ma et al. 2018 Spigler et al. 2019 Belkin
2021 Bartlett et ol.
2021 Jacot et al.
2015). While
thcoretical work in this field has led to diverse set of new overparancterized neural network architectures
Prei el al. 2020 Fang el al. 2021 Lu el al.
2020 Huang el al. 2020 Allen Zhu el al. 12119] Gu el al.
2020 Cao et al.) 2020 and training algorithms
Du et al.) 2018 Zou et al.) 2020 Soltanolkotabi et al.) 2018
Oymak & Soltanolkotabl 2019 Li et al.) 2020 Oymak 8 Soltanolkotabi 2020 most efforts fall under the
following scenario: in each iteration, we perform a gradient buddent--baedd update that involves all parameters of the neural wetwork in both the forward and backward propagation. Yet, advances in regularization techniques

---

Published in Transactions on Machine Learning Research (08/2022)
Sricostava et al.
2014
Wan et al.
2013:
Gal 8 Ghahromand 2016;
Courbariaur et al. 2015
Labach et al.
2019.
computationalll-effiie
(Shazeer el al. 2017 Fedus et al.
2021 Lepikhin el al.
2020,
LeJeune
et al. 2020
Yao et al.
2021
Yu et al.
20188 Mohtashami et al, 2021 Yuan et al. 2020 Dun et al.
2021
Wolfe
et al. 20210 and communicttion-fficiat distributed training mothods (Vogels et al. 2019
Wang
el al.
2021-
Yuan el al.
2020 favor a different narrative: one would explicitly or implicilly- train smaller
and
mandomly-scloctee
models within a large model, iteratively. This brings up the following question:
Can one meaningfully train an overparameterized ML model by iteratively training and combining together smaller versions of it?"
(b)
iteration k
Masks
m
m
Original Network
Subnetwork
Subnetwork
Subnettwrrk
A(W )
Sng w.
Smg aW
Smn (...)
(a)
Local Training
Local Training
Local Training
(n)
Aggregated Network
Local Trained
Local Trained
Local Trained
(d)
J(W Be1s
(W Ext
Smgl Su?
Sug (wa
Aggregation
Training single hidden-layer perceptron using multiple randomly masked subnetworks. Here, f(W, :) denotes the full model with W parameters, and fm; (W, denotes the surrogate model (subnetwork) with only active neurons as dictated by the mask m at the k-th iteration for subnetwork: 1. Moreover, W₁ denotes the parameter at the start of the iteration, while w is the trained parameter of subnetwork I.
Figure 1:
This question closely relates to multiple existing training algorithms, as we discuss below; goal is to work towards a unified training scheme, and seek for rigorous theoretical analysis of such framework. Focusing on this objective based on shallow feedforward neural networks, we provide positive answer, accompanied with theoretical guarantees that are supported by observations in practical scenarios.
To be more specific, the training scheme we consider is depicted in Figure Given a dense neural network, as in Fig.1(a), we sample masks within one training step; see Fib. Each of these masks deactivates a subset of the neurons in the original network's hidden layer. In way, each mask defines surrogate model, shown in Fi), based on the original network, leading to collection of subnetworks. These surrogate subnetworks independently update their own parameters (possibly on different data shards), by performing (stochastic) gradient descent ((S)GD) steps. Lastly, we aggregate the parameters of the independently trained subnetworks (Fig.1(d) to update the weights of the original network, before the next iteration starts; see (e). Note that multiple masks could share active neurons. When aggregating the updates, we take the weighted sum of the updated parameters across all subnetworks, with the aggregation weights computed on the masks of the current iteration.
2

---

Published in Transactions on Machine Learning Research (08/2022)
We mathematically illustrate the difference between traditional training (first expression below) and the considered methodology (second expression below):
W GD) (f, TR.
vs.
W k++ Reassemble GS) (fm;, Wh:r), GD) ( (Jmz, Wi:r), ,... (S) (mmg. W (i:+)) ))
Here, the acronym GD (f, W, indicates the application of (S)GD on function f for T iterations, starting from initial parameters W. Consequently, (D) (fmt. W indicates the application of (S)GD for т iterations on the surrogate function fmj based the mask m and using only the subset of parameters W The function Reassemble involves both aggregation and reassembly of the whole model WA+1-
In this work, we perform a theoretical analysis of this framework, based on a single-hidden layer perceptron with ReLU activations. This is a non-trivial, non-convex setting, that has been used extensively in studying the behavior of training algorithms on neural networks
[Du et al.) 2018; Zoo et al. 2020
Sottanolkotabi
et
2018 Oymak & Soltanolkotabl 2019 Li et al. 2020: Oymak & Soltanolkotabi
2020
Song 8 Yang
2020 Ji 8 Telgarsky 2020 Mianjy & Arora 2020).
Challenges. Much work has been devoted to analyzing the convergence of neural networks based on the
Neural Tangent Kernel (NTK) perspective [Jacot et al. 2018); see the Related Works section below. The literature in this direction notice that the NTK remains roughly stable throughout training. Therefore, the neural network output can be approximated well by the linearization defined by the NTK. Yet, training with randomly masked neurons poscs additional challenges: i) With a randomly generated mask, the NTK changes even with the same set of weights, leading to more instability of the kernel; ii) the gradient of the subnetworks introduces both randomness and bias towards optimizing the loss of the full network; and, iti) the non-linear activation makes the aggregated network function no longer linear combination of the subnetwork functions. The three challenges complicate the analysis, driving us to treat the NTK, gradient, and combined network function with special care. We will tackle these difficulties in the proof of the theorems.
Motivation and connection to existing methods. The study of partial models//ubnnetwon that reside in large dense network have drawn increasing attention.
Dropout regularization. Dropout Srivostava et al. 2014 Wan et al. 2013
Gal 8 Ghahramani 2016
Courbariaur et al. 20155) is widely-accepted technique against overfitting in deep learning. In each training step, a random mask is generated from some pre-defined distribution, and used to mask-out part of the neurons in the neural network. Later variants of dropout include the drop-connect [Wan et al. 2013), multi- sample dropout Inoud 2019,, Gaussian dropout (Wang & Manning 2015), and the variational dropout KKnggma el al. 2015). Here, we restrict our atlention ion to the vanilla dropout, and the multi-sample dropout. The vanilla dropout corresponds our framework, if in the latter we sample only one mask per iteration, and let the subnetwork perform only one gradient descent update. The multi-sample dropout extends the vanilla dropout in that it samples multiple masks per iteration. For regression tasks, our theoretical result implies convergence guarantees for these two scenarios on single hidden-layer perceptron.
Distributed ML training. Recent advances in distributed model/parallel training have led to variants of distributed gradient descent protocols [Mddonadd et al. 2009 Zinkevich et al. 2010 Zhang & Ré 2014 Zhang et al.] 2016). Yet, all training parameters are updated per outer step, which could be computationally and communication inefficient, especially in cases of high communication costs per round. The Independent Subnetwork Training (IST) protocol [Yuan et al. 2020 goes one step further: IST splits the model vertically, where each machine contains all layers of the neural network, but only with (non-overlapping) subset of neurons being active in cach layer. Multiple local SGD steps can be performed without the workers having to communicate. Methods in this line of work achieves higher communication elficiency and accuracy that is comparable to centralized training. [Wolfe et al. [022 Dun et al. 2021 Yuan et al,, 2020) Yet, the theoretical understanding of IST is currently missing. Our theoretical result implies convergence guarantees for IST for a single kidden-layer perceptron under the simplified assumption that every worker has full data access, and provides insights on how the number of compute nodes affects the performance of the overall protocol.
3

---

Published in Transactions on Machine Learning Research (08/2022)
Contributions. The present training framework naturally generalizes the approaches above. Yet, current literature -more often than not- omits any theoretical understanding for these scenarios, even for the case of shallow MLPs. While handling multiple layers is a more desirable scenario (and is, indeed, considered as future work), our presented theory illustrates how training and combining multiple randomly masked surrogate models behaves. Our findings can be summarized as follows:
We provide convergence rate guarantees for i) dropout regularization Srivostava et al. 2014 di) multi- sample dropout [Inove 2019), iii) and multi-worker IST Yuon et al. 2020 given regression task on a single-hidden layer perceptron.
We show that the NTK of surrogate models stays close to the infinite width NTK, thus being positive definite. Consequently, our work shows that training over surrogate models still enjoys linear convergence.
For subnetworks defined by Bernoulli masks with a fixed distribution parameter, we show that aggregated gradient in the first local step is a biased estimator of the desirable gradient of the whole network, with the bias term decrcasing as the number of subnetworks grows. Moreover, all aggregated gradients during local training stays close to the aggregated gradient of the first local step. This finding leads to linear convergence of the above training framework with an error term under Bernoulli masks.
For masks sampled from categorical distribution, provide tight bounds i) on the average loss increase, when sampling a subnetwork from the whole network; ii) on the loss decrease, when the independently trained subnetworks are combined into the whole model. This finding leads to linear convergence with a slightly different error term than the Bernoulli mask scenario.
Summarizing the contributions above, the main objective of our work is to provide theoretical support for the following statement:
Main statement (Informal). Consider the training scheme shown in Figure 1 and described preciscly in Algorithom 1 If the masks are generaled from Bernoulli distribution calegorical distribution, under sufficiently large over-parameterization coefficient, and sufficiently small learning rale, training the large model via surrogate subnetworks still converges linearly, up to an neighborhood around the optimal point.
2 Related Works
Convergence of Neural Network Training. Recent study on the properties of over-parameterrzee neural networks enabled their training error analysis. The NTK-based analysis studies the dynamics of the parameters in the so-called kernel regime under a particular scaling option [Jacot et al. 2018 Du
ct al.
2018 Oymak 8 Soltanolkotabi 2020 Song & Yang 2020 Ji 8 Telgarsky 2020 Su 8 Yang 2010
Arora el al. 2019 Miannj Avora 2020 Huang el al.) 2021). NTKs can be viewed as the reproducing
kernels of the function space defined by the neural network structure, and are constructed using the inner product between gradients of pairs of data points. With the observation of the NTK's stability under sufficient over- parametrrization, recent work has shown that (S)GD achieves zero training loss shallow neural networks for regression task, even if when the data-points are randomly labeled [Du et al. 2018 Oymak & Soltanolkotabi] 2020 Song & Yang [2020]. To study how the labeling of the data affects the convergence, (Arona et al. 2019) characterizes the loss update in terms of the NTK-induced inner- product of the label vector, and notices that, when the label vector aligns with the top eigenvectors of the NTK, training achieves faster convergence rate. (Su & Yang 2019) analyze the convergence of training from a functional approximation perspective, and obtains meaningful result under infinite sample size limit, where the minimum eigenvalue of the NTK matrix goes to zero.
Later works start to deviate from the NTK-based analysis and aim at reducing the over par--aaammtercization requirement. By using a more refined analysis on the evolution of the Jacobian matrix, Oynak & Soltanolkotabl 2019) reduce the required hidden-layer width to n², with n being the sample size. Nguyen 2021 leverages a property of the gradient that resembles the PL-condition, and provides convergence guar- antees for deep neural network under proper initialization. When applied to neural networks with one hidden layer, their over-parameterization requirement also reduces to n². [ong et al. 2027) further study
4

---

Published in Transactions on Machine Learning Research (08/2022)
the PL-condition along the path of the optimization and show that a subquadratic over- over- parameteerizationn is sulficient to guarantee convergence. While reducing the over-parameleriaation is ideal, the focus of our work is to extend the analysis on regular neural network training to a more general training scheme.
A different line of work explores the structure of the data data-distribution in classification tasks, by assuming separability when mapped to the Hilbert space induced by the partial application of the NTK [Ji & Telgarsky 2020 Mianjy & Aroro 2020). Rather than depending on the stability of NTK, the crux of these works relies on the small change in the lincarization of the network function. This line of work requires milder overparameterization, and can be easily extended to training stochastic gradient descent without changing the over-paramcteriaatio requirement. The above literature assumes all parameters are updated per iteration.
A Analysis of Dropout. There is literature devoted to the analysis of dropout training. For shallow linear neural networks, (Senen- Cerda & Sanders 20200/ give asymptotic convergence rate by carefully character- izing the local minima. For deep neural networks with ReLU activations, (Senen-Cerda 8 Sanders 20200 shows that the training dynamics of dropout converge to unique stationary set of a projected system of differential equations. Under NTK assumptions, Mianjy & Arora (020) shows sublinear convergence rate for an online version for dropout in classification tasks. Recently, [eJeune et al. 2021 study the duality of Dropout in a linear regression problem and transform the dropout into a penalty term in the loss. Our main theorem implies linear convergence rule of the training loss dynamic for the regression lask on a shallow neural network with ReLU activations.
Federated Learning and Distributed Training. Traditional analysis on distributed training methods assumes that the objective can be written as sum or average of sequence sub-objectives (Stich 2018
Li ct al. 2019: Haddadpour & Mahdavi 2010 Khaled et al.) 2019a B. Each worker performs some variant of local (stochastic) gradient descent on a subset of the sub-objectives. This line of work deviates greatly from the scenario we are considering. First, the assumption that the objective can be broken down into lincar combination of sub-objectives corresponds to the data-parallel training, in stark contrast to our proposed scheme: the loss computed on the whole network not necessarily equals to the mean of the losses computed on the subnetworks. Second, this line of work assumes that the objective is smooth and usually convex or strongly convex. Moreover, there is recent work on training partially masked neural networks that deviates from the assumption that the objective can be written as linear combination of sub-objectives [Mohlashami el al. 20270. In particular, in this work [Mohlahhmmi el al. 2027), the authors consider optimizing more general class of differentiable objective functions. However, assumptions including Lipschitzness and bounded perturbation made in their work cannot be easily checked for a concrete neural network, especially for the problem of minimizing mean squared error. Lastly, recent advances in the NTK theory also facilitated the theoretical work on Federated Learning (FI.) on neural network training. FI.NNTK [Huang et al.} (202)) characterize the asymmetry of the NTK matrix due to the partial data knowledge. For non-i.idd data distribution, [Deng & Mahdavi 2021) proves convergence for shallow neural network by analyzing the semi Linihhiischitznesss of the hidden layer. Our work differs since we consider training a partial model with the whole dataset. We consider the more frequently used setting of one-hidden layer perceptron with non-differentiable activation.
3 Training with Randomly Masked Neurons
We use bold lower-case letters (e.g., a) to denote vectors, bold upper-case letters (e.g., A) to denote matrices, and standard letters (e.g., a) for scalars. ||a]|| stands for the E (Euclidean) vector norm, || ll2 stands for the spectral matrix norm, and || stands for the Frobenius norm. For an integer a, we use [a] to denote the enumeration set (1, ,a}. Unless otherwise stated, p denotes the number of subnetworks, and ( E [p] its index; K denotes the number of global iterations and k e [K] its index; 7 is used for the number of local iterations and t E [r] its index. We use M, to denote the mask at global iteration k, and E| = EM,,...,, to denote the total expectation over masks Mo...., M₁- We use P(-) to denote the probability of an event, and I(-} to denote the indicator function of an event. For distributions, we use Niu, Σ) to denote the Gaussian distribution with mean μ and variance Σ. We use Bern(() to denote the Bernoulli distribution with mean ç, and we use Un) to denote the uniform distribution over the set S. For complete list of notation, see Table 1 in the Appendix.
5

---

Published in Transactions on Machine Learning Research (08/2022)
3.1 Single Hidden-Layer Neural Network with ReLU activations
We consider the single hidden-layer neural network with ReLU activations, as in:
Algorithm 1 Randomly Masked Training
Input: Mask Distribution D, local step-size
m
global aggregation weight Tarr
I(W, x) Vm
a.o(( x)) == xW.
1: Initialize Wo, a
2: for k k= == .......... 0..... K do
Here, = = Wi...... Wm E R.xd is the weight ma-
3: Sample mask M D
trix of the first layer, and a = [ai...., am] E I is the weight vector of the second layer. We assume that each Wr is initialized based on N/ x²I). Each weight entry or in the second layer is initialized uni- formly at random from 1-1, As in
4: for (= ,...., do
5:
W k. - W
6:
for t= = 0, T do
OL mg (Win.)
Du et al. 2018 Zou el. al. 2020 Sollanolkotabi el. al. 2018 Oy- mak & Soltanolkotabi
7:
W A.141 t W ki
ow
end for AW W RT W
2019: Li et al. 2020
Oymak g
Soltanolkotabi 2020).
is fixed.
10:
end for
11:
for L....,, m do
Consider a subnetwork computing scheme with workers. In the k-th global iteration, we consider each binary mask MA {0. mm to be composed of subnetwork masks
12:
Wk+1r + Whrr + 1J,r Ep=1 Awk,
13: end for
14:
end for
m' (0, for I E [p]. The r-th entry of m is denoted mi with m = indicating that neuron r is active in subnetwork 1 in the kth global iteration, and m = 0 otherwise. We assume that the sampling of the masks for each neuron is independent of other neurons, and further impose the condition that the event m = 1 happens with fixed probability for all r and L. We denote this probability with & P (mk_--1) =1). The surrogate function defined by a subnetwork mask m is given by:
1 m
fmi (W,))
m Σ a a
With colored text, we highlight the differences between the full model and the surrogate functions. Consider the dataset given by (X.y) = {xx, y.)r We make the following assumption on the dataset:
Assumption 1. For aney i E [n], a holds that [|x:|l2 1 and 1ya S C - for some consland Moreover, for any j # holds that the points x,, Xj are not co-aligned, i.e., x; # (x, for any Ç E TR.
1.
This assumption is quite standard as in previous literature Du et al. 2018 Arora et al. 2019 Song & Yang 2020 We consider training the neural network using the regression loss. Given dataset (X, y), the function output on the whole dataset is denoted as S(W, X) [J(W, x:),..., J(W, (W.x.)) Then, the (scaled) mean squared error (MSE) of surrogate model is given by:
L m (W) y- Smm (W, (W.
The surrogate gradient is computed as:
Th (W) Ow,
n2 (fmg(W, (W, xx) ) y/ xxI(((w,)) x.) 2 0}.
Vm Σ aamke (fmt
Let be a constant subnetwork training learning rate, and let the aggregation weight 7/krr be zero neuron r is active in no subnetwork in the kth iteration; otherwise Thkr is set to the inverse of the number of subnets in which it is active. Within this sctting, the general training algorithm is given by Algorithm
4 Convergence on Two-Layer ReLU Neural Network
We assumee that my,- happens with a fixed probability for all k, r and I, and such probability is denoted by E. Consequently, the forward pass of the surrogate function is linear combination of E-proportion of
6

---

Published in Transactions on Machine Learning Research (08/2022)
the neurons' output. To keep the pre-activation of the hidden layer at the same scale for both the whole network and the subnetwork, we multiply the weight of the whole network with a factor of &. Due to the homogeneity of the ReLU activation, this is equivalent to scaling the output of each neuron, as in Mianjy
Arora 2020 For notation clarity, we define:
(4)
1
m
a a.Co( x.))
mm Σ
m
a (a.o((wr.:
Adding this scaling factor gives the property that EM, [mi (W x:) = U meaning that the sampled subnetworks in the global iteration are unbiased estimators of the aggregated network the global iteration A --1 1. Here ()) is both the initial whole network output in global iteration and the aggregated network output in global iteration 1. We focus on the behavior of the following loss, computed on the scaled whole network over iterations k:
La = Fly us 12 where us = 4
This is the regression loss over iterations & between observations y and the learned model ug.
Properties of subnetwork NTK. Recent works analyzing the convergence of gradient descent for
neural networks consider approximating the function output ug with the first order Taylor expansion Du
et al. 2018 Arora et al.] 2019 Song & Yang 2020). For constant step size q, taking the gradient descent's ss
(i.e., W W4+1 = W V w L(W))) first-order Taylor expansion, we get:
k+1 u + ( W Ws)
uu
(4)
En Σ H()), 4j).
(1)
where H(k) E _ax" is the finie-widthl NTK matrix of iteration k given by
H(k); (X., Σ Il x.) 2 O, (Wa.r, x²) 2 0}.
(2)
Compared with the previous definition of finite-width NTK, we have an additional scaling factor &. This is because, based on our later definition of masked-NTK, we would like the masked-NTK to be an unbiased estimator of the finite-width NTK. In the overparameterized regime, the change of the network's weights is controlled in small region around initialization. Therefore, the change of H(k) is small, staying close to the NTK at initialization. Moreover, the latter can be well approximated by the infinite-width NTK:
H (-Ew>N(U.)) [(x, ( x;) 2 0, x;) 2 0}}.
Du et al, 2018 shows that H is positive definite.
Theorem 1. Du et al. 2018
Denote do == Amin (H*), the ménimum eigenvalue of H Then we have
do > O as long as assumption
1 holds.
With H(k) staying sufficiently close to H Du et al. 2018 Arora et al, 2019 Song
Yang 2020 show
that Amin (H())) Z da > 0. Moreover, Equation implies that
2
u&+1 ug = -knH())((u= - y),
that further leads to linear convergence rate:
L = LA + (V un Le uk+1 - up) = Lb - En/ux - y. H(k)(u: y)) = (1 Endo) LR-
In NTK analysis, the Taylor expansion for both un and Ls produces an error term that improves the convergence rate from TiAg to 7m with 7 E (0, being a constant.
For our scenario, the randomly sampled subnetworks bring a trickier situation onto the table: in each iteration, due to the different masks, the NTK changes even when the weights stay the same. To tackle this difficulty, we provide generalization of the definition of the finite-width NTK that takes both the mask and the weight into consideration:
7

---

Published in Transactions on Machine Learning Research (08/2022)
Definition 1. Let my be the mask of subnetwork I in iteration W. We define the masked-NTi NTK in global ileration &s and local ileralion t. induced by m as:
m
(mb o H)k, m (x,, xj)
Σ
" xi) 2 O, (Wi,t,r x.) 2 0}.
Here, with colored text we highlight the main differences to the common NTK definition. Although we are only interested in the masked-NTK with = k', to facilitate our analysis on the minimum cigenvalue of masked-NTK, we also allow + k. We point out two connections between our masked-VTK and the vanilla NTK: i) the masked-NTK is an unbiased estimator of the whole network's NTK; ii when & { 1, the masked-NTK reduce to the vanilla NTK as in equation ( Throughout iterations of the algorithm, the following theorem shows that all masked-NTKs stay sufficiently close to the infinite-width NTK.
Theorem 2. Suppose the number of hidden nodes satisfies m = S2 (n² (os(KKm//)//AA)) If for all k, t it holds that [Wa.rr - Wu.rl12 & R
kdg then with probability at least 1 - 6, for all k. k E [K] we have: 8n
Amin(mi o H(k, t)) 2 Ag
The above theorem relies on the small weight change in iteration (k,t). Such assumption also made in previous work Duu et all 2018; Song & Yang 2020 to show the positive definiteness of the NTK matrix. In order to guarantee each subnetwork's loss decrease, need to ensure that the the weight change is bounded up to global iteration & (this implies that when subnetwork is sampled from the whole network, its weights do not deviate much from the initialization); and, úi) the weight change during the local training of the subnetwork also bounded. The following hypothesis establishes these two conditions, and sets up the "skeleton" " to construct different theorems, based on problems considered. The aim of this work is to prove this hypothesis for several cases.
Hypothesis 1. Fir the rumber of global iterations K. Suppose the number of hidden nodes satisfies m Q (n2 ppo((//////)))) and suppose we use a constant step size O (A0/u²). If the weight perturbation before ileration k is bounded by
|l - -o,rll2 + 2nTT
V
aK má
EiM 11) - u + (K k)N E)pn S R,
(3)
then, for all e [r], with probability at least 1 -8, we have:
y fmi (Wk W..++ x)
x)| 112 (1 1 77
) |1- 1 (W (Whe: X)
(4)
2
and the local weight perturbation satisfies:
|WEEr Wa. S
/r Vmn
Wo. [lly uxll2] + 2
2E(1-OpK md
(5)
The hypothesis above states that, in given global step k, given small weight perturbation guarantee (Equation 3) up to the current global iterations, each subnetwork's local loss also decreases linearly (Equa- tion D. as well as the weight perturbation remains bounded (Equation 5 Yet, the above hypothesis does not connect the subnetwork's 's loss with the whole network's loss through the sampling and aggregation process. Our aim is to turn Hypothesis [ into a series of specific theorems that cover different cases. In particular, we prove using induction the condition for which Hypothesis (1) holds under: i) masks with i.i.d. Bernoulli; and ii) masks with i.id categorical rows. Utilizing these results, we provide convergence results for the two scenarios. This is the goal in the following section.
4.1 Generic Convergence Result under Bernoulli Mask
While the local gradient descent for each subnetwork is guaranteed to make progress with high probability, when large network is split into small subnetworks, the expected loss on the dataset increases. Since EM, [mg xW), =uk simply expanding the MSE reveals that:
(4)
EMe -11x- - -mm (W&, x, x)][[] = lly gy-umll:
uul1 + EM [1/mmm (W,,xi) - mn
27
8

---

Published in Transactions on Machine Learning Research (08/2022)
When analyzing the convergence, the second term on the right-hand side needs to be carefully dealt with. It is non-trivial to show that, when combining the updated network of the local steps, the loss computed on the whole network is smaller than or equal to the error of each sub-network. We will solve these technical difficultics for the training procedure with subnetworks created using masks sampled from two types of In this section, we locus on masks satisfies the following Bernoulli assumption:
Assumption 2. (Bernoulli Mask) Each mask entry m is independently from Bernoulli distribution with meon &, i.e,, m Bern(E).
Masks sampled in this fashion allow a neuron to be active in more than one subnetworks, or none of the subnetworks. For convenience, we denote the probability that a neuron is active in at least one subnetwork with 0 = 1 - (1 g) In the meantime, subnetworks created using Bernoulli masks enjoy full independence, and thus have nice concentration properties. By carefully analyzing the aggregated gradient of each local step, we arrive at the following generic convergence theorem, under the Bernoulli mask assumption.
Theorem 3. Let assumptions i and 2 hold. Then do 0. Fir the number of global iterations to K and the number of local ileralions to T. Let the number of hädden neurons salisfy:
12
K max
n²
nK²B1
K² ko0))
(6)
RR2E0x0"
k²2A2
Then Algorithm according to:
An with a constant step-size 1) = maxx(n.pjur
converges with probability at least 5,
E M [lly
us S (1 110T220 Ily uoll 112 B1.
(7)
for some error region level B1 > 0, defined as:
B1 = (1- mX + (0 60-E2mmn + (1 (1= (1 - E)mn3)
(8)
Overall, given the overparameterization requirement in Equation 6 the neural network training error, as expressed in Equation 7 drops linearly up to a neighborhood around the optimal point, defined by B₁ in Equation We notice that B1 has three terms that all reduce to zero when & =1 In the first term of B1, m appears in the denominator, implying that this term can be arbitrarily decreased as the cost of increasing the number of hidden neurons. The second term is kept at constant scale, long the initialization scale is small enough. The third term disappears when the number of local steps is one. However, the loss decreases more in cach global itcration, when 7 is larger, since the convergence rate is 1 Oo) In the case of { = 1,p = 1 and 7 = 1, the proposed framework reduces to the whole network training. Choosing K = 1, Theorem reduces to a form similar to Song & Yang 2020 with the same convergence rate and wwe-pparametcrization requirement.
Remark. Compared Lo .uu el al. 2018 Song & Yang 2020,, the scenario considered in our work involves additional randomness introduced by the mask. Thus, convergence result is based the expectation of the loss: we derive the bound of the loss from the bound of its expectation, using concentration incqualities, and apply a union bound over all iterations k [K]. Therefore, the required prraattteeetiiiioo on m grows as we increase the number of global iterations, meaning that, under a fixed m, the convergence is only guaranteed for a bounded number of iterations. This is not a concern in general since to guarantee small training error only need K to be log +logn This is termed as early-stopping, and is used in previous literature [Su & Yang 2019 Allen- Zhuu al. 2018).
The complete proof of this theorem is defered to Appendix E and sketch the proof below:
1. Let Xa. = Ed mp, denote the number of subnetworks that update neuron in global iteration k. Let Nk.r max max(Xar: 1} to be the normalizer of the aggregated gradient, N min(X 1}
9

---

Published in Transactions on Machine Learning Research (08/2022)
to be the indicator of whether neuron is selected by at least one subnetwork. Then, the update of each weight vector can be written as:
N
k.r OL.
7-1 P
(
k,t.r)
(9)
Wk+1, = Wk.r 17 Nkrr
Ow,
ml (W (Wi...))
No. 2. We first focus on the aggregated gradient of the first local step Nu.- that this aggregated gradient satisfies concentration property around point near the ideal gradient Ot().) and such concentration closer if p larger. Ow,
Ei=
and show
dw.
3. We notice that the difference between the aggregated gradient in the later local steps and the aggregated gradient the first local step depends on how much the local weight of each subnetwork in the later local step deviates from the weight of the first local step. We then show that the local weight change is bounded, implying that the aggregated gradient in all local steps lie near to the aggregated gradient in the first local step.
4. Lastly, use the standard NTK technique to show that the aggregated gradient update in equation 9 leads to linear convergence per each global step, with additional error term.
To interpret the theorem, we choose = n make mild assumptions and simplify the key messages of the form in Theorem (3. Note that this choice of is the same as (Arora et al. 2019
Assumption 3. For the simplicily of our theorem, assume that max { [K, d, S and Ao S 1.
Notice that for all m that satisfies equation 6), the first term in B₁ is upper bounded by O(1). Moreover, n 2, the second and third term are also upper bounded by 0(1). Therefore, Bi is upper bounded by O(1). Moreover, since Xo 2 I and max(K, n, we have that both nk*B B, and K2p are smaller than <²g0x; so the over-parameterization requirement equation (6) x20x3
since p 1 and 2 1, by choosing
n²
reduces to =²K' For different choice of and p, our considered scenario reduces to different existing algorithms. In the following, we provide convergence results of these algorithms, as corollaries of Theorem by considering different and p values.
5geX3
Dropout. The dropout algorithm Srivastava et al. 2014 corresponds to the case 1,p 1. For this assignment, we arrive at the following corollary.
Corollary 1. Let assumptions D and 3 holds. Fir the number of dropout ilerations K, the step size to 17 O (No/n²), and let the number of hidden neurons satisfies θ (n² K K/<² Then, the dropout algorithm on two-layer ReLU neural network converyes with probubility at least 1 8, according to:
M [lly -u-- 1111 < (1-
1 (1E0) ly- uoll 1 +0
Typically, 1- is usually referred to the "dropout rate". In our result, & approaches 0, which corresponds to the scenario that neurons are selected, the convergence rate approaches 1, meaning that the loss hardly decreases. In the mcan time, the error term remains constant. On the contrary, as & approaches 1, which corresponds to the scenario that all neurons are selected, we get the same convergence rale of 1 - O (1)A0) as in previous literature (Du et al. 2018 Song & Yang 2020 , and the error term decreases to 0. Moreover, we should note that the over-parameeerrzztion requirement also depends on ç. In particular, as g becomes smaller, we need a larger number of hidden neurons to guarantee convergence.
Multi-Sample Dropout. The multi-sample dropout [nooc 2019 corresponds to the scenario where 1,p 2 1. Our corollary below indicates how increasing p helps the convergence.
Corollary 2. Let assumptions D. D and 3 hold. Fix the number of dropout iterations K, the step size to 7 O (No/22), and let the number of hidden newrons satisfy θ (n (n2K/E0x6)). Then the p-sample
10

---

Published in Transactions on Machine Learning Research (08/2022)
dropout algorithm on two-layer U neural network converges with probability at least 1 - 6, according to:
1 E M [ly up I 1 1444" noxuu)
10--5 nk P
g)
0 5²) )
I uoll +0
Recall that 0 == 1 (1 E)2 denotes the probability that neuron is selected by at least one subnetwork. Based on this corollary, increasing the number of subnetworks improve the convergence rate and the over-parammelrization requirement since increases as p increases. Moreover, increasing the number of subnetworks help decreasing the error term even when the dropout rate & is fixed. After p is as large as nK, the error term stops decreasing, dominated by the termn O (1=s)²), Lastly, compared with the result of dropout, the wet-paramettrriztion depends not only on E, but also on 8.
,
Multi- IST. The multi-worker IST algorithm Yuan et al. 2020 is very similar to the general scheme with 2 1 and t 2 1, but with the additional assumption that max d, p} & n, and a special choice of initialization k k= 2.
Corollary 3. Let assumptions i and 3 hold. Fir the namber of dropoul ilerations K, the step size to 1j O (Nafee max(n.p)). and let the number of hidden neurons satisfy m θ (n' /g0xg5)). Then the IST algorithm on two-layer ReLU neural network converyes with probobility at least Ö, according to:
) byy woll +0 [lly - 1 S (1- 1 4 - 100Tho 1I
(
(1-E²- nk
8 E
(1-2) 0 - ((()))
P
While IST with subnetworks constructed using Bernoulli masks allows a neuron to be active more than one or none of the subnetworks. In the next section, we consider another mask sampling approach that fits better into the scenario of the original IST, where each hidden neuron is distributed to and only one subnetwork with uniform probability.
4.2 Multi-Subnetwork Convergence Result for Categorical Mask
We consider masks sampled from categorical distribution, explained by the assumption below: Assumption We assume that M2 Categorical(p). To be specific, for each r E [m], let let Unif(pp)), and we define mh_r 1 if,=!, and m otherwise.
In this way, the masks endorsed by each worker are non-overlapping (as stated in Yuan el al. 2020 and the union of the masks covers the whole set of hidden neurons. However, we note that the subnetworks created by the masks sampled according to this fashion are no longer independent. The following theorem presents the convergence result under this setting.
Theorem 4. Let assumplions T and (4) hold. Then No > 0. Moreover, let Amax denote the maximum eigencalue of H Fir the number of global iterations to K and the number of local iterations to T. Let the rumber of hidden neurons be m Ω n = and choose the inilialization scale K = vnAmmax251
X
Lel 7 = (1-p-1) Then, Algoritlom 10 with a consland slep-size O probability at least 6, according to:
An min 11))
converges with
E M [|ly uall < (,+(1-2) (1-2)22))22)))) I1y 111- +0
n700 ))
?Tnk? Amax
X
We defer the proof of this theorem to Appendix F This theorem has couple noticeable properties. First, when the number of workers p = 1, i.e., the scenario of multi-worker IST reduces to the full-network training, we achiove ? = 0, which implics that the error term disappears, driving further connections between regular and IST training. Second, when the number of subnetworks p increases, y also increases, leading slower decreasing of the training MSE and convergence to bigger error neighborhood. In particular, as increase
11

---

Published in Transactions on Machine Learning Research (08/2022)
the number of subnetworks, the number of active neurons in each subnetwork becomes smaller, which makes the subnetworks both harder to train and harder to synchronize. We defer the complete proof of this theorem to Appendix F and sketch the proof below:
1. Let 1 Smi W k X
We notice that Ei= fmt. Using this property, show that
Lk+1 = Ei= lly ú 2 Ei=1 ux+1 starting from an initial value of lyy o020
// 12. The first term here enjoys linear convergence
It then follows that EM. 1 Ei by ú | La L Ei=1 1
úf. Putting things together,
we have that EM, [L++1] $ (1-a1 (1-a) L + 4k, where = EMA 1 Ei= |
ú 12 / ú 12)
where 1) is convergence rate achieved by invoking Hypothesis
3. We then use the small weight perturbation induced by over-paramederiiation,, which means that 1 o ll2 is small. This allows us to bound the term 4. and arrive at the final convergence.
5
Experiments
Probability:
Probablitty 0.5
0175
Probability:
Training Erro
Submetwores
0125
10-1
0100
0.075
0.050
0.025
(2)
10-
0.4
0.5
0.8
0.
10
Total Training Steps (a)
Probebility
(b)
Probabilty
Probabilit: 0.5
2nd Global Step
10
3rd Global Step
Ineining Chaagge
Training Erro
Network Width
Subnetwores
(c)
(d)
Figure& 2: Validation experiments on a single hidden layer perceptron.
We empirically validate main theorems one- one-hidden-layer. ReLU- activated neural network. With the purpose of performing the experiments a task that is both widely used and representative, and possess some simplicity to be solved by the one-hidden-layer MLP, we use the features/embin /embeddings extracted from a convolutional--esee neural network Krizheusky 2009). This is common practice: E.g., the recent work in Chowdhury et al. (022 proposes the use of a library of pre-trained networks to extract useful features, which later on are processod by added topmost layers used for classification. This results a procodure that takes an image, creates an embedding, and then uses that embedding build classifier, by feeding the embedding into multi-layer perceptron with single/ multiple hidden layers. In this work, take
12

---

Published in Transactions on Machine Learning Research (08/2022)
ResNet-50 model [He et al. 2015 pretrained on ImageNet as our feature extractor and concatenate it with two fully-connected layers. We then train this combined models on the CIFAR-10 dataset, and take the outputs of the re- re-trained ResNet-50 model as the input features, and use the logits output of the combined model as the labols. The obtained input feature has a dimension of 2048, and we choose a constant learning rale and a sample size of 1000.
In Figure Za we plot the logarithm of the mean and variance of the training error dynamic with respect to the K (for clarity, we only plot the first 160 iterations), which includes the sampling step, local training steps, as well as the gradient aggregation step. Since the algorithm converges with a stable/simiae manner across all trials, the variance is too small to be observed from the figure. Notice that there are three types of dynamics, as annolated in the figure: (1) A smooth decrense of training error: This corresponds to subnetworks' local training, which is supported by our theory that each subnetwork makes local progress. (2) A sudden decrease of training error: This corresponds to the aggregation of locally-trained subnetworks, and is consistent with our proof in Theorem 4 (3) A sudden increase of training error: This corresponds to re-sampling subnetworks; according to our theory, the expected average training error increases after sampling.
Figure 2b provides heatmap results that demonstrate the change of the error term as we vary the number of subnetworks, and the selection probability. In Figure 26 the subnetworks are generated using Bernoulli masks, and the training process assumes a fixed number of local steps. Note that, as we fix the number of subnetworks and increase the selection probability, the error decreases (lighter colors in heatmap). Moreover, if we fix the number of selection probability and increase the number of subnetworks, the training error also decreases. This is consistent with Theorem 3
Figure 2c studies how the error term changes as we increase the width of the neural network. In Theorem we notice that the error term decreases as we choose a smaller initialization scale K, while a smaller K would require a larger wen-parameteriizaionn. This is consistent with our experiments in figure 2 as we increase the number of hidden neurons and adjust the initialization scale, we observe that the training converges to a smaller error.
Figure 2d shows how the convergence rate changes we increase the number of subnetworks under the categorical mask assumption. In particular, the y-axis denotes the training error improvement in the first, second, and third global step, respectively. The training error improvement is defined to be the training crror in last step-trainn training error in current step, We observe that the training error improvement decreases training
in last, step
consistently across the first three global steps, as we increase the number of subnetworks. This corresponds to what we have shown in Theorem 44, where - 7 decreases as we increase p.
6 Conclusion
We prove linear convergence up to a neighborhood around the optimal point when training and combining subnetworks in a single hidden-layer perceptron scenario. Our work extends results on dropout, multi-sample dropout, and the Independent Subnet Training, and has broad implications on how the sampling method, the number of subnetworks, and the number of local steps affect the convergence rate and the size of the neighborhood around the optimal point. While our work focus on the single hidden-layer perceptron, we consider multi-layer perceptrons as an interesting direction: we conjecture that a more refined analysis of each layer's output is required (Du et al. [209 Allen- Allen-Zhu et al. [20196). Moreover, focusing on the convergence of a stochastic algorithm for our framework, as well as considcring different losses (c.g., classification tasks or even generic generalization losses) are interesting future research directions. Lastly, training with randomly sampled subnetwork may result into a regularization benefit. Theoretically studying how the proposed training scheme affects the generalization ability of the neural network is interesting next step.
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers, 2018. URL
04918
13

---

Published in Transactions on Machine Learning Research (08/2022)
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural informalion processing systems, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th Inter- national Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 242-252. PMLR, 09-15 Jun 2019b. URL http:://ppoceeddiggg/ mlr http:/mmmrr.r.ss/e/eeeee ppres/v97/allee 2hu1la. html
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In International Conference Machine Learning, pp. 254-263. PMLR, 2018.
Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks, 2019.
Devansh Arpit, Slanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kan- wal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization in deep networks. In International Conference on Machine Learning, pp. 233-242. PMLR, 2017.
Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spertaally-normalized margin bounds for neural networks. Advances in Neural Information Processing Systems, 30:6240-6244 2017.
Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: statistical viewpoint. ar.ivv preprint arX:::2108.091 2021.
Mikhail Belkin. Fit without fear: remarkable mathematical phenomena of deep learning through the prism of interpolation. arXiv preprint 1005: 2021.
Mikhail Belkin, Daniel Hsu, and Partha Mitra. Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate. Advances in Neural Information Processing Systems, 31:2300- -111100-231,,
2018.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias-variance trade-off. Proceedings of the National Academy of Sciences, 116(32):15849- 15854, 2019.
Jinming Cao, Yangyan Li, Mingchao Sun, Ying Chen, Dani Lischinski, Daniel Cohon-Or, Baoquan Chen, and Changhe Tu. Do-conv: Depthwise over-paramelerized convolutional layer. arXio preprinl arXiv:2006.1900 2020.
Arkabandhu Chowdhury, Mingchao Jiang, Swarat Chaudhuri, and Chris Jermaine. Few-shot image classi- fication: Just use a library of pre-trained feature extractors and a simple classifier. In Proceedings of the IEEE/CVP International Conference on Computer Vision, pp. 9445 9454, 2021.
Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. Binaryconnect: Training deep neural net- works with binary weights during propagations. In Advances in neural information processing systems, pp. 3123-3122211 0031,,1,,2110 2015.
Yuyang Deng and Mehrdad Mahdavi. Local sgd optimizes overparameterizee neural networks in polynomial time. arXiv preprint arXiv:2107.10860 2021.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 1675 1685. PMLR, 09 15 Jun 2019. URL
mlr wwwpsp//9/euullcccc.mm html
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over- parameterized neural networks. In International Conference on Learning Representalions, 2018.
Chen Dun, Cameron R. Wolfe, Christopher M. Jermaine, and Anastasios Kyrillidis. decomposition of resnets for distributed training, 2021.
Resist: Layer-wise
14

---

Published in Transactions on Machine Learning Research (08/2022)
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint WrXii:1703.11000 2017.
Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang. Modcling from features: a mcan-field framework for over-parameterrzeed deep neural networks. Conference on Learning Theory, pp. 1887 1936. PMLR, 2021.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, arXiv preprint arXiv:2101.0390 2021.
Vitaly Feldman. Does learning require memorization? a short tale about long tail. In Proceedings of the 52nd Anseual ACM SIGACT Symposium on Theory of Computing, pp. 954 959, 2020.
Spencer Frei, Yuan Cao, and Quanquan Gu. Algorithm-dependent generalization bounds for overparame- terized deep residual networks. Advances in neural information processing systems, 2020.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050 1059. PMLR, 2016.
Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size-independent sample complexity of neural net- works. In Conference On Learning Theory, pp. 297-299. PMLR, 2018.
lan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Yihong Gu, Weizhong Zhang, Cong Fang, Jason D. Lee, and Tong Zhang 0001. How to characterize the landscape of overparameterized convolutional neural networks. In Hugo Larochelle, Mare' Aurelio Ran- zato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Newral Information Processing Systems 33: Annual Conference Newral Information Processing Systems 2020, NewrIPS 2020, December 6-, 2020, virtual, 2020. URL htttps:/pproceeddngg/ I70+5B22D00655540022111047774 Awstrdd-bbbttt.....
Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning, 2019. URL https://arxiv. wwtp:/ivw.rgg////111..//. 14425
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. URL hbttps///arxxivoorr//////1/2
Baihe Huang, Xisoxiao Li, Zhao Song, and Xin Yang. Fl-ntk: A neural tangent kernel-based framework for federated learning convergence analysis, 2021.
Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks generalize better than deep feedforward networks? a neural tangent kernel perspective. Advances in Neural Information Processing Systems, 33, 2020.
Hiroshi Inoue. sulti-sample dropout for accelerated training and better generalization. arXie preprint arXiv:1905.09788 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural works In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 8589, 2018.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks, 2020.
Kenji Kawaguchi, Leslie Pack Kaelbling, and Yoshua Bengio. Generalization in deep learning. orXiv preprint arXiv:1710.055 2017.
15

---

Published in Transactions on Machine Learning Research (08/2022)
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First analysis of local gd heterogeneous data, 2019a. URL htttp:://aarxtw.orr//b////////// 04715.
Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local sgd on identical and heterogencous data, 2019b. URL httpp::/arxx...or//ab////// 04476
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems, 28:2575-2558 2015.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Alex Labach, Hojjat Salehinejad, and Shahrokh Valaee. Survey of dropout methods for deep neural networks. arXiv preprint ar 1904. 13310, 2019.
Daniel LeJeune, Hamid Javadi, and Richard Baraniuk. The implicit regularization of ordinary least squares ensembles. In International Conference on Artificial Intelligence and Statistics, pp. 2020.
PMLR,
Daniel LeJeune, Hamid Javadi, and Richard Baraniuk. The flip side of the reweighted coin: Dual- ity of adaptive dropout and regularization. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.). Advances in Neural Information Processing Systems, volume 34, pp. 23401-23412. Curran Associates, Inc., 2021. URL https://ppoocedinngg.eerrrpppppp nwwrine F111://888002557000r000011166682 .pff
Dmitry Lepikhin, HyoukJoong Lec, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint crXiv:2006.1663 2020.
Mingchen Li, Mahdi Soltanolkotabi, and Samet Oymak. Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks. In International conference on artificial intelligence and slalislics, pp. 4313 4324. PMLR, 2020.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg on non-iid data, 2019. URL http:::t////:te /t///rxxiiva// wttp://wgggbbbb/199....
Tengyuan Liang, Tomaso Poggio, Alexander Rakhlin, and James Stokes. Fisher-Rao metric, geometry, and complexity of neural networks. In The 22nd International Conference on Artificial Intelligence and Slatislics, pp. 888 896. PMLR, 2019.
Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean field analysis of deep ResNet and beyond: Towards provably optimization via overparameterization from depth. In International Conference on Machine Learning, pp. 6426 6436. PMLR, 2020.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD modern over-parametrized learning. In International Conference on Machine Learning, pp. 5325-3334 PMLR, 2018.
Ryan Medonald, Mchryar Mohri, Nathan Silberman, Dan Walker, and Gidcon S Mann. Efficient large- scale distributed training of conditional maximum entropy models. In Advances in Neural Information Processing Systems, pp. 1231-1239 2009.
Poorya Mianjy and Raman Arora. On convergence and generalization of dropout training, 2020.
Amirkeivan Mohtashami, Martin Jaggi, and Sebastian U Stich. Simultancous training of partially masked neural networks. Xriee preprint arXiv:2106. 08895, 2021.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring generalization in deep learning. Advances in Newral Information Processing Systems, 30:5947-599 5956, 2017.
16

---

Published in Transactions on Machine Learning Research (08/2022)
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-Bayesian approach to spectrally- normalized margin bounds for neural networks. In International Conference on Learning Represenlations. 2018.
Quynh Nguyen. On the proof of global convergence of gradient descent for deep relu networks with lincar widths, 2021. URL httpp/s///xxx..orrr//////.1....0
Samet Oymak and Mahdi Soltanolkotabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning, pp. 4951-4960. PMLR, 2019.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization:: Global convergence
guarantoes for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, 1(1):84 105, 2020.
Benjamin Recht, Rebecca Roclofs, Ludwig Schmidt, and Vaishaal Shankar. Do ImageNet classifiers generalize to ImageNet? In International Conference on Machine Learning, pp.
PMLR, 2019.
Albert Senen-Cerda and Jaron Sanders. Asymptolic convergence rate of dropout on shallow linear neural networks, 2020a.
Albert Senen-Cerda and Jaron Sanders. Almost sure convergence of dropout algorithms for neural networks, 2020b.
Vatsal Shah, Soumya Basu, Anastasios Kyrillidis, and Sujay Sanghavi. On generalization adaptive methods for over-parameteriee linear regression. arXiv preprint orXiv:2011.1006, 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXie preprint
arXiv:1701.06538 2017.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Leo. Thcoretical insights into the optimization land- scape of over- aver-parameterizee shallow neural networks. IEEE Transactions on Information Theory, 65()): 742 769, 2018.
Chaehwan Song, Ali Hamezani-Kebrye Thomas Pethick, Armin Eftekhari, and Volkan Cevher. Subquadratic overparameterization for shallow neural networks, 2021. URL http://awwvv..rr//bb11111.....1/
Zhao Song and Xin Yang. Quadratic suffices for over-parametrization via matrix chernoff bound, 2020.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on scparable data. The Journal of Machine Learning Research, 19(1):2822-2878, 2018.
Stefano Spigler, Mario Geiger, Stéphanc d'Ascoli, Levent Sagun, Giulio Biroli, and Matthicu Wyart. A jamminn transition from under-to over-parametrriatioon affects generalization deep learning. Journal of Physics A: Mathematical and Theoretical, 52(47):47,001, 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1): 2014.
Sebastian U. Stich. Local sgd converges fast and communicates little, 2018. URL https://arxiv.org/ass
09767
Lili Su and Pengkun Yang. On learning over-parameterized neural networks: A functional approximation perspective, 2019.
Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Ge- offrey Gordon. An empirical study of example forgetting during deep neural network learning. In International Conference on Learning Representations, 2018.
17

---

Published in Transactions on Machine Learning Research (08/2022)
Thijs Vogels, Sai Praneeth Karinireddy, and Martin Jaggi. PowerSGD: Practical low-rank gradient compres- sion for distributed optimization. Advances In Neural Information Processing Systems 32 (Nips 2019), 32 (CONF), 2019.
Li Wan, Matthew Zeiler, Sixin Zhang, Yann Le Cun, and Rob Fergus. Regularization of neural networks using dropconncct. In International conference on machine learning, pp. 1058-1066. PMLR, 2013.
Hongyi Wang, Saurabh Agarwal, and Dimitris Papailiopoulos. Pufferfish: Communccation-effcceet models no extra cost. Proceedings of Machine Learning and Systems, 3, 2021.
Sida Wang and Christopher Manning. Fast dropout training. In international conference on machine learning, pp. 118-126, 126. PMLR, 2013.
Cameron R Wolfe, Jingkang Yang, Arindam Chowdhury, Chen Dun, Artun Bayer, Santiago Segarra, and Anastasios Kyrillidis. GIST: Distributed training for large-scale graph convolutional networks. arXiv preprint 1r11c:2102.10424 2021.
Tianyi Yao, Daniel LeJeune, Hamid Javadi, Richard G Baraniuk, and Genevera Allen. Minipatch learning implicit ridge-like regularization. In 2021 IEEE International Conference on Big Data and Smart Compating (BigComp), pp. 65 68. IEEE, 2021.
Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable neural networks. International Conference on Leorning Representations, 2018.
Binhang Yuan, Cameron R. Wolfe, Chen Dun, Yuxin Tang, Anastasios Kyrillidis, and Christopher M.
Jermaine. Distributed learning of deep neural networks using independent subnet training, 2020.
Ce Zhang and Christopher Ré. Dimmwitted: A study of main-memory statistical analytics. Proceedings of the VLDB Endoument, 7112):1283-12M)) 2014.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking gencralization. Communications of the ACM, 64(3):107-115, 2021.
Jian Zhang, Christopher De Sa, Ioannis Mitliagkas, and Christopher Ré. Parallel SGD: When does averaging help? Xixiv preprint ar Xiv::11606. 07365, 2016.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P Adams, and Peter Orbanz. Non-vacuous generaliza- tion bounds at the ImageNet scale: PAC-Bayesian compression approach. In International Conference on Learning Representations, 2018.
Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In Advances in neural information processing systems, pp. 2595-2603, 2010.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over-parameterized deep ReLU networks. Machine Learning, 109(3):467 492, 2020.
18

---

Published in Transactions on Machine Learning Research (08/2022)
A
Notation
Table 1: Notations
SYMBOL
DESCRIPTION
MATHEMATICAL DEFINITION
K
Number of global iterations
K K KN
Index of global iterations
k E [K]
Number of local iterations
E N+
Index of local iterations
t E [T]
Number of subnetworks
p E N+
Index of subnetworks
I E [p]
Probability of selecting neuron
& E (0, 1]
Vector probability selection a
9 E (0, 1l2
neuron by each worker
Constant step size for local gradient update
7 E R
Binary mask in iteration &
M, E {0, 1
ma.
Binary mask for neuron in
marr E {0, 1J", the vector of rth column of M,
iteration k
m
Binary mask for subnetwork in iteration
m E 1)m the vector of Ith row of M
m
Binary mask for neuron in subnetwork iteration & Number of subnetworks
m {0, 13 the (1,r)th entry of M2
selecting neuron r in iteration Xk.
X = Ei= m
Ne.
Aggregated gradient normalizer for neuron r in itcration
NR.- max (N 1}
Indicator of gradient existing for neuron in iteration &
N min (X 13
Global gradient aggregation step size for neuron r in iteration Output of the whole network at
Ma.r NT.//Ner
uk
global iteration & for sample Output of the whole network at
Ub
(i)
vm
Em ((Wk(W
(1) (1)
(11)
Ug
global iteration & for all X Output of subnetwork I at iteration (k:, 1) for sample Output of subnetwork 1 at iteration (A:, 1) for all X Output of subnetwork i at itcration (ke, 0) for sample Output of subnetwork at iteration (k, 0) for all X
u
Ub ,...,,
üg,t
.((()
1(6) 0
Er aym' (wh
(wi))
vm
uf,
ūk,r
0x,t
((1)
..... Ukkt
[(x)
(())
1(())
ūg.0
ú/.
Li
Global loss at iteration
L = l|y us 1
Lmo (Wl (W))) k
Local loss for subnetwork at iteration (k, t)
m (WL k.)) lly -Smi (W.))
19

---

Published in Transactions on Machine Learning Research (08/2022)
B Preliminary and Definition
In the proofs of our theorems, we use extensively the following tools. Definition Gaub-Gaussinn Random Variable) A random variable X 122 -ub Gaub-Gaassinn if
S +
Definition 3. Sub--aaponential Random Variable) A random variable with meon E[X] = is (N', a)-sub- exponential if there exists non-negative (r'', a) such that for all S a
E[e4x-H1] s 22"
Property 1. (Sub-Eapor Eob-Eerpoeenial Tail Bound) For (k' sub we have
(lsb eaponential random variable X with E μ
Z
if
P(X > ++t)s
2
20
ift>
Property 2. (Markov's Inequality) For non-negalive random variable X, we have
P(X o) E[X
Property 3. (Hoeffding's 's Inequality for Bounded Random Variahles) Let XI...., Xa be independent random variables bounded by [Xi] S 1 for all E [n]]. Then we have
Se 21
-
22
x 21
2n12
Property 4. (Berstein's Inequalily) Let XI,..., Xn be random variables with E[[,]] 0 for all E [n]. If [&J] S M almost surely, then
2/ Se E a
P
Σ
X
Property 5. (Jensen's Inequality for Expectation) For a non-negative random variable X, we have
x S (B[X))
Apart from the properties above, we also need the following definitions to facilitate our analysis. First, we note that, in the following proofs, we let R
xλo
Define
192n
A (3w B(Wa,r: R) I[(WW x,) 2 03 * If x,) 0)}
to denote the event that for sample x,, the activation pattern of neuron may change through training if the weight vector change is bounded in the R-ball centered at initialization. Moreover, let
S = fr E [m] -
S = [m] \ St
to the set of neurons whose activation pattern does not change for sample x; the weight vector change is bounded by the R-ball centered at initialization. Moreover, since we are interested the loss dynamic computed on the following function
m
Σ a (WW,,,
mm
we denote the full gradient of loss with respect to each weight vector w,
(T()
{ Σ
(u
--- ( 2 20
Ow.
Vm
20

---

Published in Transactions on Machine Learning Research (08/2022)
C
Proof of Theorem 2
Recall the definition of masked-NTK
(m H(;, (x,, xz) mi If (Wk,nr; xx) 2 0, (We.er x;) 2 0}
To start with, we fix k' E [K] [p], and i E [nt]. In this case, let
hr ma. (x,,,)) (wo,. xx) 2 0, (Wo,r, x;) 20}
Then we have
(ml HHH,, 0)),) m Σ hy
Also we have
H Note that for all r we have | S 1. Thus we apply Hoefiding's inequality for bounded random variables and get
EM w [hr]] (o [M,, [h+]]
_ hr H
P (I H(O 0)),; H 2 P
2 1 S 2e 2mmt
Apply union bound over i gives that with probability at least 1 e 2mmt? it holds that
[mb (0,00)); H <
for all i,j E [n]. Therefore,
m o H(0.0) 0) H < m H(,,0) - 2
(ml H(0,0)), H 2
1=
< 2
Let Ag 3n gives
Au
m H(0,0) - H
holds with probability at least 1 1 2n a.f
Next we show that for all A [k] and t [r], as long
|Wkeer Wo I R for all E [m], then it holds that
| H(k, t) m H ((O,0)]l, s 2nx R
Following the argument of Song 6 Yang 2020
lemma 3.2, we have
2
m H(k, m H(000)
Σ
Σ Sp.id
i-=1
with
Sr.i.j my (If(Wo,r, x;) 2 0: (Wo,r: x;) 2 0} I[(Wr,,...;) Xj) 2 0; (Wr,k,, 2 0))
21

---

Published in Transactions on Machine Learning Research (08/2022)
Then $y,i. 0 if -Ai and -Ajr happend. In other cases have |sr,agl $ 1. 1.J e [n]
Thus have that for all
4ER kv2n
MM, Wo = =P (Air
S 2Êk 2 -R R
and
4ER 2Ek
BM, Wo (Srg.) E. www.woolsss. woldrs)))))
EM, (s [
KV 2m
Thus applying Bernstein inequality with = Ex -1-- gives
m
S cxp (-
m
P
Sta.j 2 3 -R R
m
10k
Therefore, taking a union bound gives that, with probability least n²² meR
we have that
m H(k,t) m H(O 0) S m Hk: t) m ((0,0)]l S 3[nx-1p R
ed Kdo gives || o H(k,t) m H(0 0) - T Using R S 12n Therefore, we have
S
da with probability at least 1 n²²
mEAn
12m
Ao 2
m H(k,t) H 1|
which implies that Amin (mp H(k, 1))
2 A 2
holds with probability at least
n²
(
e
m5A0
2
mn
on
for a fixed K' E [K] and I E [p]. Taking a union bound over all and I and plugging in the requirement Ω x² log Kpn/s) gives the desired result.
εAo
22

---

Published in Transactions on Machine Learning Research (08/2022)
D Proof of Hypothesis 1
In this proof, we follow the idea of [Du et al. NTK during the analysis, and ii) we use different technique for bounding the weight perturbation. We repeat the key requirement stated in the theorem here: for all r E [m]]
2018 However, the difference is that i) we use our masked-
nk E P Wo.a Illy - uxlla] + (K k)N E(1 E)pn S R
|w.r wor ll2 + 2nT
(10)
mo
To start, we notice that, using the required over- over-parameteriaatinn lemma 23 holds with probability at least 0(6). We use induction on the following two conditions to prove the theorem:
I up 2..4
(1
n7d0
||y uk.1112
(11)
2
- 17T 2n
/ - E)pK
|wwer - wkrll
E S Wo,a [lly U1 [x + 217 /
(12)
vmd
||wlerr - worrl S R
(13)
Base Case: For the case of O, we notice that equation I1 and 12 naturally holds. Moreover, equation implics equation 112
10
Inductive Case: the inductive case is divided into three parts.
12
13): assume that equation 12 holds in local iteration t. Combine the result with equation
10
gives that equation 13h holds in iteration t.
13 :: assume that equation 13 holds local iteration We are going to prove that equation
holds. In particular, we are interested in
ly 1 = lly 11 - 2 2(y wh. ah.44 - akee) + 110h, a/
We define 0 2,,+1
11s) 0
=T((1) + 2,,, with
((0)
1 Σ ammlrr (a ((wkett), -a0--a((ww ((wker x:))
1,k.t
Vm
10)
a,mb (o (( ((Wke ((wk,,,):::))
2
m
rESt
and notice that, with the 1-Lipchitzness of ReLU,
1(0)
Σ lo ((wh.- x:)) ((wherr:))))
2,
mm
rESt
|h,, -k.erll2
v m
rES.
W ket aL
Σ mm
Owr
rES!
n n
||y
m
rESt
< 4mp
VnRlly
0/ 12
23

---

Published in Transactions on Machine Learning Research (08/2022)
where the last incquality uses |.S//| | x 4mx ' R from Lemma 16 Therefore,
K u/e.0 I2 k.,, S v
max ie(m)
f(()) 2,k,t
Ily - o2,cll2 S 41/ss "nlelly -0f,c112
Similarly, we have
((0)
(((0)
1
Σ whetr.r wharlls
2R,4+1
ukit
772
whi,tir we.sorl2 11
2
Tl.m! (Wh..
k.t)
S n7² Σ
Ow,
$ n7"n>1ly fly
11,,112
Lastly, we define m H(: 1)- with
1 xx) rES!
(mk t() m
Σ
( xi) 2 0, (,
x,) 202 0}
and we have
100 =
mm Σ
ammier (WE,tLY - wherr xx)
20
1,
rES:
0Lm! (Whs))
7
ar
xx I((WW::: (x)
20 20)
m
Ow,
rES.
1
Σ
mkr Uj - ((11)) (xi, x1) 11 (Wkir) 2 0, (Wiir x1) 0}
m
rES, j=1
= 17 Σ
(m/ tkk - -m a H(k))))) (us -ah.))
Therefore,
(y ahe 1 =1 Σ
I(1)) y/ ux.²) ) (ml. H(k, 1) mn H(k:
(16))
(1s
i,j=1
17 (y ûc.t (mk H(k, () m² H(k, 1) 1) (y an.e))
17A0
2 |ly -i - mm| H((k () 112lly -/,,112
11200 Ayp 'nn] ||y I o2.111
where the last incquality follows from the fact that
|m o (kk < m H(k
2
(x)))))
I[ x;) 200 (Wk,,, x,) 2 20} rEst
Σ
m 1==1
n²
| 12
m
16n k 22
24

---

Published in Transactions on Machine Learning Research (08/2022)
Putting things together gives
by úf k1+1
12 (1 n7o 167K
'nR + n²n2² 1
lly ,1,
Choose S
AdA 64m
and 17 S
da 4n
gives
by u k,t41 I (1 (1- 1100 lly - ,/
11
12
Assume that equation 11 holds for local itcration O..... We are going to prove equation
12 for local iteration & + 1. We start by noticing that equation 11 implies that for all local ileration
E [t], we have that lyy u I Iyy a 22. Moreover, we notice that
amm
W k.t
1(() UR5
xi)
Vi) a,
Owr
.l,(0) Uktt
yE
m
n
y uf,ell2 ll2
m
Therefore,
HLm (W/ kitt Owr
WkAr - Wk.r 12 -
12 1n
8==0
vn
t-1
$ 17
y u//,e l12
==
vn Vm
lly
ú/ l|2
n (Ily 1lll122 + ||u -6(l12
Applying Markov's inequality to the global convergence, with probabiltiy at least 2k holds that
lly unlla S 2K. [lly 11ly-ulla)
uxl1a]
By Lemma 25
we have
EM, [||u: 11 S 4E(1 g)
Thus with probability at least 1
2pk it holds that
Iu úill2 5 2x 2 E)mpK/8
Plugging in gives
2nk Vmd
2 E)pK
Wk.tr Wa. <
B 1). Wa.a llly uull2l + 2ntran
m8
which completes the proof.
25

---

Published in Transactions on Machine Learning Research (08/2022)
E
Proof of Theorem 3
Before start the proof, we introduce several notations. Define
1
Σ
ar
(o( o((wkrr x,),,
Vm
rES;
(1)
{
Σ
ar (o( xx)) ( x.)))
2,k
m rESt
Let
(1) 1,.-.... 1,k
[n)
1,
and similarly,
2, 2.k
2,
Then we have 4+1 u600 =0=1p ((00 + 1 and
U₁ 11,6 12 Also, we define H())+ to be
u
H
Σ (xi, xx, I(WW,) 2 (we.r, 0}
m rEST
For k-th global iteration, first local iteration, we define the mixing gradient as
aLmg
Wh. K.D
gRr
Th.r
Owr
1=1
T. (Ws)
1k.r
Ow,
The
m (ūk 111) 1 - 2 0}
mm
1=1 im
( kr
- N
[I 20
where we define the mixing function as
1(1 mgge 0k
TJk.r
As usual, we let ferr
per.....
f(m)
We note that S has the form
(1)
(1())
fkrr
K,
m
I=1
1
Σ
p mb.
ol(wam)x))
ar
T]krr
Vm
Let Verr"" = TJrr Ei mh The mixing function reduce to the form
(1)
kk.r
m
26

---

Published in Transactions on Machine Learning Research (08/2022)
Also, note that if N O, we have Vkrr." 0.
We prove Theorem 3 by fashion of induction, with the
two conditions we consider stated below:
1
&
M [lly 1x
S
4
- 20 l1y uoll + B
(14)
2nK md
4
[w.r - Wo.rl12 + 2
BlM 1) Illy usl22 + (K - k)) SR
(15)
nOT2O
|eerr Worrl2 S R
(16)
with
4B1
B =
+ < p)
100
Base Case: Note that equation (14] and equation
16 holds naturally for t = 0. To show that equation
15p holds, need to use the over pre-pparamettriiation property. In particular, want to show that
2nK
4
RAA
2
Ewa. llly uolle] KBB 100T>0
RS
md
144n
Apply lemma 26 and move the factor on the left hand side of the equation to the right. Then we equivalently want
RλO
md
4C
LBKB kB})
n3K
Ω
maxx
100
Plugging in the value of B and the requirement of 17, and solve for get that
K
n
nk²B,
K2p) ke0)
m = Ω
5
max
K
R2OX?
Inductive Case: again the inductive case is divided into three parts.
15
16) Observe that B 2 kvE((1-)pm) Thus, if equation 15 is satisfied and the over-
parameterize terization requirement holds, then Hypothesis holds. Thus equation 16 holds naturally.
15). Assume that equation 14 holds, and 15 holds for global iteration k, we want show that holds for global iteration 1. In particular, we would like to show that
oquation (15
2nK 4 md
Wk+1r - -w.rll-Worrllz + 2nT
10T B,M.I.W... lly u 1//2 + +(K - - 1-)B S R
it suffice to show that
2nk md
2nK md
8nK B1
| Wk.rlla V
E8 11 llly -u- uxl12] + 21/T
B 27T
Recall the definition of B as
B
=
+ E(1 pm
It then suffice to show that
2nk mō
2E(1 E)pK
| - War S mT
El Wo.a Illy - u4ll2] + 2
md
27

---

Published in Transactions on Machine Learning Research (08/2022)
Note that under these two conditions and the over-parameteriiation requirement, Hypothesis Thus, we have
holds.
2nK B
2E(1 E)pK
Wwker Wk. 112 S V
B 1.. Wo.a [lly 1en--ullla u1l12] + 2
md
md
NE. No.r"
for all I E (p) and E [r]- Using the definition that "err
have
Wk+1, Warll2 S Tkrr
mier /w
W
2nK mo
2 E)pK
S W/T
E
11. Illy uuslaa + 2
md
16
15
14
Assume that equation
16
and
15 holds for iteration k. We want to show
14 for up
to iteration +1. Under these conditions, WC have that with probability at least - 46, Hypothesis holds. Throughout this proof, we assume that
m
(Wo,r, x,) & 2mnn - mnR
and that
Wolle S 22md /mR
Note that Lemma 22 and Lemma 23 shows that, as long as m 2 (log 7) the above assumption holds with
probability at least 1 - over initialization. Moreover, Lemma 16 shows that as long as -EAA EAo
log with
probability least 5 over initialization we have
|St| S 4mx R
To start, expanding the loss al iteration k + gives
EM.
|lly us |2 = ly ux12 2 2 (y- - U1 EM, [ug+1 uz]) + EMA [| - ux 112]
y 1 2 (y- un EM, [I,,]) 2 (y u EM [I,, +
EM. ||
1
Following previous work, we bound the second, third, and fourth term separately. However, the second term requires more detailed analysis. In particular, we let
1 I1, ,)0TH(()(y - ux)
Then the loss at iteration & + 1 has the form
EM. [lly U
= lly uxl12 112 2107 (y Ug H(k)(y u&)) + EM, [lle41 - ux
2 2(y us. EM, [T,ee)) 2(y - us, EMA [I,,]]]
< (1 - 10 )o))o - - uu 11 + 2 ui EMA [T],e]>] ]
2 2 Uk. EM, [2,,])| + EM, [| Us 112]
where in the last inequality use Amin (H()) 2 Au from Du et al. 2z
2018 Assumption 3.1. Moreover,
Lemma 6 Lemma and Lemma 10 shows that under the given assumption, with 0 nT unxx [ have
1 n - un (1 ndd
2n 2m
1)2
Ky uv EM [HAI)| <
mdo
0ho
28

---

Published in Transactions on Machine Learning Research (08/2022)
Ky ub, EMA, [2,x]>] - - + n7052 -=))n
+
24pr
9678
EM, [||u++1 U2 [ s d -ull2 17m252220(0 - kn² = p)²pn(
p
Putting things together gives
EM ||ly ux+1ll $ 1
4 n 11y 11y- -uxl
1
16 ( ndd 2T 1) pCr
+
mAo
A
nA 2 - =n²
(
- 1)2pC1
17n)25""00
E =2)nnk²
24rp
9678
>0=2 - 1)pnC1
(1 1 11y - uus 12 + 7
1
n10 A1
Therefore, we have
EM, 1) [lyy ux] I 311 1
x 11
uoll B
This completes the prool.
29

---

Published in Transactions on Machine Learning Research (08/2022)
F Proof of Theorem 4
Again we use induction to prove the theorem. Consider the following conditions
EM, [lly u I (1 - a) ly uxl12 + B1
(17)
2nK md
Wk.r wo + 2n/
PIMMell. - uxllal (K (55-k)) -k)BB)
R
(18)
a
with a and B defined as below
a= -- (1-))))-) (1-
)) (1-1-)))))1-11111))) 1700
B + +xv E)pn: B = 01=0((1-01) nx²1p
Amax Ao
B =
with = For notation clarity we will not plug in the value of & for now. Notice that equation implics the convergence result as in the thcorem statement as long as y = and that max K, d, p} Thus, as long as we eslablish the inductive relations as above we are done.
17 2 n.
Base Case: Notice by lemma 26 equation
17 holds.
We use the over-parametterization requirement
show equation 18 In particular, show
2nK
(- llly uoll2] + KB S R =
2
mo
As before, using Lemma 26
we have
Ew llly - uolle] & (Ew..a [lly
un uoll2)) CVn
Plugging in this value and the value of n, B and a, we arrive at the over paraametttttiinnoonnn requirement
KNaax Amax
Ω
188
Inductive Case:
the proof is again divided into two parts.
18 Assume that equation 17 holds, and equation 18 holds for global iteration k. Similar to the
proof
Theorem
3 it suffice show that
2nK mo
2nK
2nKB, māa
W Wo. & /
EiM, W.W. Illy uella + 2n/7
V md
B 2nT
Plugging the value of B and a, then suffice to show that
2nK
2E(1 =)pK md
| Wo S V
md 1IM Wo,a llly uxll2] + 2ny
By Hypothesis we have
2nK E Wo.a llly -uel2l + 2nyTan
2 p)p
I Wk.ll2 nTv
mo
m8
for all 1 (p] and E [7]. Then we have
| - Wk. S Ther
m Il Wa.ll2
[=]
2nK md
{(1 E)p
< 17T
E E Wa.a W.. [lly - us [zlla] + 21yTkn
md
30

---

Published in Transactions on Machine Learning Research (08/2022)
18 17 Now, assume that equation [88] holds. Then the result of Hypothesis holds. Our target
is lo show equation 17). As in previous theorem, we start by studying lly u In the case of a categorical mask, we have the nice property that the average of the sub-networks equals to the full network
u&+1 = > - u/r.. Using this property, Lemma 13 characterize lly uk+ill as
ly u 12
y
p
We start by assuming the condition of Hypothesis holds. We proceed by proving the convergence, then we prove the weight perturbation bound with a fashion of induction. Hypothesis implics that
lly a S (1 (1- 17:00
lly
0
12
2
= lly u/y 11
7
Σ
7-1
1
n7
yyy - 0//12
2
1=0
Using the fact that EM, [o]]] Uk, we have that
EM [lly us [e] lly
us
11
+
EM
[] a
Therefore, we have
11.0
19:00
EM, [lly- Us++
EM, [|ly
EM. [|ly - [
2p
2
1=1
|10b - a 112
EM.
p2
1=
P=1
11
11100
lly
11
1
EM. [|ly uf +
2p
2
1=0 1=1
p
EM
[llu-
- 0/
EM || uf.-113
p2
Lemma 14 studics the error term | út 112 and gives
Uug 0//12
1
Plugging in we have
EM, [lly U 1
< Ily
-uel1
A7Ao
Σ
T-1
1
17d0
Σ
EM,
[lly - 01. +
2p
2
1mo
1=1
EM,
u/! 113
Denote the last term in the right hand side as 4.. respect to the initialization. In particular, if lemma 23 holds for some R 2 and the weight perturbation
Lemma 15 shows bound of the expectation of 4 with
31

---

Published in Transactions on Machine Learning Research (08/2022)
is bounded by lIwE, 1o.r S R for all [m], then we have
(1 (1-1-1)
p
[||
u a/ccll2 24 (1- 1--1) nu.
[1
Amaa
4k S
EM,
2p
1=1 t=0
Using this result, we have that
EM, [|ly- u
by 1y--ul1
17A00
1
n7
Σ
BM: [|y ( [ly-0k1(] +
2p
2
1=0
(=1
17:00
p
r-1
1770
lly
uxl3
EM: [|ly 111y-u111))
2p
2
(1-p-1)
17
17/00
EMe [lly d':13] + 48 (1 p 1-1) mxx2
Amax
1
r A
4p
1=1 1=0
p 1)
7700
p T-1
1110
= Ily 1y=-111 1x113
EM, [|ly- - 10%y--1,.2)
4p
2
1
48 (1 - p- 1-1)
3
Amax
Ao
< Ily mn 1 (1 p
77
7700
11y 1p-uu16.6 1
4
2
1=
2 Amax
48 (1 - (1-p-1) nx
Ao
= -(L-p-)) + (1- (1 - p1 (1 (220) )) Ily ux11 48 (1- -r-1) nnnn nTT
Amax
do
Therefore,
EM, ||ly uR+1 < (1 (1-a) a)) lly uu 112 B₁
This is the same as the form in equation 17
Thus we arrive the convergence
B
B(M,) [|ly
u
I (1 - a)* Ily uo +
This shows the convergence result in the theorem with and B1 plugged in.
32

---

Published in Transactions on Machine Learning Research (08/2022)
G
Lemmas for Theorem 3
Lemma 1. The expectation of the mixing function satisfies
EM.
[ Ou
0(1 -
a (WWk(()
vm
Proof. Note that if NE, 0, then we have Ske for all [n]. Thus
E
= 0 =0
kr
Er
Moreover, if NET = 1, the expectation can be computed as
m - ke ( xi)) No. =1
BM, [60 Kyr = 1
EM,
Ther vm
r==1
m m Kir 1 apo((We
EM
TKKr
v m
r'=l
1m1
1-
ar ((ww xx))
+
((
v
by using Lemma 20
Combining the two conditions above gives that
EM.
[kr
P(N,, = 1)EM,
Skrr NE, =1 + P(NE, O)EM, [sRe No.,- 0]
OC
apo(( x;)) + 8(1 - 5)
a, ((Wk.", xx1)
Vm
r'=1
8(1 E)
= Buk
(8) +
a, ((oa
mm
Lemma The expectation of the mixing gradient satisfies
9 (L((W))
8(1 5) n
EM. gkrl E Ow,
+
Σ x ((Wk,s
m
Proof. With the result from Lemma
we have
n
EM [gk.r]
EM
((0) krr
- y.E (Avtt.) xxx; (Wk.:; 20
m
0
n
(7)
1- E a (( a.xiI((We.r. 2 0}
U
Vi
m
Vm
0 [LL k)
0(1 - 5
x,o x;oo (wo((())))))))
& Ow,
m
1=1
Lemma 3. Suppose m 2 p. If for some R > O and all r E [m] the inilialization salisfies
Wu,r: xz) S 2mnx2 mnRR
33

---

Published in Transactions on Machine Learning Research (08/2022)
and for all [m], it holds that |war Worr R. the expected norm of the difference between the mizing function and ux satisfies
EMs [[lfkrr u No, 1] 5 8(0 m2)
P
Proof. Since EM, | Nir 1 E for r' o r, we have for rı # r2, there is at least one of 71,72 7,72 that is not r. Thus
EM, [Urknr - (Vs,,rr - E) Nor = 1] =0
and for + rl
Varm, (Veyyor Ni. = =1 EM, [[Vk, - E)² NE. =
Morcover, for r= r', Lemma 19
EM. [(KKITr.- -5) 5²
Therefore, using Lemma 21 we have
- ()o(((srr) ) Nerr =1
EMa [1 ux )00) No.r = 1
EME Σ
VarMA (Vkmer |N, = (( xx))+
171
EM, [(Kkinn. E) o(iWe. x,))
m
0 E
0 - E2 ((WWirr xi))
(Wi, x:)²
pm
2(0 <2)
m
2(6 - (2)
(Wo.r. x/2 + mR²
((Wer) + +R)
+
pm
8(0 E2) k22
p
Plugging this in gives
EM, [[]fe,r 1 No. 1] x 8(0 m2)
p
Lemma 4. Under the condition of Lemma is bounded by
the expected norm and squared-norm of the miring gradienl.
2n0 lly uu 112
160 g2)n?>²
EM, ||gtr
112] 5
m
pm
[|grrr112] 5 fa
0(0 - 52) pm
EM.
/m Ily uu
122 + 4nng/
Proof. Using Lemma
3 we have
80(0 g2)
EM, INe, 111 P(NE) 1)EM, [[]fer 11.r--+1331
u11l S
p
34

---

Published in Transactions on Machine Learning Research (08/2022)
According to Jensen's inequality, we also have
20(0 E²)n
EM. [N,, uxl12] S 2x V
Moreover, we have
1
kir 3i) a, x.) 2 03
(()
gkrr
Vm
NE. OL(W&)
krr u )
(1)
20 g Ow.
Therefore,
EM, ggtr [2]
2EM, 5²
Mir
(L(W&)
2
2 EM
/krr
(1)
ux (1) No x; x ((WK,, xi) 2 20]
27
Ow,
2 m
2n
EM.
No.
uk (FLe)
(2)
2
2nf
Ily
u 112
m
m
2n
v
(EMe
[Na.r forr -u 11 + Blly 12x112)
2nd
160(0 5² )n?n? R²
Ily
1l122 +
m
pm
This shows the first inequality. To show the second, similarly we have
EM, 1igtr 11a)
S
EM
N
OL(W&))
+
1
FM.
[0 - -(a
NT x;) 2 0}
Ow,
vm
vno m
EM.
Ni.
(()
y
uzll2
Jrr
V m
vn0
FM.
[N llferr
uxllz] + vmm lly
ug 1|2
m
vnl
0(0 0(0- - 22)
<
Ily - uull2 + =nnyg
Vm
pm
Lemma 5. Under the condition of Theorem
3 we have
((1) @
a S nntnllyy-
- -u/lly
and therefore,
(1(33)
s /t (lly 1zll2 + | u/(l2)
k 0x
((0)
1(1)
2r n (lly
m 1 a út l12)
35

---

Published in Transactions on Machine Learning Research (08/2022)
Proof. We have
((6)
_10)
ar m, (o ((W) - ((Warr, x,)))
uk
vmm
(W v. ((Wkr)
m
|Witer Whell2
at (Wh.) k.t Ow,
m
"
(-1 Ily i
S n/n
-==0
S nt/nnly
Therefore,
1())
1((1)
I mt (lly uxll2 |lue ú/. l12)
Moreover,
1(0 k
((1))
[(()
f 1 217² (lly uxll3+ + llu - d/113)
Uk
Lemma 6. Under the condition of Lemma g with
Ag 1om7-1ynz we have
Ky us uz EM, [T,N) S 8 n uxl1 16 - 5)2x2nd
+ 20 2 - 1)2npC1
mAo
QNo
Proof. We start by analyzing Wk+1. Wk. Taking expectation, we have
ÖL 0L
W k.t
EM, [We+11 - Wk.rl -1/BM,
Ow,
1=1
ÖL (Wkr))
θT mi
(Wx)
Tm (Wa) dw.
=-nEMM.
+ 1k.r
Ow.
Ow.
1m1
Iml
W k.t
al (W+)
=-nyEM,
[g
7/EMI
Mk.r
Ow,
dwr
nθT (L(WW)
-10(1-E) - T)T
Σ xio
xio((Wkirs x;))=
3 Ow.
m
i=1
p ΣΣ
ÖLm!
(W/) kie))
Ohmy (W+) Ow,
17MM,
1,krr
Ow,
1=1
36

---

Published in Transactions on Machine Learning Research (08/2022)
Therefore
EM.
70
m Σ
a, E, (o((w++,r;x))) x:)) -w((wr::
1,k
=
rES;
Σ rES;
ar (EM, |W++1, Werl xz) (w 20}
mm
_nOT
(L(W
Ow, xx) I((W 2 03 17 (Ef,f + +E2,
ar
vm
PES;
nEOT
1
(Us u ) (x, x,) 2 0, (Werr x,) 2 03 17 'f +E
m
rES: j=1
(H()), H()))t (w3 - (()) 17 ( + +
= n10
j=1
where
e(t)
AE(1 7)
Or (x,, X,) o (W),,,
1,
m²
rES: j=1
iL (Whs) Ow,
àL (Wx)
)
c(t)
2,k
vm
Or
EM.
Thir
Owr
X2 I( [((Warr.x)) 20)
rES:
1=1 t=1
Let E1, =
((1))
11
c(n) 1,
and E2,k
[((1) 22,
c(a) 2,
Then we have
EM, [11,k] = n10 (H(k) H(*)²) (y -u:) -7 + E22)
Thus,
EM [[A]]] = EM, [1,a]] - m)TH(()(( uu)
-(0 (y uc) (E1,6 + E2,6)
According to Lemma
and Lemma
we have the bound of & and (l)) as
2d
c(() 11
I Ag(1 E)rR
c(i) 2,
TTT(T
1
(0lly uAllz + vpCi
<
2
Moreover, according to Lemma. 17
we have
H()) 1|2 < 45nxx-11 R
Let R
we have
125n
[[(())+12 < Ao
32
37

---

Published in Transactions on Machine Learning Research (08/2022)
Therefore, we have
Ky us, EM [T,e)D] & n0T (y u4. (y uz)}m t Σ (1(v: [ ux E + (3: - 2 ) E
im1
,10 [H()) 12|ly us
max
(11..) | Σ
y/ ukk
(6)
icim]
2=1
+ max 32 32"00 100 u
[
+ + ) //
n/llly
us 1|
tE[n]
TT 1)n²
$
lly 1x1
32
2
an
17 gT(T - 1)m²
n - F)Tm
+
11 11u-uul1a
2
Using the general inequality that ab S (a² + (2)), and S
16(7-1)u7 we get
An
Ky - uR EM,, [T]]]|| S 8 n - us -) 2n - -1)2n"pO,
1
mdo
Oda
Lemma 7. Under the assumption of Theorem we have that for all k E [K] i E [n], it holds that
2d
0E E)TRKy
m
Proof. We have
OE(1 - E)T
e(0) 1,k
<
Or (x,, ( ((We.))
m
rES. j=1
BE(1 E)T
(Wir x,yl
m 3
rES,
AC(1 - Tm
Σ
/Wwrrlla
m
rES.
OE(1 E)
(l/ + R)
m
rES.
0E(1 - Tn
BE(1 TITR
W WolF
m
Vm
2d
0E(1 E)TAA
m
where for the bound of [ use Lemma 22
Lemma 8. Suppose W < R for all r E [m]. Then we have
c < TET(T 1)
(flly us 12l + VpCi)
2
38

---

Published in Transactions on Machine Learning Research (08/2022)
Proof. Since r Si, the difference between the surrogate gradients of sub-network has the form
mL (Whsr k.tt) Ow,
T mi
(W&)
ar m. k Xj
=1(6)) ((( (ug
If [[(w..,) xz) 2 0}
Ow,
Vm
jml
((j) a
m
k
Therefore, using the convexity of norm,
ŌLmi (Wax)
dLm! (
p
aLmi (Wx) Ow,
ŌLmi (Wu) Sw,
EM.
Aker
EM.
Ah,r
Ow,
Ow,
dLm!
(W're)
abmn (Wa)
< IBMs
1hK
Ow,
Ow,
1=1
S EM.
Tker
(()) RRt
m
1=1
By Lemma 5 we have
21(9)
S ntvn (Ily uel12 |lue - 0%. 12)
k.t
Therefore,
OLmj (Whet) Ow,
aL (Wx) Ow,
2,2
EM,
ak.r
m
rES,
Obmb (W're)
OLmi, (Ww) Owr
$
EM.
Mh.r
/m
Ow,
1=1 rES;
Uker
(()
EM
Vm
1=1 rES,
1=1
1=1
EM
T
mkrr (lly - uull2 l + |ly at. (12))))
1=1 rES.
TTTTT 1)
10lly uxll2 + EM..
m [ng 0b/ll2
mk.r 1=1
TT 1)nž
(Blly
(olly-uad - uxll2 + pCi
where the last inequality follows from Lemma 24
Lemma 9. Under the condition of Theorem we have
2 =
nA -T)
Ky ug EM, [2])] < noT 1
24pt
96T8
39

---

Published in Transactions on Machine Learning Research (08/2022)
Proof. To start, we notice that Using the 1-Lipschitzness of ReLU, we have
M
(())
EM,,
Σ
(a (o((we+1.))))) o
2.
Vm
reSt
EM, |1 (lo((Wet),)) Xx) -((
m
res!
Ma [k W.),,
v m
rESt
M [1 Werll2
m
rES!
Olim W k.c) Ow,
n
EM
Thr
rES!
t=
HLmr (Wh.t) Owr
T-1 M,krr
EM, [|gg,r 112] + EM.
m
rES!
1=1
n0EVn
4n
0(0 52)
1St llly
uslla
m
m
p
v
r-1
EM
Ther
m lly 0/
rest
1=1 1=1
v
4nEkn]S St
0(0 5²)
1SS |ly
us ll +
m
ne Σ EM
r-1 p
m m lly u
Thkr
m rESt
0(0 52)
|sst Ily
us 2|2
p
ne
EM
Ther
m 11u
u//l22 ll
rESt
where in the seventh incquality we use the bound on EM, [1 from Lemma 4
Morcover, using Lemma
24
we have
E
Tjkrr Σ
mh.rluu
0//l2
S vpCr
Then we have
EM.
(()
nBETVT.
1Stt Illy uu 20 + 4ngkn| S
0 {2)
VnpCr (sstt
2,
m
m
p
m
0(0 (
< SnAET Vnn" -1 BRly u4l12 + 16nEnR
+ AnE(T - 1)x nn/pp
p
40

---

Published in Transactions on Machine Learning Research (08/2022)
where in the last inequality use |S I 4mx RR Therefore,
Ky - us EMA [I,,,>]]] =
Σ
Ui
uk
EMM
n
(1)
|M
[
Vi
uk
2.k
max =E/n|
EM,
7 2,
Σ
y/
uk
(1)
i=1
V max rE
EMA;
1,50
||y
u&l12
0(0 =²)m²
- 8105TH "nRlly
us 12 + R
Iyy - uu 1
p
AnE(T 1)x0 INR pC1ll> uzll2
1 Jolly A ux113 + A 2 -g))mx²
A0 - 11)ppc
8
24pr
96TQ
where in the last inequality use R 5
RAo and ab - (a² + 1²).
Lemma 10. Under the condition of Theorem with l S zsn7 max [a.pp
we have
1 17 - )n)=
EM, [| us [ S 4
=02 1)?
p
Proof. As in previous lemma, we use the Lipschitzness of ReLU to get
EM, (
27
5² EM
Σ
(a((Wa-L,)))))) -olp-oooowwaao -0(
m
< 8²
EM, (o((WW+),r;x))) - -o((w,, x,)))
r=1
m
2
EM (WetLr Wk. x;
<3
EM. |[Wttrr warll
r=1
_² (Di,t D2,6)
where
D1.k =
Σ EM || - wk.rl
reS;
D2. EM.
rest
Using Lemma 11 and Lemma 12 we have
16n =2)mi2²²
D1,E S 7n0 + 4n 2n - 1)p) lly - 11 +
+ +1
An 1)2
000000
4n7 0
T2)
1)AopC1
D2, S 18
(1 + (T 1)p) Ily ue 1
+
9p
18
41

---

Published in Transactions on Machine Learning Research (08/2022)
Therefore we have
EM, [||u++1 ux12 S < + D2,x)
1f@rnAo
$ An 2 41 ( 1)p
(1+ (T (111)) ly ux1
18
161)252 - )2)
4n2 20030000 =z)Tn²n²
( (1-1))p
P
9p
52 - 1)ndopCss
18
Ao With 77 S 48n7 max (n.p)
we have
17m2522 - =2)n²k²
EM, [| - us [3] S > 10lly-ulll
uxl12
+12g2>o(T 1)²
p
Lemma 11.
Di,e S 7 + 4n/ 'nn0 1)p) bly un uxl12+ 16np²=0(00 - )2)
4n n (T-1)=1
p
Proof. We have
27
OL
W! Owr
D1.6 17 Σ
EM.
k.t
rES:
1=0 1=1
Ölmij
W! k.t
a. (Wk) Ow,
$ 2 Σ EM.
Tgk.r + Uetr
Ow,
rES:
1=1 Iml
OL
W kit
aL (W*) Ow,
< 21 7
EM,
|1 21 1)p
Σ
2 EM,
1k.r
Ow,
rES,
rES. 1=1
1=1
Note that for r E S₁, we have
dLm; (W'rr)
2
dmmn (Wn) dw,
maker
(20())
100) T(( 2
Ow,
==1
nmlker
I(i)
1(1))
u k.t
0x
m
n n'm m
(2n
uxll + 2:72 u 1111
where the last inequality we use Lemma from Lemma gives
5 Plugging in the bound above and the bound on EM, [[]g:
16n yn
Di. S 4 112
4n An/0
1)elly u=ll
P
n - 1)plEM
nrr Σ
| - al/12
m
1=1
16n7-=2000 =g²)m²
< 24n + An 2n 1)p) Ily
u
un
11 p1)
p
42

---

Published in Transactions on Machine Learning Research (08/2022)
Lemma 12.
0/0020
(1+ (T 1)p) byy un uzl12 + 4n²² >00(0 - -TR
,
1)AopC
D2,x S 18
9p
18
Proof.
dbm! (Wirr) Ow,
D2,k = 1/ Σ
EM.
1k.r
rESt
1=
1=1
W! R.1)
S n²² Σ
EM |||gar
EM
T
Owr
rES.
12
27
T-1
(W
< n Σ
EM. [|lgar [2]
+p
EM
Ow,
rES,
2n7² Tnf
So72 - g2)TT????
Tpp
1x112+ + 1/1 11
18:
EM
1irr
mx. lly - i//
pm
1=1
21/² Tno
8172 - E²) Tn?K' R2
2n 0T(T
1)mp
|/11 -uul1
|8
| Us 112+
pm
m11
2ny² T(T - 1)mp
S [EM
Tkrr
m ||u - a//12
Using |S1| S 4mn 1 R with R
Emin 144n gives
1²0
An0200(0 - TTx)
D2.k
(1 + (r (1)) 11y -u 1ul12
18
9p
T²T(T - 1)AuP
EM, nhrr Σ
m |u - i1//12
18
1=1
1//07>0
4n Ag - Tnxx
T
1)Aop(1
(1 + (r 1)p) 11y -u +
18
9p
18
43

---

Published in Transactions on Machine Learning Research (08/2022)
H Lemmas for Theorem 4
Lemma 13. The kth global step produce the squared error satisfying
p
ly u
út
11
p²
Proof. We have
1 P
yy - uk+1
P 1=1
y
p
Iyy
Iy
R y uk
p
Σ
Iyy
11
P 1=1
|ly
0/ 11 + lly
o 11
p
0 y alt.r) ))
y
2p
1=1
1=1
p
1-1
lyy
11
| -r
p2
Im
Lemma 14. We have
112
Σ
ukkr 1/6,-13
1=1
Proof. Using u:
1 p
ΣP 0k have
|
1=1
P
|
11
W Kol uif
P
P
(iat 11 + 11
Σ Σ
(ak ukrr)
2p
1=1
|
p
T=
44

---

Published in Transactions on Machine Learning Research (08/2022)
Lemma 15. Suppose the condition of lemma 25 holds and the step size satisfies 17
0
(n2(1-p-1)
dop
then
with probability at least 1 - it holds for all k E [K] that
(1 - 11-p-1) a1
P
EM [|x ui 1l2,
24 (1 - p-1))= nu:21
Amax
4k -
2p
Proof- Recall the definition of lks expanding the quadratic form gives
p
[uk u//12 | - 11,
Lk
EM2
u/ u//12 1 -Il(ux a//) + (ux us+1 0f. + -/,,r)11
MM
Uk
P
EM 2 2(u
u1g. uu us+1 if ( > + llue u8+1 - Oif. + 1//e.r [2]
p
Σ EM, (ux u1e 6/e,r utt)]
|uu uk+1 ui + o 11
EM,
where in the third inequality, use the fact that
u/ un - ux+1)
pu
uig
un
u6+t1 =
1=1
1=1
For convenience, we denote
(1) k,r
= ((Wr.
(0) k = a ((wk
Noticing that 1
Eim (ik nnk.r)
= us
Uk+1, apply a trick similar to lemma 14 to get that
p
P
EM;
||ue:
Uk++1
0/. 0 2
EM,
[1
ap.r 0% 0 [2]
a (ohir
((i)
mh (oho on.t.ol (i)
(8)
mp
m
m (a(0) (i)
1(0) 2
+ mm (a oR, (i) 2
mp
Tm² 2
ΣΣΣΣ
mkrr ||y - 11. m |||- ob/.ll1)
mpp
122 Tp
8
ΣΣΣ mp. |ly 1/..112
2mp
Notice that fixing an E [p], we have that m 's are independent for all r E [m]. Apply Hoeffding's inequality to get that
IP
m 2mp-1 <
scxp exp (-2mp
45

---

Published in Transactions on Machine Learning Research (08/2022)
Apply the union bound over all & [K] and Z [p], and apply the wer-paramettrrization requirement to get that, with probability at least 5, it holds that
pfrn2(p 1)
P
EM
uu
uk+1
ut ú 122
I
d/fe l1
p²
-
1-pp1)
11
0/,.012 1=1
||y
4p
0=
by choosing 17 0
(n((1-p-))tz
App
The first term can be bounded by
A = (ux uig a/e,r - -u/)) - L uf
a/ell
We study the term 0f (/ - o l12 Pollowing analysis in the proof of hypothesis 1, we write 02/,11 02,,
I i,k,t 1g 2.6, We first study the magnitude of I I,K,t- Its ith entry has
mm) o Hk, m H(k V/ Σ
-0
((6)
I.ket
2
As we have shown in the proof of theorem 2
Ao 2
m/ H(k,t) H
Therefore, Amax (ml o H(k,t)) & Amax + da S 2Amax Moreover, in the proof of hypothesis we have shown that
m H(k,t)+11 I 4nk-
Therefore, we have
| = (m) o H(ktt) - m H(k,t) t) (y - ale..)
I (21)Amax + 41p 4mp-=''n1 ||y 0 l12
Using the bound of k,
(()
in the proof of hypothesis 1, we have that
12 I
1(0)) 2,000
2 Amp=-12 ||y -1/eel
Therefore,
0 uk.1112 S [| + 1222 S (21) Smps --11511177) nR)) Ily- a/4,oll 212 < 37 ||y 0/2,012 by our choice of R. Thus, we have that
A < Amax Σ y dh.rl1 Imu ukll2 <
11
n7
7-1 ||y
uheell2 + 12
17T ardmax | u//l
8
(1 p-1) Ao
t=0
1=0
Therefore, by applying lemma
25 we have that
4k 5
(1 (1-
11-p-1) nnao
Exa, [|lx o'c.l12]
+ 24 (1 - ---)
nk. 17T An
2
Amax
2p
1=1 1=
46

---

Published in Transactions on Machine Learning Research (08/2022)
Auxiliary Results
Lemma 16. With probability least 1 ne
ne-mx mm R we have |St] S 4mx R for all e [n]].
Proof. Note that I[r E 5 [((AAr) 0) I(Arr). Therefore, we have
| Σ T(r E 8t) I(Air)
Since Eworr [T(A++]] P ((A+r) S
28
1-1 R. also have
E wo
(T(Air)
For. [T(Air)]) S Ewo.r [TAiir]]]
2R
S K --
RV 2m
Again apply Bernstein incquality over the random variable I(Air} - -Ewor [I(Air]] with t == 3mx- 1 R gives
P (( S 4mx -R) P.
Σ mArt 2 4mk 4mx - S Sexp (-mx- -10)
Lemma 17. Define H Ruax such that
H 5 (x,
Σ
If 2 0; (wr 2 0}
rESt
If [S S 4mm-1 R, then we have
H
||2 S 4nEx R
Proof. We note that
H
Σ |
i,==1
For each i,j pair we have
H | 46K= R
Thus
H | s (| - < - 4n R
Lemma 18. For i.i.d Bernoulli masks with parameter &, No.r Bern(O) with
0 = P(N 1) 1 (1- 1--$i E
Proof. We have
P
P(N,, = 1) P(N, =0 =1 II
P( = 1-11 E)
47

---

Published in Transactions on Machine Learning Research (08/2022)
Lemma 19. We have
EM, (vk - S 50 -
Proof. To start, we notice that Pkrr, = 1,,r Ei=1 mkr = Th.r Ei= m Myr Therefore EM, [pana
EM, N 0. Moreover, since No.2 k,r, have EM, 122 [krrrr =. Thus, using 0 02 &, have
EM [VRRT. E) L8Me [Ph - -2EBMM [Mwerll + +
0 2E0 5² s =²
Lemma 20. For i.idd Bernoulli masks with parameter &, we have
{ fr
EM, [ | =
{ ifr=
Proof. r = we have
P m |N
Xkrr No.
EM. |krtyr NE, 1| EM
TKKr I=1
EMa
New
EE [N
NE, = 1] =1
If r' Fr, then we have that m/. is independent from m/,r and Nkr- Therefore,
EMC |krrrr N =1 -EM.
A,rr Σ
EMa, [mkerd
Xkir Nk.r
EEMe
Ner = =
Lemma 21. The variance follows
g-
if # r
Varm, (Veny | = 1)
p
if r = r
Proof. For * r', the expectation of v given No.- 1 is
Ei Ei=n my m m
EM, [2 | = 1] EM,
No. 1
X
mk. X
mk.r
EM (mk,))) BM (mm BM,,
Nohr
EM,
No.rr = +
Xx.r
1=1 1'
m x2
EM [m, |EM
N =
mm
mkrr
Em:
N
EM
| =
+
Xkrr
Xb.r
1=1 11
mm krr
EM
N
X2
1
m
-
= {BEM
XR.
N
kr
EM
X
N kr
1
48

---

Published in Transactions on Machine Learning Research (08/2022)
Therefore, the variance of Vk,r,r given N,,, 1 has the form
Var (vkyr 1gN 1) EM, [VR Nitrr = 1]
EM, [VReT
No. =]
m
-2
RiT = 1 S²
EM
N R,r =
Xk,
Xk.r
1=1
Let ((p) = Ei=1 m. ~ Bl,,E), then we have
1
EM,
[N krr
EMs
Xx.r
1+ X(p 1)
m
E EM.
EM₃
N
P(mmhr
NE.r 1)
Xy
1+ X(p 1)
1+ X - 1)
Moreover, using reciprocal moments we have
0
EM:
=
1 + X(p 1)
De
Therefore
1
VarM, (pene Nir =1) {BM=
1 X(p - 1)
1 xp
0 - 22
If r== r', the variance
Varm. (Veir Ir =) VarM, (gr gr 1)
Lemma 22. Suppose $1, xV 32- With probability at least 1 - edd/32 we have that
[olle I xv2mdd /mR
Proof. For all r [m], [d], have EM. [w]d]] n²2 Moreover, each whs, is (2x² 2x -su-- exponential random variable
(win, -x3)
1
2 e (wwid xx²
E
zx dw
KV 2m
1
2 Dnuie ux? dwad, 600
Kv 2m
Ix²
KV 2m 2m
V (2x)
e fi?
2t2tn n
VI 2tx
with t S 2
1. Thus, using independence between entries of Wo gives
(llloll -mdr3)] <
Awoee -x²)] S 2mdi
E
c
d==1
49

---

Published in Transactions on Machine Learning Research (08/2022)
Invoking the tail bound of sub-exponential random variable gives
P (1 2 mdk? + S
e
Smant
if 0 t S 2mdk
if t> 2mdk2
Let / mdkx SmnRy2d mRC Then
|Wll s 2mdk mRl -2mxRV22 (NV2md vmll)
with probability at least 1 amdde
Using R S mi 32 we have / 2 mmdd? Thus with probability at least
we have
W xv2mdd vmR
Lemma 23. Assume K. S and S 8 all i [n] that
With probability least ne 32 over inilialization, it holds for
"2
m
(wo,,, x1)² I 22mx> mR²
r=1
(Wor, x1)² < 2mnk> 2 mnR?
Proof. It sulfice to prove the first inequality, and the second follows by summing over To begin, show that each (wo.,,) are Gaussian with zero mean and variance x?? Using independence between entries of Wo.r, we have
E e xx)))
=
t
E
e
two
j=1
1=1
3x²
Σj
o7
c Ar2
where the last equality follows from our assumption that ||x:||2 1. Next, we treat each wi. (wo,r, xi) as random variable. First, we compute the mean of Wr.i
27
E [[wr. = Ewo (worr, x,)² = Rwo
Worrd
<d1=1
Ew W
x²
2
x2
di
Then, we show that each why is sub-exponential with parameter (2x² 2n²). 2x22,kx?)
(unna-x)) e zz d /
E
2V 2T
1 1)( 2 -nn= en² d Wwi
2
π
e tk²
(2x2)-1--
tx2
< 20x²
1 2
50

---

Published in Transactions on Machine Learning Research (08/2022)
1 for t S 2xc2-
Since each Wu,r is independent, have that each Wrys is independent for fixed i. Thus
E
Σ TM
(wne-
II R
e (000)--22)) x2)
S 2
Thus we have
12 Sunk
if St I 2mk
P
Σ Whi 2 mn² +8
2x
if 2mx
We choose t = mn² mRR Since R S x
we have that mx.
S t < mk²
Thus
V2
P
Σ Whi 2 2mx
mR²
Apply union bound over all [ni] gives that with probability at least 1 that
32, holds for all [n]]
Wr,s S 2mx² - -mR²
Lemma 24. If for some R> 0 and all E [m] the initiolization satisfies
(wrrr xi)² < 2maxx -nRR²
and for all [m], holds that l|Wer wo, S R. Then with C₁ (1-(1n) have
FM
M
||
1il12 SCi
EM Tkrr
mk,r ||ue dille
VPCa
1=1
Proof. Using reciprocal moments, we have
0²
EM, [Mwr]] = P (N = EM [nher N =1]
pc
To start, we compute that for # r'. Using the independence mber and mb, we have
P
EM
Mkrr 1=1
m (E m/. )2
EMa me mar")" )7]
Σ EM, [Tte mi EM, [[E m2
{(1 B)EMr Th.r mk.r
= E(1 5) [Mr.r]
For r = =r, we use the idempotent
EM
Mkr
mk (5 mm
= (1 - EE)
mkrr
2
Σ mkr
1=1
(1 - [Mw.r]
51

---

Published in Transactions on Machine Learning Research (08/2022)
Therefore
EM,
1 mikr
(0) "k
(6)
Ok
1(8))
2
MM
mkr
Σ ar(5 (
m
EM
2 Tk.r
m
(E m ) (o(((w
1=1
r==1
EM
2 Thkr
m m m )
o o((wa,r(,))) xx))?
171
5(1
E) EM. [[/e.]]
o ((WA,r',, xx))
p=1
(1 - E)(1 - -2E)
EM [nner] o o((Wkr; xi))
m
0 5)
(W,r' x,)² + (1 g)
(Wi,r,
mp
mpf
r'=l
20 (1 k)K²
20 (1 - =)2
p
mpf
40 (1 K)K²
p
where in the last inequality m _ Thus, we have
2
EM
Mh.r
Σ
marr
ll - 1/.12
Σ EM
Th.r
2
Σ
m
(1)
(((0)
C
1=1
I=
A Also, we have
27
EMa
m llu
0//122
< ENa
mh.r
mi,r
|
0f/l2
1=1
P Th.r
S /pEM,
2
m
[u
0//1
1=1
Plugging in the previous bound gives the desired result.
Lemma 25. If for some and all E [m]] the inilialization salisfies
(Wurr: xi)² < 2mnn2 mnR
and for all E [m)) it holds that ||war Wa.rl22 5 R.
Then we have
EM, [||u - a < 4g(1 x)x
52

---

Published in Transactions on Machine Learning Research (08/2022)
Proof. To start, we have
EM4 (1/)
2/
((3)
EM
Σ ar(E m
m
EM, [[E m2 o((Wkrr:
5(1 E
(Wirrs xx)
m
2E(1 - E)
Σ
(Wo,r: x:)² 2E(11 R?
45(1 E)R²
Therefore,
EM. [||ux 0/.
EM, ("h
a.( )
2
I 4g(1 E)nx?
Lemma 26. Assume that for all i [n], ys satisfies lyil S 1 for some C² 1. Then, we have
E W [lly
uo [ S C²n
(2)
Proof. It is casy to see that Ew..a
u
= O. Now, note that
27
E.o.a
u10
(1) 27
E2
Ewo.a
Σ
m
a (Worr xi)
m
Ewo
Xi))
d
27
EwD
Wo
Ewo
Tw
22
d
Σ
Σ
x
m
r= d'
g²
Therefore,
Ewo
(u-- - ubb)) )
27
2y, Wo..
(1)
u10
+ Ew
110
(1) 2]
vt g²
Thus,
(u: (ws-ulh -16)))
E Ew
[|ly
Upl
Ew...
Σ 1=1
3 + n C
1=1
E
53


Reviews:
[{"review_id": "9GwqgZ0zAc", "reviewer": "Reviewer_1me7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\nThis paper develops a theoretical analysis for a unified training scheme, which can be adapted to many practical scenarios.\n\nWeakness:\n\n1. More motivations are needed to support the goal of this work. Currently, the authors only mention that regularization techniques and distributed training methods are widely studied in recent years. However, this may not be enough to motivate the study of convergence rates of training neural networks with random masks (in the NTK regime). The authors may need to (1) mention the advantage/improvements of using distributed training methods; (2) summarize the existing convergence results (not limited to NN models); and (3) point out why studying convergence is important in the introduction section (before the stated question).\n\n2. I do not quite see whether studying the convergence in the NTK regime is a good approach for training the random masked network (or dropout regularization). Take the dropout regularization as an example, the regularization typically encourages the algorithm to find a solution with lower complexity from a global perspective, while NTK analysis requires the solution to be extremely close to the initialization, which contradicts the spirit of the regularization.\n\n3. I am also not satisfied with the overparameterization condition stated in Theorem 3. Particularly, the number of hidden neurons needs to be larger than at least $O(K)$, where $K$ is the number of outer iterations. This implies the theory cannot support the case that the neural network width is fixed while the training iteration number becomes extremely large.\n\n4. I believe the benefit of using dropout or other random masking approaches is to get a better \"generalization\" ability. The authors may also need to discuss the generalization performance of the solution found by using random masked neurons. Otherwise, as shown in Corollary 1, the optimization guarantee of using dropout is definitely worse than that without using dropout. ", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the convergence of two-layer neural network training with randomly masked neurons. In particular, several contributions are claimed by the authors:\n\nThis paper shows that with randomly masked neurons, the training loss decreases at a linear rate to a certain error region level.\n\nThe developed convergence guarantees imply the convergence rate for a class of empirical training approaches, such as dropout, multi-sample dropout, and multi-worker IST.\n\nThe authors further compare the convergence performance when using categorical masks and Bernoulli masks and show that using  the categorical mask can lead to a more desirable error term.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "The authors may need to make the modification according to the aforementioned weakness. In particular, I would like to see \n\n1. More rigorous motivations.\n2. Why study the convergence in the NTK regime.\n3. If it's possible to remove/improve the dependency on $K$.\n4. Discussion on the generalization.\n", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "4EXQ8lj9X9", "reviewer": "Reviewer_uQXy", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nThe paper’s primary contributions are Thms. 3 and 4, which involve some generalization of the standard NTK analysis. The authors do a good job of explaining the steps involved in this proof at a high level and discussing the different terms that emerge in the results. For example, the discussion of B_1 after Thm. 3 is appreciated.\n\nIn summary, the main claims outlined in the contributions are supported by the theorems. Although, since the bias of the surrogate gradient is identified in the contributions section, it might make sense to move some of the analysis of this into the main text from the supplement.\n\n\n**Weaknesses**\n\nOne point of confusion I have is whether the paper investigates Dropout in a regime that is practically relevant. Setting aside the assumptions required for the NTK analysis, it seems like showing that the surrogate gradient concentrates around the ideal gradient implies that the updates to the weights are the same with and without Dropout. However, Dropout is commonly used as a regularizer, which would require that the weights learned depend on the Dropout rate. This concern does not matter for IST, where the motivation is to reduce computation time.\n\nThe experimental section could do a better job of supporting the paper’s main contributions. While the results are consistent with Thm. 3, it would be better to understand the assumptions of Thms. 3 and 4.\n", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study different algorithms that iteratively train subnetworks of a neural network as a way of training the full network. Examples of this include Dropout and the Independent Subnet Training protocol. Using the neural tangent kernel approach, the main contributions of the analysis are to prove convergence guarantees resulting in linear convergence of the training loss. The other contributions are in service of this main result but are of some independent interest also. For example, the authors identify the NTK corresponding to the subnetworks and show it is close to the infinite-width limit. They also derive the bias in the surrogate gradient compared to the ideal gradient, and show that this bias decreases to zero as the number of subnetworks grows.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "1. Can the authors discuss the differences between their work and Mohtashami et al., 2021? Is the only difference in the differentiability of the activation function? Is there any difference in technique?\n\n2. Can the authors comment on the scale of \\kappa? I believe the standard NTK initialization would set \\kappa=1/sqrt(d), but the discussion around B_1 indicates \\kappa=1/sqrt(n) is the relevant scale. \n\n3. As mentioned above, the experimental section should include some analysis of the effect of width of the results. In particular, there should be finite-size effects from the width. \n\n4. Another claim the paper makes is that categorical sampling has a more desirable error term. The paper would be improved by an experiment supporting this claim.\n\n5. What is the purpose of using a pre-trained model as a feature extractor in the experimental section? \n\n**Minor**\n- Is “error region” a standard term? It was not clear to me what this meant in the abstract.\n- In Fig. 1(c), f_{m_k^1} should be f_{m_k^2} \n- Page 4, “minima” is already plural\n- a_r should be a_m in the vector on page 5\n- “right hand side” should be hyphenated on page 8\n- In Eq 7, why is there an expectation over M_{k-1}? I thought u_k and y did not depend on the mask.\n- I think the y-axis in Fig. 2(a) is incorrect and should just be Training Error, as the axis’ ticks are already in log-scale.\n- \"Conbine” -> Combine on page 21", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JvqVj40ARL", "reviewer": "Reviewer_M1SE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\nThis paper is sound and provides a rigorous analysis for a challenging problem. The authors clearly defined the problem and training algorithm. The assumptions are stated explicitly and notations are mostly clear. The paper is mostly well-written. In particular, proof sketches are helpful to help the readers get the main ideas. The related work are mostly provided though there is room for improvement. \n\nWeakness: \n\nThe setting should be better motivated, in particular, the design of masks and the fact that all workers have access to the entire training samples. The sufficient number of neurons for convergence is another concern. It is also unclear whether local updates are useful given the convergence analysis. Finally, related work and experimental results can be improved.  ", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper focuses on analyzing training error and establishing convergence rates when training a perceptron with a single-hidden layer and multiple randomly masked subnetworks. Masks are randomly sampled at each iteration. The authors used techniques based on the NTK analysis. However, unlike standard NTK analysis, they handled the additional randomness due to mask sampling. They defined a new notion of masked-NTK induced by another mask and show that it becomes positive definite under some assumptions in Theorem 2. Then they consider masked drawn from Bernoulli and Categorical distributions and established convergence guarantees when the number of hidden neurons is sufficiently large and the step size is sufficiently small. They provided results on important and interesting special cases including standard Dropout, Multi-Sample Dropout, and  Multi-Worker Independent Subnetwork Training. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "- Problem formulation: all compute nodes have access to the entire training samples. It is important that the authors further motivate and provide practical settings that fit to the model considered in this paper.\n\n- It seems underlying distributions over which masks are sampled directly impact training and maybe test error. The authors considered two well-known distributions, which is certainly interesting. Beyond those distributions, are there principles to design optimal masks?\n\n- How is the aggregation of weights  performed? Does it depend on the construction of masks and the number of local updates?  \n\n- Page 3 \"Distributed ML .... communication inefficient\" This is partially true. If the tuning/convergence of IST takes a lot of time, the overall speedup gain may disappear.\n\n- Loss at iteration $k$ $L_k$: This loss is in the form of $||\\mathrm{E}[\\dots] ||^2$ where the expectation insider inside the quadratic function is over randomess of masks. It is more common to analyze the expected loss, i.e.,  $\\mathrm{E}[||\\dots ||^2]$. Is it possible to extend the analysis? \n\n- Theorem 2: This boundedness assumption 2 is a bit strange. Could you please clarify when this assumption holds in general? \n\n- Theorem 3: The second term in $B_1$ increases with $n$. Is it possible to improve the bound? In addition, Theorem 3 shows that increasing the number of local steps does not help. Why do they authors consider multiple local steps? \n\n- Section 5: It is not clear how the variance is shown in Figure 2a. \n\nMinor comments: \n\n- $\\eta$ and $\\eta_{k,r}$ in Algorithm 1 should be defined.  \n\n- Please specify the precise definition of co-alignment in Assumption 1. \n\n-  $\\xi$ should be red in two equations right after Theorem 1. \n\n- Corollary 1: Please provide citations for the convergence rate under $\\zeta=1$. Do other work also require  $m=\\Omega(n^5)$?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Dear Authors,

thank you submitting your manuscript, your response and revision to the reviewers questions.  After your response, the reviewers have all agreed that, comparing the papers contents to its claims, it should be accepted, and I agree. 

May I only ask one very minor modifications. Two of the reviewers raised questions about if the assumptions needed for the NTK analysis are well suited to the setting of dropout/random masking. I find the authors response satisfactory, but I would also encourage the authors to include in the abstract a small sentence on the assumptions under which their analysis holds. For instance, perhaps a comment on the need for having a small weight change for the NTK analysis to hold. Or something else similar, just in the spirit of further matching claims of the paper to the results.

The reviewers also asked for more background material and better motivations as relating to the distributed setting. But with the much increased "Related Work" section, they agree that it is now clear and satisfactory.
Decision: Accept (poster)