On the Trade-off between Flatness and Optimization
in Distributed Learning

Understanding the optimization behavior of mini-batch gradient descent (GD) algorithms and their influence on the generalization of models has emerged as a prominent topic of interest in recent years. Large-batch training is increasingly important in distributed learning for its potential to enhance the training speed and scalability of modern neural networks. However, it has already been observed in the literature that small-batch GD tends to flatter minima than the large-batch version [30], and that flat minima usually generalize better than sharp ones [31, 32, 33, 34, 35, 36]. Intuitively, the loss values around flat minima change slowly when the model parameters are adjusted, thereby reducing the disparity between the training and test data [30]. The works by [37, 38, 39, 40, 34, 41] utilized the stochastic differential equation (SDE) approach as a fundamental tool to analyze the regularization effects of mini-batch GD. This type of analysis inherently introduces an extra layer of approximation error due to substituting the discrete-time update by a different equation. Moreover, to enable the analysis with SDE, it is necessary to make some extra assumptions on the gradient noise, such as being Gaussian distributed [34, 39], parameter-independent [37] or heavy-tailed distributed [42]. To overcome these limitations, another line of works [43, 44, 33] focused on the dynamical stability of mini-batch GD with discrete-time analysis. Specifically, they determined the conditions on the Hessian matrices to ensure that the distance to a local minimum does not increase, thereby stabilizing the iterates in the vicinity of a local minimum.

Contrary to the aforementioned works focusing on the single-agent case, we will study in this work the optimization bias of distributed algorithms in the multi-agent setting, which allows us to compare the performance of various distributed strategies. Inspired by [34], we start from investigating the esca** efficiency of algorithms from local minima, and then relate it to the trade-off between flatness and optimization. Our contributions are listed as follows:

(1) We propose a general framework to examine the esca** efficiency of distributed algorithms from local minima by staying in the discrete-time domain. Since the dependence of the Hessian matrix on the immediate iterate in the original recursion makes the optimization analysis intractable, it is necessary to resort to a short-term model where the original Hessian is replaced by one evaluated at some local minima. We rigorously verify that the approximation error between the short-term and true models is negligible, ensuring that the short-term model represents the true model accurately enough. Afterwards, we obtain closed-form expressions of the short-term excess risk which quantifies the esca** efficiency. Note that we follow a different discrete-time analysis from the dynamical stability approach used in [43, 44, 33] and which focused on studying the stability of algorithms. Similar to [34], our emphasis is on quantifying the extent to which algorithms can escape from local minima.

(2) According to the results obtained from (1), we compare the esca** efficiency of the centralized, consensus and diffusion methods. We show that decentralized approaches, i.e., consensus and diffusion, gain additional efficacy from the network heterogeneity and graph structure, enabling them to escape faster from local minimum than the centralized strategy. We further verify that consensus outperforms diffusion in terms of esca** ability from local minima. We also show that higher esca** efficiency encourages algorithms to favor flatter minimum by relating esca** efficiency to the flatness metric.

(3) If the additional power generated by the network heterogeneity and graph structure are not sufficiently strong to allow decentralized methods to successfully leave the current basin, then both decentralized and centralized methods will be stuck within the basin of the current local minimum. A similar rule also holds for the comparison between consensus and diffusion. This motivates us to pursue next a long-term analysis, which corresponds to the optimization performance. In this context, we verify that the extra power boosting the esca** efficiency could inversely deteriorate the optimization performance. This reveals the inherent trade-off between flatness and optimization, with diffusion showing superior optimization performance over consensus in terms of smaller estimation error.

(4) We finally illustrate the performance of centralized, consensus and diffusion training strategies on real data. In addition to the flatness and optimization performance of the three distributed methods, we show that diffusion achieves a favorable balance between flatness and optimization, thereby exhibiting superior test accuracy than the other two methods (i.e., centralized and consensus).

Consider a collection of $K$ agents (or nodes) linked by a graph topology. Each agent receives a streaming sequence of data realizations arising from independently distributed observations. The agents wish to collaborate to solve a distributed learning problem of the following form:

