Accelerated Stochastic ExtraGradient: Mixing Hessian and Gradient Similarity to Reduce Communication in Distributed and Federated Learning

Empirical risk minimization is a key concept in supervised machine learning (Shalev-Shwartz et al.,, 2010). This paradigm is used to train a wide range of models, such as linear and logistic regression (Tripepi et al.,, 2008), support vector machines (Moguerza and Muñoz,, 2006), tree-based models (Hehn et al.,, 2020), neural networks (Sannakki and Kulkarni,, 2016), and many more. In essence, this is an optimization of the parameters based on the average losses in data. By minimizing the empirical risk, the model aims to generalize well to unseen samples and make accurate predictions. In order to enhance the overall predictive performance and effectively tackle advanced tasks, it is imperative to train on large datasets. By leveraging such ones during the training process, we aim to maximize the generalization capabilities of models and equip them with the robustness and adaptability required to excel in challenging real-world scenarios (Feng et al.,, 2014). The richness and variety of samples in large datasets provides the depth needed to learn complex patterns, relationships, and nuances, enabling them to make accurate and reliable predictions across a wide range of inputs and conditions. As datasets grow in size and complexity, traditional optimization algorithms may struggle to converge in a reasonable amount of time. This increases the interest in developing distributed methods (Verbraeken et al.,, 2020). Formally, the objective function in the minimization problem is a finite sum over nodes, each one containing its own possibly private dataset. In the most general form, we have:

where $M$ refers to number of nodes/clients/machines/devices, $n_{m}$ is the size of the $m$-th local dataset, $x$ is the vector representation of a model using $d$ features, $y_{j}^{m}$ is the $j$-th data point on the $m$-th node and $\ell$ is the loss function. We consider the most computationally powerful device that can communicate with all others to be a server. Without loss of generality, we assume that $r_{1}$ is responsible for the server, and the rest $r_{m}$ for other nodes.

This formulation of the problem entails a number of challenges. In particular, transferring information between devices incurs communication overhead (Sreenivasan et al.,, 2006). This can result in increased resource utilization, which affects the overall efficiency of the training process. There are different ways to deal with communication costs (McMahan et al.,, 2017; Koloskova et al.,, 2020; Alistarh et al.,, 2017; Beznosikov et al.,, 2023). One of the key approaches is the exploiting the similarity of local data and hence the corresponding losses. One way to formally express it is to use Hessian similarity (Tian et al.,, 2022):

where $\delta>0$ measures the degree of similarity between the corresponding Hessian matrices. The goal is to build methods with communication complexity depends on $\delta$ because it is known that $\delta\thicksim\nicefrac{{1}}{{\sqrt{n}}}$ for non-quadratic losses and $\delta\thicksim\nicefrac{{1}}{{n}}$ for quadratic ones (Hendrikx et al.,, 2020). Here $n$ is the size of the training set. Hence, the more data on nodes, the more statistically "similar" the losses are and the less communication is needed. The basic method in this case is Mirror Descent with the appropriate Bregman divergence (Shamir et al., 2014a, ; Hendrikx et al.,, 2020):

The computation of $\nabla r$ requires communication with clients. Our goal is to reduce it. The key idea is that by offloading nodes, we move the most computationally demanding work to the server. Each node performs one computation per iteration (to collect $\nabla r$ on the server), while server solves the local subproblem and accesses the local oracle $\nabla r_{1}$ much more frequently. While this approach provides improvements in communication complexity over simple distributed version of Gradient Descent, there is room for improvement in various aspects. In particular, the problem of constructing an optimal algorithm has been open for a long time (Shamir et al., 2014b, ; Arjevani and Shamir,, 2015; Yuan and Li,, 2020; Sun et al., 2022a, ). This has recently been solved by the introduction of Accelerated ExtraGradient (Kovalev et al.,, 2022).

