summaries: The paper proposes a new method to solve the finite sum problem in federated learning where each local model is an expectation function. The STEM algorithm uses momentum variance-reduced estimator in local updates (similar to STORM estimator in stochastic optimization) and also send these estimators to the server when performing server-side update. Therefore, STEM includes momentum updates for both local and server updates. Convergence analysis reveals that STEM nearly match the best known sample complexity in centralized setting and it also matches the lower bound communication complexity up to a log factor. Numerical experiments on neural network training not only show advantage of STEM compared with FedAvg and SCAFFOLD but also presenting the tradeoff between communication frequency and minibatch size.

summaries: This paper introduced a two-step momentum method for federated learning. The proposed method achieves the best achievable sample and communication complexities for first-order stochastic FL algorithms. The theoretical analysis provides guidance on selecting synchronization interval and batch size for achieving the nearly optimal complexities.

summaries: The authors presented a momentum extension of the classic FedAvg method for non-convex federated learning, along with a unified framework for convergence analysis. The core result shows that the proposed method can achieve near-tight sample and communication complexities under mild conditions, whilst also revealing a trade-off between local update frequencies and local minibatch sizes. The actual performance of the proposed method was evaluated in CIFAR-10 and MNIST image classification tasks.

summaries: Authors propose Stochastic Two-Sided Momentum algorithm, that utilizes certain momentum-assisted stochastic gradient directions for both the WNs and SN updates. Authors show that this method achieves near optimal sample complexity $\tilde{\mathcal{O}}\left(\epsilon^{-3 / 2}\right)$ and $\tilde{\mathcal{O}}\left(\epsilon^{-1}\right)$ communication complexity.  Also, authors provide interesting observations about trade-off between the minibatch sizes and the local update frequency.

summaries: The paper proposes STEM for federated learning under *sample gradient smoothness condition*. By using momentum-based techniques, the proposed STEM algorithm achieves the optimal sample complexity of $O(\epsilon^{-3/2})$ and communication complexity of $O(\epsilon^{-1})$ for non-convex functions. The paper also shows how to balance the batch size and the number of local steps.

abstractText: Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs’ and the server’s update directions, the minibatch sizes, and the number of local updates, so that the WNs use the minimum number of samples and communication rounds to achieve the desired solution. This work addresses the above question and considers a class of stochastic algorithms where the WNs perform a few local updates before communication. We show that when both the WN’s and the server’s directions are chosen based on certain stochastic momentum estimator, the algorithm requires Õ(✏ 3/2) samples and Õ(✏ 1) communication rounds to compute an ✏-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such near-optimal sample and communication complexities simultaneously. Further, we show that there is a trade-off curve between the number of local updates and the minibatch sizes, on which the above sample and communication complexities can be maintained. Finally, we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the number of local updates, WNs’ and server’s update directions, and minibatch sizes to achieve the best performance.

authors:
J. Konečnỳ
H.B. McMahan
D. Ramage
P. Richtárik

title: Federated optimization: Distributed machine learning for on-device intelligence

venue: arXiv preprint arXiv:1610.02527, 2016.

authors:
M. Li
D.G. Andersen
A.J. Smola
K. Yu

title: Communication efficient distributed machine learning with the parameter server

venue: Advances in Neural Information Processing Systems 27, Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger, Eds. Curran Associates, Inc., 2014, pp. 19–27.

authors:
J. Dean
G. Corrado
R. Monga
K. Chen
M. Devin
M. Mao
M. Ranzato
A. Senior
P. Tucker
K. Yang

title: Large scale distributed deep networks

venue: Advances in neural information processing systems, 2012, pp. 1223–1231.

authors:
T. Léauté
B. Faltings

title: Protecting privacy through distributed computation in multi-agent decision making

venue: Journal of Artificial Intelligence Research, vol. 47, pp. 649–695, 2013.

authors:
C. Fang
C.J. Li
Z. Lin
T. Zhang

title: Spider: Near-optimal non-convex optimization via stochastic path-integrated differential estimator

venue: Advances in Neural Information Processing Systems, 2018, pp. 689–699.

authors:
D. Zhou
P. Xu
Q. Gu

title: Stochastic nested variance reduction for nonconvex optimization

venue: arXiv preprint arXiv:1806.07811, 2018.

authors:
A. Cutkosky
F. Orabona

title: Momentum-based variance reduction in non-convex SGD

venue: Advances in Neural Information Processing Systems 32. Curran Associates, Inc., 2019, pp. 15 236–15 245.

authors:
Q. Tran-Dinh
N.H. Pham
D.T. Phan
L.M. Nguyen

title: Hybrid stochastic gradient descent algorithms for stochastic nonconvex optimization

venue: arXiv preprint arXiv:1905.05920, 2019.

authors:
X. Zhang
M. Hong
S. Dhople
W. Yin
Y. Liu

title: Fedpd: A federated learning framework with adaptivity to non-iid data

venue: IEEE Transactions on Signal Processing, pp. 1–1, 2021.

authors:
T. Li
A.K. Sahu
M. Zaheer
M. Sanjabi
A. Talwalkar
V. Smith

title: Federated optimization in heterogeneous networks

venue: arXiv preprint arXiv:1812.06127, 2018.

authors:
B. McMahan
E. Moore
D. Ramage
S. Hampson
B.A. y Arcas

title: Communication-efficient learning of deep networks from decentralized data

venue: Artificial Intelligence and Statistics. PMLR, 2017, pp. 1273–1282.

authors:
H. Yu
S. Yang
S. Zhu

title: Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning

venue: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 33, no. 01, 2019, pp. 5693–5700.

authors:
B. Woodworth
K.K. Patel
N. Srebro

title: Minibatch vs local sgd for heterogeneous distributed learning

venue: arXiv preprint arXiv:2006.04735, 2020.

authors:
H. Yu
R. Jin
S. Yang

title: On the linear speedup analysis of communication efficient momentum sgd for distributed non-convex optimization

venue: International Conference on Machine Learning. PMLR, 2019, pp. 7184–7193.

authors:
S.P. Karimireddy
S. Kale
M. Mohri
S. Reddi
S. Stich
A.T. Suresh

title: Scaffold: Stochastic controlled averaging for federated learning

venue: International Conference on Machine Learning. PMLR, 2020, pp. 5132–5143.

authors:
H. Yang
M. Fang
J. Liu

title: Achieving linear speedup with partial worker participation in non-iid federated learning

venue: arXiv preprint arXiv:2101.11203, 2021.

authors:
S.P. Karimireddy
M. Jaggi
S. Kale
M. Mohri
S.J. Reddi
S.U. Stich
A.T. Suresh

title: Mime: Mimicking centralized stochastic algorithms in federated learning

venue: arXiv preprint arXiv:2008.03606, 2020.

authors:
R. Das
A. Hashemi
S. Sanghavi
I.S. Dhillon

title: Improved convergence rates for non-convex federated learning with compression

authors:
B. Woodworth
K.K. Patel
S.U. Stich
Z. Dai
B. Bullins
H.B. McMahan
O. Shamir
N. Srebro

title: Is local sgd better than minibatch sgd?

authors:
H. Yu
R. Jin

title: On the computation and communication complexity of parallel sgd with dynamic batch sizes for stochastic non-convex optimization

venue: International Conference on Machine Learning. PMLR, 2019, pp. 7174–7183.

authors:
J. Wang
G. Joshi

title: Cooperative sgd: A unified framework for the design and analysis of local-update sgd algorithms

venue: Journal of Machine Learning Research, vol. 22, no. 213, pp. 1–50, 2021.

authors:
A. Khaled
K. Mishchenko
P. Richtárik

title: Better communication complexity for local sgd

title: Local sgd converges fast and communicates little

venue: arXiv preprint arXiv:1805.09767, 2018.

authors:
T. Lin
S.U. Stich
K.K. Patel
M. Jaggi

title: Don’t use large mini-batches, use local sgd

venue: International Conference on Learning Representations, 2020.

authors:
F. Zhou
G. Cong

title: On the convergence properties of a k-step averaging stochastic gradient descent algorithm for nonconvex optimization

venue: Proceedings of the Twenty-Seventh International Joint Conference on Artificial Intelligence, IJCAI-18, 7 2018, pp. 3219–3227.

authors:
F. Sattler
S. Wiedemann
K.-R. Müller
W. Samek

title: Robust and communication-efficient federated learning from non-iid data

venue: IEEE transactions on neural networks and learning systems, vol. 31, no. 9, pp. 3400–3413, 2019.

authors:
Y. Zhao
M. Li
L. Lai
N. Suda
D. Civin
V. Chandra

title: Federated learning with non-iid data

venue: arXiv preprint arXiv:1806.00582, 2018.

authors:
J. Wang
V. Tantia
N. Ballas
M. Rabbat

title: Slowmo: Improving communication-efficient distributed sgd with slow momentum

venue: arXiv preprint arXiv:1910.00643, 2019.

authors:
X. Liang
S. Shen
J. Liu
Z. Pan
E. Chen
Y. Cheng

title: Variance reduced local sgd with lower communication complexity

venue: arXiv preprint arXiv:1912.12844, 2019.

authors:
P. Sharma
P. Khanduri
S. Bulusu
K. Rajawat
P.K. Varshney

title: Parallel restarted SPI- DER – communication efficient distributed nonconvex optimization with optimal computation complexity

venue: arXiv preprint arXiv:1912.06036, 2019.

authors:
S.J. Reddi
S. Kale
S. Kumar

title: On the convergence of adam and beyond

venue: arXiv preprint arXiv:1904.09237, 2019.

authors:
A. Koloskova
N. Loizou
S. Boreiri
M. Jaggi
S. Stich

title: A unified theory of decentralized sgd with changing topology and local updates

venue: International Conference on Machine Learning. PMLR, 2020, pp. 5381–5393.

authors:
R. Johnson
T. Zhang

title: Accelerating stochastic gradient descent using predictive variance reduction

venue: Advances in Neural Information Processing Systems 26. Curran Associates, Inc., 2013, pp. 315–323.

authors:
L. Bottou
F.E. Curtis
J. Nocedal

title: Optimization methods for large-scale machine learning

authors:
Y. Drori
O. Shamir

title: The complexity of finding stationary points with stochastic gradient descent

