FeDXL: Provable Federated Learning for Deep X-Risk Optimization

We consider the following FL objective for DXO:

To highlight the challenge and motivate FeDXL, we decompose the gradient of the objective function into:

Let $\nabla F_{i}(\mathbf{w}):=\!\Delta_{i,1}\!+\!\Delta_{i,2}$. Then $\nabla F(\mathbf{w})=\frac{1}{N}\sum\limits_{i=1}^{N}\nabla F_{i}(\mathbf{w})$.

With the above decomposition, we can see that the main task at the local client $i$ is to estimate the gradient terms $\Delta_{i1}$ and $\Delta_{i2}$. Due to the symmetry between $\Delta_{i1}$ and $\Delta_{i2}$, below, we only use $\Delta_{i1}$ as an illustration for explaining the proposed algorithm. The difficulty in computing $\Delta_{i1}$ lies at it relies on data in other machines due to the presence of $\mathbb{E}_{\mathbf{z}^{\prime}\in\mathcal{S}_{2}^{j}}$ for all $j$. To overcome this difficulty, we decouple the data-dependent factors in $\Delta_{i1}$ into two types marked by green and blue shown below:

It is notable that the three green terms can be estimated or computed based the local data. In particular, local1 can be estimated by sampling data from $\mathcal{S}_{1}^{i}$ and local2 and local3 can be computed based on the sampled data $\mathbf{z}$ and the local model parameter. The difficulty springs from estimating and computing the two blue terms that depend on data on all machines. We would like to avoid communicating $h(\mathbf{w};\mathbf{z}^{\prime})$ at every iteration for estimating the blue terms as each communication would incur additional communication overhead. To tackle this, we propose to leverage the historical information computed in the previous round ²²2A round is defined as a sequence of local updates between two consecutive communications.. To put this into context of optimization, we consider the update at the $k$-th iteration during the $r$-th round, where $k=0,\ldots,K-1$. Let $\mathbf{w}^{r}_{i,k}$ denote the local model in $i$-th client at the $k$-th iteration within $r$-th round. Let $\mathbf{z}^{r}_{i,k,1}\in\mathcal{S}_{1}^{i},\mathbf{z}^{r}_{i,k,2}\in\mathcal{S}_{2}^{i}$ denote the data sampled at the $k$-th iteration from $\mathcal{S}_{1}^{i}$ and $\mathcal{S}_{2}^{i}$, respectively. Each local machine will compute $h(\mathbf{w}^{r}_{i,k},\mathbf{z}^{r}_{i,k,1})$ and $h({\mathbf{w}^{r}_{i,k}},\mathbf{z}^{r}_{i,k,2})$, which will be used for computing the active parts. Across all iterations $k=0,\ldots,K-1$, we will accumulate the computed prediction outputs over sampled data and stored in two sets $\mathcal{H}^{r}_{i,1}=\{h(\mathbf{w}^{r}_{i,k},\mathbf{z}^{r}_{i,k,1}),k=0,\ldots,K-1\}$ and $\mathcal{H}^{r}_{i,2}=\{h(\mathbf{w}^{r}_{i,k},\mathbf{z}^{r}_{i,k,2}),k=0,\ldots,K-1\}$. At the end of round $r$, we will communicate $\mathbf{w}^{r}_{i,K}$ and $\mathcal{H}^{r}_{i,1}$ and $\mathcal{H}^{r}_{i,2}$ to the central server, which will average the local models to get a global model $\mathbf{w}_{r}$ and also merge $\mathcal{H}^{r}_{1}=\mathcal{H}^{r}_{1,1}\cup\mathcal{H}^{r}_{2,1}\ldots\cup\mathcal{H}^{r}_{N,1}$ and $\mathcal{H}^{r}_{2}=\mathcal{H}^{r}_{1,2}\cup\mathcal{H}^{r}_{2,2}\ldots\cup\mathcal{H}^{r}_{N,2}$. These merged information will be broadcast to each individual client. Then, at the $k$-th iteration in the $r$-th round, we estimate the blue term by sampling $h^{r-1}_{2,\xi}\in\mathcal{H}^{r-1}_{2}$ without replacement and compute an estimator of $\Delta_{i1}$ by