where $M$ is the dimension of the unknown vector $w$, $J_{k}(w)$ is the risk function for the $k$-th agent, and $J(w)$ is the aggregate risk across the graph. Each individual risk is defined as the empirical average of a loss function over a collection of training data points, namely,

where $Q_{k}(\cdot)$ is the possibly nonconvex loss function and the $\{x_{k,i}=\{\gamma_{k,i},h_{k,i}\}\}$ refers to $N_{k}$ training samples with feature vector $\gamma_{k,i}$ and true label $h_{k,i}$ arising as realizations from a random source $\boldsymbol{x}_{k}$ associated with agent $k$.

Since we aim to understand the optimization behavior of algorithms around local minima of nonconvex risk functions, we need to distinguish between the local minima of $J_{k}(w)$ and $J(w)$. Thus, we let $w^{\star}$ denote a local minimizer for the aggregate risk $J(w)$, and let $w_{k}^{\star}$ denote a local minimizer for the individual risk $J_{k}(w)$. If all agents in the network share the same risk function, i.e., $J_{k}(w)=J(w)$, then $w^{\star}=w_{k}^{\star}$, and we refer to the network as being homogeneous. In this case, all local minimizers of $J(w)$ are also local minima of $J_{k}(w)$. Otherwise, we refer to the network as being heterogeneous, where the $w_{k}^{\star}$ of all agents can be distinct among themselves and also in relation to $w^{\star}$. In this paper, we focus on this latter more general case.

We consider three popular classes of algorithms that can be used to seek a solution for (1). The first algorithm is the mini-batch centralized method, in which all data from across the graph are shared with a central processor. At every instant $n$, $B$ samples denoted by $\{\boldsymbol{x}_{k,n}^{b}\}$ are selected uniformly at random with replacement from the training data available for each agent $k$. Starting from a random initial condition $\boldsymbol{w}_{-1}$, the central processor updates the estimate by using:

Here, $B$ is the batch size and $\mu$ is a small positive step-size parameter. Observe that we are using boldface letters to refer to the data samples and iterates in (3); it is our convention in this paper to use the boldface notation to refer to random quantities.

The other two algorithms are of the decentralized type, where the data remains local and agents interact locally with their neighbors to solve (1) through a collaborative process. We consider two types of decentralized methods, which have been studied extensively in the literature, namely, the consensus and diffusion strategies [4, 5, 6, 7, 20, 24].

Before listing the algorithms, we describe the graph structure that drives their operation. The agents are assumed to be linked by a weighted graph topology. The weight on the link from agent $\ell$ to agent $k$ is denoted by $a_{\ell k}$; this value is used to scale information sent from $\ell$ to $k$. Each $a_{\ell k}$ is non-negative and lies within $[0,1]$; it will be strictly positive if there exists a link from $\ell$ to $k$ over which information can be shared. We collect the $\{a_{\ell k}\}$ into a $K\times K$ matrix $A$. In this paper, we consider a symmetric doubly-stochastic matrix $A$: the entries on each column of $A$ are normalized to add up to $1$.

It follows from the Perron-Frobenius theorem [6] that $A$ has a single eigenvalue at 1. Moreover, if we let $\pi=\{\pi_{k}\}_{k=1}^{K}$ denote the corresponding right eigenvector, then all its entries are positive and they can be normalized to add up to $1$:

For the doubly-stochastic matrix $A$, we have $\pi=\frac{1}{K}\mathbbm{1}$. This means that the entries $\{\pi_{k}\}$ of the Perron vector $\pi$ are identical.

The diffusion strategy consists of the following two steps:

At every iteration $n$, every agent $k$ samples $B$ data points and uses $(\ref{x_e})$ to update its iterate $\boldsymbol{w}_{k,n-1}$ to the intermediate value $\boldsymbol{\phi}_{k,n}$. Subsequently, the same agent combines the intermediate iterates from across its neighbors using (5b). The symbol $\mathcal{N}_{k}$ denotes the collection of neighbors of $k$. It is useful to remark that the centralized implementation (3) can be viewed as a special case of (5a)–(5b) if we select the combination matrix as $A=\pi\mathbbm{1}^{\sf T}$.

In comparison, the consensus strategy involves the following steps:

