Fast and Secure Distributed
Nonnegative Matrix Factorization

In the sequel, we briefly review three research areas which are related to this paper.

Assume there are $N$ computing nodes in the cluster. We partition the row indices $\{1,2,\dots,m\}$ of the input matrix $M$ into $N$ disjoint sets $I_{1},I_{2},\dots,I_{N}$, where $I_{r}\subset\{1,2,\dots,m\}$ is the subset of rows assigned to node $r$, as in [21]. Similarly, we partition the column indices $\{1,2,\dots,n\}$ into disjoint sets $J_{1},J_{2},\dots,J_{N}$ and assign column set $J_{r}$ to node $r$. The number of rows and columns in each node are near the same in order to achieve load balancing, i.e., $\left|I_{r}\right|\approx m/N$ and $\left|J_{r}\right|\approx n/N$ for each node $r$. The factor matrices $U$ and $V$ are also assigned to nodes accordingly, i.e., node $r$ stores and updates $U_{I_{r}:}$ and $V_{J_{r}:}$ as shown in Fig. 1(a).

Data partitioning in distributed NMF differs from that in parallel NMF. Previous works on parallel NMF [19, 20] choose to partition $U$ and $V$ along the long dimension, but we adopt the row-partitioning of $U$ and $V$ as in [21]. To see why, take the $U$-subproblem (2) as an example and observe that it is row-independent in nature, i.e., the $r$-th row block of its solution $U_{I_{r}:}$ is given by:

and thus can be solved independently without referring to any other row blocks of $U$. The same holds for the $V$-subproblem. In addition, no communication is needed concerning $M$ when solving (5) because $M_{I_{r}:}$ is already present in node $r$.

On the other hand, solving (5) requires the entire $V$ of size $n\times k$, meaning that every node needs to gather $V$ from all other nodes. This process can easily be the bottleneck of a naive distributed ANLS implementation. As we will explain shortly, our DSALNS algorithm alleviates this problem, since we use a sketched matrix of reduced size instead of the original complete matrix $V$.

To better understand DSANLS, we first introduce the Sketched ANLS (SANLS), i.e., a centralized version of our algorithm. Recall that in Sec. 2.1.1, at each step of ANLS, either $U$ or $V$ is fixed and we solve a nonnegative least square problem (2) or (3) over the other variable. Intuitively, it is unnecessary to solve this subproblem with high accuracy, because we may not have reached the optimal solution for the fixed variable so far. Hence, when the fixed variable changes in the next step, this accurate solution from the previous step will not be optimal anymore and will have to be re-computed. Our idea is to apply matrix sketching for each subproblem, in order to obtain an approximate solution for it at a much lower computational and communication cost.

Specifically, suppose we are at the $t$-th iteration of ANLS, and our current estimations for $U$ and $V$ are $U^{t}$ and $V^{t}$ respectively. We must solve subproblem (2) in order to update $U^{t}$ to a new matrix $U^{t+1}$. We apply matrix sketching to the residual term of subproblem (2). The subproblem now becomes:

where $S^{t}\in\mathbb{R}^{n\times d}$ is a randomly-generated matrix. Hence, the problem size decreases from $n\times k$ to $d\times k$. $d$ is chosen to be much smaller than $n$, in order to sufficiently reduce the computational cost¹¹1However, we should not choose an extremely small $d$, otherwise the the size of sketched subproblem would become so small that it can hardly represent the original subproblem, preventing NMF from converging to a good result. In practice, we can set $d=0.1n$ for medium-sized matrices and $d=0.01n$ for large matrices if $m\approx n$. When $m$ and $n$ differ a lot, e.g., $m\ll n$ without loss of generality, we should not apply sketching technique to the $V$ subproblem (since solving the $U$ subproblem is much more expensive) and simply choose $d=m\ll n$. . Similarly, we transform the $V$-subproblem into:

where $S^{\prime t}\in\mathbb{R}^{m\times d^{\prime}}$ is also a random matrix with $d^{\prime}\ll m$.