where $\xi=(j,t,\mathbf{z}^{r-1}_{j,t,2})$ represents a random variable that captures the randomness in the sampled client $j\in\{1,\ldots,N\}$, iteration index $k\in\{0,\ldots,K-1\}$ and data sample $\mathbf{z}^{r-1}_{j,t,2}\in\mathcal{S}_{2}^{j}$, which is used for estimating the global1 in (4). We refer to the green factors in $G_{i,k,1}$ as the active parts and the blue factor in $G_{i,k,1}$ as the passive part. Similarly, we can estimate $\Delta_{i2}$ by $G_{i,k,2}$

where $h^{r-1}_{1,\zeta}\in\mathcal{H}^{r-1}_{1}$ is a randomly sampled prediction output in the previous round with $\zeta=(j^{\prime},t^{\prime},\mathbf{z}^{r-1}_{j^{\prime},t^{\prime},1})$ representing a random variable including a client sample $j^{\prime}$ and iteration sample $t^{\prime}$ and the data sample $\mathbf{z}^{r-1}_{j^{\prime},t^{\prime},1}$. Then we will update the local model parameter $\mathbf{w}^{r}_{i,k}$ by using a gradient estimator $G^{r}_{i,k,1}+G^{r}_{i,k,2}$.

We present the detailed steps of the proposed algorithm FeDXL1 in Algorithm 1. Several remarks are following: (i) at every round, the algorithm needs to communicate both the model parameters $\mathbf{w}^{r}_{i,K}$ and the historical prediction outputs $\mathcal{H}^{r-1}_{i,1}$ and $\mathcal{H}^{r-1}_{i,2}$, where $\mathcal{H}^{r-1}_{i,*}$ is constructed by collecting all or sub-sampled computed predictions in the $(r-1)$-th round. The bottom line for constructing $\mathcal{H}^{r-1}_{i,*}$ is to ensure that $\mathcal{H}^{r-1}_{*}$ contains at least $K$ independently sampled predictions that are from the previous round on all machines such that the corresponding data samples involved in $\mathcal{H}^{r-1}_{*}$ can be used to approximate $\frac{1}{N}\sum_{i=1}^{N}\mathbb{E}_{\mathbf{z}\in\mathcal{S}^{i}_{*}}$ $K$ times. Hence, to keep the communication costs minimal, each client at least needs to sample $O(\lceil K/N\rceil)$ sampled predictions from all iterations $k=0,1,\ldots,K-1$ and send them to the server for constructing $\mathcal{H}^{r-1}_{*}$, which is then broadcast to all clients for computing the passive parts in the round $r$. As a result, the minimal communication costs per-round per-client is $O(d+Kd_{o}/N)$. Nevertheless, for simplicity in Algorithm 1 we simply put all historical predictions into $\mathcal{H}^{r-1}_{i,*}$.

Similar to all other FL algorithms, FeDXL1 does not require communicating the raw input data, hence protects the privacy of the data. However, compared with most FL algorithms for ERM, FeDXL1 for DXO has an additional communication overhead at least $O(d_{o}K/N)$ which depends on the dimensionality of prediction output $d_{o}$. For learning a high-dimensional model (e.g. deep neural network with $d\gg 1$) with score-based pairwise losses ($d_{o}=1$), the additional communication cost $O(K/N)$ could be marginal. For updating the buffer $\mathcal{B}_{i,1}$ and $\mathcal{B}_{i,2}$, we can simply flush the history and add the newly received $\mathcal{R}^{r-1}_{i,1}$ and $\mathcal{R}^{r-1}_{i,2}$ with random shuffling to $\mathcal{B}_{i,1}$ and $\mathcal{B}_{i,2}$, respectively.

For analysis, we make the following assumptions regarding the DXO with linear $f$ problem, i.e., problem (3).

The first three assumptions are standard in the optimization of DXO problems (Wang & Yang, 2022). The last assumption embodies the data heterogeneity that is also common in federated learning (Yu et al., 2019a; Karimireddy et al., 2020b). Next, we present the theoretical results on the convergence of FeDXL1.