Despite the fact that there is an optimal algorithm for communication, there are still open challenges: since each dataset is stored on its own node, it is expensive to compute the full gradient. Moreover, by transmitting it, we lose the ability to keep procedure private (Weeraddana and Fischione,, 2017), since in this case it is possible to recover local data. One of the best known cheap ways to protect privacy is to add noise to communicated information (Abadi et al.,, 2016). To summarize the above, in this paper, we are interested in stochastic modifications of Accelerated ExtraGradient (Kovalev et al.,, 2022), which would be particularly resistant to the imposition of noise on the transmitted packages.

Since our goal is to investigate stochastic versions of algorithms that exploit data similarity (in particular, to speed up the computation of local gradients and to protect privacy), we give a brief literature review of related areas.

$\bullet$ New method. We propose a new method based on the deterministic Accelerated ExtraGradient (AEG) (Kovalev et al.,, 2022). Unlike AEG, we add the ability to sample computational nodes. This saves a significant amount of runtime. As mentioned above, sampling was already implemented in Three Pillars Algorithm (Beznosikov et al.,, 2024), Catalyzed SVRP (Khaled and Jin,, 2022) and AccSVRS (Lin et al.,, 2024). However, this was done using SVRG-like approaches. We rely instead on SGD, which does not require decreasing the step size as the number of computational nodes increases.

$\bullet$ Privacy via noise. We consider stochasticity not only in the choice of nodes, but also in the additive noise in the transmission of information from client to server. This makes it much harder to steal private data in the case of an attack on the algorithm.

$\bullet$ Theoretical analysis. We show that the Hessian similarity implies gradient one with an additional term. Therefore, it is possible to estimate the variance associated with sampling nodes more optimistically. Thus, in some cases ASEG converges quite accurately before getting "stuck" near the solution.

$\bullet$ Analysis of subproblem. We consider different approaches to solving a subproblem that occurs on a server. Our analysis allows us to estimate the number of iterations required to ensure convergence of the method.

$\bullet$ Experiments. We validate the constructed theory with numerical experiments on several real datasets.

Algorithm 1 (ASEG) is the modification of AEG (Kovalev et al.,, 2022). In one iteration, server communicates with clients twice. In the first round, random machines are selected (Line 6). They perform a computation and send the gradient to the server (Line 15). Then the server solves the local subproblem at Line 16 and the solution is used to compute the stochastic (extra)gradient at the second round of communication (Line 17 and Line 26). This is followed by a gradient step with momentum at Line 27.

Let us briefly summarize the main differences with AEG. Line 16 of Algorithm 1 uses a gradient over the local nodes’ batches, rather than all of them, as is done in the baseline method. Moreover, the case where local gradients are noisy is supported. This is needed for privacy purposes. Similarly in Line 27. Note that we consider the most general case where different devices are involved in solving the subproblem and performing the gradient step.

ASEG requires $\Theta(2B)$ of communication rounds (Line 15 and Line 26), where $B$ is the number of nodes involved in the iteration. Unlike our method, AEG does $\Theta(2M-2)$ per iteration. AccSVRS does $\Theta(2M)$ on average. Thus, in experiments, either the difference between the methods is invisible, or AEG loses. ASEG deals with a local gradient $\nabla r_{m}(x,\xi)$, which can have the structure:

where, for example, $\xi\in N(0,\sigma^{2})$ or $\xi\in U(-c,c)$. This is the previously mentioned additive noise overlay technique, popular in the federated learning.

Before we delve into analysis, let us introduce some definitions and assumptions to build the convergence theory.

In this paper, complexity is understood in the sense of number of server-node communications. This means that we consider how many times a server exchange vectors with clients. In this case, a single contact is counted as one communication.

We work with functions of a class widely used in the literature.

If $\mu=0$, we call $f$ a convex function.

We assume a strongly convex optimization problem (1) with possibly non-convex components.

This assumption does not require convexity of local functions. Indeed, $\delta$-relatedness is a strong enough property to allow them to be non-convex. Within this paper we are interested in the case $\mu<\delta$, since $\mu\ll\delta\ll L$ for large datasets in practice (Sun et al., 2022b, ).

