Coded Matrix Computations for D2D-enabled Linearized Federated Learning

TL;DR

This paper introduces a novel coded matrix computation method for federated learning that reduces communication delays, enhances privacy, and improves local computation speed, especially with sparse data matrices.

Contribution

It presents a straggler-optimal coded matrix computation approach tailored for D2D-enabled federated learning, addressing privacy and efficiency issues.

Findings

Reduces communication delay in federated learning.

Enhances privacy by minimizing D2D data transmissions.

Improves local computation speed with sparse data matrices.

Abstract

Federated learning (FL) is a popular technique for training a global model on data distributed across client devices. Like other distributed training techniques, FL is susceptible to straggler (slower or failed) clients. Recent work has proposed to address this through device-to-device (D2D) offloading, which introduces privacy concerns. In this paper, we propose a novel straggler-optimal approach for coded matrix computations which can significantly reduce the communication delay and privacy issues introduced from D2D data transmissions in FL. Moreover, our proposed approach leads to a considerable improvement of the local computation speed when the generated data matrix is sparse. Numerical evaluations confirm the superiority of our proposed method over baseline approaches.

Tables

Table 1. Table 1 : Comparison among different approaches in terms of communication delay for a system with n = 20 𝑛 20 n=20 , k A = 18 subscript 𝑘 𝐴 18 k_{A}=18 and s = 2 𝑠 2 s=2 .

Poly	Ortho-	RKRP	Conv.	Prop.
Code [7]	Poly[9]	Code[13]	Code[17]	Sch.
$14.13s$	$14.02s$	$2.49s$	$2.56s$	$2.21𝐬$

Table 2. Table 2 : Per client product computation time where n = 30 𝑛 30 n=30 , k A = 28 , s = 2 formulae-sequence subscript 𝑘 𝐴 28 𝑠 2 k_{A}=28,s=2 and ζ = 95 % 𝜁 percent 95 \zeta=95\% , 98 % percent 98 98\% or 99 % percent 99 99\% entries of 𝐀 𝐀 \mathbf{A} are zero.

Methods	Product Comp. Time (in ms)
	$\zeta =99\%$	$\zeta =98\%$	$\zeta =95\%$
Poly Code [7]	$54.7$	$55.2$	$53.7$
Ortho-Poly [9]	$54.3$	$54.8$	$55.2$
RKRP Code [13]	$55.1$	$53.4$	$53.7$
Conv. Code [17]	$56.2$	$55.8$	$56.8$
Prop. Scheme	$14.9$	$21.1$	$29.6$

Equations

$$\mathbf{A}=\left[\begin{array}{cccc}{\mathbf{A}}_0& {\mathbf{A}}_1& \dots & {\mathbf{A}}_{k_A-1}\end{array}\right].$$

$$|\mathcal{N}(\bar{\mathcal{C}})|$$

$$=\text{min}\left({\omega }_A+p-1,k_A\right)=\text{min}\left(s+p,k_A\right)\ge m.$$

$$|\mathcal{N}(\bar{\mathcal{C}})|$$

$$=\text{min}$$

Full text

**Index Terms— ** Distributed Computing, Federated Learning, Stragglers, Heterogeneous Edge Computing, Privacy.

1 Introduction

Contemporary computing platforms are hard-pressed to support the growing demands for AI/ML model training at the network edge. While advances in hardware serve as part of the solution, the increasing complexity of data tasks and volumes of data will continue impeding scalability. In this regard, federated learning (FL) has become a popular technique for training machine learning models in a distributed manner [1, 2, 3]. In FL, the edge devices carry out the local computations, and the server collects, aggregates and updates the global model.

