Title: Towards a Better Theoretical Understanding of Independent Subnetwork Training
Venue: NeurIPS 2023
Accepted: no
Keywords: Optimization

Distributed Training

Federated Learning

Independent Subnetwork Training
Abstract: Modern advancements in large-scale machine learning would be impossible without the paradigm of data-parallel distributed computing. Since distributed computing with large-scale models imparts excessive pressure on communication channels, a lot of recent research was directed towards co-designing communication compression strategies and training algorithms with the goal of reducing communication costs. While pure data parallelism allows better data scaling, it suffers from poor model scaling properties. Indeed, compute nodes are severely limited by memory constraints, preventing further increases in model size. For this reason, the latest achievements in training giant neural network models rely on some form of model parallelism as well. In this work, we take a closer theoretical look at Independent Subnetwork Training (IST), which is a recently proposed and highly effective technique for solving the aforementioned problems. We identify fundamental differences between IST and alternative approaches, such as distributed methods with compressed communication, and provide a precise analysis of its optimization performance on a quadratic model.

1

Abstract

Modern advancements in large-scale machine learning would be impossible without the paradigm of data-parallel distributed computing. Since distributed computing with large-scale models imparts excessive pressure on communication channels, a lot of recent research was directed towards co-designing communication compression strategies and training algorithms with the goal of reducing communication costs. While pure data parallelism allows better data scaling, it suffers from poor model scaling properties. Indeed, compute nodes are severely limited by memory constraints, preventing further increases in model size. For this reason, the latest achievements in training giant neural network models rely on some form of model parallelism as well. In this work, we take a closer theoretical look at Independent Subnetwork Training (IST), which is a recently proposed and highly effective technique for solving the aforementioned problems. We identify fundamental differences between IST and alternative approaches, such as distributed methods with compressed communication, and provide a precise analysis of its optimization performance on a quadratic model.

1 Introduction

A huge part of today's machine learning success drives from the possibility to build more and more complex models and train them on increasingly larger datasets. This fast progress has become feasible due to advancements in distributed optimization, which is necessary for proper scaling when the training data sizes grow [50]. In a typical scenario data parallelism is used for efficiency which consists of sharding the dataset across computing devices. This allowed very efficient scaling and accelerating of training moderately sized models by using additional hardware [19]. Though, such data parallel approach can suffer from communication bottleneck, which sparked a lot of research on distributed optimization with compressed communication of the parameters between nodes [3, 27, 38].

1.1 The need for model parallel

Despite the efficiency gains of data parallelism, it has some fundamental limitations when it comes to scaling up the model size. As the model dimension grows, the amount of memory required to store and update the parameters also increases, which becomes problematic due to resource constraints on individual devices. This has led to the development of model parallelism [11, 37], which splits a large model across multiple nodes, with each node responsible for computations of model parts [15, 47]. However, naive model parallelism also poses challenges because each node can only update its portion of the model based on the data it has access to. This creates a need for a very careful management of communication between devices. Thus, a combination of both data and model parallelism is often necessary to achieve efficient and scalable training of huge models.

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

Algorithm 1 Distributed Submodel (Stochastic) Gradient Descent

IST. Independent Subnetwork Training (IST) is a technique which suggests dividing the neural network into smaller independent sub-parts, training them in a distributed parallel fashion and then aggregating the results to update the weights of the whole model. According to IST, every subnetwork is operational on its own, has fewer parameters than the full model, and this not only reduces the load on computing nodes but also results in faster synchronization. A generalized analog of the described method is formalized as an iterative procedure in Algorithm 1. This paradigm was pioneered by [45] for networks with fully-connected layers and was later extended to ResNets [14] and Graph architectures [43]. Previous experimental studies have shown that IST is a very promising approach for various applications as it allows to effectively combine data with model parallelism and train larger models with limited compute. In addition, [28] performed theoretical analysis of IST for overparameterized single hidden layer neural networks with ReLU activations. The idea of IST was also recently extended to the federated setting via an asynchronous distributed dropout [13] technique.

Federated Learning. Another important setting when the data is distributed (due to privacy reasons) is Federated Learning [22, 27, 31]. In this scenario computing devices are often heterogeneous and more resource-constrained [5] (e.g. mobile phones) in comparison to data-center setting. Such challenges prompted extensive research efforts into selecting smaller and more efficient submodels for local on-device training [2, 6, 8, 12, 20, 21, 29, 35, 42, 44]. Many of these works propose approaches to adapt submodels, often tailored to specific neural network architectures, based on the capabilities of individual clients for various machine learning tasks. However, there is a lack of comprehension regarding the theoretical properties of these methods.

1.2 Summary of contributions

When reviewing the literature, we have found that a rigorous understanding of IST convergence virtually does not exist, which motivates our work. The main contributions of this paper include

A novel approach to analyzing distributed methods that combine data and model parallelism by operating with sparse submodels for a quadratic model.

The first analysis of independent subnetwork training in homogeneous and heterogeneous scenarios without restrictive assumptions on gradient estimators.

Identification of settings when IST can optimize very efficiently or converge not to the optimal solution but only to an irreducible neighborhood which is also tightly characterized.