Remark. To get $\mathbb{E}[\frac{1}{R}\sum_{r=1}^{R}\|\nabla F(\bar{\mathbf{w}}^{r-1})\|^{2}]\leq\epsilon^{2}$, we just need to set $R=O(\frac{1}{\epsilon^{3}})$, $\eta=N\epsilon^{2}$ and $K=\frac{1}{N\epsilon}$. The number of communications is much less than the total number of iterations i.e., $O(\frac{1}{N\epsilon^{4}})$ as long as $N\leq O(\frac{1}{\epsilon})$. And the sample complexity on each machine is $\frac{1}{N\epsilon^{4}}$, which is linearly reduced by the number of machines $N$.

Novelty of Analysis. As the passive parts are computed in different machines in a previous round, the gradient estimators $G^{r}_{i,k,1}$ and $G^{r}_{i,k,2}$ will involve the dependency between the local model parameter $\mathbf{w}^{r}_{i,k}$ and the historical data contained in $\xi,\zeta$ used for computing $G^{r}_{i,k,1}$ and $G^{r}_{i,k,2}$, which makes the analysis more involved. We need to make sure that using the gradient estimator based on them can still result in “good” results. To this end, we borrow an analysis technique in (Yang et al., 2021b) to decouple the dependence between the current model parameter and the data used for computing the current gradient estimator, in which they used data in previous iteration to couple the data in the current iteration in order to compute a gradient of the pairwise loss $\ell(h(\mathbf{w}_{t};\mathbf{z}_{t}),h(\mathbf{w}_{t};\mathbf{z}_{t-1}))$. Nevertheless, in federated DXO controlling the error brought by the passive parts is more challenging since the delay is much longer and they were computed on different machines. In our analysis, we replace $\mathbf{w}^{r}_{i,k}$ with $\bar{\mathbf{w}}^{r-1}$ to decouple the dependence between the model parameter $\bar{\mathbf{w}}^{r-1}$ and the historical data $\xi,\zeta$, then we need to control the latency error $\|\bar{\mathbf{w}}^{r-1}-\bar{\mathbf{w}}^{r}\|^{2}$ and the gap error between different machines $\sum_{i}\sum_{k}\mathbb{E}\|\bar{\mathbf{w}}^{r}-\mathbf{w}^{r}_{i,k}\|^{2}$ such that the complexities are not compromised.

With nonlinear $f$, we consider the following FL problem of DXO minimization,

We compute the gradient and decompose it into:

where

Let $\nabla F_{i}(\mathbf{w})=\Delta_{i,1}+\Delta_{i,2}$. Then we have $\nabla F(\mathbf{w})=\frac{1}{N}\sum\limits_{i=1}^{N}\nabla F_{i}(\mathbf{w})$.

Compared to that in (4) for DXO with linear $f$, the $\Delta_{i1}$ term above involves another factor $\nabla f(g(\mathbf{w},\mathbf{z},\mathcal{S}_{2}))$, which cannot be computed locally as it depends on $\mathcal{S}_{2}$ distributed over all machines. Similarly, the $\Delta_{i2}$ term above involves another non-locally computable factor $\nabla f(g(\mathbf{w},\mathbf{z},\mathcal{S}_{2}))$. To address the challenge of estimating $g(\mathbf{w},\mathbf{z},\mathcal{S}_{2})$, we leverage the similar technique in the centralized setting (Wang & Yang, 2022) by tracking it using a moving average estimator based on random samples. In a centralized setting, one can maintain and update $\mathbf{u}(\mathbf{z})$ for estimating $g(\mathbf{w},\mathbf{z},\mathcal{S}_{2})$ by

where $\mathbf{z}^{\prime}$ is a random sample from $\mathcal{S}_{2}$. However, this is not possible in an FL setting as $\mathcal{S}_{2}$ is distributed over many machines. To tackle this, we leverage the delay communication technique used in the last subsection. At the $k$-th iteration in the $r$-th round, we update $\mathbf{u}(\mathbf{z}^{r}_{i,k,1})$ for a sampled $\mathbf{z}^{r}_{i,k,1}$ by