Recent approaches have looked at linearizing the training operations in FL [1, 4]. This is advantageous as it opens the possibility for coded matrix computing techniques that can improve operating efficiency. Specifically, in distributed settings like FL, the overall job execution time is often dominated by slower (or failed) worker nodes, which are referred to as stragglers. Recently, a number of coding theory techniques [5, 6, 7, 8, 9, 10, 11, 12, 13, 14] have been proposed to mitigate stragglers in distributed matrix multiplications. A toy example [5] of such a technique for computing $\mathbf{A}^{T}\mathbf{x}$ across three clients is to partition $\mathbf{A}$ as $\mathbf{A}=[\mathbf{A}_{0}\leavevmode\nobreak\ |\leavevmode\nobreak\ \mathbf{A}_{1}]$, and to assign them the job of computing $\mathbf{A}^{T}_{0}\mathbf{x}$, $\mathbf{A}^{T}_{1}\mathbf{x}$ and $\left(\mathbf{A}_{0}+\mathbf{A}_{1}\right)^{T}\mathbf{x}$, respectively. In a linearized FL setting, $\mathbf{A}\in\mathbb{R}^{t\times r}$ is the data matrix and $\mathbf{x}\in\mathbb{R}^{t}$ is the model parameter vector. While each client has half of the total computational load, the server can recover $\mathbf{A}^{T}\mathbf{x}$ if any two clients return their results, i.e., the system is resilient to one straggler. If each of $n$ clients computes $1/k_{A}$ fraction of the whole job of computing $\mathbf{A}^{T}\mathbf{x}$, the number of stragglers that the system can be resilient to is upper bounded by $n-k_{A}$ [7].

In contemporary edge computing systems, task offloading via device-to-device (D2D) communications has also been proposed for straggler mitigation. D2D-enabled FL has recently been studied [15, 16, 2], but can add considerable communication overhead as well as compromise data privacy. In this work, we exploit matrix coding in linearized FL to mitigate these challenges. Our straggler-optimal matrix computation scheme reduces the communication delay significantly compared to the techniques in [7, 9, 12]. Moreover, unlike [7, 9, 12, 13, 17], our scheme allows a client to access a limited fraction of matrix $\mathbf{A}$, and provides a considerable protection against information leakage. In addition, our scheme is specifically suited to sparse matrices with a significant gain in computation speed.

2 Network and Learning Architecture

We consider a D2D-enabled FL architecture consisting of $n=k_{A}+s$ clients, denoted as $W_{i}$ for $i=0,1,\dots,n-1$. The first $k_{A}$ of them are active clients (responsible for both data generation and local computation) and the next $s<k_{A}$ are passive clients (responsible for local computation only).

Assume that the $i$-th device has local data $(\mathbf{D}_{i},\mathbf{y}_{i})$, where $\mathbf{D}_{i}$ and $\mathbf{y}_{i}$ are the block-rows of full system dataset $(\mathbf{D},\mathbf{y})$. Under a linear regression-based ML model, the global loss function is quadratic, i.e., $f(\beta_{\ell})=||\mathbf{D}\beta_{\ell}-\mathbf{y}||^{2}$, where the model parameter after iteration $\ell$ is obtained through gradient methods as $\beta_{\ell}=\beta_{\ell-1}-\mu_{\ell}\nabla_{\beta}f(\beta_{\ell-1})$ and $\mu_{\ell}$ is the stepsize. Based on the form of $\nabla_{\beta}f(\beta_{\ell})$, the FL local model update at each device includes multiplying the local data matrix $\mathbf{D}_{i}$ with parameter $\beta_{\ell}$. For this reason, recent work has also investigated linearizing non-linear models for FL by leveraging kernel embedding techniques [1]. Thus, our aim is to compute $\mathbf{A}^{T}\mathbf{x}$ – an arbitrary matrix operation during FL training – in a distributed fashion such that the system is resilient to $s$ stragglers. Our assumption is that any active client $W_{i}$ generates a block-column of matrix $\mathbf{A}$, denoted as $\mathbf{A}_{i}$, $i=0,1,\dots,k_{A}-1$, such that

[TABLE]

In our approach, every client is responsible to compute the product of a coded submatrix (linear combinations of some block-columns of $\mathbf{A}$) and the vector $\mathbf{x}$. Stragglers will arise in practice from computing speed variations or failures experienced by the clients at particular times [18, 8, 17]. Now, similar to [15, 16, 19], we assume that there is a set of trusted neighbor clients for every device to transmit its data via D2D communications. The passive clients receive coded submatrices only from active clients. Unlike the approaches in [4, 20, 1, 3], we assume that the server cannot access to any uncoded/coded local data generated in the edge devices and is only responsible for transmission of vector $\mathbf{x}$ and for decoding $\mathbf{A}^{T}\mathbf{x}$ once the fastest clients return the computed submatrix-vector products.