Thus, from the similarity of the Hessians follows the similarity of the gradients with some correction. Let us introduce a more general assumption on the gradient similarity:

Assumption 2 is a generalization of Proposition 1, where the correction term is multiplied by some constant $\zeta$. Formally, it can take any non-negative value, but it is important to analyze only two cases corresponding to different degrees of similarity: $\zeta=0$ and $\zeta=1$. Indeed, any non-negative values of $\zeta$ give asymptotically the same complexity.

A stochastic oracle is available for the function at each node. We impose the standard assumption:

We want to obtain a criterion that guarantees convergence of Algorithm 1. The obvious idea is to try to get one for the new algorithm by analogy with how it was done in the non-stochastic case (Kovalev et al.,, 2022):

See the proof in Appendix A. We first deal with the case under the gradient similarity condition ($\zeta=0$), since it removes two additional constraints on the parameters and makes the analysis simpler.
The following theorem is obtained by rewriting the result of Lemma 1 with a finer tuning of the parameters. We take $\mu<\delta$ into account and find that $1-\frac{\sqrt{\mu\theta}}{3}\geq 1-\frac{1}{3}$.

The theorem asserts convergence only to some neighborhood of the solution. To guarantee convergence with arbitrary accuracy, a suitable parameter $\theta$ must be chosen. We propose a corollary of Theorem 1:

See the proof in Appendix C. Note that the linear part of the obtained estimate reproduces the convergence of AEG. However, for example, with $B=1$, an iteration of ASEG costs $\Theta(2)$ communications instead of $\Theta(2M-2)$. Thus, our method requires significantly less communication to achieve the value determined by the sublinear term. Moreover, as mentioned earlier, in the Catalyzed SVRP and AccSVRS estimates, an additional $M^{-\nicefrac{{1}}{{4}}}$ multiplier arises due to the use of variance reduction. Comparing $M^{-\nicefrac{{1}}{{4}}}\cdot\sqrt{\nicefrac{{\delta}}{{\mu}}}$ to $BM^{-1}\cdot\sqrt{\nicefrac{{\delta}}{{\mu}}}$, we notice the superiority of ASEG in terms of chosen approach to measure communication complexity.

Note that if the noise of the gradients sent to the server is weak ($\sigma_{noise}=0$) and the functions are absolutely similar in the gradient sense ($\sigma_{sim}=0$), then ASEG has $\nicefrac{{M^{\nicefrac{{3}}{{4}}}}}{{B}}$ times less complexity in terms of number of communications.

In the case of a convex objective, it is not possible to implement the method under Assumption 2 with $\zeta\neq 0$, since the arising correction terms cannot be eliminated. Under Assumption 2 with $\zeta=0$, we propose a modification of ASEG for convex objective. In this case, we need

to reach an arbitrary $\varepsilon$-solution. See Appendix D.

The case $\zeta=1$ is more complex and requires more fine-tuned analysis.

See the proof in Appendix B.

See the proof in Appendix C. It can be seen from the proof that Assumption 2 with $\zeta=1$ gives worse estimates. Nevertheless, optimal communication complexity can be achieved if we choose large enough $B$, which equalizes two linear terms (at least $72\frac{\delta^{\nicefrac{{3}}{{2}}}}{\mu^{\nicefrac{{3}}{{2}}}}$). It comes from equality

Thus, even on tasks with a poor ratio of $\delta$ to $\mu$, if the number of available nodes is large enough, ASEG outperforms AEG, Catalyzed SVRP and AccSVRS.

In this section, we are focused on

defined in Line 16 of Algorithm 1. Note that the problem (7) is solved on the server and does not affect the communication complexity. Obviously:

Moreover:

Thus, subproblem is $L_{A}$-smooth with $L_{A}=\frac{1}{\theta}+L_{1}$ and $\frac{1}{\theta}$ strongly convex.

Despite the fact that the problem (7) does not affect the number of communications, we would still like to spend less server time solving it. The obvious solution is to use stochastic approaches.