Now, we come to our proposal: the distributed version of SANLS called DSANLS. Since the $U$-subproblem (6) is the same as the $V$-subproblem (7) in nature, here we restrict our attention to the $U$-subproblem. The first observation about subproblem (6) is that it is still row-independent, thus node $r$ only needs to solve:

For simplicity, we denote:

and the subproblem (8) can be written as:

Thus, node $r$ needs to know matrices $A^{t}_{r}$ and $B^{t}$ in order to solve the subproblem.

For $A^{t}_{r}$, by applying matrix multiplication rules, we get $A^{t}_{r}=\left(MS^{t}\right)_{I_{r}:}=M_{I_{r}:}S^{t}$. Therefore, if $S^{t}$ is stored at node $r$, $A^{t}_{r}$ can be computed without any communication. On the other hand, computing $B^{t}=\left(V^{t\top}S^{t}\right)$ requires communication across the whole cluster, since the rows of $V^{t}$ are distributed across different nodes. Fortunately, if we assume that $S^{t}$ is stored at all nodes again, we can compute $B^{t}$ in a much cheaper way. Following block matrix multiplication rules, we can rewrite $B^{t}$ as:

Note that the summand $\bar{B}^{t}_{r}\triangleq\left(V^{t}_{J_{r}:}\right)^{\top}S^{t}_{J_{r}:}$ is a matrix of size $k\times d$ and can be computed locally. As a result, communication is only needed for summing up the matrices $\bar{B}^{t}_{r}$ of size $k\times d$ by using MPI all-reduce operation, which is much cheaper than transmitting the whole $V_{t}$ of size $n\times k$.

Now, the only remaining problem is the transmission of $S^{t}$. Since $S^{t}$ can be dense, even larger than $V^{t}$, broadcasting it across the whole cluster can be quite expensive. However, it turns out that we can avoid this. Recall that $S^{t}$ is a randomly-generated matrix; each node can generate exactly the same matrix, if we use the same pseudo-random generator and the same seed. Therefore, we only need to broadcast the random seed, which is just an integer, at the beginning of the whole program. This ensures that each node generates exactly the same random number sequence and hence the same random matrices $S^{t}$ at each iteration.

In short, the communication cost of each node is reduced from $\mathcal{O}(nk)$ to $\mathcal{O}(dk)$ by adopting our sketching technique for the $U$-subproblem. Likewise, the communication cost of each $V$-subproblem is decreased from $\mathcal{O}\left(mk\right)$ to $\mathcal{O}\left(d^{\prime}k\right)$. The general framework of our DSANLS algorithm is listed in Alg. 2.

A key problem in Alg. 2 is how to generate random matrices $S^{t}\in\mathbb{R}^{n\times d}$ and $S^{\prime t}\in\mathbb{R}^{m\times d^{\prime}}$. Here we focus on generating a random $S^{t}\in\mathbb{R}^{d\times n}$ satisfying Assumption 1. The reason for choosing such a random matrix is that the corresponding sketched problem would be equivalent to the original problem on expectation; we will prove this in Sec. 3.5.

Different options exist for such matrices, which have different computation costs in forming sketched matrices $A^{t}_{r}=M_{I_{r}:}S^{t}$ and $\bar{B}^{t}_{r}=\left(V^{t}_{J_{r}:}\right)^{\top}S^{t}_{J_{r}:}$. Since $M_{I_{r}:}$ is much larger than $V^{t}_{J_{r}:}$ and thus computing $A^{t}_{r}$ is more expensive, we only consider the cost of constructing $A^{t}_{r}$ here.

The most classical choice for a random matrix is one with i.i.d. Gaussian entries having mean 0 and variance $1/d$. It is easy to show that $\mathbb{E}\left[S^{t}S^{t\top}\right]=I$. Besides, Gaussian random matrix has bounded variance because Gaussian distribution has finite fourth-order moment. However, since each entry of such a matrix is totally random and thus no special structure exists in $S^{t}$, matrix multiplication will be expensive. That is, when given $M_{I_{r}:}$ of size $|I_{r}|\times n$, computing its sketched matrix $A^{t}_{r}=M_{I_{r}:}S^{t}$ requires $\mathcal{O}(|I_{r}|nd)$ basic operations.