3 Homogeneous Edge Computing

Here we assume that each active client generates equal number of columns of $\mathbf{A}$ (i.e. all $\mathbf{A}_{i}$’s have the same size in (1)) and all the clients are rated with the same computation speed. In this scenario, we propose a distributed matrix-vector multiplication scheme in Alg. 1 which is resilient to any $s$ stragglers.

The main idea is that any active client $W_{j}$ generates $\mathbf{A}_{j}$, for $0\leq j\leq k_{A}-1$ and sends it to another active client $W_{i}$, if $j=i+1,i+2,\dots,i+\omega_{A}-1$ ( modulo $\,k_{A})$. Here we set $\omega_{A}=s+1$, thus, any data matrix $\mathbf{A}_{j}$ needs to be sent to only $\omega_{A}-1=s$ other clients. Then, active client $W_{j}$ computes a linear combination of $\mathbf{A}_{i},\mathbf{A}_{i+1},\dots,\mathbf{A}_{i+\omega_{A}-1}\,$ (indices modulo $\,k_{A}$) where the coefficients are chosen randomly from a continuous distribution. Next, active client $W_{i}$ sends another random linear combination of the same submatrices to $W_{i+k_{A}}$ (a passive client), when $i=0,1,\dots,s-1$. Note that all $n$ clients receive the vector $\mathbf{x}$ from the server. Now the job of each client is to compute the product of their respective coded submatrix and the vector $\mathbf{x}$. Once the fastest $k_{A}$ clients finish and send their computation results to the server, it decodes $\mathbf{A}^{T}\mathbf{x}$ using the corresponding random coefficients. The following theorem establishes the resiliency of Alg. 1 to stragglers.

Theorem 1**.**

Assume that a system has $n$ clients including $k_{A}$ active and $s$ passive clients. If we assign the jobs according to Alg. 1, we achieve resilience to any $s=n-k_{A}$ stragglers.

Proof.

In order to recover $\mathbf{A}^{T}\mathbf{x}$, according to (1), we need to decode all $k_{A}$ vector unknowns, $\mathbf{A}^{T}_{0}\mathbf{x},\mathbf{A}^{T}_{1}\mathbf{x},\dots,\mathbf{A}^{T}_{k_{A}-1}\mathbf{x}$; we denote the set of these unknowns as $\mathcal{B}$. Now we choose an arbitrary set of $k_{A}$ clients each of which corresponds to an equation in terms of $\omega_{A}$ of those $k_{A}$ unknowns. Denoting the set of $k_{A}$ equations as $\mathcal{C}$, we have $|\mathcal{B}|=|\mathcal{C}|=k_{A}$.

Now we consider a bipartite graph $\mathcal{G}=\mathcal{C}\cup\mathcal{B}$, where any vertex (equation) in $\mathcal{C}$ is connected to some vertices (unknowns) in $\mathcal{B}$ which have participated in the corresponding equation. Thus, each vertex in $\mathcal{C}$ has a neighborhood of cardinality $\omega_{A}$ in $\mathcal{B}$. Our goal is to show that there exists a perfect matching among the vertices of $\mathcal{C}$ and $\mathcal{B}$. We argue this according to Hall’s marriage theorem [21] for which we need to show that for any $\bar{\mathcal{C}}\subseteq\mathcal{C}$, the cardinality of the neighbourhood of $\bar{\mathcal{C}}$, denoted as $\mathcal{N}(\bar{\mathcal{C}})\subseteq\mathcal{B}$, is at least as large as $|\bar{\mathcal{C}}|$. Thus, for $|\bar{\mathcal{C}}|=m\leq k_{A}$, we need to show that $|\mathcal{N}(\bar{\mathcal{C}})|\geq m$.

Case 1: First we consider the case that $m\leq 2s$. We assume that $m=2p,2p-1$ where $1\leq p\leq s$. Now according to Alg. 1, the participating unknowns are shifted in a cyclic manner among the equations. If we choose any $\delta$ clients out of the first $k_{A}$ clients $\left(W_{0},W_{1},W_{2},\dots,W_{k_{A}-1}\right)$, according to the proof of cyclic scheme in Appendix C in [8], the minimum number of total participating unknowns is $\textrm{min}(\omega_{A}+\delta-1,k_{A})$, where $\omega_{A}=s+1$. Now according to Alg. 1, same unknowns participate in two different equations corresponding to two different clients, $W_{j}$ and $W_{k_{A}+j}$, where $j=0,1,\dots,s-1$. Thus, for any $|\bar{\mathcal{C}}|=m=2p,2p-1\leq 2s$, we have

