GP-FL: Model-Based Hessian Estimation for Second-Order Over-the-Air Federated Learning

Early efforts to improve the communication efficiency in the FL framework can be categorized roughly into two approaches: ($i$) the first approach tries to reduce the communication overhead per round. One way to achieve this involves utilizing quantization [7] and sparsification [8] techniques. These techniques aim to diminish the transmitted bits and eliminate redundant parameter updates. In general, these approaches come with the trade-off of potentially degrading the model performance, and they must take into account the compatibility for the aggregation operation in FL [9]. ($ii$) The second approach intends to minimize the total communication rounds. The most well-known example is the federated averaging (FedAvg) algorithm, in which clients execute multiple local iterations before sharing their models with the PS [1], which can significantly reduce the total number of rounds. Some variants of the gradient descent (GD) method have further demonstrated the ability to substantially reduce the overall communication rounds compared to the naive GD [10, 11].

These previous approaches ignore the properties of the communication network, and look at the communication links as rate-limited orthogonal noiseless channels. Nevertheless, later studies have shown that using the properties of the underlying network, further gain in terms of communication efficiency can be achieved [2]. A well-known case is wireless FL, in which clients can use the superposition property of multiple access channel (MAC) to perform the computation directly over the air and significantly reduce the communication costs [12]. In fact, in wireless networks the abstract view on communication links corresponds to the conventional digital transmission, where coded symbols are sent using techniques such as orthogonal frequency-division multiplexing (OFDM). For wireless FL, however, the clients can send their local parameters simultaneously via analog transmission, and the PS can apply the concept of over-the-air computation (AirComp) to aggregate the global parameter. This method allows aggregation to be performed directly over the air, significantly reducing the required bandwidth [12, 13]. This approach has gained traction in wireless FL as an effective strategy for mitigating communication costs [14, 15, 16, 17, 18, 19].

A standard framework for wireless FL consider first-order algorithms, where the clients transmit their local gradients, and the PS aggregates them gradient directly over the air using the AirComp technique. We refer to this framework as first-order AirComp. Similar to other first-order methods, first-order AirComp at best achieves linear convergence, which leads to a relatively large number of iteration rounds to reach the desired accuracy. We further note that the analog nature of AirComp-aided computation leads to extra aggregation error from the channel. This can be seen as zero-mean fluctuations, which adds to the noiseless aggregated gradient, i.e., the estimator that PS computes in perfect FL settings. As a result, the estimator of global gradient computed in first-order AirComp has higher variance as compared with FL with noiseless aggregation. This leads to further degradation in terms of convergence behavior.

One potential solution is to incorporate second-order methods, e.g., Newton-type methods, into the wireless FL setting, as explored in [20]. Unlike first-order methods, second-order approaches benefit from a quadratic convergence rate by computing the update direction using both first and second-order derivatives of the loss function. The update direction computed in these methods is often called Newton direction.

Integrating second-order information into FL requires further access to the second-order information at the PS. This information is obtained by clients sharing their local Hessian matrices with the PS, which imposes a significant communication burden, especially in wireless settings with strictly-restricted links. To mitigate this issue, various studies have explored ways to approximate the Hessian matrix from the first-order information and/or reduced second-order information. For instance, GIANT [21] estimates the Newton direction by utilizing the global gradient and local Hessians on each client, combined with a global backtracking line search across all clients. Although this approach reduces the communication overhead, it still requires partial sharing of second-order information resulting in higher communication load as compared to first-order methods.

The study in [22] suggests an alternative approach, in which each client the computation of Newton direction is offloaded to the clients: the PS and devices invoke the conjugate gradient descent algorithm to compute iteratively an estimate of the global Newton direction. Although the information shared per communication round in this approach is the same as that of the first-order methods, a single iteration of gradient descent (on the model parameters) requires at least two rounds of communication. To circumvent further this overhead, the recent study in [23] proposes a novel second-order method that eliminates the aggregation of local gradients, enabling a single communication round per iteration. Motivated by this work, the authors in [20] develop a second-order wireless FL method, in which the clients compute their Newton directions locally and share them with the PS. The PS then estimates the global Newton direction from the aggregation of these local directions. This approach enables devices to communicate with the PS only once per iteration, reducing the communication overhead to that of the first-order methods.

As a natural extension to available proposals, some recent studies have proposed integration of communication-efficient second-order FL algorithms into the analog computation framework [21, 20]. Experiments however show that, unlike their first-order counterparts, these second-order AirComp algorithms struggle to perform adequately. This observation can be intuitively explained as follows: the noisy aggregation of the second-order information in AirComp-aided approaches introduces significant biases to the estimate of the Newton direction, leading to suboptimal performance. The investigations reveal that such biases can severely deteriorates the convergence, and hence the overall efficiency, of these methods [21, 20]. This study aims to develop a novel second-order AirComp algorithm that addresses this challenge in a communication-efficient way.