In this case, the existing iterates $\boldsymbol{w}_{\ell,n-1}$ are first combined to generate the intermediate value $\boldsymbol{\phi}_{k,n}$, after which (6b) is applied. Observe the asymmetry on the right-hand side in (6b). The starting iterate is $\boldsymbol{\phi}_{k,n}$, while the loss functions are evaluated at the different iterates $\boldsymbol{w}_{k,n-1}$. In contrast, the same iterate $\boldsymbol{w}_{k,n-1}$ appears in both terms on the RHS of (5b). This symmetry has been shown to enlarge the stability range of diffusion implementations over its consensus counterparts in the convex case, namely, diffusion is mean-square stable for a wider range of step-sizes $\mu$ [45, 1, 7].

We first characterize the basin (or valley) of a local minimum:

We illustrate the basin of attraction related to $w^{\star}$ in Figure 2, where $\Omega(w^{\star})=(w_{1},w_{2})$ is a bounded open set. Also, we use $\partial\Omega(w^{\star})$ to denote the boundary of $\Omega(w^{\star})$, according to which $\partial\Omega(w^{\star})=\{w_{1},w_{2}\}$ in Figure 2. If ${w}_{k,n}$ or $w_{n}\notin\Omega(w^{\star})$, then we say the algorithm escapes from $w^{\star}$ exactly.

To analyze the esca** behavior of learning algorithms from local minima, it is necessary to quantify their esca** ability. Inspired by the continuous-time definition of esca** efficiency for the single-agent case from [34], and by the metrics used to assess the optimization performance of stochastic learning algorithms [6, 7], we use the following metric to measure the discrete-time esca** efficiency for general non-convex environments.

$\mathrm{ER}_{n}$ assesses the average excess-risk value across the network, and it quantifies the deviation between the current model at iteration $n$ and the local minimum. For a fixed $n$, a larger value of $\mathrm{ER_{n}}$ indicates an expected faster escape from the local minimum (since the iterate will be further away from it). To justify whether an algorithm escapes from a local basin or not, earlier studies [46, 47, 48] have established various criteria based on the analysis methods they use. They nevertheless generally adhered to the principle of the minimum effort to reach the boundary of the local basin. In a similar vein, we introduce the risk barrier, which is defined as the infimum of the risk gap between the boundary points and the local minimum:

For example, the risk barrier of $w^{\star}$ in Figure 2 is $J(w_{2})-J(w^{\star})$. When $\mathrm{ER}_{n}$ exceeds the corresponding risk barrier defined in (8), namely,

then we say that the network escapes from $w^{\star}$ on average. If, on the other hand, $\boldsymbol{w}_{k,n}$ gets stuck around (or converges to) $w^{\star}$, then $\mathrm{ER}_{n}$ in that case would serve as a measure of the resulting optimization performance[6, 7].

Next, we formally characterize the notion of flat minima. Consider a local minimum $w^{\star}$ and a small drift $\Delta w$ around it. The change in the risk value can be approximated by

where $H^{\star}$ refers to the(positive definite) Hessian matrix of $J(w)$ at $w^{\star}$. If the above change in the risk value is small then we refer to $w^{\star}$ as a flat minimum. Otherwise, if the change in the risk value is large, then we refer to $w^{\star}$ as a sharp minimum. Motivated by (10), various metrics related to the Hessian matrix are applied in the literature to measure the flatness of local minima, e.g., the spectral norm $\|H^{\star}\|_{2}$[33, 43], the Frobenius norm $\|H^{\star}\|_{F}$[33, 44] and the trace $\mathrm{Tr}(H^{\star})$[49, 33, 50]. In this paper, we use the trace of the Hessian as a flatness measure, but the other two metrics may also be applied due to the norm equivalence, namely:

Note that for highly ill-conditioned minima, which is commonly observed in over-parameterized neural networks [34, 51], the first $d$ largest eigenvalues of $H^{\star}$ dominate the remaining $M-d$ eigenvalues where $d\ll M$. In this case, the norm equivalence in (11) can be more tightly guaranteed since now the $M$ in the upper and lower bounds can be approximately replaced by $d$. More insights related to the flatness measure $\mathrm{Tr}(H^{\star})$ can be found in [49, 50].