[TABLE]

Case 2: Now we consider the case where $m=2s+q$, $1\leq q\leq k_{A}-2s$. We need to find the minimum number of unknowns which participate in any set of $m$ equations. Now, the same unknowns participate in two different equations corresponding to two different clients, $W_{j}$ and $W_{k_{A}+j}$, where $j=0,1,\dots,s-1$. Thus, the additional $q$ equations correspond to at least $q$ additional unknowns until the total number of participating unknowns is $k_{A}$. Therefore, in this case

[TABLE]

Thus, for any $m\leq k_{A}$ (where $|\bar{\mathcal{C}}|=m$), we have shown that $|\mathcal{N}(\bar{\mathcal{C}})|\geq|\bar{\mathcal{C}}|$. So, there exists a perfect matching among the vertices of $\mathcal{C}$ and $\mathcal{B}$ according to Hall’s marriage theorem.

Now we consider the largest matching where vertex $c_{i}\in\mathcal{C}$ is matched to vertex $b_{j}\in\mathcal{B}$, which indicates that $b_{j}$ participates in the equation corresponding to $c_{i}$. Let us consider a $k_{A}\times k_{A}$ system matrix where row $i$ corresponds to the equation associated to $c_{i}$. Now we replace this row $i$ by $\mathbf{e}_{j}$ which is a unit row-vector of length $k_{A}$ with $j$-th entry being $1$, and [math] otherwise. Thus we have a $k_{A}\times k_{A}$ matrix where each row has only one non-zero entry which is $1$. Since we have a perfect matching, this $k_{A}\times k_{A}$ matrix has only one non-zero entry in every column. This is a permutation of the identity matrix, and thus, is full rank. Since the matrix is full rank for a choice of definite values, according to Schwartz-Zippel lemma [22], it will be full rank for random choices of non-zero entries. Thus, the server can recover all $k_{A}$ unknowns from any $k_{A}$ clients, hence the system is resilient to any $s=n-k_{A}$ stragglers. ∎

Example 1**.**

Consider a homogeneous system of $k_{A}=10$ active clients and $s=2$ passive clients. According to Alg. 1, $\omega_{A}=s+1=3$, and client $W_{i}$ ($0\leq i\leq 11$) has a random linear combination of $\mathbf{A}_{i},\mathbf{A}_{i+1}$ and $\mathbf{A}_{i+2}\,\left(\textrm{indices modulo}\,10\right)$ as shown in Fig. 1. Thus, according to Theorem 1, this system is resilient to $s=2$ stragglers. Note that our scheme requires any active client to send its local data matrix to only up to $s+1=3$ other clients, thus involves a significantly lower communication cost in comparison to the approaches in [7, 9].

Remark 1**.**

In comparison to [7, 9, 13], our proposed approach is specifically suited to sparse data matrices, i.e., most of the entries of $\mathbf{A}$ are zero. The approaches in [7, 9, 13] assign dense linear combinations of the submatrices which can destroy the inherent sparsity of $\mathbf{A}$, leading to slower computation speed for the clients. On the other hand, our approach assigns linear combinations of limited number of submatrices which preserve the sparsity up to certain level that leads to faster computation.

4 Heterogeneous Edge Computing

In this section, we extend our approach in Alg. 1 to heterogenous system where the clients may have different data generation capability and different computation speeds. We assume that we have $\lambda$ different types of devices in the system, with client type $j=0,1,\dots,\lambda-1$. Moreover, we assume that any active client $W_{i}$ generates $\alpha_{i}=c_{ij}\alpha$ columns of data matrix $\mathbf{A}$ and any client $W_{i}$ has a computation speed $\beta_{i}=c_{ij}\beta$, where $W_{i}$ is of client type $j$ and $c_{ij}\geq 1$ is an integer. Thus, a higher $c_{ij}$ indicates a “stronger” type client $W_{i}$ which can process at a $c_{ij}$ times higher computation speed than the “weakest” type device, where $\alpha$ is the number of the assigned columns and $\beta$ is the number of processed columns per unit time in the “weakest” type device. Note that $\lambda=1$ and all $c_{ij}=1$ lead us to the homogeneous system discussed in Sec. 3 where $0\leq i\leq n-1$ and $j=0$.