This work proposes a novel second-order AirComp scheme for FL which estimates the Newton direction directly from the aggregated first-order information. It is hence extremely communication-efficient, in the sense that it does not require the exchange of local Hessians or a functions of them: in each round, the clients leverage AirComp to share only their local gradients. The PS then estimates the global Hessian matrix from a finite sequence of its noisy aggregations. For Hessian estimation, we deviate from the classical deterministic schemes and develop a nonparametric model-based estimator. The proposed estimator assumes a Gaussian prior for the Hessian matrix and determines its posterior distribution conditional to a window of most recent noisy aggregations. It then computes an unbiased estimator of the Newton direction by sampling from the posterior distribution. We refer to this method as Gaussian process-based Hessian modeling for FL (GP-FL). Our analysis and numerical experiments demonstrate that the scheme can efficiently suppress the undesired directional bias.

In summary, the contributions of this paper are three-fold:

The FedAvg scheme relies solely on first-order gradients for updates [1], offering significantly faster computation compared to Hessian-based methods, especially in settings with fast communication. Adaptive step-size methods like AdaGrad [28] and ADAM [29] have also been adapted to distributed settings, showing improved convergence [30]. The adaptive step-size can be represented as a diagonal matrix $\mathbb{D}$, providing an alternative preconditioning to the Newton direction $\mathbb{H}^{-1}\nabla f(\boldsymbol{\theta})$, where the direction is computed as $\mathbb{D}^{-1}\nabla f(\boldsymbol{\theta})$.

The second-order methods studied in the literature can be categorized into two groups: ($i$) those that utilize second-order information implicitly [31, 32, 33], and ($ii$) those that explicitly compute them [22, 21, 34]. DANE [31] computes a mirror descent update on the local loss functions, which is equivalent to the GIANT update for a quadratic function. The study in [32] proposes FedDANE as a variant of DANE, specifically tailored for FL, where it uses FedAvg as a baseline with 20 local epochs and observes no improvement with their proposed method. AIDE, proposed in [33], is an alternative accelerated inexact version of DANE. Another line of study in the first group, i.e., group $i$, considers employing distributed quasi-Newton methods [35]. CoCoA [36] and its trust-region extension [37] also perform local steps on a second-order local subproblem, but they specifically address the special case of generalized linear model objectives.

The second group of studies employ the Hessian by computing it indirectly through the use of the so-called Hessian-free optimization approach. In DiSCO [22], the clients compute the Hessian-vector products, and subsequently the PS executes the conjugate gradient method [38]. This process entails one communication round for each iteration of the conjugate gradient method. GIANT [21] and LocalNewton [34] both employ the conjugate gradient method at the clients: GIANT utilizes the global gradient, while LocalNewton uses the local gradients. The study in [39] and its FL extension [40] iteratively approximate the global Hessian, requiring a similar number of communication rounds as GIANT but achieving better convergence rates. However, they lack experimental comparisons with FedAvg using multiple local steps.

In general, the performance of mentioned methods tends to degrade when integrated into the AirComp framework, in which the PS receives a distorted and noisy aggregation. The proposed scheme in this work, GP-FL, addresses this issue by explicitly accounting for the channel imperfections.

The gradient and Hessian of a function $f(\cdot)$ are denoted by $\nabla f(\cdot)$ and $\nabla^{2}f(\cdot)$, respectively. The operators $(\cdot)^{\mathsf{T}}$ and $(\cdot)^{\mathsf{H}}$ denote the transpose and Hermitian transpose, respectively; and $\mathbb{E}(\cdot)$ represents the mathematical expectation. The notation $\mathbf{x}[j]$ refers to the $j$-th entry of vector $\mathbf{x}$. The notation $\mathcal{GP}(\mu,\kappa)$ denotes a Gaussian process with mean $\mu$ and covariance $\kappa$.