To support theoretical analysis and insights, we solve both logistic regression problem:

and quadratic regression problem:

Here $x$ refers to parameters of the model, $a_{j}^{m}$ is the feature description and $b_{j}^{m}$ is the corresponding label. To avoid overfitting, we use ${l_{2}}$-regularization. The penalty parameter $\lambda$ is set to $\nicefrac{{L}}{{100}}$, where $L$ is a smoothness constant (2). We use SVRG and SARAH as the subtask solvers in Algorithm 1. In all tests, the original datasets are merged into batches $\left\{r_{b}\right\}_{i=1}^{M+N}$. A function $r_{1}(x)$ is created from part of the acquired ones:

Due to the Hessian similarity, we have the following expression for the smoothness constant

where

Then for the quadratic task we have

For the logistic regression, $\delta$ is estimated by enumerating 100 points in the neighborhood of the solution. In both cases, obtained value is multiplied by $\nicefrac{{3}}{{2}}$. Experiments are performed on three datasets: Mushrooms, W8A, A9A, (Chang and Lin,, 2011).

$\bullet$ We investigate how the size of the batches affects the convergence of ASEG on tasks with different parameters. Figure 1 shows the dependence of untuned ASEG convergence on the batch size. In the first experiment the method on $50$ nodes solves the problem with $\mu=0.105$ and $\delta=1.45$. According to Corollary 2 one should choose $B>72\cdot(\nicefrac{{\delta}}{{\mu}})^{\nicefrac{{3}}{{2}}}\approx 3695$ to obtain optimal complexity. That is, for any size of the batch, there is non-optimal rate. In the second experiment, the method on $200$ nodes solves the problem with $\mu=0.06$, $\delta=0.07$. In this case, ASEG has guaranteed optimal complexity for $B>57$ and converges to rather small values of the criterion.

$\bullet$ We compare ASEG to AccSVRS (Figure 2 and Figure 3). 50 nodes are available for Mushrooms and 200 nodes are available for A9A. If the ratio between $\mu$, $\delta$, and $B$ yields the optimal complexity according to Corollary 2, then due to similarity there is an optimistic estimate on the batching noise provides a significant gain over AccSVRS. Otherwise, it is worth using full batching and obtaining a less appreciable gain.

$\bullet$ In ASEG stability experiments (Figure 4, Figure 5, Figure 6, Figure 7) it is tested how white noise in each $\nabla r_{m}(x)$ affects the convergence. It can be seen that increasing the randomness does not worsen the convergence. This can be explained by the fact that the convergence is mainly determined by accuracy of chosen solver.

$\bullet$ ASEG experiments test how changing the epoch size changes the ASEG and subproblem (Line 16 of Algorithm 1) convergence with SVRG and SARAH solvers (Figure 8, Figure 9, Figure 10, Figure 11, Figure 12, Figure 13, Figure 14, Figure 15).

Based on the SGD-like approach, we have adapted Accelerated ExtraGradient to solve Federated Learning problems by adding two stochasticities: sampling of computational nodes and application of noise to local gradients. Numerical experiments show that for small values of $\nicefrac{{\delta}}{{\mu}}$ or a sufficiently large number of available nodes, Algorithm 1 (ASEG) is significantly superior to accelerated variance reduction methods (Figure 3), which is consistent with the theory (Corollary 1). If the problem parameters are not "good" enough, there is a $\nicefrac{{\delta^{2}}}{{(\mu^{2}M)}}\cdot\log\nicefrac{{1}}{{\varepsilon}}$ term in the communication complexity of batched ASEG (Corollary 2). Nevertheless, the superiority in the number of communications remains even in this case (Figure 2, Figure 3). Experiments also show the robustness of our method against noise (Figure 4, Figure 5, Figure 6, Figure 7). The server subproblem has also been studied numerically (Figure 10, Figure 11, Figure 14, Figure 15) and theoretically. The analysis shows the inefficiency of randomly choosing the number of solver iterations for the internal subtask (Proposition 2).