Now, we have $n=k_{A}+s$ clients including $k_{A}$ active and $s$ passive clients in the heterogeneous system. Aligned to the homogeneous system, we assume that the number of passive clients of any type $j$ is less than the number of active clients of the same type. Next, without loss of generality, we sort the indices of active clients in such a way so that, $c_{ij}\geq c_{kj}$ if $i\leq k$, for $0\leq i,k\leq k_{A}-1$. We do the similar sorting for the passive clients too so that $c_{ij}\geq c_{kj}$ if $i\leq k$, for $k_{A}\leq i,k\leq n-1$. Now if a client $W_{i}$ is of client type $j$, it requires the same time to process $c_{ij}\geq 1$ block-columns (each consisting of $\alpha$ columns) of $\mathbf{A}$ as the “weakest” device to process $c_{ij}=1$ such block-column. Moreover, if it is an active client, it also generates $\alpha_{i}=c_{ij}\alpha$ columns of data matrix $\mathbf{A}$. Thus, client $W_{i}$ can be thought as a collection of $c_{ij}$ homogeneous clients of “weakest” types where each of the active “weakest” clients generates equally $\alpha$ columns of $\mathbf{A}$ and each of the “weakest” clients processes equally $\alpha$ columns.

Theorem 2**.**

(a) A heterogeneous system of $k_{A}$ active and $s$ passive clients of different types can be considered as a homogeneous system of $\bar{k}_{A}=\sum_{i=0}^{k_{A}-1}c_{ij}$ active and $\bar{s}=\sum_{i=k_{A}}^{n-1}c_{ij}$ passive clients of the “weakest” type. Next (b) if the jobs are assigned according to Alg. 1 in the modified homogeneous system of $\bar{n}=\bar{k}_{A}+\bar{s}$ “weakest” clients, the system can be resilient to $\bar{s}$ such clients.

Proof.

Each $\mathbf{A}_{k}$ (generated in $W_{k}$) in (1) is a block-column consisting of $c_{kj}\alpha$ columns of $\mathbf{A}$ when client $W_{k}$ is of client type $j$. Thus, for any $k=0,1,\dots,k_{A}-1$, we can partition $\mathbf{A}_{k}$ as $\mathbf{A}_{k}=\begin{bmatrix}\bar{\mathbf{A}}_{m}&\bar{\mathbf{A}}_{m+1}&\dots&\bar{\mathbf{A}}_{m+c_{kj}-1}\end{bmatrix}$, where $m=\sum_{i=0}^{k-1}c_{ij}$ and each $\bar{\mathbf{A}}_{\ell}$ is a block-column consisting of $\alpha$ columns of $\mathbf{A}$, $m\leq\ell\leq m+c_{kj}-1$. Thus using (1), we can write $\mathbf{A}=\begin{bmatrix}\mathbf{A}_{0}&\mathbf{A}_{1}&\dots&\mathbf{A}_{\bar{k}_{A}-1}\end{bmatrix}$, where $\bar{k}_{A}=\sum_{i=0}^{k_{A}-1}c_{ij}$. Now from the matrix generation perspective, $k_{A}$ active clients in a heterogeneous system generating $\bar{k}_{A}$ block-columns can be considered as the same as $\bar{k}_{A}$ active clients in a homogeneous system generating one block-column each.

Similarly, any client $W_{i}$ of type $j$ can process $c_{ij}\alpha$ columns in the same time when the “weakest” type device can process $\alpha$ columns. Thus, from the computation speed perspective, $k_{A}$ active clients and $s$ passive clients in the heterogeneous system can be thought as $\bar{k}_{A}=\sum_{i=0}^{k_{A}-1}c_{ij}$ active clients and $\bar{s}=\sum_{i=k_{A}}^{n-1}c_{ij}$ passive clients, respectively, in a homogeneous system by assigning $\alpha$ coded block-columns to each client. Hence, we are done with the proof of part (a). Moreover, part (b) of the proof is straight-forward from Theorem 1 when we have $\bar{k}_{A}$ active and $\bar{s}$ passive clients. ∎