For an integer $K$, $[K]$ denotes $\{1,\cdots,K\}$. For a scalar or function $f$, $\{f_{k}\}_{k\in[K]}=\{f_{1},\dots,f_{K}\}$, and $[f_{k}]_{k\in[K]}=[f_{1},\dots,f_{K}]^{\mathsf{T}}$. Scalars are denoted by non-boldface letters (e.g. $a$), vectors by boldface lowercase letters (e.g. $\mathbb{a}$), and matrices by boldface uppercase letters (e.g. $\mathbb{A}$). The $j$-th entry of vector $\mathbb{a}$ is denoted by $\mathbb{a}[j]$. The matrix $\mathbb{I}$ represents the identity matrix. For two matrices $\mathbb{A}$ and $\mathbb{B}$ of the same size, $\mathbb{A}\preceq\mathbb{B}$ implies that $\mathbb{B}-\mathbb{A}$ is positive semi-definite.

The remainder of this paper is organized as follows: Section II provides an in-depth discussion of the preliminary knowledge on Newton and quasi-Newton methods. Section III formulates the problem. Section IV presents the proposed stochastic approach for estimating the quasi-Newton direction. Section V provides the convergence analysis for GP-FL. Numerical results and comparison with baselines are given in Section VI. Finally, Section VII concludes the paper.

The de-facto solution to problem (1) is to employ the distributed (stochastic) gradient descent method (DGD/DSGD). In the $t$-th round of DGD, the PS shares parameter $\boldsymbol{\theta}_{t}$ with all clients. Each client computes its local gradient at $\boldsymbol{\theta}_{t}$, i.e., $\nabla f_{k}(\boldsymbol{\theta}_{t})$, and returns it to the PS. The PS then aggregates these local gradients into a global gradient as

and performs one step of gradient descent, i.e., it computes $\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_{t}-\eta_{t}\nabla f(\boldsymbol{\theta}_{t})$, where $\eta_{t}>0$ is the global learning rate at the $t$-th round. First-order methods often converge slowly with linear (or sub-linear) convergence rates [41, 42], requiring many iterations to approach a local minimum. Since the number of communication rounds in FL corresponds to iterations, this increases communication overhead.

A method to improve FL communication efficiency is to use faster-converging optimizers, such as second-order methods. These methods estimate the loss landscape’s local curvature, enabling faster and more adaptive updates. While computationally intensive per iteration, they require fewer iterations to converge. Second-order methods are Newton-type techniques. Specifically, Newton’s method updates the model using the second-order Taylor series approximation of the empirical loss. For a small $\mathfrak{d}\in\mathbb{R}^{d}$, the Taylor series expansion of $f(\boldsymbol{\theta}+\mathfrak{d})$ around $\boldsymbol{\theta}$ up to the quadratic term is given by $f(\boldsymbol{\theta}+\mathfrak{d})\approx f(\boldsymbol{\theta})+\mathfrak{d}^{T}\nabla f(\boldsymbol{\theta})+\frac{1}{2}\mathfrak{d}^{T}\nabla^{2}f(\boldsymbol{\theta})\mathfrak{d}$, with $\nabla^{2}f(\boldsymbol{\theta})$ being the Hessian matrix of $f$ computed at $\boldsymbol{\theta}$. Assuming the Hessian is positive definite, the Newton direction that minimizes this quadratic approximation is given by $\mathfrak{d}^{\star}=-(\nabla^{2}f(\boldsymbol{\theta}))^{-1}\nabla f(\boldsymbol{\theta})$. Newton’s method hence updates the model parameters by stepping proportional to $\mathfrak{d}^{\star}$.

Deploying Newton’s method for the distributed training task in (1), the PS needs to update the global model as

This requires transmitting local Hessians to the PS, creating a trade-off: second-order methods reduce communication rounds via faster convergence but increase communication costs per round due to larger transmission. Standard analyses show that the naive use of Newton’s method in distributed settings often leads to inefficiency, as communication overhead dominates.

To address issues with Newton’s method in distributed settings, quasi-Newton methods provide an alternative by avoiding explicit computation of local Hessians. Instead of $\nabla^{2}f(\boldsymbol{\theta})$ in (3a), an approximation matrix $\mathbb{B}_{\boldsymbol{\theta}}$ is built using the sequence of local gradients from previous iterations. This matrix is updated iteratively to include new information. To understand the quasi-Newton method, assume $f(\cdot)$ is twice continuously differentiable. We can hence write

By adding and subtracting the term $\nabla^{2}f(\boldsymbol{\theta})\mathfrak{d}$ to the right-hand side (r.h.s.) of (4), we have

Note that since the gradient function $\nabla f(\cdot)$ is continuous, the magnitude of the integral on the r.h.s. of (II-C) is $\mathcal{O}(\|\mathfrak{d}\|)$. We now set $\boldsymbol{\theta}=\boldsymbol{\theta}_{t-1}$ and $\mathfrak{d}=\boldsymbol{\theta}_{t}-\boldsymbol{\theta}_{t-1}$. This leads to