Experimental validation of the proposed theory through carefully designed illustrative experiments. Due to space limitations, the results (and proofs) are provided in the Appendix.

2 Formalism and Setup

We consider the standard optimization formulation of distributed/federated learning problem [41], 68

where n is the number of clients/workers, each f i : R d → R d represents the loss of the model 69 parameterized by vector x ∈ R d on the data of client i . 70

A typical Stochastic Gradient Descent (SGD) type method for solving this problem has the form 71

where γ > 0 is a stepsize and g k i is a suitably constructed estimator of ∇ f i ( x k ) . In the distributed 72 setting, computation of gradient estimators g k i is typically performed by clients, sent to the server, 73 which subsequently performs aggregation via averaging g k = 1 n ∑ n i =1 g k i . The result is then used to 74 update the model x k +1 via a gradient-type method (2), and at the next iteration the model is broadcast 75 back to the clients. The process is repeated iteratively until a model of suitable qualities is found. 76

77

78

79

80

81

82

One of the main techniques used to accelerate distributed training is lossy communication compression [3, 27, 38]. It suggests applying a (possibly randomized) lossy compression mapping C to a vector/matrix/tensor x before it is transmitted. This saves bits sent per every communication round at the cost of transmitting a less accurate estimate C ( x ) of x . The error caused by this routine also causes convergence issues, and to the best of our knowledge, convergence of IST-based techniques is for this reason not yet understood.

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

Definition 1 (Unbiased compressor) . A randomized mapping C : R d → R d is an unbiased compression operator ( C ∈ U ( ω ) for brevity) if for some ω ≥ 0 and ∀ x ∈ R d

A notable example of a mapping from this class is the random sparsification ( Rand-q for q ∈ { 1 , . . . , d } ) operator defined by

where e 1 , . . . , e d ∈ R d are standard unit basis vectors in R d , and S is a random subset of [ d ] := { 1 , . . . , d } sampled from the uniform distribution on the all subsets of [ d ] with cardinality q . Rand-q belongs to U ( d/q -1) , which means that the more elements are 'dropped' (lower q ), the higher is the variance ω of the compressor.

In this work, we are mainly interested in a somewhat more general class of operators than mere sparsifiers. In particular, we are interested in compressing via the application of random matrices, i.e., via sketching . A sketch C k i ∈ R d × d can be used to represent submodel computations in the following way:

where we require C k i to be a symmetric positive semidefinite matrix. Such gradient estimate corresponds to computing the local gradient with respect to a sparse submodel model C k i x k , and additionally sketching the resulting gradient with the same matrix C k i to guarantee that the resulting update lies in the lower-dimensional subspace.

Using this notion, Algorithm 1 (with one local gradient step) can be represented in the following form 99

which is equivalent to the SGD-type update (2) when perfect reconstruction property holds

where I is the identity matrix (with probability one). This property holds for a specific class of compressors that are particularly useful for capturing the concept of an independent subnetwork partition.

Definition 2 (Permutation sketch) . Assume that model size is greater than number of clients d ≥ n and d = qn , where q ≥ 1 is an integer 1 . Let π = ( π 1 , . . . , π d ) be a random permutation of [ d ] . Then for all x ∈ R d and each i ∈ [ n ] we define Perm-q operator

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

Perm-q is unbiased and can be conveniently used for representing (non-overlapping) structured decomposition of the model such that every client i is responsible for computations over a submodel C i x k .

Our convergence analysis relies on assumption previously used for coordinate descent type methods.

Assumption 1 (Matrix smoothness) . A differentiable function f : R d → R is L -smooth, if there exists a positive semi-definite matrix L ∈ R d × d such that

Standard L -smoothness condition is obtained as a special case of (8) for L = L · I .

2.1 Issues with existing approaches

Consider the simplest gradient type method with compressed model in the single node setting

Algorithms belonging to this family require a different analysis in comparison to SGD [16, 18], Distributed Compressed Gradient Descent [3, 26] and Randomized Coordinate Descent [34, 36] type methods because the gradient estimator is no longer unbiased

̸

That is why such kind of algorithms are harder to analyze. So, prior results for unbiased SGD [25] can not be directly reused. Furthermore, the nature of the bias in this type of gradient estimator does not exhibit additive (zero-mean) noise, thereby preventing the application of previous analyses for biased SGD [1].

An assumption like bounded stochastic gradient norm extensively used in previous works [30, 48] hinders an accurate understanding of such methods. This assumption hides the fundamental difficulty of analyzing biased gradient estimator:

and may not hold even for quadratic functions f ( x ) = x ⊤ A x . In addition, in the distributed setting such condition can result in vacuous bounds [23] as it does not allow to accurately capture heterogeneity.

3 Results in the Interpolation Case

To conduct a thorough theoretical analysis of methods that combine data with model parallelism, we simplify the algorithm and problem setting to isolate the unique effects of this approach. The following considerations are made:

We assume that every node i computes the true gradient at the submodel C i ∇ f i ( C i x k ) .

A notable difference from the original IST algorithm 1 is that workers perform single gradient descent step (or just gradient computation).

Finally, we consider a special case of quadratic model (12) as a loss function (1).