where $h^{r-1}_{\xi,2}$ is a random sample from $\mathcal{H}^{r-1}_{2}$ where $\xi=(j^{\prime},t^{\prime},\hat{\mathbf{z}}^{r-1}_{j^{\prime},t^{\prime},2})$ captures the randomness in client, iteration index and data sample in the last round. Then, we can use $\nabla f(\mathbf{u}^{r}_{i,k}(\mathbf{z}^{r}_{i,k,1}))$ in place of $\nabla f(g(\mathbf{w}^{r}_{i,k},\mathbf{z}^{r}_{i,k,1},\mathcal{S}_{2}))$ for estimating $\Delta_{i1}$. However, it is more nuanced for estimating $\nabla f(g(\mathbf{w},\mathbf{z},\mathcal{S}_{2}))$ in $\Delta_{2i}$ since $\mathbf{z}\in\mathcal{S}^{2}_{j}$ is not local random data. To address this, we propose to communicate $\mathcal{U}^{r-1}=\{\mathbf{u}^{r-1}_{i,k}(\hat{\mathbf{z}}^{r-1}_{i,k,1}),i\in[N],k\in[K]-1\}$. Then at the $k$-iteration in the $r$-th round of the $i$-th client, we can estimate $\nabla f(g(\mathbf{w},\mathbf{z},\mathcal{S}_{2}))$ with a random sample from $\mathcal{U}^{r-1}$ denoted by $u^{r-1}_{\zeta}$, where $\zeta=(j^{\prime},t^{\prime},\hat{\mathbf{z}}^{r-1}_{j^{\prime},t^{\prime},1})$, i.e., by using $\nabla f(\mathbf{u}_{\zeta}^{r-1})$. Then we estimate $\Delta_{1i}$ and $\Delta_{2i}$ by

where $j,\xi,j^{\prime},\zeta$ are random variables. Another difference from DXO with linear $f$ is that even in the centralized setting directly using $G^{r}_{i,k,1}+G^{r}_{i,k,2}$ will lead to a worse complexity due to that non-linear $f$ make the stochastic gradient estimator biased (Wang et al., 2017). Hence, in order to improve the convergence, we follow existing state-of-the-art algorithms for stochastic compositional optimization (Ghadimi et al., 2020; Wang & Yang, 2022) to compute a moving average estimator for the gradient at local machines, i.e., Step 17 in Algorithm 2. With these changes, we present the detailed steps of FeDXL2 for solving DXO with non-linear $f$ in Algorithm 2. The buffers $\mathcal{B}_{i,*}$ and $\mathcal{C}_{i}$ are updated similar to that for FeDXL1. Different from FeDXL1, there is an additional communication cost for communicating $\mathcal{U}_{i}^{r-1}$ and an additional buffer $\mathcal{C}_{i}$ at each local machine to store the received $\mathcal{P}^{r-1}_{i}$ from aggregated $\mathcal{U}^{r-1}$. Nevertheless, these additional costs are marginal compared with communicating $\mathcal{H}^{r-1}_{*}$ and maintaining the buffer $\mathcal{B}_{i,*}$.

We make the following assumptions regarding problem (8).

We present the convergence result of FeDXL2 below.

Remark. To get $\mathbb{E}[\frac{1}{R}\sum_{r=1}^{R}\|\nabla F(\bar{\mathbf{w}}^{r})\|^{2}]\leq\epsilon^{2}$, we just set $R=O(\frac{M^{1/2}}{\epsilon^{3}})$, $\eta=O(\frac{\epsilon^{2}}{M})$, $\gamma=O(\epsilon^{2})$, $\beta=\frac{\epsilon^{2}}{\sqrt{M}}$ and $K=\frac{M^{1/2}}{\epsilon}$. The number of communications $R=O(\frac{M^{1/2}}{\epsilon^{3}})$ is less than the total number of iterations i.e., $O(\frac{M}{\epsilon^{4}})$ by a factor of $O(M^{1/2}/\epsilon)$. And the sample complexity on each machine is $\frac{M}{\epsilon^{4}}$, which is less than that in (Wang & Yang, 2022) which has a sample complexity of $O(\sum\nolimits_{i=1}^{N}|\mathcal{S}^{1}_{i}|/\epsilon^{4})$. When the data are evenly distributed on different machines, we have achieved a linear speedup property. And in an extreme case where all data are on one machine, the sample complexity of FeDXL2 matches that in (Wang & Yang, 2022), which is expected. Compared with FeDXL1, the analysis of FeDXL2 has to deal with extra difficulties. First, with non-linear $f$, the coupling between the inner function and outer function adds to the complexity of interdependence between different rounds and machines. Second, we have to deal with the error for the passive part related to $\mathbf{u}$.