with $\boldsymbol{\varepsilon}=\mathcal{O}(\|\boldsymbol{\theta}_{t}-\boldsymbol{\theta}_{t-1}\|)$. Note that when $\boldsymbol{\theta}_{t-1}$ and $\boldsymbol{\theta}_{t}$ reside in a region close to the minimizer, the r.h.s. of (6) is dominated by the first term. We can hence write

The approximation in (7) suggests estimating the Hessian matrix using a matrix $\mathbb{B}_{t}$ that satisfies the relation in (7) as an identity. Specifically, by defining $\mathbb{w}_{t}=\boldsymbol{\theta}_{t}-\boldsymbol{\theta}_{t-1}$ and $\mathbb{y}_{t}=\nabla f(\boldsymbol{\theta}_{t})-\nabla f(\boldsymbol{\theta}_{t-1})$, the quasi-Newton method computes Newton’s direction using $\mathbb{B}_{t}$ that satisfies the following equation, commonly referred to as the secant equation:

which is an estimator of the Hessian matrix. A classic solution to (8) given by Broyden–Fletcher–Goldfarb–Shanno (BFGS) method [43], which iteratively updates matrix $\mathbb{B}_{t+1}$ using information from ${\mathbb{B}_{t},\mathbb{w}_{t},\mathbb{y}_{t}}$ according to $\mathbb{B}_{t+1}=\mathbb{B}_{t}-\frac{\mathbb{B}_{t}\mathbb{w}_{t}\mathbb{w}_{t}^{\mathsf{T}}\mathbb{B}_{t}}{\mathbb{w}_{t}^{\mathsf{T}}\mathbb{B}_{t}\mathbb{w}_{t}}+\frac{\mathbb{y}_{t}\mathbb{y}_{t}^{\mathsf{T}}}{\mathbb{w}_{t}^{\mathsf{T}}\mathbb{y}_{t}}$, ensuring the positive-definiteness of $\mathbb{B}_{t}$ and consequently making the following quasi-Newton search direction a descent direction:

Unlike Newton’s method, the quasi-Newton approach relies only on computed gradients, reducing communication requirements. However, this comes with higher variance in estimating the optimal direction. Studies show that the quasi-Newton method strikes a good trade-off, achieving faster convergence than first-order methods with similar communication costs.

We now develop the second-order algorithm GP-FL for FL in wireless networks based on the quasi-Newton method. The algorithm uses two key notions to adapt the quasi-Newton search directions approach to wireless networks: ($i$) it uses AirComp to realize the aggregation directly over the air, and ($ii$) it models the quasi-Newton estimator of the Hessian matrix as a Gaussian process and employs the maximum-likelihood method to estimate it from noisy aggregations.

In the next sections, we present GP-FL, sketch its derivation and analyze its convergence. For the sake of simplicity, we use the following notation hereafter in the paper: we denote the local gradient of client $k$ in iteration $t$ as $\mathbb{g}_{t,k}\triangleq\nabla f_{k}(\boldsymbol{\theta}_{t})$ and global gradient as $\mathbb{g}_{t}\triangleq\nabla f(\boldsymbol{\theta}_{t})$, i.e.,

Using this notation, we can summarize the vanilla quasi-Newton method as follows: in iteration $t$,

1. 1.

   The PS aggregates local gradients as per (10).

2. 2.

   It computes the difference

   $\displaystyle\mathbb{y}_{t}=\mathbb{g}_{t}-\mathbb{g}_{t-1}.$

   (11)

3. 3.

   It finds a quasi-Newton matrix by solving (8) for $\mathbb{B}_{t}$.

4. 4.

   It updates the global model parameter as $\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_{t}-\eta_{t}\mathbb{B}_{t}^{-1}\mathbb{g}_{t}$, and sets $\mathbb{w}_{t+1}=\boldsymbol{\theta}_{t+1}-\boldsymbol{\theta}_{t}$.

We consider a Gaussian fading multiple access channel with slow frequency-flat fading, where the coherence time exceeds $d$ symbol intervals. Clients synchronously transmit $d$-dimensional local gradient entries over $d$ consecutive intervals within a single channel coherence interval. The PS is assumed to have an array of $N$ antennas, while each client has a single antenna.

We denote the channel coefficient vector between client $k$ and the PS in round $t$ as $\mathbb{h}_{t,k}\in\mathbb{C}^{N}$. Assuming perfect channel state information (CSI) at both ends, clients can adjust their transmitted signals based on the channel coefficients.