We introduce the following commonly-used assumptions. First, since our purpose is to understand the behavior of the algorithm around the same local minima $w^{\star}$, we assume all agents start from points near $w^{\star}$.

By using Jensen’s inequality on (12), we have

Since we will establish that mini-batch gradient descent algorithms approach an $O(\frac{\mu}{B})$ neighborhood of $w^{\star}$, it is reasonable to assume all agents start from the $o(\frac{\mu}{B})$ neighborhood, which is dominated by $O(\frac{\mu}{B})$. This assumption is weaker than the one used in [34] where models should exactly start from $w^{\star}$.

We further require the loss function $Q_{k}$ of all agents to be smooth [11, 21].

Next, consider the Hessian matrix of agent $k$ at $w^{\star}$ denoted by:

and the global Hessian matrix of $J(w)$ at $w^{\star}$:

Assumption III.3 can be satisfied when the data heterogeneity among agents is sufficiently small. For example, if all agents observe independently and identically distributed data, and consider the stochastic Hessian defined by:

Then by resorting to the matrix Bernstein inequality [52], for any $\epsilon>0$, we have:

Thus, when each agent collects sufficient amount of data, all local Hessian matrices are guaranteed to be close to the stochastic Hessian with high probability, from which we get:

Intuitively, even when all agents independently collect data from different distributions, if the data distribution among agents are sufficiently close to each other, then (III-A) can still be satisfied.

To enable the analysis, we describe properties associated with the gradient noise process for later use. For any $w\in\mathbbm{R}^{M}$, the stochastic gradient noise at agent $k$ at iteration $n$ is defined by the difference:

where the term stochastic refers to the case $B=1$ (i.e., batch size equal to 1). In the mini-batch case, the gradient noise is instead given by:

We denote the covariance matrix of the gradient noise by:

which is symmetric and non-negative definite. Likewise, for the mini-batch case,

We can further relate the noise covariance matrix at $w^{\star}$ and the Hessian matrix $\bar{H}$. In the single-agent case, when $B=1$ and negative log-likelihood losses are used, it is well-known that there is an exact equivalence between the Hessian and the gradient covariance matrices at local minima in stochastic risk minimization [34], and this relationship can also be approximately used in the context of empirical risk minimization [34, 53, 39]. By log-likelihood losses we mean choosing $Q_{k}(w;\boldsymbol{x}_{k})$ to be of the form

where $\boldsymbol{h}_{k}$ and $\boldsymbol{\gamma}_{k}$ represent the feature vector and label associated with $\boldsymbol{x}_{k}$. In the multi-agent case, we define the gradient covariance matrix of the global risk function at $w^{\star}$ for the case $B=1$ as:

Lemma III.5 establishes the approximate equivalence between $\bar{R}_{s}$ and $\frac{1}{K}\bar{H}$ in the context of empirical risk minimization. Since negative log-likelihood losses, e.g., mean-square loss and cross-entropy loss, are popular in the machine learning community, we mainly focus on this category in this paper, but our framework can also be applied to other losses.

We now proceed with the network analysis, which will enable us to derive expressions for the value of $\mathrm{ER}_{n}$. We will then use these expressions to deduce properties about the behavior of the centralized and decentralized algorithms in relation to flat minima and esca** efficiency. To begin with, we follow the decomposition from [6] and note that the combination matrix $A$ admits an eigen-decomposition of the form:

where the matrices $V$ and $P$ are:

Here, $P_{\alpha}\in\mathbbm{R}^{(K-1)\times(K-1)}$ is a diagonal matrix with elements from the second largest-magnitude eigenvalue $\lambda_{2}$ to the smallest-magnitude eigenvalue $\lambda_{K}$ of $A$ appearing on the diagonal, and $V_{\alpha}\in\mathbbm{R}^{K\times(K-1)}$. Consider the extended network version policy:

where $\otimes$ denotes the Kronecker product, and $I_{M}$ is the identity matrix of size $M$. Then, $\mathcal{A}$ satisfies

where

We further collect quantities from across the network into the block variables:

where $\mathrm{col}$ denotes a block column vector, and each Hessian matrix $H_{k,n}(\boldsymbol{w}_{k,n})$ is defined by