A seemingly better choice for $S^{t}$ would be a subsampling random matrix. Each column of such random matrix is uniformly sampled from $\{e_{1},e_{2},\dots,e_{n}\}$ without replacement, where $e_{i}\in\mathbb{R}^{n}$ is the $i$-th canonical basis vector (i.e., a vector having its $i$-th element 1 and all others 0). We can easily show that such an $S^{t}$ also satisfies $\mathbb{E}\left[S^{t}S^{t\top}\right]=I$ and the variance $\mathbb{V}\left[S^{t}S^{t\top}\right]$ is bounded, but this time constructing the sketched matrix $A^{t}_{r}=M_{I_{r}:}S^{t}$ only requires $\mathcal{O}\left(|I_{r}|d\right)$. Besides, subsampling random matrix can preserve the sparsity of original matrix. Hence, a subsampling random matrix would be favored over a Gaussian random matrix by most applications, especially for very large-scale or sparse problems. On the other hand, we observed in our experiments that a Gaussian random matrix can result in a faster per-iteration convergence rate, because each column of the sketched matrix $A^{t}_{r}$ contains entries from multiple columns of the original matrix and thus is more informative. Hence, it would be better to use a Gaussian matrix when the sketch size $d$ is small and thus a $\mathcal{O}(|I_{r}|nd)$ complexity is acceptable, or when the network speed of the cluster is poor, hence we should trade more local computation cost for less communication cost.

Although we only test two representative types of random matrices (i.e., Gaussian and subsampling random matrices), our framework is readily applicable for other choices, such as subsampled randomized Hadamard transform (SRHT) [42, 43] and count sketch [44, 45]. The choice of random matrices is not the focus of this paper and left for future investigation.

Before describing how to solve subproblem (10), let us make an important observation. As discussed in Sec. 2.2.2, the sketching technique has been applied in solving linear systems before. However, the situation is different in matrix factorization. Note that for the distributed matrix factorization problem we usually have:

So, for the sketched subproblem (10), which can be equivalently written as:

where the non-zero entries of the residual matrix $\left(M_{I_{r}:}-U_{I_{r}:}V^{t\top}\right)$ will be scaled by the matrix $S^{t}$ at different levels. As a consequence, the optimal solution will be shifted because of sketching. This fact alerts us that for SANLS, we need to update $U^{t+1}$ by exploiting the sketched subproblem (10) to step towards the true optimal solution and avoid convergence to the solution of the sketched subproblem.

DSANLS and all lines of works discussed in Sec. 2.2.1 store copies of $M$ across two-dimensional (shown in Fig. 1(a)), and exploit the independence of local update computation for rows of $U$ and $V$ to apply communication-optimal matrix multiplication. They cannot be applied directly to secure distributed NMF setting. The reason is that, in secure distributed NMF setting (shown in Fig. 1(b)), only one column copy is stored in each node, while the others cannot be disclosed.

Nevertheless, DSANLS can be adapted to this secure setting with modification, but only for one or limited iterations. The reason is illustrated in Theorem 2. In modified DSANLS algorithm, each node still takes charge of updating $U_{I_{r}:}$ and $V_{J_{r}:}$ as before, but only one copy $M_{:J_{r}}$ of $M=[M_{1},M_{2},...,M_{N}]$ will be stored in node $r$. Thus, $V$-subproblem is exactly the same as in DSANLS. Differently, we need to use MPI-AllReduce function to gather $M_{:J_{r}}S^{t}$ from all nodes before each iteration of $U$-subproblem, so that each node has access to fully sketched matrix $MS^{t}$ to solve sketched $U$-subproblem. Note that here random matrix $S^{t}$ not only helps reduce the communication cost from $\mathcal{O}(mn)$ to $\mathcal{O}(md)$ with a smaller NLS problem, but also conceals the full matrix $M$ in each iteration.

However, NMF is an iterative algorithm (shown in Alg. 1). Secure computation in limited iterations cannot guarantee an acceptable accuracy for practical use due to the following reason:

Theorem 3 suggests that DSANLS algorithm suffers from the dilemma of choosing between information disclosure and unacceptable accuracy, making it impractical to real applications. Therefore, we need to propose new practical solutions to secure distributed NMF.

A straightforward solution to secure distributed NMF is that each node solves a local NMF problem with a local copy of $U$ (denoted as $U_{(r)}$ for node $r$). Periodically, nodes communicate with each other, and update local copy of $U$ to the aggregation of all local copies $U_{(j)},j\in\{1,\cdots,N\}$ by All-Reduce operation. We name this method as Syn-SD under synchronous setting. The detailed algorithm is shown in Alg. 4. Within inner iterations, every node maintains its own copy of $U$ (i.e., $U_{(r)}$) by solving the regular NMF problem. Every $T_{2}$ rounds, different local copies of $U$ will be averaged through nodes by using $\sum_{j=1}^{N}U_{(j)}/N$. Note that, $U_{(r)}$ is one copy of the whole matrix $U$ stored locally in node $r$, while $V_{J_{r}:}$ is the corresponding part of the matrix $V=[V_{J_{1}:},V_{J_{2}:},...,V_{J_{N}:}]$ stored in node $r$.

In Syn-SD, the local copy $U_{(r)}$ in node $r$ will be updated to a uniform aggregation of local copies from all nodes periodically. Small number of inner iteration $T_{2}$ incurs large communication cost caused by All-Reduce. Larger $T_{2}$ may lead to slow convergence, since each node does not share any information of its local copy $U_{(r)}$ inside the inner iterations.

To improve the efficiency of data exchange, we incorporate matrix sketching to Syn-SD, and propose an improved version called Syn-SSD. In Syn-SSD, information of local copies is shared across cluster nodes more frequently, with communication overhead roughly the same as Syn-SD. As shown in Alg. 5, the sketched version $S^{t}U_{(r)}$ of the local copy $U_{(r)}$ is exchanged within each inner iteration. There are two advantages of applying matrix sketching: (1) Since the sketched matrix has a much smaller size, All-Reduce operation causes much less communication cost, making it affordable with higher frequency. (2) Solving a sketched NLS problem can also reduce the computation cost due to a reduced problem size of solving $U_{(r)}$ and $V_{J_{r}:}$ for each node. It is worth noting that $S_{1}^{t}$ is exactly the same for each node by using the same seed and generator. The same for $S_{2}^{t}$. But $S_{1}^{t}$ and $S_{2}^{t}$ are not necessarily equivalent. With such a constraint, the algorithm is equivalent to NMF in single-machine environment and the convergence can be guaranteed.

It is straightforward to see that Syn-SD and Syn-SSD satisfy Definition 1 and they are $(N-1)$-private protocols, since $V_{J_{r}:}$ and $M_{:J_{r}}$ are only seen by node $r$.

In Syn-SD and Syn-SSD, each node must stall until all participating nodes reach the synchronization barrier before the All-Reduce operation. However, highly imbalanced data in real scenario of federated data mining may cause severe workload imbalance problem. The synchronization barrier will force nodes with low workload to halt, making synchronous algorithms less efficient. In this section, we study secure distributed NMF in an asynchronous (i.e., server/client architecture) setting and propose corresponding asynchronous algorithms.

First of all, we extend the idea of Syn-SD to asynchronous setting and name the new method Asyn-SD. In Asyn-SD, the server (in Alg. 6) takes full charge of updating and broadcasting $U^{t}$. Once received $U^{t}_{(r)}$ from the client node $r$, the server would update $U^{t}$ locally, and return the latest version of $U^{t}$ back to the client node $r$ for further computing. Note that the server may receive local copies of $U^{t}$ from clients in an arbitrary order. Consequently, we cannot use the same operation of All-Reduce as Syn-SD any more. Instead, $U^{t}$ in server side is updated by the weighted sum of current $U^{t}$ and newly received local copy $U^{t}_{(r)}$ from client node $r$. Here the relaxation weight $\omega^{t}$ asymptotically converges to 0. Thus a converged $U^{t}$ is guaranteed on server side. Our experiments in Sec. 5 suggest that this relaxation has no harm to factorization convergence.