The AirComp-based model aggregation works as follows: client $k$ sends $\phi_{k}(\mathbb{g}_{t,k})$ for some pre-processing function $\phi_{k}(\cdot)$ simultaneously with all other clients over the channel. The PS then receives a superimposed version of these transmissions. Denoting the received signal by $\mathbb{r}_{t}$, the PS estimates the aggregated gradient as $\hat{\mathbb{g}}_{t}=\psi(\mathbb{r}_{t})$ for some post-processing function $\psi(\cdot)$. The pre- and post-processing functions serve as analog filters aiming to fulfill the transmission constraints, e.g, transmit power constraint, and minimize the estimation error.

Due to the diversity of $\mathbb{g}_{t,k}$ among devices, a universal pre- and post-processing function cannot guarantee the joint stationarity of the information-bearing symbols. i.e., local gradients. We hence opt for a data-and-CSI-aware design: prior to uplink transmission, client $k$ normalizes its local gradient as

As the normalized gradients are unit-norm, we can model them as jointly stationary processes. This means that $\mathbb{E}(\|\mathbb{s}_{t,k}[j]\|^{2})={1}/{d}$ for entry $j\in[d]$. After normalization, client $k$ sends

over its uplink channel, where $b_{t,k}\in\mathbb{R}$ is a power scaling factor satisfying

with $\mathsf{P_{0}}$ denoting maximum transmit power of devices.

The PS receives the superimposed version of transmitted signals: let $\mathbb{r}_{t,j}\in\mathbb{C}^{N}$ denotes the received signal of the PS in the $j$-th symbol interval of iteration $t$. Using the channel model, we can write

where $\hat{\mathbb{h}}_{t,k}\triangleq{\mathbb{h}_{t,k}}/{\|\mathbb{g}_{t,k}\|_{2}}$ is the effective channel, and $\mathbb{n}_{t,j}\in\mathbb{C}^{N}$ is complex additive white Gaussian noise (AWGN) with mean zero and variance $\sigma^{2}$, i.e., $\mathbb{n}_{t,j}\sim\mathcal{N}(\boldsymbol{0},\sigma^{2}\mathbb{I})$.

The PS uses the received signals $\mathbb{r}_{t,j}$ for $j\in[d]$ to determine an estimator of the global gradient $\mathbf{g}_{t}$. Let $\mathbb{d}_{t}[j]$ denote the estimator of $\mathbf{g}_{t}[j]$. Using linear post-processing, the PS determines $\mathbb{d}_{t}[j]$ from $\mathbb{r}_{t,j}$ as

for some linear receiver $\mathbb{c}_{t}\in\mathbb{C}^{N}$ and the power factor $\alpha_{t}\neq 0$.

Then, the PS employs a zero-forcing approach to estimate $\mathbb{g}_{t}$. It first computes a linear receiver $\mathbb{c}_{t}$ using the CSI (details on computing $\mathbb{c}_{t}$ will be discussed later) and determines the power scaling factor $\alpha^{\text{ZF}}_{t}$ as:

This way, the PS guarantees that all clients satisfy their transmit power constraint, and then it broadcasts $\alpha^{\text{ZF}}_{t}$ to the clients. Client $k$ upon receiving $\alpha^{\text{ZF}}_{t}$ determines its scaling as

which satisfies the transmit power constraint (14).

The above coordination of clients realizes the desired superposition over the air in the absence of noise during uplink transmission. Substituting (17) and (18) in (15), it is concluded that the PS aggregates $\mathbb{d}_{t}[j]=\sum_{k\in\mathcal{S}_{t}}|\mathcal{D}_{k}|\mathbb{g}_{t,k}[j]+\frac{1}{\sqrt{\alpha^{\text{ZF}}_{t}}}\mathbb{c}_{t}^{\mathsf{H}}\mathbb{n}_{t,j}$, which after scaling by ${1}/{|\mathcal{D}|}$ concludes the following

The estimator of the gradient in iteration $t$ is hence given by

where, for notational simplicity, we defined the $d$-dimensional column vector $\tilde{\mathbb{n}}_{t}$ as

The variance of $\tilde{\mathbb{n}}_{t}$ which specifies the variance of the gradient estimator is given by $\tilde{\sigma}^{2}(\mathbb{c}_{t})=\frac{\sigma^{2}\|\mathbb{c}_{t}\|^{2}}{|\mathcal{D}|^{2}{\alpha^{\text{ZF}}_{t}}}$. Substituting (17) in this term, we have

Thus, for a given set of selected devices $\mathcal{S}_{t}$, the optimal receiver that minimizes the estimation variance is given by a solution to the following optimization problem:

Using a similar analysis as in [45], the optimization problem in (23) can be reformulated as