Condition (1) is mainly for the sake of simplicity and clarity of exposition and can be potentially generalized to stochastic gradient computations. (2) is imposed because local steps did not bring any theoretical efficiency improvements for heterogeneous settings until very recently [32]. And even then, only with the introduction of additional control variables, which goes against resourceconstrained device setting. The reason behind (3) is that despite the seeming simplicity quadratic problem has been used extensively to study properties of neural networks [46, 49]. Moreover, it is a non-trivial model which allows to understand complex optimization algorithms [4, 10, 17]. It serves as a suitable problem for observing complex phenomena and providing theoretical insights, which can also be observed in practical scenarios.

Having said that we consider a special case of problem (1) 145

In this case, f ( x ) is L -smooth, and ∇ f ( x ) = L x -b , where L = 1 n ∑ n i =1 L i and b := 1 n ∑ n i =1 b i .

3.1 No linear term: problems and solutions

First, let us examine the case of b i ≡ 0 , which we call interpolation for quadratics, and perform the analysis for general sketches C k i . In this case the gradient estimator (2) takes the form

where B k := 1 n ∑ n i =1 C k i L i C k i . We prove the following result for a method with such an estimator. Theorem 1. Consider the method (2) with estimator (13) for a quadratic problem (12) with L ≻ 0 and b i ≡ 0 . Then if W := 1 2 E [ LB k + B k L ] ⪰ 0 and there exists constant θ > 0 :

and the step size is chosen as 0 < γ ≤ 1 θ , the iterates satisfy

146

147

148

149

150

151

152

153

and 154

155

156

157

158

159

This theorem establishes an O (1 /K ) convergence rate with constant step size up to a stationary point and linear convergence for the expected distance to the optimum. Note that we employ weighted norms in our analysis, as the considered class of loss functions satisfies the matrix L -smoothness Assumption 1. The use of standard Euclidean distance may result in loose bounds that do not recover correct rates for special cases like Gradient Descent.

160

161

162

163

164

165

166

167

168

169

170

It is important to highlight that inequality (14) may not hold (for any θ > 0 ) in the general case as the matrix W is not guaranteed to be positive (semi-)definite in the case of general sampling. The intuition behind it is that arbitrary sketches C k i can result in gradient estimator g k , which is misaligned with the true gradient ∇ f ( x k ) . Specifically, the inner product 〈 ∇ f ( x k ) , g k 〉 can be negative, and there is no expected descent after one step.

Next, we give examples of samplings for which the inequality (14) can be satisfied.

1. Identity. Consider C i ≡ I . Then B k = L , B k LB k = L 3 , W = L 2 ≻ 0 and hence (14) is satisfied for θ = λ max ( L ) . So, (15) says that if we choose γ = 1 θ , then

which exactly matches the rate of Gradient Descent in the non-convex setting. As for iterates convergence, the rate in (16) is λ max ( L ) / λ min ( L ) corresponding to precise Gradient Descent result for strongly convex functions.

2. Permutation. Assume n = d 2 and the use of Perm-1 (special case of Definition 2) sketch C k i = ne π k i e ⊤ π k i , where π k = ( π k 1 , . . . , π k n ) is a random permutation of [ n ] . Then

̸

171

172

173

174

175

176

177

Then 178

where D := 1 n ∑ n i =1 D i , D i := Diag( L i ) . Then inequality (14) leads to

which may not always hold as LD + DL is not guaranteed to be positive definite even in case of L ≻ 0 . However, such kind of condition can be enforced via a slight modification of permutation sketches { ˜ C i } n i =1 , which is done in Section 3.1.2. The limitation of such an approach is that compressors ˜ C i become no longer unbiased.

Remark 1. Matrix W in case of permutation sketches may not be positive-definite. Consider the following homogeneous ( L i ≡ L ) two-dimensional problem example

179

which for c > 2 ab a + b has det( W ) < 0 , and thus W ⊁ 0 according to Sylvester's criterion.

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

Next, we focus on the particular case of Permutation sketches, which are the most suitable for model partitioning according to Independent Subnetwork Training (IST). At the rest of the section, we discuss how the condition (14) can be enforced via a specially designed preconditioning of the problem (12) or modification of sketch mechanism (7).

3.1.1 Homogeneous problem preconditioning

To start consider a homogeneous setting f i ( x ) = 1 2 x ⊤ L x , so L i ≡ L . Now define D = Diag( L ) -diagonal matrix with elements equal to diagonal of L . Then problem can be converted to

which is equivalent to the original problem after a change of variables ˜ x := D -1 2 x . Note that D = Diag( L ) is positive definite as L ≻ 0 , and therefore ˜ L ≻ 0 . Moreover, the preconditioned matrix ˜ L has all ones on the diagonal: Diag( ˜ L ) = I . If we now combine it with Perm-1 sketches

Therefore, inequality (14) takes the form ˜ W = n ˜ L ⪰ 1 θ n 2 ˜ L , which holds for θ ≥ n , and left hand side of (15) can be transformed the following way

for an accurate comparison to standard methods. The resulting convergence guarantee

which matches classical Gradient Descent.

3.1.2 Heterogeneous sketch preconditioning