Remark 2**.**

The heterogeneous system is resilient to $\bar{s}$ block-column processing. The number of straggler clients that the system is resilient to can vary depending on the client types.

Example 2**.**

Consider the example in Fig. 2 consisting of $n=7$ clients. There are $k_{A}=5$ active clients which are responsible for data matrix generation. Let us assume, $W_{0}$ and $W_{1}$ are of type $1$ clients which generate twice as many columns of $\mathbf{A}$ than $W_{2},W_{3}$ and $W_{4}$ which are of type [math] clients. The jobs are assigned to all clients (including $s=2$ passive clients) according to Fig. 2(b). It can be verified that this scheme is resilient to two type [math] clients or one type $1$ client.

5 Numerical Evaluation

In this section, we compare the performance of our proposed approach against different competing methods [7, 9, 13] in terms of different metrics for distributed matrix computations from the federated learning aspect. Note that the approaches in [1, 4] require the edge devices to transmit some coded columns of matrix $\mathbf{A}$ to the server which is not aligned with our assumptions. In addition, the approaches in [8] and [11] do not follow the same network learning architecture as ours. Therefore, we did not include them in our comparison.

Communication Delay: We consider a homogeneous system of $n=20$ clients each of which is a t2.small machine in AWS (Amazon Web Services) Cluster. Here, each of $k_{A}=18$ active clients generates $\mathbf{A}_{i}$ of size $12000\times 1000$, thus the size of $\mathbf{A}$ is $12000\times 18000$. The server sends the parameter vector $\mathbf{x}$ of length $12000$ to all $20$ clients including $s=2$ passive clients. Once the preprocessing and computations are carried out according to Alg. 1, the server recovers $\mathbf{A}^{T}\mathbf{x}$ as soon as it receives results from the fastest $k_{A}=18$ clients, thus the system is resilient to any $s=2$ stragglers.

Table 1 shows the comparison of the corresponding communication delays (caused by data matrix transmission) among different approaches. The approaches in [7, 9] require all active clients to transmit their generated submatrices to all other edge devices. Thus, they lead to much more communication delay than our proposed method which needs an edge device to transmit data to only up to $s+1=3$ other devices. Note that the methods in [13, 17] involve similar amounts of communication delay as ours, however, they have other limitations in terms of privacy and computation time as discussed next.

Privacy: Information leakage is introduced in FL when we consider the transmission of local data matrices to other edge devices. To protect against privacy leakage, any particular client should have access to a limited portion of the whole data matrix. Consider the heterogeneous system in example 2 where the clients are honest but curious. In this scenario, the approaches in [7, 9, 13, 17] would allow clients to access the whole matrix $\mathbf{A}$. In our approach, as shown in Fig. 2, clients $W_{0}$ and $W_{1}$ only have access to $4/7$-th fraction of $\mathbf{A}$ and clients $W_{2}$, $W_{3}$ and $W_{4}$ have access to $3/7$-th fraction of $\mathbf{A}$. This provides significant protection against privacy leakage.

Product Computation Time for Sparse Matrices: Consider a system with $n=30$ clients where $k_{A}=28$ and $s=2$. We assume that $\mathbf{A}$ is sparse, where each active client generates a sparse submatrix of size $40000\times 1125$. We consider three different scenarios with three different sparsity levels for $\mathbf{A}$ where randomly chosen $95\%$, $98\%$ and $99\%$ entries of $\mathbf{A}$ are zero. Now we compare our proposed Alg. 1 against different methods in terms of per client product computation time (the required time for a client to compute its assigned submatrix-vector product) in Table 2. The methods in [7, 9, 13, 17] assign linear combinations of $k_{A}=28$ submatrices to the clients. Hence, the inherent sparsity of $\mathbf{A}$ is destroyed in the encoded submatrices. On the other hand, our approach combines only $s+1=3$ submatrices to obtain the coded submatrices. Thus, the clients require a significantly less amount of time to finish the respective tasks in comparison to [7, 9, 13, 17].