which is a quadratically constrained quadratic programming problem that is computationally challenging. However, it can be transformed into a difference-of-convex-function (DC) program, as shown in Appendix A.

The device selection is a combinatorial problem, which lies int the set of nondeterministic polynomial (NP) hard problems. We therefore invoke the sub-optimal approach developed in [20] to approximate its solution via Gibbs sampling (GS) method. The key idea in the GS method is to iteratively sample a device set from the neighboring devices based on a suitable distribution. This iterative process allows the selected devices to gradually converge towards an optimal set. We omit the details of the algorithm due to lack of space and refer the interested readers to [20] for more details.

In model-based estimation approach, the stochastic nature of the solution to (8) is described through a stochastic process, which approximates the original process. In the sequel, we adapt the Gaussian model for this problem. This choice is particularly useful as Gaussian processes offer non-parametric probabilistic models to capture complexities of nonlinear functions [47].¹¹1Modeling Hessians as Gaussian processes was proposed in [46] for use in noisy settings but has not yet been applied to tasks like FL.

To understand the idea of probabilistic modeling of Hessian, let us look more precisely into the noisy quasi-Newton matrix computed from the over-the-air aggregations: at round $t$, the PS uses its latest two aggregations to compute

The quasi-Newton matrix $\tilde{\mathbb{B}}_{t}$, which is an estimator of Hessian, can then be computed from $\tilde{\mathbb{y}}_{t}$ by solving

With ideal links, $\tilde{\mathbb{y}}_{t}=\mathbb{y}_{t}$, and using a deterministic scheme like BFGS, the solution to (26) matches that of (8). However, with over-the-air computation, deterministic methods yield a Hessian estimate with potentially higher bias and variance compared to the standard case with perfect aggregation²²2This is demonstrated in the numerical results.. To address this, we deviate from the classical approach of using deterministic schemes to solve (26). Instead, we model $\tilde{\mathbb{B}}_{t}$ and $\tilde{\mathbb{y}}_{t}$ as jointly Gaussian random processes. We approximate their statistics from the observation sequence and use these to derive a more robust Hessian estimator, effectively mitigating the impact of aggregation noise.

We start by modeling the prior: let $\tilde{\mathbb{B}}_{t}$ be a Gaussian matrix with mean $\mathbf{M}$ and covariance $\mathbb{C}$. For $i,j\in[d]$, entry $(i,j)$ of $\tilde{\mathbb{B}}_{t}$ is an independent Gaussian process with mean $\mu_{i,j}=[\mathbf{M}]_{i,j}$ and variance $\psi_{i,j}$ given by $\mathbb{C}$. For simplicity, we focus on a single entry of $\tilde{\mathbb{B}}_{t}$, denoting it as $\tilde{b}_{t}\sim\mathcal{N}(\mu,\psi)$, dropping the index $(i,j)$. This does not affect the generality of the analysis, as entries of $\tilde{\mathbb{B}}_{t}$ are statistically independent and follow the same estimation procedure.

Our ultimate goal is to model $\tilde{b}_{t}$ from the observations $\tilde{\mathbb{y}}_{t}$ collected through time $t$. To this end, we collect the last $r$ gradient difference in a set $\mathcal{Y}_{t}$, i.e.,

and model it jointly with entry $\tilde{b}_{t}$ as a Gaussian processes $\mathcal{GP}$. More precisely, let us define $\mathbb{z}=[\mathbb{o}_{t}^{\sf T},\tilde{b}_{t}]^{\sf T}$, where $\mathbb{o}_{t}$ is defined as the concatenation of entries in $\mathcal{Y}_{t}$, i.e.,

We model $\mathbb{z}_{t}$ as a Gaussian vector whose mean is $\boldsymbol{\xi}_{t}\in\mathbb{R}^{rd+1}$ and whose covariance matrix $\boldsymbol{\Sigma}_{t}\in\mathbb{R}^{(rd+1)\times(rd+1)}$. Considering this statistical model, our goal is to use the sample data collected in the last $r$ communication rounds, i.e., $\mathcal{Y}_{t}$, to find an estimate of the mean and covariance. We then use the assumed statistical model to find an alternative estimator for $\tilde{b}_{t}$.

To estimate the covariance matrix, we follow the conventional approach in the literature [48]: we estimate the covariance of $\mathbb{z}_{t}$ from its samples using the kernel function $\kappa$. The choice of an appropriate kernel is based on assumptions such as smoothness and likely patterns that are expected in the data. Given the nature of data in this problem, i.e., the fact that $\tilde{\mathbb{y}}_{t}$ is the gradient difference, one popular choice of the kernel is the radial basis function [48].