In contrast to homogeneous case the heterogeneous problem f i ( x ) = 1 2 x ⊤ L i x can not be so easily preconditioned by a simple change of variables ˜ x := D -1 2 x , as every client i has its own matrix L i . However, this problem can be fixed via the following modification of Perm-1 , which scales the output according to the diagonal elements of local smoothness matrix L i :

199

In this case E [ ˜ C i L i ˜ C i ] = I , E [ B k ] = I , and W = L . Then inequality (14) is satisfied for θ ≥ 1 .

If one plugs these results into (15), such convergence guarantee can be obtained 200

which matches the Gradient Descent result as well. Thus we can conclude that heterogeneity does not bring such a fundamental challenge in this scenario. In addition, a method with Perm-1 is significantly better in terms of computational and communication complexity as it requires calculating the local gradients with respect to much smaller submodels and transmits only sparse updates.

This construction also shows that for γ = 1 /θ = 1

which after plugging into the bound for the iterates (16) shows that the method basically converges in 1 iteration. This observation that sketch preconditioning can be extremely efficient, although it uses only the diagonal elements of matrices L i .

Now when we understand that the method can perform very well in the special case of ˜ b i ≡ 0 we can move on to a more complicated situation.

4 Irreducible Bias in the General Case

Now we look at the most general heterogeneous case with different matrices and linear terms f i ( x ) ≡ 1 2 x ⊤ L i x -x ⊤ b i . In this instance gradient estimator (2) takes the form

where C b = 1 n ∑ n i =1 C k i b i . Herewith let us use a heterogeneous permutation sketch preconditioner (23) like in Section 3.1.2 Then E [ B k ] = I and E [ C b ] = 1 √ n ˜ D b , where ˜ D b := 1 n ∑ n i =1 D -1 2 i b i . Furthermore expected gradient estimator (26) results in E [ g k ] = x k -1 √ n ˜ D b and can be transformed the following way

which reflects the decomposition of the estimator into optimally preconditioned true gradient and a bias, depending on the linear terms b i .

4.1 Bias of the method

Estimator (27) can be directly plugged (with proper conditioning) into general SGD update (2)

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

The resulting recursion (28) is exact, and its asymptotic limit can be analyzed. Thus for constant 222 γ < 1 by using the formula for the sum of the first k terms of a geometric series, one gets 223

which shows that in the limit, the first initialization term (with x 0 ) vanishes while the second converges to 1 √ n ˜ D b . This reasoning shows that the method does not converge to the exact solution

̸

which for the positive-definite L can be defined as x ⋆ = L -1 b , while x ∞ = 1 n √ n ∑ n i =1 D -1 2 i b i . So, 224 in general, there is an unavoidable bias. However, in the limit case: n = d →∞ , the bias diminishes. 225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

4.2 Generic convergence analysis

While the analysis in Section 4.1 is precise, it does not allow us to compare the convergence of IST to standard optimization methods. Due to this, we also analyze the non-asymptotic behavior of the method to understand the convergence speed. Our result is formalized in the following theorem.

Theorem 2. Consider the method (2) with estimator (26) for a quadratic problem (12) with the positive definite matrix L ≻ 0 . Assume that for every D i := Diag( L i ) matrices D -1 2 i exist, scaled permutation sketches (23) are used and heterogeneity is bounded as E [ ∥ ∥ g k -E [ g k ]∥ ∥ 2 L ] ≤ σ 2 . Then for step size is chosen as

where γ c,β ∈ (0 , 1] for β ∈ (0 , 1 / 2) , the iterates satisfy

Note that the derived convergence upper bound has a neighborhood proportional to the bias of the gradient estimator h and level of heterogeneity σ 2 . Some of these terms with factor γ can be eliminated via decreasing learning rate schedule (e.g., ∼ 1 / √ k ). However, such a strategy does not diminish the term with a multiplier 2 β -1 (1 -γ ) , making the neighborhood irreducible. Moreover, this term can be eliminated for γ = 1 , which also minimizes the first term that decreases as 1 /K . Though, such step size choice maximizes the terms with factor γ . Furthermore, there exists an inherent trade-off between convergence speed and the size of the neighborhood.

In addition, convergence to the stationary point is measured in the weighted by L -1 squared norm of the gradient. At the same time, the neighborhood term depends on the weighted by L norm of h . This fine-grained decoupling is achieved by carefully applying Fenchel-Young inequality and provides a tighter characterization of the convergence compared to using standard Euclidean distances.

Homogeneous case. In this scenario, every worker has access to the all data f i ( x ) ≡ 1 2 x ⊤ L x -x ⊤ b . Then diagonal preconditioning of the problem can be used as in the previous Section 3.1.1. This results in a gradient ∇ f ( x ) = ˜ L x -˜ b for ˜ L = D -1 2 LD -1 2 and ˜ b = D -1 2 b . If it is further combined with a scaled by 1 / √ n Permutation sketch C i := √ ne π i e ⊤ π i , the resulting gradient estimator is

for ˜ h = ˜ L -1 ˜ b -1 √ n ˜ b . In this case heterogeneity term σ 2 from upper bound (30) disappears as E [ ∥ ∥ g k -E [ g k ]∥ ∥ 2 L ] = 0 , thus the neighborhood size can significantly decrease. However, the bias term depending on ˜ h still remains as the method does not converge to the exact solution x k → x ∞ = x ⋆ = ˜ L -1 ˜ b for positive-definite ˜ L . Nevertheless the method's fixed point x ∞ = ˜ b / √ n and solution x ⋆ can coincide when ˜ L -1 ˜ b = 1 √ n ˜ b , which means that ˜ b is the right eigenvector of matrix ˜ L -1 with eigenvalue 1 √ n .