venue: International Conference on Machine Learning. PMLR, 2020, pp. 2658–2667.

authors:
D.A.E. Acar
Y. Zhao
R. Matas
M. Mattina
P. Whatmough
V. Saligrama

title: Federated learning based on dynamic regularization

authors:
S. Caldas
S.M.K. Duddu
P. Wu
T. Li
J. Konečný
H.B. McMahan
V. Smith
A. Talwalkar

title: Leaf: A benchmark for federated settings

heading: 1 Introduction

text:
In Federated Learning (FL), multiple worker nodes (WNs) collaborate with the goal of learning a joint model, by only using local data. Therefore it has become popular for machine learning problems where datasets are massively distributed [1]. In FL, the data is often collected at or off-loaded to multiple WNs which in collaboration with a server node (SN) jointly aim to learn a centralized model
⇤Corresponding Author: Mingyi Hong.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
[2, 3]. The local WNs share the computational load and since the data is local to each WN, FL also provides some level of data privacy [4]. A classical distributed optimization problem that K WNs aim to solve:
min x2Rd
⇢ f(x) := 1
K
KX
k=1
f (k)(x) := 1
K
KX
k=1
E⇠(k)⇠D(k) [f (k)(x; ⇠(k))] . (1)
where f (k) : Rd ! R denotes the smooth (possibly non-convex) objective function and ⇠(k) ⇠ D(k) represents the sample/s drawn from distribution D(k) at the kth WN with k 2 [K]. When the distributions D(k) are different across the WNs, it is referred to as the heterogeneous data setting.
The optimization performance of non-convex FL algorithms is typically measured by the total number of samples accessed (cf. Definition 2.2) and the total rounds of communication (cf. Definition 2.3) required by each WN to achieve an ✏-stationary solution (cf. Definition 2.1). To minimize the sample and the communication complexities, FL algorithms rely on the following four key design elements: (i) the WNs’ local model update directions, (ii) Minibatch size to compute each local direction, (iii) the number of local updates before WNs share their parameters, and (iv) the SN’s update direction. How to find effective FL algorithms by (optimally) designing these parameters has received significant research interest recently.
Contributions. The main contributions of this work are listed below: 1) We propose the Stochastic Two-Sided Momentum (STEM) algorithm, that utilizes certain momentum-assisted stochastic gradient directions for both the WNs and SN updates. We show that there exists an optimal trade off between the minibatch sizes and number of local updates, such that on the trade-off curve STEM requires Õ(✏ 3/2)2 samples and Õ(✏ 1) communication rounds to reach an ✏-stationary solution; see Figure 1 for an illustration. These complexity results are the best achievable for first-order stochastic FL algorithms (under certain assumptions, cf. Assumption 1); see [5–8] and [9, 10], as well as Remark 1 of this paper for discussions regarding optimality. To the best of our knowledge, STEM is the first algorithm which – (i) simultaneously achieves the optimal sample and communication complexities for FL and (ii) can optimally trade off the minibatch sizes and the number of local updates.
2) A momentum-less special case of our STEM result further reveals some interesting insights of the classical FedAvg algorithm (a.k.a. the Local SGD) [11–13]. Specifically, we show that for FedAvg, there also exists a trade-off between the minibatch sizes and the number of local updates, such that it requires O(✏ 2) samples and O(✏ 3/2) communication rounds to achieve an ✏-stationary solution.
2The notation Õ(·) hides the logarithmic factors.
Collectively, our insights on the trade-offs provide practical guidelines for choosing different design elements for FL algorithms.
Related Works. FL algorithms were first proposed in the form of FedAvg [11], where the local update directions at each WN were chosen to be the SGD updates. Earlier works analyzed these algorithms in the homogeneous data setting [19–25], while many recent studies have focused on designing new algorithms to deal with heterogeneous data settings, as well as problems where the local loss functions are non-convex [9, 10, 12–16, 18, 26–32]. In [12], the authors showed that Parallel Restarted SGD (Local SGD or FedAvg [11]) achieves linear speed up while requiring O(✏ 2) samples and O(✏ 3/2) rounds of communication to reach an ✏-stationary solution. In [14], a Momentum SGD was proposed, which achieved the same sample and communication complexities as Parallel Restarted SGD [12], without requiring that the second moments of the gradients be bounded. Further, it was shown that under the homogeneous data setting, the communication complexity can be improved to O(✏ 1) while maintaining the same sample complexity. The works in [15, 16] conducted tighter analysis for FedAvg with partial WN participation with O(1) local updates and batch sizes. Their analysis showed that FedAvg’s sample and communication complexities are both O(✏ 2). Additionally, SCAFFOLD was proposed in [15], which utilized variance reduction based local update directions [33] to achieve the same sample and communication complexities as FedAvg. Similarly, VRL-SGD proposed in [29] also utilized variance reduction and showed improved communication complexity of O(✏ 1), while requiring the same computations as FedAvg. Importantly, both SCAFFOLD and VRL-SGD’s guarantees were independent of the data heterogeneity. The FedProx proposed in [10] used a penalty based method to improve the communication complexity of FedAvg (i.e., the Parallel Restarted and Momentum SGD [14, 12]) to O(✏ 1). FedProx used a gradient similarity assumption to model data heterogeneity which can be stringent for many practical applications. This assumption was relaxed by FedPD proposed in [9].
Recently, the works [17, 18] proposed to utilize hybrid momentum gradient estimators [7, 8]. The MIME algorithm [17] matched the optimal sample complexity (under certain smoothness assumptions) of O(✏ 3/2) of the centralized non-convex stochastic optimization algorithms [5–8]. Similarly, Fed-GLOMO [18] achieved the same sample complexity while employing compression to further reduce communication. Both MIME and Fed-GLOMO required O(✏ 3/2) communication rounds to achieve an ✏-stationary solution. Please see Table 1 for a summary of the above discussion.
The comparison of Local SGD (FedAvg) to Minibatch SGD for convex and strongly convex problems with homogeneous data setting was first conducted in [19] and later extended to heterogeneous setting in [13]. It was shown that Minibatch SGD almost always dominates the Local SGD. In contrast, it was shown in [24] that Local SGD dominates Minibatch SGD in terms of generalization performance. Although existing FL results are rich, but they are somehow ad hoc and there is a lack of principled
understanding of the algorithms. We note that the proposed STEM algorithmic framework provides a theoretical framework that unifies all existing FL results on sample and communication complexities.
Notations. The expected value of a random variable X is denoted by E[X] and its expectation conditioned on an Event A is denoted as E[X|Event A]. We denote by R (and Rd) the real line (and the d-dimensional Euclidean space). The set of natural numbers is denoted by N. Given a positive integer K 2 N, we denote [K] , {1, 2, . . . ,K}. Notation k · k denotes the `2-norm and h·, ·i the Euclidean inner product. For a discrete set B, |B| denotes the cardinality of the set. Uniform distribution over a discrete set {1, . . . , T} is denoted as U{1, . . . , T}.

heading: 2 Preliminaries

text:
Before we proceed to the algorithms, we make the following assumptions about problem (1). Assumption 1 (Sample Gradient Lipschitz Smoothness). The stochastic functions f (k)(·, ⇠(k)) with ⇠(k) ⇠ D(k) for all k 2 [K], satisfy the mean squared smoothness property, i.e, we have
Ekrf (k)(x; ⇠(k)) rf (k)(y; ⇠(k))k2  L2kx yk2 for all x, y 2 Rd. Assumption 2 (Unbiased gradient and Variance Bounds). (i) Unbiased Gradient. The stochastic gradients computed at each WN are unbiased
E[rf (k)(x; ⇠(k))] = rf (k)(x), 8 ⇠(k) ⇠ D(k), 8 k 2 [K]. (ii) Intra- and inter- node Variance Bound. The following bounds hold:
Ekrf (k)(x; ⇠(k)) rf (k)(x)k2  2, krf (k)(x) rf (`)(x)k2  ⇣2, 8 ⇠(k) ⇠ D(k), 8k, ` 2 [K].
Note that Assumption 1 is stronger than directly assuming f (k)’s are Lipschitz smooth (which we will refer to as the averaged gradient Lipschitz smooth condition), but it is still a rather standard assumption in SGD analysis. For example it has been used in analyzing centralized SGD algorithms such as SPIDER [5], SNVRG [6], STORM [7] (and many others) as well as in FL algorithms such as MIME [17] and Fed-GLOMO [18]. The second relation in Assumption 2-(ii) quantifies the data heterogeneity, and we call ⇣ > 0 as the heterogeneity parameter. This is a typical assumption required to evaluate the performance of FL algorithms. If data distributions across individual WNs are identical, i.e., D(k) = D(`) for all k, ` 2 [K] then we have ⇣ = 0.
Next, we define the ✏-stationary solution for non-convex optimization problems, as well as quantify the computation and communication complexities to achieve an ✏-stationary point. Definition 2.1 (✏-Stationary Point). A point x is called ✏-stationary if krf(x)k2  ✏. Moreover, a stochastic algorithm is said to achieve an ✏-stationary point in t iterations if E[krf(xt)k2]  ✏, where the expectation is over the stochasticity of the algorithm until time instant t. Definition 2.2 (Sample complexity). We assume an Incremental First-order Oracle (IFO) framework [34], where, given a sample ⇠(k) ⇠ D(k) at the kth node and iterate x, the oracle returns (f (k)(x; ⇠(k)),rf (k)(x; ⇠(k))). Each access to the oracle is counted as a single IFO operation. We measure the sample (and computational) complexity in terms of the total number of calls to the IFO by all WNs to achieve an ✏-stationary point given in Definition 2.1. Definition 2.3 (Communication complexity). We define a communication round as a one back-andforth sharing of parameters between the WNs and the SN. Then the communication complexity is defined to be the total number of communication rounds between any WN and the SN required to achieve an ✏-stationary point given in Definition 2.1.
3 The STEM algorithm and the trade-off analysis
In this section, we discuss the proposed algorithm and present the main results. The key in the algorithm design is to carefully balance all the four design elements mentioned in Sec. 1, so that sufficient and useful progress can be made between two rounds of communication.
Let us discuss the key steps of STEM, listed in Algorithm 1. In Step 10, each node locally updates its model parameters using the local direction dkt , computed by using b stochastic gradients at two
Algorithm 1 The Stochastic Two-Sided Momemtum (STEM) Algorithm 1: Input: Parameters: c > 0, the number of local updates I , batch size b, stepsizes {⌘t}. 2: Initialize: Iterate x(k)1 = x̄1 = 1K PK k=1 x (k) 1 , descent direction d (k) 1 = d̄1 = 1 K PK k=1 d (k) 1
with d(k)1 = 1 B P ⇠(k)1 2B (k) 1 rf (k)(x(k)1 ; ⇠ (k) 1 ) and |B (k) 1 | = B for k 2 [K].
3: Perform: x(k)2 = xk1 ⌘1d (k) 1 , 8 k 2 [K] 4: for t = 1 to T do 5: for k = 1 to K do #at the WN 6: d(k)t+1 = 1
b
X
⇠(k)t+12B (k) t+1
rf (k)(x(k)t+1; ⇠ (k) t+1)+(1 at+1) ⇣ d(k)t 1
b
X
⇠(k)t+12B (k) t+1
rf (k)(x(k)t ; ⇠ (k) t+1)
⌘
where we choose |B(k)t+1| = b, and at+1 = c · ⌘2t ; 7: if t mod I = 0 then #at the SN 8: d(k)t+1 = d̄t+1 := 1 K PK k=1 d (k) t+1 9: x(k)t+2 := x̄t+1 ⌘t+1d̄t+1 = 1 K PK k=1 x (k) t+1 ⌘t+1d̄t+1 #server-side momentum
10: else x(k)t+2 = x (k) t+1 ⌘t+1d (k) t+1 #worker-side momentum 11: end if 12: end for 13: end for 14: Return: x̄a where a ⇠ U{1, ..., T}.
consecutive iterates x(k)t+1 and x (k) t . After every I local steps, the WNs share their current local models {x(k)t+1} K k=1 and directions {d (k) t+1} K k=1 with the SN. The SN aggregates these quantities, and performs a server-side momentum step, before returning x̄t+1 and d̄t+1 to all the WNs. Because both the WNs and the SN perform momentum based updates, we call the algorithm a stochastic two-sided momentum algorithm. The key parameters are: b the minibatch size, I the local update steps between two communication rounds, ⌘t the stepsizes, and at the momentum parameters.
One key technical innovation of our algorithm design is to identify the most suitable way to incorporate momentum based directions in FL algorithms. Although the momentum-based gradient estimator itself is not new and has been used in the literature before (see e.g., in [7, 8] and [17, 18] to improve the sample complexities of centralized and decentralized stochastic optimization problems, respectively), it is by no means clear if and how it can contribute to improve the communication complexity of FL algorithms. We show that in the FL setting, the local directions together with the local models have to be aggregated by the SN so to avoid being influenced too much by the local data. More importantly, besides the WNs, the SN also needs to perform updates using the (aggregated) momentum directions. Finally, such two-sided momentum updates have to be done carefully with the correct choice of minibatch size b, and the number of local updates I . Overall, it is the judicious choice of all these design elements that results in the optimal sample and communication complexities.
Next, we present the convergence guarantees of the STEM algorithm.
3.1 Main results: convergence guarantees for STEM
In this section, we analyze the performance of STEM. We first present our main result, and then provide discussions about a few parameter choices. In the next subsection, we discuss a special case of STEM related to the classical FedAvg and minibatch SGD algorithms. Theorem 3.1. Under the Assumptions 1 and 2, suppose the stepsize sequence is chosen as:
⌘t = ̄
(wt + 2t)1/3 , (2)
where we define :
̄ = (bK)2/3 2/3
L , wt = max
⇢ 2 2, 4096L3I3̄3 2t, c3̄3
4096L3I3
.
Further, let us set c = 64L 2 bK + 2 24̄3LI = L 2 ✓ 64 bK + 1 24(bK)2I ◆ , and set the initial batch size as
B = bI; set the local updates I and minibatch size b as follows:
I = O (T/K2)⌫/3 , b = O (T/K2)1/2 ⌫/2
(3)
where ⌫ satisfies ⌫ 2 [0, 1]. Then for STEM the following holds:
(i) For x̄a chosen according to Algorithm 1, we have:
Ekrf(x̄a)k2 = O ✓ f(x̄1) f⇤
K2⌫/3T 1 ⌫/3
◆ + Õ ✓ 2
K2⌫/3T 1 ⌫/3
◆ + Õ ✓ ⇣2
K2⌫/3T 1 ⌫/3
◆ . (4)
(ii) For any ⌫ 2 [0, 1], we have Sample Complexity: The sample complexity of STEM is Õ(✏ 3/2). This implies that each WN requires at most Õ(K 1✏ 3/2) gradient computations, thereby achieving linear speedup with the number of WNs present in the network. Communication Complexity: The communication complexity of STEM is Õ(✏ 1).
The proof of this result is relegated to the Supplemental Material. A few remarks are in order. Remark 1 (Near-Optimal sample and communication complexities). Theorem 3.1 suggests that when I and b are selected appropriately, then STEM achieves Õ(✏ 3/2) and Õ(✏ 1) sample and communication complexities. Taking them separately, these complexity bounds are the best achievable by the existing FL algorithms (upto logarithmic factors regardless of sample or batch Lipschitz smooth assumption) [35]; see Table 1. We note that the O(✏ 3/2) complexity is the best possible that can be achieved by centralized SGD with the sample Lipschitz gradient assumption; see [5]. On the other hand, the O(✏ 1) complexity bound is also likely to be the optimal, since in [9] the authors showed that even when the local steps use a class of (deterministic) first-order algorithms, O(✏ 1) is the best achievable communication complexity. The only difference is that [9] does not explicitly assume the inter-node variance bound (i.e., the second relation in Assumption 2-(ii)). We leave the precise characterization of the communication lower bound with inter-node variance as future work. Remark 2 (Large Batch Sizes and/or Local Updates). At first glance, it may seem that the requirement of STEM to compute large mini-batches and/or local updates (cf. Table 1) to achieve this (near) optimal performance is a drawback, however, we note that it is in fact an advantage of STEM that it allows the WNs to perform larger number of local updates (or compute large minibatches) without communicating often. This follows from the fact that irrespective of the number of local updates (or batch sizes) STEM achieves near-optimal communication complexity while attaining optimal overall sample complexity. Moreover, note that even with b = I = O(1) (i.e., b and I are chosen as constants), STEM achieves the same (optimal) sample and communication complexities as achieved by FedGLOMO [18] and MIME [17]. We further note that to the best of our knowledge the algorithms that achieve the communication complexity of O(✏ 1) either require the number of local updates or the batch-sizes that depend on the solution accuracy ✏. For example, FedProx [10], FedPD [9], and FedDyn [36] rely on solving the “local problems" to achieve an ✏-accuracy, which implies that the number of local updates (or the batch sizes) implicitly depends on the desired solution accuracy ✏, as is the case for STEM. Similarly, as shown in [12] and [14] the communication complexity of FedAvg and its momentum version can be improved from O(✏ 2) to O(✏ 3/2) when the number of local updates (or batch size) is chosen as O(✏ 1/2) (cf. Section 3.2 for a more detailed discussion). Remark 3 (The Optimal Batch Sizes and Local Updates Trade-off). The parameter ⌫ 2 [0, 1] is used to balance the local minibatch sizes b, and the number of local updates I . Eqs. in (3) suggest that when ⌫ increases from 0 to 1, b decreases and I increases. Specifically, if ⌫ = 1, then b is a constant but I = O(T 1/3/K2/3). In this case, each WN chooses a small minibatch while executing multiple local updates, and STEM resembles a FedAvg (a.k.a. Local SGD) algorithm but with double-sided momentum update directions, and is referred to as Fed STEM. In contrast, if ⌫ = 0, then b = O(T 1/2/K) but I is a constant. In this case, each WN chooses a large batch size while executing only a few, or even one, local updates, and STEM resembles the Minibatch SGD, but again with different update directions, and is referred to as Minibatch STEM. Such a trade-off can be seen in Fig. 1b. Due to space limitation, these two special cases will be precisely stated in the supplementary materials as corollaries of Theorem 3.1.
Algorithm 2 The FedAvg Algorithm 1: Input: {⌘t}Tt=0; I , the # of local updates per communication round; b, the minibatch sizes. 2: for t = 1 to T do 3: for k = 1 to K do 4: d(k)t = 1 b P ⇠(k)t 2B (k) t rf (k)(x(k)t ; ⇠ (k) t ) with |B (k) t | = b
5: x(k)t+1 = x (k) t ⌘td (k) t 6: if t mod I = 0 then 7: x(k)t+1 = x̄t+1 = 1 K PK k=1 x (k) t+1 8: end if 9: end for
10: end for 11: Return: x̄a where a ⇠ U{1, ..., T}.
Remark 4 (The Sub-Optimal Batch Sizes and Local Updates Trade-off). From our proof (Theorem C.10 included in the supplemental material), we can see that STEM requires Õ max (b ·
I)✏ 1,K 1✏ 3/2 samples and Õ max ✏ 1, (b · I) 1K 1✏ 3/2 and communication rounds.
According to the above expressions, if b · I increases beyond O(K 1✏ 1/2), then the sample complexity will increase from the optimal Õ(✏ 3/2); otherwise, the optimal sample complexity Õ(✏ 3/2) is maintained. On the other hand, if b · I decreases beyond O(K 1✏ 1/2), the communication complexity increases from Õ(✏ 1). For instance, if we choose b = O(1) and I = O(1) the communication complexity becomes Õ(✏ 3/2) while the optimal sample complexity Õ(✏ 3/2) is maintained. This trade-off is illustrated in Figure 1a, where we maintain the optimal sample complexity, while changing b and I to generate the trade-off surface. Remark 5 (Data Heterogeneity). The term Õ ⇣ ⇣2
K2⌫/3T 1 ⌫/3
⌘ in the gradient bound (4) captures the
effect of the heterogeneity of data across WNs, where ⇣ is the parameter characterizing the intra-node variance and has been defined in Assumption 2-(ii). Highly heterogeneous data with large ⇣2 can adversely impact the performance of STEM. Note that such a dependency on ⇣ also appears in other existing FL algorithms, such as [9, 14, 18]. However, there is one special case of STEM that does not depend on the parameter ⇣. This is the case where I = 1, i.e., the minibatch SGD counterpart of STEM where only a single local iteration is performed between two communication rounds. We have the following corollary.
Corollary 1 (Minibatch STEM). Under Assumptions 1 and 2 , and choose the algorithm parameters as in Theorem 3.1. At each WN, choose I = 1, b = (T/K2)1/2, and the initial batch size B = b · I . Then STEM satisfies:
(i) For x̄a chosen according to Algorithm 1, we have
Ekrf(x̄a)k2 = O ⇣f(x̄1) f⇤
T
⌘ + Õ ⇣ 2
T
⌘ .
(ii) Minibatch STEM achieves Õ(✏ 3/2) sample and Õ(✏ 1) communication complexity.
Next, we show that FedAvg also exhibits a trade-off similar to that of STEM but with worse sample and communication complexities.

heading: 3.2 Special cases: The FedAvg algorithm

text:
We briefly discuss another interesting special case of STEM, where the local momentum update is replaced by the conventional SGD (i.e., at = 1, 8 t), while the server does not perform the momentum update (i.e., d̄t = 0, 8 t). This is essentially the classical FedAvg algorithm, just that it balances the number of local updates I and the minibatch size b. We show that this algorithm also exhibits a trade-off between b and I and on the trade-off curve it achieves O(✏ 2) sample complexity and O(✏ 3/2) communication complexity.
Theorem 3.2 (The FedAvg Algorithm). Under Assumptions 1 and 2, suppose the stepsize is chosen as: ⌘ = q bK T ; Let us set:
I = O (T/K3)⌫/4 , b = O (T/K3)1/3 ⌫/3
(5)
where ⌫ 2 [0, 1] is a constant. Then for FedAvg with T 81L2I2bK, the following holds
(i) For x̄a chosen according to Algorithm 2, we have
Ekrf(x̄a)k2 = O ✓ f(x̄1) f⇤
K⌫/2T 2/3 ⌫/6
◆ +O ✓ 2
K⌫/2T 2/3 ⌫/6
◆ +O ✓ ⇣2
K⌫/2T 2/3 ⌫/6
◆ .
(ii) For any choice of ⌫ 2 [0, 1] we have: Sample Complexity: The sample complexity of FedAvg is O(✏ 2). This implies that each WN requires at most O(K 1✏ 2) gradient computations, thereby achieving linear speedup with the number of WNs in the network. Communication Complexity: The communication complexity of FedAvg is O(✏ 3/2).
Note that the requirement on T being lower bounded is only relevant for theoretical purposes, a similar requirement was also imposed in [14] to prove convergence. Again, the parameter ⌫ 2 [0, 1] in the statement of Theorem 3.2 balances I and b at each WN while maintaining state-of-the-art sample and communication complexities; please see Table 1 for a comparison of those bounds with existing FedAvg bounds. For ⌫ = 1, FedAvg (cf. Theorem 3.2) reduces to FedAvg proposed in [12, 14] and for ⌫ = 0, the algorithm can be viewed as a large batch FedAvg with constant local updates [15, 16]. Note that similar to STEM, it is known that for I = 1, the Minibatch SGD’s performance is independent of the heterogeneity parameter, ⇣ [13]. We also point out that if Algorithm 1 uses Nesterov’s or Polyak’s momentum [14] at local WNs instead of the recursive momentum estimator we get the same guarantees as in Theorem 3.2.
In summary, this section established that once the WN’s and the SN’s update directions (SGD in FedAvg and momentum based directions in STEM) are fixed, there exists a sequence of optimal choices of the number of local updates I , and the batch sizes b, which guarantees the best possible sample and communication complexities for the particular algorithm. The trade-off analysis presented in this section provides some useful guidelines for how to best select b and I in practice. Our subsequent numerical results will also verify that if b or I are not chosen judiciously, then the practical performance of the algorithms can degrade significantly.

heading: 4 Numerical results

text:
In this section, we validate the proposed STEM algorithm and compare its performance with the de facto standard FedAvg [11], and the algorithms stated in Table 1. Note that instead of FedPD we include the performance comparison with FedDyn [36] since they are known to be very closely
Algorithm Training Acc. Testing Acc. FedAvg 40.1 39.2 FedProx 43.5 43.2 FedDyn 43.7 43.2 SCAFFOLD 40.3 41.3 MIME 32.1 32.1 FedGLOMO 40.3 40.1 STEM 44.5 43.8
Table 4: Training and testing accuracy on Shakespeare dataset.
related. The goal of our experiments are three-fold: (1) To show that STEM performs on par, if not better, compared to other algorithms in different heterogeneity settings, (2) there are multiple ways to reach the desired solution accuracy, one can either choose a large batch size and perform only a few local updates or select a smaller batch size and perform multiple local updates, and finally, (3) if the local updates and the batch sizes are not chosen appropriately, the WNs might need to perform excessive computations to achieve the desired solution accuracy, thereby slowing down convergence.
Data and Parameter Settings: We compare the algorithms for image classification tasks on CIFAR10 and MNIST data sets with 100 WNs, and for next character prediction task on Shakespeare dataset [37] with 143 WNs in the network. For both CIFAR-10 and MNIST, each WN implements a two-hidden-layer convolutional neural network (CNN) architecture followed by three linear layers for CIFAR-10 and two for MNIST. For CIFAR-10 (and MNIST) datset, we consider three settings with mild, moderate and high heterogeneity. For all the three settings, the data is partitioned into disjoint sets among the WNs. In the mild heterogeneity setting, the WNs have access to partitioned data from all the classes. In the moderate (resp. high) heterogeneity setting the data is partitioned such that each WN can access data from only 5 (resp. 2) out of 10 classes. For CIFAR-10 (resp. MNIST), each WN has access to 490 (resp. 540) samples for training and 90 (resp. 80) samples for testing purposes.
We also compare the performance of algorithms on a popular FL benchmarking dataset, Shakespeare dataset [37]. For this task, we adopt the settings from [10] and utilize a 2-Layer LSTM network with 100 hidden units and an 8-D embedding layer at each WN. Each WN has access to 3616 samples on average, and the samples are randomly split into an 80% training set and a 20% testing set. We randomly sample 10 nodes out of 143 for the training purpose. All the experiments are implemented on a single NVIDIA Quadro RTX 5000 GPU. More details are provided in appendix.
For the proposed STEM algorithm, recall that the step-size is ⌘t = ̄/(wt + 2t)1/3 with momentum parameter defined as at = c⌘2t . The step-size is used to update the iterates while the momentum
parameter is used to construct the stochastic gradient estimate (cf. Algorithm 1 and Theorem 3.1). For the experiments, we set wt = 2 = 1 and c = c̄/̄2 and tune for ̄ 2 [10 1, 10 2] for the CIFAR-10 dataset and for ̄ 2 {101, 100, 10 1, 10 2} for the Shakespeare dataset. For both the datasets we tune for c̄ in the range [1, 10]. For FedProx [10] and FedDyn [36] we choose the regularization constant to be 0.1. The momentum parameters for FedGLOMO [18] and MIME [17] are set based on the choices given in the respective papers. Specifically, for FedGLOMO we choose the parameter k = 0.2 and design the momentum gradient using a damping factor given in Appendix A.4 of FedGLOMO [18]. Moreover, for MIME we choose the momentum parameter as 0.9. For the rest of the algorithms (including FedAvg and SCAFFOLD), the step-size is tuned from the set {101, 100, 10 1, 10 2}.
Discussion: We evaluate the training and testing performance of STEM against multiple algorithms for different heterogeneity settings, minibatch sizes, and number of local updates. In Tables 2a, 2b and 3, we compare the training and testing accuracy of STEM to that of other algorithms on the CIFAR-10 dataset. Specific, heterogeneity settings, the choices of minibatches, and number of local updates are stated along with the tables. Note that STEM performs uniformly well under all the conditions. Moreover, note from Table 2b that FedGLOMO diverges once the number of local updates are high. Also, note from Table 3 that FedProx and STEM adapt well to high heterogeneity. Finally, with the next set of experiments we emphasize the importance of choosing b and I carefully. In Figure 2, we compare the training and testing performance of STEM, FedAvg and SCAFFOLD, against the number of samples accessed at each WN for the classification task on MNIST dataset with moderate heterogeneity. We fix b = 8 and conduct experiments under two settings, one with I = 67, and the other with I = 536 local updates at each WN. Note that although a large number of local updates might lead to fewer communication rounds but it can make the sample complexity extremely high as is demonstrated by Figure 2. For example, Figure 2 shows that to reach testing accuracy of 96 97% with I = 67, STEM requires approximately 5000 6000 samples, in contrast with I = 536 it requires more than 25000 samples at each WN. Similar behavior can be observed if we fix I > 1 and increase the local batch sizes. This implies not choosing the local updates and the batch sizes judiciously might lead to increased sample complexity. Additional experiments are included in the supplementary material to further evaluate the performance of the proposed algorithms.
Conclusion
In this work, we proposed a novel algorithm STEM, for distributed stochastic non-convex optimization with applications to FL. We showed that STEM reaches an ✏-stationary point with Õ(✏ 3/2) sample complexity while achieving linear speed-up with the number of WNs. Moreover, the algorithm achieves a communication complexity of Õ(✏ 1). We established a (optimal) trade-off that allows interpolation between varying choices of local updates and the batch sizes at each WN while maintaining (near optimal) sample and communication complexities. We showed that FedAvg (a.k.a LocalSGD) also exhibits a similar trade-off while achieving worse complexities. Our results provide guidelines to carefully choose the number of local updates, update directions, and minibatch sizes to achieve the best performance. The future directions of this work include developing lower bounds on communication complexity that establishes the tightness of the analysis conducted in this work.

heading: Acknowledgement

text: We thank the anonymous reviewers for their valuable comments and suggestions. The work of Prashant Khanduri and Mingyi Hong was supported by NSF grant CMMI-1727757, AFOSR grant 19RT0424 and ARO grant W911NF-19-1-0247. The work of Mingyi Hong was also supported by an IBM Faculty Research award. The work of Jia Liu has been supported in part by NSF grants CAREER CNS-2110259, CNS-2112471, CNS-2102233, CCF-2110252, ECCS-2140277, and a Google Faculty Research Award.

title: STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning


summaries: This paper study the interplay between generalization bound and the covariance of perturbed gradient descent. The authors invoke a recent (information-theoretic) bound for stochastic optimization methods. Then, they optimize the upper bound with respect to trace constraint on the covariance of noise in each step and experimentally show that this optimization may be effective for training neural networks.

summaries: This paper considers the problem of understanding the generalization of the SGLD algorithms using information-theoretic techniques. This line of work initiates by [27],[24], and [13]. The goal in this paper is to "modify" the SGLD algorithms such that it has a good generalization properties and good performance on the training set. Basically, the modification that the authors considered is to make the covariance of the Gaussian noise in the SGLD update to depend on the trajectory. Note that in the usual setting the noise is isotropic. Then, authors claim that controlling the trace of the SGLD ensures we have good performance on training set and good generalization. Then the authors provide a generalization bound which uses reverse KL between posterior and the prior distribution. I think this bound is not true as stated in my review. Since this bound has been used in the paper, I think the authors should comment on the validity of the theorem.

summaries: This interesting paper studies the connection between generalization abilities of of SGLD and the covariance structure of its noise term. It first proves that with constraint to guarantee low empirical risk and commonly used data-dependent priors, the optimal noise covariance of SGLD in terms of the bound is similar to the empirical gradient covariance. The paper further proves that if the generalization bound is jointly optimized with respect to both prior and posterior, the optimal noise covariance is the square root of the expected gradient covariance. These facts can be used to further support the belief of the superiority of SGD noise over isotropic noise.

summaries: From usual isentropic noise to data-dependent priors, the authors extend the existing information-theoretic analysis of SGLD. Additionally, they identify the optimal choice of prior distribution covariance. Last but not least, the paper includes empirical observations to validate its technical findings.

summaries: This paper studies optimization of the information theoretical generalization bounds of Stochastic Gradient Langevin Dynamics (SGLD) with respect to the noise covariance. Here are the main findings: * When updating rule is fixed and only optimizing noise covariance under the constraint that the trace of the covariance is fixed, the noise covariance of SGLD is similar to the empirical gradient covariance * When both updating rule and the noise covariance are jointly minimized, the noise covariance of SGLD is the square root of the expected gradient covariance.  The two results can potentially be used to improve information theoretical generalization bounds obtained in the literature. The theoretical results are validated by some numerical experiments.

abstractText: Recently, the information-theoretical framework has been proven to be able to obtain non-vacuous generalization bounds for large models trained by Stochastic Gradient Langevin Dynamics (SGLD) with isotropic noise. In this paper, we optimize the information-theoretical generalization bound by manipulating the noise structure in SGLD. We prove that with constraint to guarantee low empirical risk, the optimal noise covariance is the square root of the expected gradient covariance if both the prior and the posterior are jointly optimized. This validates that the optimal noise is quite close to the empirical gradient covariance. Technically, we develop a new information-theoretical bound that enables such an optimization analysis. We then apply matrix analysis to derive the form of optimal noise covariance. Presented constraint and results are validated by the empirical observations.

authors:
S. Arora
R. Ge
B. Neyshabur
Y. Zhang

title: Stronger generalization bounds for deep nets via a compression approach

venue: International Conference on Machine Learning, pages 254–263. PMLR

authors:
A.R. Asadi
E. Abbe
S. Verdú

title: Chaining mutual information and tightening generalization bounds

venue: Advances in Neural Information Processing Systems (NeurIPS)

authors:
S. Boucheron
G. Lugosi
P. Massart

title: Concentration inequalities: A nonasymptotic theory of independence

venue: Oxford university press

authors:
O. Bousquet
S. Boucheron
G. Lugosi

title: Introduction to statistical learning theory

venue: Summer School on Machine Learning, pages 169–207. Springer

authors:
Y. Bu
S. Zou
V.V. Veeravalli

title: Tightening mutual information-based bounds on generalization error

venue: IEEE Journal on Selected Areas in Information Theory, 1(1):121–130

authors:
Z. Charles
D. Papailiopoulos

title: Stability and generalization of learning algorithms that converge to global optima

venue: International Conference on Machine Learning, pages 745–754. PMLR

authors:
T.M. Cover
J.A. Thomas

title: Elements of information theory

venue: Wiley

authors:
C. Dwork
V. Feldman
M. Hardt
T. Pitassi
O. Reingold
A. Roth

title: Generalization in adaptive data analysis and holdout reuse

venue: Advances in Neural Information Processing Systems

authors:
G.K. Dziugaite
D.M. Roy

title: Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data

venue: arXiv preprint arXiv:1703.11008

authors:
V. Feldman
T. Steinke

title: Calibrating noise to variance in adaptive data analysis

venue: Conference On Learning Theory, pages 535–544. PMLR

authors:
S.B. Gelfand
S.K. Mitter

title: Recursive stochastic algorithms for global optimization in rˆd

venue: SIAM Journal on Control and Optimization, 29(5):999–1018

authors:
N. Gergely
G.K. Dziugaite
M. Haghifam
D.M. Roy

title: Information-theoretic generalization bounds for stochastic gradient descent

venue: COLT

authors:
H. Hafez-Kolahi
Z. Golgooni
S. Kasaei
M. Soleymani

title: Conditioning and processing: Techniques to improve information-theoretic generalization bounds

venue: Advances in Neural Information Processing Systems (NeurIPS), volume 33

authors:
M. Haghifam
J. Negrea
A. Khisti
D.M. Roy
G.K. Dziugaite

title: Sharpened generalization bounds based on conditional mutual information and an application to noisy

venue: iterative algorithms. arXiv preprint arXiv:2004.12983

authors:
M. Hardt
B. Recht
Y. Singer

title: Train faster

venue: generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pages 1225–1234. PMLR

authors:
I. Issa
A.R. Esposito
M. Gastpar

title: Strengthened information-theoretic bounds on the generalization error

venue: 2019 IEEE International Symposium on Information Theory (ISIT), pages 582–586. IEEE

authors:
Y. Jiang
B. Neyshabur
H. Mobahi
D. Krishnan
S. Bengio

title: Fantastic generalization measures and where to find them

venue: International Conference on Learning Representations

authors:
N.S. Keskar
D. Mudigere
J. Nocedal
M. Smelyanskiy
P.T.P. Tang

title: On large-batch training for deep learning: Generalization gap and sharp minima

venue: arXiv preprint arXiv:1609.04836

authors:
I. Kuzborskij
C. Lampert

title: Data-dependent stability of stochastic gradient descent

venue: International Conference on Machine Learning, pages 2815–2824. PMLR

authors:
Y. Lei
Y. Ying

title: Sharper generalization bounds for learning with gradient-dominated objective functions

authors:
J. Li
X. Luo
M. Qiao

title: On generalization error bounds of noisy gradient methods for non-convex learning

authors:
S. Mandt
M.D. Hoffman
D.M. Blei

title: Stochastic gradient descent as approximate Bayesian inference

venue: The Journal of Machine Learning Research, 18(1):4873–4907

authors:
Q. Meng
S. Gong
W. Chen
Z.-M. Ma
T.-Y. Liu

title: Dynamic of stochastic gradient descent with state-dependent noise

venue: arXiv preprint arXiv:2006.13719

authors:
W. Mou
L. Wang
X. Zhai
K. Zheng

title: Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints

venue: Conference on Learning Theory, pages 605–638. PMLR

authors:
J. Negrea
M. Haghifam
G.K. Dziugaite
A. Khisti
D.M. Roy

title: Information-theoretic generalization bounds for SGLD via data-dependent estimates

venue: arXiv preprint arXiv:1911.02151

authors:
B. Neyshabur
S. Bhojanapalli
N. Srebro

title: A PAC-Bayesian approach to spectrallynormalized margin bounds for neural networks

authors:
E. Parrado-Hernández
A. Ambroladze
J. Shawe-Taylor
S. Sun

title: PAC-Bayes bounds with data dependent priors

venue: The Journal of Machine Learning Research, 13(1):3507–3531

authors:
A. Pensia
V. Jog
P.-L. Loh

title: Generalization error bounds for noisy

venue: iterative algorithms. In 2018 IEEE International Symposium on Information Theory (ISIT), pages 546–550. IEEE

authors:
M. Raginsky
A. Rakhlin
M. Telgarsky

title: Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis

venue: Conference on Learning Theory, pages 1674–1703. PMLR

authors:
D. Russo
J. Zou

title: Controlling bias in adaptive data analysis using information theory

venue: Artificial Intelligence and Statistics, pages 1232–1240. PMLR

authors:
T. Steinke
L. Zakynthinou

title: Reasoning about generalization via conditional mutual information

venue: Conference on Learning Theory, pages 3437–3452. PMLR

authors:
G. Strang
G. Strang
G. Strang
G. Strang

title: Introduction to linear algebra

venue: volume 3. Wellesley-Cambridge Press Wellesley, MA

authors:
M. Welling
Y.W. Teh

title: Bayesian learning via stochastic gradient langevin dynamics

venue: Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688. Citeseer

authors:
Y. Wen
K. Luk
M. Gazeau
G. Zhang
H. Chan
J. Ba

title: An empirical study of large-batch stochastic gradient descent with structured covariance noise

venue: arXiv e-prints, pages arXiv–1902

authors:
H. Xiao
K. Rasul
R. Vollgraf

title: Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms

venue: arXiv preprint arXiv:1708.07747

authors:
Z. Xie
I. Sato
M. Sugiyama

title: A diffusion theory for deep learning dynamics: Stochastic gradient descent escapes from sharp minima exponentially fast

venue: arXiv preprint arXiv:2002.03495

authors:
A. Xu
M. Raginsky

title: Information-theoretic analysis of generalization capability of learning algorithms

venue: arXiv preprint arXiv:1705.07809

authors:
C. Zhang
S. Bengio
M. Hardt
B. Recht
O. Vinyals

title: Understanding deep learning requires rethinking generalization

authors:
W. Zhou
V. Veitch
M. Austern
R.P. Adams
P. Orbanz

title: Non-vacuous generalization bounds at the imagenet scale: a PAC-Bayesian compression approach

authors:
Z. Zhu
J. Wu
B. Yu
L. Wu
J. Ma

title: The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects

venue: arXiv preprint arXiv:1803.00195

heading: 1 Introduction

text:
Generalization ability is one of the core questions in learning theory [4], but remains unclear for deep learning models [17, 38]. Existing generalization bounds based on the capacity control become vacuous for practical deep learning models due to the over-parameterization property [1, 26, 39].
From the information theoretical perspective, recent works [30, 37] bound the generalization error by the mutual information between the dataset and the learned parameters. This result reveals that good generalization occurs when the learned parameters do not depend on a specific dataset too much, which is intuitively reasonable and closely related with the idea of algorithm stability [6, 15, 19, 24, 20] and differential privacy [8, 10]. Meanwhile, based on information-theoretic metrics, one can analyze general classes of updates and models, e.g., stochastic iterative algorithms for non-convex objectives, hence applicable to deep learning. It has been shown that the information theoretical bounds are non-vacuous and closely related with the real generalization error even in deep learning [14, 26, 9, 39].
Notably, the information theoretical bound is realized in [28] via decomposing the mutual information across iterations. This framework can perform a step-wise analysis of Stochastic Gradient Langevin Dynamics (SGLD) [33, 29], by evaluating the mutual information conditional on the previous learned
∗Corresponding Author
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
parameters at each step. We note that the noise added in SGLD is critical for both the mutual information evaluation and the empirical risk minimization. Most existing work focus on SGLD with constant and isotropic noise covariance [22, 25, 14] due to the technical difficulty. However, it is observed that the test accuracy of SGLD with isotropic noise has a considerable gap compared to that of the widely-used Stochastic Gradient Descent (SGD) [40]. This empirical gap motivates us to consider the following question:
Can we find the optimal noise added in SGLD in terms of generalization?
An affirmative answer will lead to an algorithm imitating SGD better, which helps us to better understand the generalization behavior of SGD.
Specifically, we propose to optimize the structure of the noise in SGLD such that the generalization bound is minimized while a low empirical risk is guaranteed. To this end, we first show that the trace of the noise covariance in SGLD is a valid constraint that governs the empirical risk behavior both theoretically and empirically. Then we devise a new information theoretical generalization bound that are parallel to those bounds in [25], but facilitate the derivation of the optimal noise. With these technical preparations, we prove that when jointly optimizing both the prior and the posterior, the optimal noise covariance is the square root of the expected gradient covariance, i.e., their eigenvectors are the same and the corresponding eigenvalues of the former are the square root of the latter, and the optimal prior recovers the prior in [25]. This indicates the optimal noise covariance of SGLD would be quite close to the empirical gradient covariance, i.e., the noise covariance of SGD, because of the concentration of measure.
Our result lends support to the belief that the noise introduced by Stochastic Gradient Descent (SGD) is superior to the isotropic noise, which has been widely observed [18, 36, 40, 38]. As an illustrative example, we plot the generalization errors of SGD, SGLD with the isotropic noise and SGLD with the optimal noise in Figure 1, where their training curves behave almost the same (do not show here). We can see the optimal noise captures the behavior of SGD much better than the isotropic noise.
Specifically, our contribution can be summarized as follows:

heading: 1.1 Related Works

text:
Information-theoretical bounds. Recently, researchers [37, 30, 28] propose to bound the generalization error by the mutual information between output hypothesis and input samples. Negrea et al. [25] further tighten the bound by designing a data-dependent prior. Following [25], Haghifam et al. [14] obtain comparable results through conditional mutual information [31]. Other related work [2, 16, 5, 13] tightens the information theoretical generalization bounds from different perspectives. There are also high probability generalization bound [24, 26, 21] obtained by combining information theory [7] and PAC Bayesian framework [27, 33].
Effects of the noise covariance. SGD achieves excellent performance in terms of generalization error in deep learning. In contrast, one recent work [40] demonstrates empirically that even with well-tuned scaling, isotropic SGLD still achieves much worse generalization error than mini-batch SGD. They obtain an optimal covariance matrix (low-rank approximation of Hessian) in terms of escaping efficiency at local minima. Furthermore, [34] show empirically large-batch SGD with diagonal Fisher Gaussian noise can recover similar validation performance as small-batch SGD. However, both the noise covariance investigated by [40] and [34] has gap with that of SGD. It has also been proven [23] that the stationary distribution of state-dependent SGLD with the empirical loss approximated by quadratic function obeys the power law, and has a better escaping efficiency than SGLD with the isotropic noise.

heading: 2 Preliminaries

text:
Notations. Here we briefly introduce the notations which will be used throughout this paper.
• (Set operation) For a positive integer N , we use [N ] to denote the index set {1, 2, · · · , N}, and use V[T ] to denote the set {Vt}Tt=1. For a finite set S = {z1, · · · , zN}, |S| is the cardinality of S, and SJ △ = {zi}i∈J is a subset of S with indices in set J ⊂ [N ].
• (Probability) z ∼ S denotes that z is uniformly sampled from the set S. For two random variables X and Y , we denote the conditional distribution of X given Y as PY (X) and denote the conditional expectation of X given Y as EY (X).
• (Matrix) We use Id×d for the d× d identity matrix, abbreviated as I when dimension is clear. For a differentiable function f , we denote the gradient of f at point W as∇f(W ). For a positive semi-definite matrix A ∈ Rd×d, we say B = A 12 ∈ Rd×d if B ·B = A.
Supervised Learning. In this paper, we focus on the supervised learning. It conducts the empirical risk minimization (ERM) over training data: MinimizeW RS(W ) △ = 1|S| ∑ z∈S ℓ(z;W ), where W ∈ Rd is the parameter of the model, S is the training set with each data point i.i.d. sampled from a distribution D, ℓ is the (individual) loss function. A stochastic algorithm solves the ERM by outputting a distribution QS of parameter W . The generalization error measures the gap between the population risk and the training risk as
Gerr △ = ES,QS [RD(W )−RS(W )] . (1)
SGD and SGLD. The update rule of Stochastic Gradient Descent (SGD) at step t is defined as Wt ←Wt−1 − ηt∇RSVt (Wt−1), (2)
where Vt is sampled uniformly without replacement from [N ] with size bt. Given Wt−1, Wt is random with
E [Wt|Wt−1] = Wt−1 − ηt∇RS (Wt−1) , Cov [Wt|Wt−1] = N − bt
bt(N − 1) ΣsdS,Wt−1 ,
where ΣsdS,Wt−1 △ = Covz∼S [∇Rz(Wt−1)] is the covariance of∇Rz with z single drawn from S.
The superscript sd means “single draw”. Similarly, we define the population covariance of gradient as ΣpopW △ = Covz∼D [∇Rz(W )]. We define the Hessian ofRS(W ) at point W asHS,W .
The Stochastic Gradient Langevine Dynamics (SGLD) is given by
Wt ←Wt−1 − ηt∇RSVt (Wt−1) +N (0,Σt(S,Wt−1)), (3)
where Σt(S,Wt−1) ∈ Rd×d is a positive semi-definite matrix with dependence on S and Wt−1. SGLD (Eq. 3) with Σt(S,Wt−1) = ctI where ct a positive constant, is called isotropic SGLD, and SGLD with Σt(S,Wt−1) dependent on Wt−1 is called state-dependent SGLD.
Statistics for iterative algorithms. Both SGD and SGLD are stochastic iterative algorithms [11]. Specifically, the update rule at step t of a stochastic iterative algorithm A can be generally characterized as Wt =Mt(Wt−1,S,Vt), where Vt is the auxiliary random variable at step t, e.g., Vt in Eq.(2) and Eq.(3). We use QS,V[T ] to denote the joint distribution of (Wt)Tt=0 conditional on {S,V[T ]}, Q S,V[T ] i:j to denote the joint distribution of (Wt) j t=i conditional on {S,V[T ]}, and Q S,V[T ] j|j−1 to denote the distribution of Wj conditional on {Wj−1,S,V[T ]}. Decomposition of KL divergence. The following Lemma is extensively used throughout this paper. Lemma 1 (Proposition 2.6, [25]). Let Q0:T and P0:T are two probability measures on Rd×(T+1) with Q0 = P0. Then, the KL divergence between Q0:T and P0:T can be decomposed into
KL(Q0:T ||P0:T ) = T∑
t=1
EQ0:t−1 [ KL ( Qt|[t−1]∥Pt|[t−1] )] .
Information theoretical generalization bound. Several existing information-theoretical bounds [14, 25] share similar framework. We state one representative proposed in [25].
Proposition 1 (Theorem 2.5 in [25]). Let S be the data set i.i.d. sampled from D, and let the loss function ℓ be [a1, a2] bounded. Let A be an algorithm with update rule
Wt ←Mt(Wt−1,S,Vt),W0 ∼ W0.
Let B be another stochastic iterative optimization algorithm with update rule
Wt ← M̃t(Wt−1,SJ ,Vt,J),W0 ∼ W0. where J is sampled uniformly without replacement from [N ] with size N − 1. Given S, J , and V[T ], denote the joint distribution of (Wt)Tt=1 of A as QS,V[T ] , and the joint distribution of (Wt)Tt=1 of B as PJ,SJ ,V[T ] . Then, for any M̃[T ], the generalization error of A can be bounded as:
ES,V[T ] [ RD ( Q S,V[T ] T ) − R̂S ( Q S,V[T ] T )] ≤ ES,V[T ],J √ (a2 − a1)2
2 KL
( QS,V[T ] ∥∥PJ,SJ ,V[T ] ). This proposition is used to obtain the generalization error bound for SGLD [25, Theorem 3.1] by further combing Lemma 1. As M̃[T ] can be arbitrarily picked, P works as an "Auxiliary Line" and is called the prior distribution, while Q is the real distribution of parameters called the posterior distribution. In Section 3.2.1, we will argue the difficulty of applying Proposition 1 to analyze the effect of noise structures.

heading: 3 Formulate the Problem: Proper Constraints and Optimization Target

text: In this part, we formulate the optimization problem, i.e., finding the optimal noise covariance of SGLD in terms of the information-theoretical generalization bound, by selecting the proper optimization constraint and optimization target. Specifically, in Section 3.1, we argue that the trace of the noise covariance is a proper constraint to ensure the same optimization error; in Section 3.2, we argue that existing generalization bounds are not proper candidates as the optimization target, and propose a new information-theoretical bound parallel to existing ones but easier to analyze.

heading: 3.1 Constraint on the Covariance to Control the Empirical Risk

text:
We first derive the constraint of noise covariance Σt in Eq. (3) from the perspective of training performance, under which we optimize the generalization bounds in the rest of this paper.
Without any constraint on the noise covariance, optimizing generalization error is trivial but meaningless: a direct combination of Theorem 1 of [37] and Theorem 1 of [28] shows that for isotropic SGLD with Σt = σtI, the generalization error after T step satisfies Gerr ≤ O(( ∑T t=1 log(1 + 1/σt))
1/2). Therefore, as σt →∞ for t ∈ [T ], we have Gerr→ 0, but then the update of SGLD is dominated by the noise, leading to arbitrary bad empirical risk. Hence, we need constraints on the covariance in order to control the empirical risk when minimizing the generalization error. Specifically, the expected decrease of the empirical risk for one iteration can be bounded as follows. Lemma 2. Let empirical riskRS(W ) be β-smooth w.r.t. W . Let W[T ] be given by state-dependent SGLD (Eq. (3)). Then, Et+1|tRS(Wt+1)−RS(Wt) ≤ − ( 1 − βηt+1
2
) ηt+1 ∥∇RS (Wt)∥2+ β
2 tr ( η2t+1(N − bt) (N − 1)bt Σ sd S,Wt + Σt+1(S,Wt) ) .
The proof can be obtained by a standard analysis in optimization, and we defer it to Appendix B. By Lemma 2, the noise covariance Σt affects the upper bound of the empirical risk by its trace. Therefore, it is reasonable to keep tr(Σt(S,Wt−1)) unchanged while seeking the optimal Σt(S,Wt−1) to minimize the generalization error. The constraint is given formally as follows: Constraint 1. The trace of Σt(S,Wt−1) is fixed when optimizing the generalization error. That is, there exist positive constants ct(S,Wt−1) depending on S and Wt−1, such that,
tr(Σt(S,Wt−1)) = ct(S,Wt−1).
We do not put any constraint on the value of ct in our latter analyses. Therefore, it is also possible to manipulate ct in order to jointly optimize the empirical risk and the generalization bound, which
however, is beyond the scope of this paper and we defer it to future works. Similar constraint is also proposed by [40] from the standpoint of kinetic energy when analyzing the effect of noise structure on the escaping efficiency from saddle points.
We next verify this constraint empirically. We run SGLD with different covariances on a four-layer neural network for the Fashion-MNIST classification problem. Concretely, the noise covariances are chosen respectively as Σ(1)t = Σ sd S,Wt−1 , i.e., “EC-SGLD” ( Empirical Covariance SGLD) curve in Fig.2, and Σ(2)t = 1 d tr(Σ sd S,Wt−1
)I, i.e., the “Iso-SGLD (C)” curve in Fig.2. It is easy to verify that tr(Σ
(1) t ) = tr(Σ (2) t ), which is exactly the Constraint 1. We can see from Fig.2 that the convergence curves corresponding to Σ(1)t and Σ (2) t almost coincide with each other, validating the Constraint 1.
In comparison, another quantity tr(HS,Wt−1Σt), the trace of the empirical risk’s Hessian times the empirical covariance, has also been proposed to govern the convergence behavior [34, Theorem 4.1]. We plot the convergence curve of SGLD with noise covariance
Σ (3) t =
tr(HS,Wt−1Σ sd S,Wt−1 )
tr(HS,Wt−1 ) I ("Iso-SGLD (H)" in Fig.2). While
tr(HS,Wt−1Σ (1) t ) = tr(HS,Wt−1Σ (3) t ), there is a significant gap between curves of “EC-SGLD ” and “Iso-SGLD (H)” in Fig.2. This implies that tr(HS,Wt−1Σt) is not a good constraint for the empirical risk of SGLD, validating the Constraint 1 from another side.
In the rest of this paper, we will optimize the information-theoretical bound under Constraint 1 to ensure low empirical risk.

heading: 3.2 New Information-theoretical Bounds as the Optimization Target

text: We first demonstrate that existing information-theoretical bounds are not suitable for the optimization target in Section 3.2.1. Then we propose a new information-theoretical bound as the optimization target in Section 3.2.2.

heading: 3.2.1 Difficulties When Applying Traditional Information-theoretical Bounds

text:
By Lemma 1, the generalization error bound in Proposition 1 can be rewritten as
ES,V[T ],J √√√√ (a2 − a1)2 2 T∑ s=1 E Q S,V[s−1] s−1 KL ( QS,Vss|(s−1) ∥∥∥PJ,SJ ,Vss|(s−1) ). (4) For Problem 1, there are two difficulties. For a fixed QS,V[T ] , in order to obtain the optimal PJ,SJ ,Vss|(s−1) for any s ∈ [T ], one would first calculate the outside expectation of Eq. (4) over SJc , as PSJ ,V[T ] is independent of SJc . However, for each term E
Q S,V[s−1] s−1
KL (QS,Vss|(s−1)∥P J,SJ ,Vs s|(s−1) ) where s ∈ [T ],
both the probability measure Q S,V[s−1] s−1 and the function KL(Q S,Vs s|(s−1)∥P J,SJ ,Vs s|(s−1) ) has dependency on SJc , which makes evaluating the generalization bound with respect to PSJ ,V[T ] extremely hard. Furthermore, when we come to optimize the bound w.r.t. Σ[T ], for each Σi where i ∈ [T − 1], the Eq. (4) depends on Σi via two terms A and B, where
A = E Q
S,V[T ] i−1
KL ( Q
S,V[T ] i|(i−1)∥P J,SJ ,V[T ] i|(i−1)
) , B = T∑ s=i+1 E Q S,V[s−1] s−1 KL ( QS,Vss|(s−1)∥P J,SJ ,Vs s|(s−1) ) .
While the first term is easy to deal with as Q S,V[T ] i−1 is irrelevant of Σi, the second term depends on Σi via the distribution Q S,V[s−1] s−1 for s ≥ i+ 1, and can be very complex (please refer to Appendix C.1).

heading: 3.2.2 Information-theoretical Generalization Bound for Identifying the Noise Effect

text:
We then establish a new information-theoretical generalization bound as the optimization target that are parallel to the results in [25] but more suitable for analyzing the effect of noise structure. The
basic idea is to reverse the order of prior and posterior in the KL divergence to let Σt only affect one term in the generalization bound for each t when finding the optimal point. The formal theorem is stated as follows: Theorem 1. Let sample set S, mini-batch V[T ], random subset J , the posterior distribution Q output by Algorithm A with update ruleM and the prior distribution P output by Algorithm B with update rule M̃ be defined as Proposition 1. Let [a1, a2] be the range of loss. Then, the generalization error of A can be bounded as
Gerr ≤ ES,V[T ],J
√ (a2 − a1)2
2 KL
( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). (5) Compared to Proposition 1, Theorem 1 can be viewed as a parallel version with the positions of the prior distribution P and the posterior distribution Q reversed. This reverse benefits the optimization of the generalization error bound with respect to the noise covariance under Constraint 1 in two ways. First, by Lemma 1, the KL term in Eq. (5) can be further decomposed into
KL ( PSJ ,V[T ] ∥∥QS,V[T ] ) = T∑ s=1 E P J,SJ ,V[s−1] 1:s−1 KL ( PJ,SJ ,Vss|(s−1) ||Q S,Vs s|(s−1) ) ,
in which only each KL term depends on each Σs, respectively. Secondly, for each summand in the above decomposition, only QS,Vss|(s−1) depends on SJc , making it easy to compute the optimal prior in a "greedy" sense (see details in Section 4).
Hence in the rest of this paper, we study the optimal structure of Σt in terms of the generalization bound in Theorem 1 under the Constraint 1. Specifically, our ultimate goal can be stated as follows: Problem 1. What is the optimal structure of the noise covariance Σ[T ] for SGLD in terms of the generalization bound in Theorem 1 under the Constraint 1?
In the rest of this paper, we focus on solving Problem 1 and its variants.

heading: 4 Obtain the Optimal Noise Covariance with the Greedy Prior

text:
The generalization bound in Theorem 1 depends on both the prior distribution P and the posterior distribution Q (consequently on Σ[T ]). Therefore it requires searching P and Σ[T ] jointly to optimize the generalization error.
In this section, we solve Problem 1 with greedily selected priors. Due to that the square root in the bound Eq. (5) make the dependency on Σ[T ] even more complex, we optimize a slightly different version of the bound by taking the expectation with respect to SJc into the square root, i.e.,
GenT △ = ESJ ,V[T ],J
√ (a2 − a1)2
2 ESJc KL
( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). (6) By Jensen’s Inequality, Eq. (6) is still a generalization bound, but allows the expectation with respect to SJc to interact with the KL divergence directly. One can easily observe that GenT is a mapping from P and Q (and thus from P and Σ[T ]) to a positive real. Therefore, we can reformulate Problem 1 mathematically as the following optimization problem (P1):
(P1). min P,Σ[T ] GenT (P,Σ[T ]), subject to: Constraint 1.
For simplicity, we restrict that the considered noise covariance Σ[T ] only depends on the parameter W , and is independent of the sample S. By Lemma 1, the KL term in the optimization target GenT can be decomposed into
KL ( PJ,SJ ,V[T ] ∥∥∥QS,V[T ] ) = T∑ s=1 E P J,SJ ,V[s−1] s−1 KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1) ) . (7) Therefore, when optimizing GenT with respect to prior P , both KL(P J,SJ ,Vs s|(s−1) ||Q S,Vs s|(s−1)) and P J,SJ ,V[i−1] i−1 (i > s) have dependence on P J,SJ ,Vs s|(s−1) . Similar to the discussion in Section 3.2.1, the dependence of P J,SJ ,V[i−1] i−1 with i > s on P J,SJ ,Vs s|(s−1) can be very complex. Therefore, we approximate the optimal PJ,SJ ,Vss|(s−1) by the greedy prior which is defined as follows:
Definition 1 (Greedy Prior). We say P ∗ is the optimal prior in the greedy sense, or the greedy prior for brevity, if for any 1 ≤ s ≤ T and any SJ and V[T ], P ∗J,SJ ,Vss|(s−1) = P sJ,SJ ,Vs s|(s−1) , where P
s is defined as follows:
P s △ = argmin
P ( min Σ[s] Gens(P,Σ[s]) ) , subject to: Constraint 1.
Intuitively, the conditional probability of P ∗ of the step s is the optimal one if we only consider the generalization bound for steps up to s, and a special case is that the step T conditional probability of P ∗ agrees with the step T conditional probability of PT , which is the desired optimal prior. This is why we call P ∗ "greedy" and use it to approximate the optimal prior.
With the greedy prior, we characterize the optimal noise covariance by the following theorem.
Theorem 2. Let the iteration of SGLD with state-dependent noise QS,V[T ] be given as Eq.(3). Under Constraint 1, the greedy optimal prior of step t is given by
Wt = Wt−1 − ηt ( |Vt ∩ J | |Vt| ∇RSVt∩J (Wt−1) + |Vt ∩ Jc| |Vt| ∇RD (Wt−1) ) +N (0,Σ∗t (Wt−1)) ,
while the optimal covariance of noise Σ∗[T ] for GenT with the greedy prior is given by Σ ∗ t (W ) =
λt(W ) (Σ pop W ) 1 2 (∀t ∈ [T ]), where λt(W ) = ct(W )/ tr((ΣpopW ) 1 2 ).
As the sample size is large enough, we have ΣsdS,W → Σ pop W almost surely, which demonstrates the similarity between the solution of Problem 1 and the noise covariance of SGD. Also, by the Law of Large Numbers, ∇RSJ ≈ ∇RD, and the mean of the greedy optimal prior recovers the mean of the prior used in [25] (one can also refer to Eq. (10) in this paper for the form).
We briefly state the proof skeleton of Theorem 2, with the proof details deferred to Appendix D. To obtain the final optimal noise covariance in Theorem 2, we need to first derive the greedy prior, i.e., the optimal conditional distribution Ps|(s−1) of the prior of step s in terms of the generalization bound Gens, which has the form
ESJ ,V[s],J √√√√ (a2 − a1)2 2 ESJc s∑
t=1
E P
J,SJ ,V[t−1] t−1
KL ( PJ,SJ ,Vtt|(t−1) ∥∥∥QS,Vtt|(t−1)). (8) Typically, solving the optimal Ps|(s−1) requires optimizing Gens with respect to all Pt|(t−1) t ∈ [s] and Σ[s], which is still very complex. However, we can tackle this problem in a rather elegant way. We first investigate the optimal noise covariance in terms of a single KL divergence term in Gens.
Lemma 3. Under Constraint 1, the optimal noise covariance of the following problem
min Σs  min P
J,SJ ,Vs s|(s−1)
ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1))  . (9)
is attained at Σ∗t (W ), where Σ ∗ t (W ) is defined as Theorem 2.
By Lemma 3, the optimal solution of Eq. (9) doesn’t rely on J ,SJ , or V[s]. On the other hand, by Eq. (8), ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1)) (∀J ,SJ ,V[s]) are the only terms depending on PJ,SJ ,Vss|(s−1) and Σs. Therefore, the optimal Σs is also λt(W ) (Σ pop W ) 1 2 , which is formally stated as the following lemma.
Lemma 4. The optimal Σs and Ps|(s−1) in terms of Gens are the same as Σ∗s and P ∗s|(s−1) given by Theorem 2, respectively.
With the greedy prior derived, we apply it back to the generalization bound GenT . As GenT depends on Σs also through ESJc∼D KL ( PJ,SJ ,Vss|(s−1) ∥∥∥QS,Vss|(s−1)), by applying Lemma 3 again, we derive Theorem 2.

heading: 5 Extension: Optimal Noise Covariance with Fixed Priors

text:
In existing works [25, 12], the prior distribution is set to be the SGLD with isotropic noise, which (with the notations in Theorem 1) is given by
M̃t(Wt−1,SJ ,Vt,J) = Wt−1 − ηt ( |Vt ∩ J | |Vt| ∇RSVt∩J (Wt−1) + |Vt ∩ Jc| |Vt| ∇RSJ (Wt−1) ) +N (0, σtId) , (10)
where σt > 0 is the noise scale of prior noise covariance. In our latter analysis, we generalize the prior by allowing σt depend on Wt−1. The formal description of the iteration of the prior is deferred to Appendix E.1 for completeness.
Therefore, it is also interesting to see what the optimal noise covariance looks like if the prior is fixed as the one commonly adopted in the existing analyses, e.g., Eq. (10). We still set the optimization constraint the same as Section 4, but change the optimization target a little to
G̃enT △ = ES
√ (a2 − a1)2
2 EV[T ],J KL
( PJ,SJ ,V[T ] ∥∥QS,V[T ] ). By Jensen’s Inequality and Theorem 1, G̃enT is still a generalization bound, but allows us to treat the expectation of the KL divergence with respect to V[T ] and J for given S as a whole in optimization. Similar trick is also adopted in [25] to obtain the final generalization error of SGLD. The problem can be mathematically formulated as the following optimization problem (P2):
(P2) min Σ[T ] G̃enT (P,Σ[T ]), subject to: Constraint 1,
where P is given by the update rule Eq. (10). As SJ is obtained by removing only one sample from the dataset S and the size of S is usually large in practice, it is reasonable to make the assumption that for any fixed V[T ] and S, the prior distribution is the same regardless of J .
Assumption 1. For any fixed dataset S and mini-batches V[T ], the distribution PJ,SJ ,V[T ] is invariant of J .
We also restrict that ct(S,W ) ≤ dσt in order to guarantee the noise scale of the prior comparable to that of the posterior. The optimal noise covariance with prior fixed can then be characterized by the following theorem. Theorem 3. Let prior and posterior be defined as Eq.(10) and Eq.(3), respectively. Let Assumption 1 hold. Then, the solution of (P2) is given by
Σ∗t (S,W ) = O sd S,W Diag(ω̃ S,W 1 , · · · , ω̃ S,W d ) ( OsdS,W )⊤ ,
where ω̃S,Wi ≥ 0 (i ∈ [d]) (the exact form is omitted here) and ( OsdS,W ) is the orthogonal matrix
that diagonalizes ΣsdS,W as Σ sd S,W = O sd S,W Diag(ω S,W 1 , · · · , ω S,W d ) ( OsdS,W )⊤ . Moreover, for any i ̸= j, ω̃S,Wi ≥ ω̃ S,W j if and only if ω S,W i ≥ ω S,W j .
The proof together with the the exact formula of ω̃S,Wi is deferred to Appendix 5. From Theorem 3, the optimal point (Σ∗t ) T t=1 is similar to the empirical gradient covariance matrix Σ
sd in two ways: 1) {Σ∗t }Tt=1 share the same eigenvectors with Σsd; 2) the corresponding eigenvalues of Σ∗t has the same order as Σsd. Though the value of ω̃S,Wi is not comparable to ω S,W i because Σ ∗ t is affected by the prior noise and the posterior noise covariance, which are freely chosen, it can be shown that the condition number of Σ∗t is smaller than Σ sd S,W (please refer to Appendix E.4 for details).
Theorem 3 also reveals an interesting correlation between the noise covariance matrices of the prior and posterior. While the optimal noise covariance is affected by the prior noise covariance, it also depends on the distance between the means of the prior and the posterior (see Lemma 13 in Appendix E.2), which brings the information of empirical gradient covariance into the optimal noise structure. That being said, the optimal posterior noise covariance is biased to the empirical gradient covariance from prior covariance. We note that such analysis can be easily extended to arbitrary priors.

heading: 6 Empirical Verification

text:
In this section, we support our theoretical findings with some experiments. We adopt the setting in [40] where a four-layer neural network with 11330 parameters is used to conduct the classification task on the Fashion-MNIST except that we use 10000 training samples instead of 1200 used in [40]. We defer detailed settings of the experiments to Appendix F.
We first verify if the empirical gradient covariance is far from isotropic Gaussian distribution. We plot the ratios of the 1st eigenvalue to the 500th largest eigenvalue of the empirical gradient covariance along the training trajectory of both SGD and Iso-SGLD in Figure 4. We can see that throughout the training procedure, the ratios of empirical gradient covariance stay around 107 for both SGD and Iso-SGLD. In constrast, for the isotropic Gaussian, the distribution of eigenvalues follows semi-circle law, and the ratio would be around 1.1 for dimension 11330. This demonstrates that the energy of empirical gradient covariance concentrates in a very small subspace, less than 5% of the total 11330 dimensions. Hence the empirical gradient covariance is highly anisotropic.
We next verify our main claim Theorem 2 that under Constraint 1, the generalization error for SGLD with the optimal noise covariance is much smaller than that for Iso-SGLD. Specifically, we consider the SGLD with noise covariance equal to the (scaled) square root of the empirical gradient covariance, coined “SREC-SGLD”. We compare the training loss and the generalization error of SREC-SGLD and Iso-SGLD under Constraint 1 on the noise covariance trace. From Figure 3, we can see that the generalization error of SREC-SGLD is smaller than Iso-SGLD for different noise trace scales, which supports Theorem 2. Moreover, if we look at the training loss curves, the SREC-SGLD and Iso-SGLD behave almost the same for the same noise trace scale, which supports Lemma 2.
Finally though our analysis does not cover the generalization bound for SGD, we further empirically show how SREC-SGLD performs in comparison with SGD and the Iso-SGLD in Figure 1. We can see that with the same noise trace scale, SREC indeed provides a more accurate characterization of SGD compared to Iso-SGLD.

heading: 7 Conclusion and Future Direction

text: In this paper, we study the optimal noise covariance of SGLD in terms of its generalization ability. Specifically, we first formulate the optimization problem both by deriving constraints from the optimization performance and proposing the optimization target by constructing a new information theoretical bound. We then solve the problem with both greedy optimal prior and fixed prior. Interestingly, we observe that the optimal noise covariance aligns with the empirical gradient covariance, which indicates the superiority of the noise covariance of SGD in terms of generalization.

heading: Acknowledgement

text: The authors would like to thank Mr. Ziming Liu for helpful theoretical discussions.

title: Optimizing Information-theoretical Generalization Bounds via Anisotropic Noise in SGLD