It is worth mentioning that using the radial basis kernel, the estimator of the covariance matrix computed from $\mathbb{z^{0}}$ sampled from $\mathbb{z}$, i.e., $\kappa(\mathbb{z^{0}},\mathbb{z^{0}})$, is a positive semi-definite matrix.

To use the above covariance estimator, we require a sample from $\mathbb{z}_{t}$. To this end, we use a deterministic solver to find a solution to (26) for the last $r$ communication rounds. This means that we find $\tilde{\mathbf{B}}_{i}^{0}$ by solving (26) with $\tilde{\mathbb{y}}_{i}^{0}$ for $i\in\{t-r+1,\ldots,t\}$, where $\tilde{\mathbb{y}}_{i}^{0}$ denotes to the sample gradient difference computed from the observed received signals.³³3Superscript $0$ is to distinguish random samples from stochastic processes. Let $\tilde{b}_{i}^{0}$ be an entry of $\tilde{\mathbf{B}}_{i}^{0}$. We sample $\mathbb{z}$ as $\mathbb{z}^{0}=[\mathbb{o}_{t}^{0\sf T},\tilde{b}_{t}^{0}]^{\sf T}$ with

and estimate the covariance of $\mathbb{z}_{t}$ with the sample covariance matrix $\kappa(\mathbb{z}_{t}^{0},\mathbb{z}_{t}^{0})$ computed by the radial basis kernel.

We next compute an estimator of the mean $\boldsymbol{\xi}_{t}$ using a moving average: we estimate the mean of $\tilde{b}_{t}$ by arithmetic average of $\tilde{b}_{t-r+1}^{0},\ldots,\tilde{b}_{t}^{0}$ denoted by $\hat{\mathbb{E}}\{\tilde{b}_{t}^{0}\}$, and the mean of $\mathbb{o}_{i}$ by the moving average of $\mathbb{o}_{i}^{0}$, i.e., we set the mean of $\tilde{\mathbb{y}}_{i}$ to be the average of $\{\tilde{\mathbb{y}}_{t-r}^{0},\dots,\tilde{\mathbb{y}}_{i}^{0}\}$ for $i\in\{t-r+1,\ldots,t\}$. Let us denote this moving average by $\mu[\mathbb{o}_{t}^{0}]$. We can then approximate the random process $\mathbb{z}_{t}$ as

for scalar $\beta_{t}=\kappa(\tilde{{b}}_{t}^{0},\tilde{{b}}_{t}^{0})$, vector $\boldsymbol{\phi}_{t}=\kappa(\mathbb{o}_{t}^{0},\tilde{b}_{t}^{0})$, and symmetric and positive semi-definite matrix $\mathbb{K}_{t}=\kappa(\mathbb{o}_{t}^{0},\mathbb{o}_{t}^{0})$.

The statistical model considered for the quasi-Newton matrix and the sample observations enables us to find an alternative estimator for the Hessian matrix. To this end, we derive the posterior distribution of $\tilde{\mathbb{B}}_{t}$ in the following lemma.

Using this posterior distribution, we can compute various estimators for $\tilde{b}_{t}$. Noting that this estimator is used to estimate Newton’s direction, we compute a simple unbiased estimator by sampling from the posterior distribution:⁴⁴4From the Bayesian viewpoint, one might consider setting the estimator to the posterior mean, as it describes the minimum mean squared error (MMSE) estimate. Although MMSE returns minimal variance, it is in general biased and hence can lead to directional misalignment. for every entry of $\tilde{\mathbb{B}}_{t}$, we compute the posterior mean $\zeta(\mathbb{o}_{t}^{0})$ and variance $\psi(\mathbb{o}_{t}^{0})$ using the sample observation $\mathbb{o}_{t}^{0}$. Let $\hat{b}_{i,j,t}$ be the posterior sample for entry $(i,j)$. We build the Hessian estimator $\hat{\mathbb{B}}_{t}$ from entries $\hat{b}_{i,j,t}$ and compute

The global model is then updated as $\boldsymbol{\theta}_{t+1}=\boldsymbol{\theta}_{t}+\eta_{t}\tilde{\mathfrak{d}}_{t}$.

Although validated numerically, it is insightful to explain why (34) provides a better estimate of the true Newton direction compared to deterministic approaches like BFGS. Classical deterministic methods rely only on the latest two gradient samples, which can result in highly inaccurate curvature estimates under noisy aggregation. In contrast, the proposed scheme uses information from the last $r+1$ gradient samples to estimate the loss curvature. This allows it to effectively mitigate the impact of aggregation noise and compute a more robust curvature estimate.