̸

Let us contrast obtained result (30) with non-convex rate of SGD [25] with constant step size γ for L -smooth and lower-bounded f

where constant C depends, for example, on the variance of stochastic gradient estimates. Observe that the first term in the compared upper bounds (32) and (30) is almost identical and decreases with speed 1 /K . But unlike (30) the neighborhood for SGD can be completely eliminated by reducing the step size γ . This highlights a fundamental difference of our results to unbiased methods.

The intuition behind this issue is that for SGD-type methods like Compressed Gradient Descent 263

the gradient estimate is unbiased and enjoys the property that variance

goes down to zero as the method progresses because ∇ f ( x k ) →∇ f ( x ⋆ ) = 0 in the unconstrained case. In addition, any stationary point x ⋆ ceases to be a fixed point of the iterative procedure as

̸

in the general case, unlike for Compressed Gradient Descent with both biased and unbiased compressors C . So, even if the method (computing gradient at sparse model) is initialized from the solution after one gradient step, it may get away from there.

4.3 Comparison to previous works

Independent Subnetwork Training [45]. There are several improvements over the previous works that tried to theoretically analyze the convergence of Distributed IST.

The first difference is that our results allow for an almost arbitrary level of model sparsification, i.e., work for any ω ≥ 0 as permutation sketches can be viewed as a special case of compression operators (1). This improves significantly over the work of [45], which demands 3 ω ≲ µ 2 / L 2 . Such a requirement is very restrictive as the condition number L/µ of the loss function f is typically very large for any non-trivial optimization problem. Thus, the sparsifier's (4) variance ω = d/q -1 has to be very close to 0 and q ≈ d . So, the previous theory allows almost no compression (sparsification) because it is based on the analysis of Gradient Descent with Compressed Iterates [24].

The second distinction is that the original IST work [45] considered a single node setting and thus their convergence bounds did not capture the effect of heterogeneity, which we believe is of crucial importance for distributed setting [9, 39]. Besides, they consider Lipschitz continuity of the loss function f , which is not satisfied for a simple quadratic model. A more detailed comparison including additional assumptions on the gradient estimator made in [45] is presented in the Appendix.

FL with Model Pruning. In a recent work [48] made an attempt to analyze a variant of the FedAvg algorithm with sparse local initialization and compressed gradient training (pruned local models). They considered a case of L -smooth loss and sparsification operator satisfying a similar condition to (1). However, they also assumed that the squared norm of stochastic gradient is uniformly bounded (11), which is 'pathological' [23] especially in the case of local methods as it does not allow to capture the very important effect of heterogeneity and can result in vacuous bounds.

In the Appendix we show some limitations of other relevant previous approaches to training with compressed models: too restrictive assumptions on the algorithm [33] or not applicability in our problem setting [7].

5 Conclusions and Future Work

In this study, we introduced a novel approach to understanding training with combined model and data parallelism for a quadratic model. This framework allowed to shed light on distributed submodel optimization which revealed the advantages and limitations Independent Subnetwork Training (IST). Moreover, we accurately characterized the behavior of the considered method in both homogeneous and heterogeneous scenarios without imposing restrictive assumptions on gradient estimators.

In future research, it would be valuable to explore extensions of our findings to settings that are closer to practical scenarios, such as cross-device federated learning. This could involve investigating partial participation support, leveraging local training benefits, and ensuring robustness against stragglers. Additionally, it would be interesting to generalize our results to non-quadratic scenarios without relying on pathological assumptions.

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

References

Ahmad Ajalloeian and Sebastian U Stich. On the convergence of SGD with biased gradients. arXiv preprint arXiv:2008.00051 , 2020.

Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. FedRolex: Model-heterogeneous federated learning with rolling sub-model extraction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communication-efficient SGD via gradient quantization and encoding. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Yossi Arjevani, Ohad Shamir, and Nathan Srebro. A tight convergence analysis for stochastic gradient descent with delayed updates. In Algorithmic Learning Theory , pages 111-132. PMLR, 2020.

Sebastian Caldas, Jakub Koneˇ cny, H Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.07210 , 2018.

Zachary Charles, Kallista Bonawitz, Stanislav Chiknavaryan, Brendan McMahan, et al. Federated select: A primitive for communication-and memory-efficient federated learning. arXiv preprint arXiv:2208.09432 , 2022.

El Mahdi Chayti and Sai Praneeth Karimireddy. Optimization with access to auxiliary information. arXiv preprint arXiv:2206.00395 , 2022.

Yuanyuan Chen, Zichen Chen, Pengcheng Wu, and Han Yu. Fedobd: Opportunistic block dropout for efficiently training large-scale neural networks through federated learning. arXiv preprint arXiv:2208.05174 , 2022.