The remaining of this section provides convergence analysis for GP-FL under a set of regularity assumptions, which are commonly considered in the literature [2]. We start the analysis by stating the first assumption.

For global convergence of GP-FL, as a stochastic estimator of the Newton method, under Assumption 1 the magnitude of model update should be controlled. The classical approach for update control is to resort backtracking line search methods. In this work, we deviate from this classical scheme and control the update magnitude by scheduling the global learning rate. To this end, we invoke the results of [49], which proposes an adaptive learning rate to facilitate the global convergence of Newton-type methods. Casting GP-FL on the proposed scheme in [49], it is shown that for convergence guarantee of GP-FL to the minimizer of an $L$-smooth and $\lambda$-strongly convex loss, it is sufficient to schedule the learning rate as

We hence consider this scheduling throughout the analysis.

To present the convergence guarantee for GP-FL, we first need to define the concept of $\delta$-approximate matrix.

Using this definition, we can now present our first convergence result: Theorem 1 gives an upper-bound on the gap between the converging and optimal models.

Theorem (1) implies that the GP-FL algorithm exhibits a linear-quadratic convergence rate. In fact, the quadratic term in (65) is exactly the same as the one reported in [49] for Newton-type methods. Nevertheless, our result has an extra linear term, as a result of model-based estimation of the Hessian matrix and the AirComp aggregation error.

We next present Corollaries 1 and 2, which characterize the scenarios under which the GP-FL algorithm converges quadratically and linearly.

Corollary 1 presents an intuitive result: when GP-FL is initiated in a close enough vicinity of the optimal model, it can quadratically converge to the optimal model. We next consider the case with linear convergence.

It is worth mentioning that in classical distributed methods, the number of rounds required to achieve a desired precision $\varepsilon$, follows a linear convergence rate. Specifically, for vanilla federated averaging $T_{\varepsilon}=\mathcal{O}\left(\frac{L}{\lambda}\log\frac{1}{\varepsilon}\right)$. This underscores the superiority of GP-FL in terms of convergence rate.

The convergence analysis can be extended to characterize the optimality gap, i.e., the difference between $f(\boldsymbol{\theta}_{t})$ and the minimal loss. Theorem 2 gives an upper bound on the optimality gap at round $t$ whose proof is differed to Appendix E.

First, we show results for a binary classification task using logistic regression and three datasets from the LIBSVM library [24]: a9a, w8a, and phishing. The data samples are uniformly distributed across $K=20$ clients, all of whom participate in 50 communication rounds. We use the logistic regression model, a type of generalized linear model, as described in [50]. The results for a9a, w8a, and phishing datasets are depicted in Figures 1a, 1b and 1c. As observed, GP-FL achieves higher classification accuracy than benchmark methods.

We next use a fully-connected neural network with two hidden layers and perform 300 communication rounds, with all $K=10$ clients participating in each round. The results are shown in Figure 1d. As observed, GP-FL not only achieves higher classification accuracy than the benchmark methods but also reaches 80% accuracy within just five rounds.

We evaluate CIFAR-10 classification in two setups:

• •

  Setup I: Following [51, 52], we organize the dataset by class and divide it into 200 shards. Each client randomly selects two shards without replacement. We use a feedforward neural network with two hidden layers for the FL task, with $K=100$ clients, 4000 communication rounds, and a 10% client sampling rate per round.

• •

  Setup II: We distribute the dataset among them using Dirichlet allocation with $\beta=0.5$. Using ResNet-18 with Group Normalization, we conduct 200 communication rounds, wherein all clients participate. We set $K=10$.

The test accuracy curves for our method and the benchmark methods are depicted in Figures 1e and 1f. As illustrated, compared to the benchmark methods, GP-FL (i) achieves a higher test accuracy, and (ii) demonstrates a faster rate of convergence. For example, in Figure 1e, GP-FL reaches 45% test accuracy in approximately 700 rounds, whereas the best-performing second-order benchmark methods require around 1600 rounds to reach the same accuracy.

CIFAR-100 shares the same sample size as CIFAR-10, yet it encompasses a broader diversity with 100 distinct classes. We next train the ResNet-18 model for this classification. Through training, we use Group Normalization and let all clients participate in each communication round. For this experiment, we consider the following two setups:

• •

  Setup I: We set $K=10$ and $\beta=0.5$ for Dirichlet allocation, and use 400 communication rounds.

• •

  Setup II: We set $K=50$ and $\beta=0.05$ for Dirichlet allocation, and use 200 communication rounds.