Sélim Chraibi, Ahmed Khaled, Dmitry Kovalev, Peter Richtárik, Adil Salim, and Martin Takáˇ c. Distributed fixed point methods with compressed iterates. arXiv preprint arXiv:2102.07245 , 2019.

Leonardo Cunha, Gauthier Gidel, Fabian Pedregosa, Damien Scieur, and Courtney Paquette. Only tails matter: Average-case universality and robustness in the convex regime. In International Conference on Machine Learning , pages 4474-4491. PMLR, 2022.

Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc'aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems , 25, 2012.

Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint arXiv:2010.01264 , 2020.

Chen Dun, Mirian Hipolito, Chris Jermaine, Dimitrios Dimitriadis, and Anastasios Kyrillidis. Efficient and light-weight federated learning via asynchronous distributed dropout. arXiv preprint arXiv:2210.16105 , 2022.

Chen Dun, Cameron R Wolfe, Christopher M Jermaine, and Anastasios Kyrillidis. ResIST: Layer-wise decomposition of resnets for distributed training. In Uncertainty in Artificial Intelligence , pages 610-620. PMLR, 2022.

Philipp Farber and Krste Asanovic. Parallel neural network training on multi-spert. In Proceedings of 3rd International Conference on Algorithms and Architectures for Parallel Processing , pages 659-666. IEEE, 1997.

Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. A unified theory of SGD: Variance 348 reduction, sampling, quantization and coordinate descent. In International Conference on 349 Artificial Intelligence and Statistics , pages 680-690. PMLR, 2020. 350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

Baptiste Goujaud, Damien Scieur, Aymeric Dieuleveut, Adrien B Taylor, and Fabian Pedregosa. Super-acceleration with cyclical step-sizes. In International Conference on Artificial Intelligence and Statistics , pages 3028-3065. PMLR, 2022.

Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. Proceedings of the 36th International Conference on Machine Learning, Long Beach, California , 2019.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2018.

Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems , 34:12876-12889, 2021.

Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning on edge devices. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Oliveira, Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaïd Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Koneˇ cný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Found. Trends Mach. Learn. , 14(1-2):1-210, 2021.

Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. Tighter theory for local SGD on identical and heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 4519-4529. PMLR, 2020.

Ahmed Khaled and Peter Richtárik. Gradient descent with compressed iterates. arXiv preprint arXiv:1909.04716 , 2019.

Ahmed Khaled and Peter Richtárik. Better theory for SGD in the nonconvex world. Transactions on Machine Learning Research , 2023. Survey Certification.

Sarit Khirirat, Hamid Reza Feyzmahdavian, and Mikael Johansson. Distributed learning with compressed gradients. arXiv preprint arXiv:1806.06573 , 2018.

Jakub Koneˇ cný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. NIPS Private Multi-Party Machine Learning Workshop , 2016.

Fangshuo Liao and Anastasios Kyrillidis. On the convergence of shallow neural network training with randomly masked neurons. Transactions on Machine Learning Research , 2022.

Rongmei Lin, Yonghui Xiao, Tien-Ju Yang, Ding Zhao, Li Xiong, Giovanni Motta, and Françoise Beaufays. Federated pruning: Improving neural network efficiency with federated learning. arXiv preprint arXiv:2209.06359 , 2022.

Tao Lin, Sebastian U Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. Dynamic model pruning with feedback. In International Conference on Learning Representations , 2019.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficient Learning of Deep Networks from Decentralized Data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , volume 54 of Proceedings of Machine Learning Research , pages 1273-1282, 20-22 Apr 2017.

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

Konstantin Mishchenko, Grigory Malinovsky, Sebastian Stich, and Peter Richtárik. ProxSkip: Yes! local gradient steps provably lead to communication acceleration! finally! In International Conference on Machine Learning , pages 15750-15769. PMLR, 2022.

Amirkeivan Mohtashami, Martin Jaggi, and Sebastian Stich. Masked training of neural networks with partial gradients. In International Conference on Artificial Intelligence and Statistics , pages 5876-5890. PMLR, 2022.

Yu Nesterov. Efficiency of coordinate descent methods on huge-scale optimization problems. SIAM Journal on Optimization , 22(2):341-362, 2012.

Xinchi Qiu, Javier Fernandez-Marques, Pedro PB Gusmao, Yan Gao, Titouan Parcollet, and Nicholas Donald Lane. ZeroFL: Efficient on-device training for federated learning with local sparsity. In International Conference on Learning Representations , 2022.

Peter Richtárik and Martin Takáˇ c. Iteration complexity of randomized block-coordinate descent methods for minimizing a composite function. Mathematical Programming , 144(1-2):1-38, 2014.

Peter Richtárik and Martin Takáˇ c. Distributed coordinate descent method for learning with big data. Journal of Machine Learning Research , 17(75):1-25, 2016.

Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns. In Fifteenth Annual Conference of the International Speech Communication Association , 2014.

Egor Shulgin and Peter Richtárik. Shifted compression framework: Generalizations and improvements. In The 38th Conference on Uncertainty in Artificial Intelligence , 2022.

Rafał Szlendak, Alexander Tyurin, and Peter Richtárik. Permutation compressors for provably faster distributed nonconvex optimization. In International Conference on Learning Representations , 2022.

Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan AlShedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al. A field guide to federated optimization. arXiv preprint arXiv:2107.06917 , 2021.

Dingzhu Wen, Ki-Jun Jeon, and Kaibin Huang. Federated dropout-a simple approach for enabling federated learning on resource constrained devices. IEEE Wireless Communications Letters , 11(5):923-927, 2022.

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 arXiv:2102.10424 , 2021.

Tien-Ju Yang, Dhruv Guliani, Françoise Beaufays, and Giovanni Motta. Partial variable training for efficient on-device federated learning. In ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 4348-4352. IEEE, 2022.

Binhang Yuan, Cameron R Wolfe, Chen Dun, Yuxin Tang, Anastasios Kyrillidis, and Chris Jermaine. Distributed learning of fully connected neural networks using independent subnet training. Proceedings of the VLDB Endowment , 15(8):1581-1590, 2022.

Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems , 32, 2019.

Xiru Zhang, Michael Mckenna, Jill Mesirov, and David Waltz. An efficient implementation of the back-propagation algorithm on the connection machine cm-2. Advances in neural information processing systems , 2, 1989.

Hanhan Zhou, Tian Lan, Guru Venkataramani, and Wenbo Ding. On the convergence of heterogeneous federated learning with arbitrary adaptive online model pruning. arXiv preprint arXiv:2201.11803 , 2022.

Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic 449 models for understanding neural network dynamics. arXiv preprint arXiv:2205.11787 , 2022. 450

Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient 451 descent. Advances in neural information processing systems , 23, 2010. 452

Peer reviews:
comments: SUMMARY:
The authors provided a theoretical analysis for independent subnet training and provided a convergence analysis in the case where communication compression is present.

The authors discussed the scenario when bias is not present and provided two analyses in the homogeneous and heterogeneous case, respectively before extending their theorems to the case with bias.

STRENGTHS:
1. Sections 3 and 4 are well written and clear, which break the scenarios down in an intuitive way and presented the theorems and outlines clearly.

WEAKNESSES:
1. Introduction could be more straightforward and can dive straight into the main technical contributions of the work. It was not clear why the problem is well motivated and what are the main technical hurdles until reading section 2 and onwards. A clearer presentation in the intro can make the paper much more readable and well-motivated.

QUESTIONS:
1. The authors mentioned that prior work on the convergence of IST focus "on overparameterized single hidden layer neural networks with ReLU activations". It is not entirely clear to me why the authors considered quadratic form and what is the tradeoff between their work and prior work in this regard. A more thorough explanation would be appreciated.

LIMITATIONS:
The work is theoretical and does not seem to have any potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a theoretical analysis of the convergence properties of Independent Subnetwork Training (IST), for distributed Stochastic Gradient Descent (SGD) optimization with a quadratic loss function. The analysis considers both the cases of homogeneous and heterogeneous distributed scenarios, without restrictive assumptions on the gradient estimator.\n\nThe work characterizes situations where IST converges very efficiently, and cases where it does not converge to the optimal solution but to an irreducible neighbourhood. Experimental results that validate the theory are provided in the Appendix.\n\nSTRENGTHS:\nThe paper provides a solid analytical treatment of the important problem of distributed optimization with reduced communication overhead by means of Independent Subnetwork Training (IST).\n\nCompared to previous work, the analysis of the paper does not rely on the restrictive assumption of a bounded stochastic gradient norm.\n\nThe paper is well written – the exposition is clear, and the material is well structured and well presented.\n\nWEAKNESSES:\nThe work considers distributed Stochastic Gradient Descent (SGD) training with a quadratic loss. As mentioned by the authors in Section 3, a simple quadratic loss function has been used in other work to analyze properties of neural networks. While this loss function can still provide interesting theoretical insights, it would be valuable to extend the analysis and the experimental results to more generally used loss functions.\n\nMinor comments:\n\n-- Line 17: \"drives from\" may be changed to \"derives from\".\n\n-- Equation after line 217: it seems that the first part of the equation\n\"$\\mathbb{E}[g^k] = \\bar{\\mathbf{L}}^{-1} \\bar{\\mathbf{L}} x^k \\pm \\bar{\\mathbf{L}}^{-1} \\bar{b} - \\frac{1}{\\sqrt{n}} \\tilde{\\mathbf{D} b} = $ ...\" \nshould be rewritten as\n\"$\\mathbb{E}[g^k] = \\bar{\\mathbf{L}}^{-1} \\bar{\\mathbf{L}} x^k - \\frac{1}{\\sqrt{n}} \\tilde{\\mathbf{D} b} = $ ...\".\n\n-- Equation after line 221: it seems that the first part of the equation\n\"$\\mathbb{E}[x^{k+1}] = x^k - \\gamma \\mathbb{E}[g^k] = $ ...\"\nshould be\n\"$\\mathbb{E}[x^{k+1}] = \\mathbb{E}[x^k] - \\gamma \\mathbb{E}[g^k] = $ ...\".\n\nQUESTIONS:\nIt would be valuable if the authors could consider a discussion or possible additional experiments to extend some of the results and insights presented in the paper to the case of distributed optimization with more generally used loss functions.\n\nLIMITATIONS:\nThe work relates to distributed training of large-scale models, which generally correspond to significant power consumption and CO2 emissions. However, the IST method studied in the paper aims at allowing distributed training with reduced communication overhead, corresponding to potentially reduced power consumption.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIndependent Subnetwork Training (IST) is a technique that divides the neural network into smaller independent subnetworks, trains them in a distributed parallel way, and aggregates the results of each independent subnetwork to update the weights of the whole model.\nThis paper aims to analyze the behavior of IST theoretically. Specifically, it considers a quadratic model trained by IST. It conducts convergence analysis under both homogeneous and heterogeneous scenarios and shows that IST can only converge to an irreducible neighborhood of optimal solution.\n\nSTRENGTHS:\nThis is probably the first work providing a thorough theoretical analysis of IST.\n\nWEAKNESSES:\n1. This paper should include a more comprehensive motivation for the theoretical study of IST. This could involve discussing the potential limitations of current IST architectures and how a theoretical analysis can guide future modifications to improve their performance. By doing so, reviewers will have a clearer understanding of the significance of the paper's findings and how they can be applied in practice.\n2. The main body of this paper does not have any experimental results. The authors should include some key experiments to validate their theoretical analysis.\n3. The authors should consider expanding the scope of their experiments beyond quadratic models to include other types of models that are commonly used in SOTA IST papers, e.g., ResNet, and Graph Convolutional Networks, as listed in this paper's reference. This would allow reviewers to better understand the generality of the paper's findings and how they can be applied to real-world applications.\n\nQUESTIONS:\n1. How do the findings of this work be of any help to the future design of IST architecture? In addition, the authors should include more experiments on SOTA IST applications, e.g., ResNet, and Graph Convolutional Networks, to indicate the generality and significance of this paper's findings.\n2. The authors should clarify the assumptions made in the permutation example in Section 3.1. Specifically, they should explain the case that n=d^2 and the use of Perm-1 sketch, while according to definition 2, perm-q refers to d=qn, which leads to n=d=1 and it’s a naïve configuration. By clarifying the assumptions made in this section, the reader can better understand the example and its implications for IST.\n\nLIMITATIONS:\nFuture works are discussed in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis submission presents a theoretical analysis for the independent subnetwork training (IST) algorithm for training models under both data and model parallel settings. The convergence guarantees are analyzed for quadratic loss functions, when using permutation sketches as the model compressors.\n\nSTRENGTHS:\n- Optimization using data and model parallelization is an important practical problem, for which more theoretical insights are welcome \n- The submission presents a fairly thorough analysis of the convergence guarantees for IST for the quadratic loss function analyzed\n- The authors highlight some limitations of IST, which would be useful to be aware of (e.g. the fact that in the general case of the quadratic function the algorithm does not converge to the solution)\n- The submission is overall well written, and the authors are careful to introduce the setup and all the assumptions used in their analysis\n\nWEAKNESSES:\n- Overall the analysis presented in this submission is very limited, as it only deals with a specific type of loss function, which is not a practical instance where model parallelization would be useful\n- While the authors argue that the quadratic model has been previously used in the literature for studying neural networks (lines 140-144), in my understanding this model still relies on a Taylor approximation for the loss function of a neural network (for non-linear models). It is therefore not clear how the error introduced through this approximation would translate into the convergence analysis presented in the submission\n- The authors use additional simplifying assumptions, such as the fact that each node can compute the true gradient of its submodel, which would be infeasible in the case of large datasets. Additionally, the results are only presented for Perm-1 sketches when the number of nodes matches the dimension of the model, which is again not a practical use-case. While the authors argue that their results can generalize beyond these limitations, a more general formulation is not provided in the submission.\n- Other works have analyzed the convergence guarantees for IST, notably [28] has done so for a one hidden layer network with ReLU activations, which is a more general case than the one from this submission. It is not clear what are the additional insights presented here, compared to the previous work.\n\nQUESTIONS:\n- Can the authors please detail how their choice of the quadratic loss function ties to practical applications of model parallelism (e.g. for neural networks)?\n- In Equation (12), do the matrices $L_i$ also need to be semi-positive definite?\n- For Theorems 1 and 2, can the authors please describe how the results would change when using an unbiased gradient estimator, instead of the full gradient for each node?\n- Are there any situations, without preconditioning, which would satisfy the conditions from Theorem 1, for model parallelization? The Identity example from Section 3.1 seems to only apply to data parallelization. \n- For Equation (23), I am confused by the notation for the scaling coefficient. Should it be $[(L_i)_{\\pi_i, \\pi_i}]^{-1/2 }$ instead?\n- From the analysis in Section 3.1.2 it looks like the fact that $C_i$ is biased does not have any effect on the bound. Can the authors provide more insights on why that is the case?\n- In Appendix B.4. the authors show generalizations for $d>n$ and for different sketches. I think these would be useful to highlight in the main submission, since they correspond to a more practical setting.\n\nLIMITATIONS:\nI believe the authors have properly addressed the limitations of their analysis. However, I am not convinced that the contributions presented in this submission are broad enough, which ultimately motivates my score.\n\n\n---------------------------------\n**Edited after rebuttal**\n\nAfter reading the authors' answers, and the other reviews, I decided to increase my overall rating to 5, as well as the scores for Soundness, Presentation and Contribution.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}