SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Marie Maros
mmaros@purdue.edu
Gesualdo Scutari
gscutari@purdue.edu
We explore the matrix sensing problem from near-isotropic linear measurements, distributed across a network of agents modeled as an undirected graph, with no server. We provide the first study of statistical, computational/communication guarantees for a decentralized gradient algorithm that solves the (nonconvex) Burer-Monteiro type decomposition associated to the low-rank matrix estimation. With small random initialization, the algorithm displays an approximate two-phase convergence: (i) a spectral phase that aligns the iterates' column space with the underlying low-rank matrix, mimicking centralized spectral initialization (not directly implementable over networks); and (ii) a local refinement phase that diverts the iterates from certain degenerate saddle points, while ensuring swift convergence to the underlying low-rank matrix. Central to our analysis is a novel "in-network" Restricted Isometry Property which accommodates for the decentralized nature of the optimization, revealing an intriguing interplay between sample complexity, network connectivity & topology, and communication complexity.
Matrix sensing-the estimation of a low-rank matrix from a set of linear measurements-finds applications in diverse fields such as image reconstruction (e.g., [48, 29]), object detection (e.g., [33, 52]) and array processing (e.g., [21]), to name a few. It also serves as a benchmark for determining the statistical and computational guarantees achievable in deep learning theory, since it retains many of the key phenomena in deep learning while being simpler to analyze. Despite significant progress in understanding the convergence and generalization properties of various solution methods for training such learning models, a majority of these advances focus on a centralized paradigm, aggregating data at a central location with vast computing resources-good tutorials on the topic include [7, 4]. This centralized approach, however, is increasingly unsuitable for modern applications due to server bottlenecks, inefficient communication, and power usage. Therefore, the development of statistical learning methods for massively decentralized networks without servers is timely and crucial.
This paper tackles the matrix sensing problem from data distributed over networks. We contemplate a network of  $m$  agents modeled as an undirected graph with no servers, where agents can communicate with their immediate neighbors—these architectures are also known as mesh networks. The collective objective is to estimate a ground-truth matrix  $\bar{Z}^{\star} \in \mathbb{R}^{d \times d}$ , based on  $N = m \cdot n$  total observations  $y_{1}, \ldots, y_{N}$ , equally split into  $n$ -sized, disjoint datasets  $\mathcal{D}_1, \mathcal{D}_2, \ldots, \mathcal{D}_m$ . Each agent's signal model is thus given by
$$
y _ {j} = \left\langle A _ {j}, \bar {Z} ^ {\star} \right\rangle := \operatorname {t r a c e} \left(A _ {j} \bar {Z} ^ {\star}\right), \quad \text {f o r} j \in \mathcal {D} _ {i}. \tag {1}
$$
Here,  $A_{j}\in \mathbb{R}^{d\times d}$ , with  $j\in \mathcal{D}_i$ , are the known symmetric measurement matrices to agent  $i$ ;  $\bar{Z}^{\star}$  is assumed symmetric, positive semidefinite, and low-rank, i.e.,  $r^{\star}\coloneqq \mathrm{rank}(\bar{Z}^{\star}) <   <   d$
To minimize communication overhead and avoid the need for  $d \times d$  matrix transmission, we employ the Burer-Monteiro-type decomposition of the estimate  $\bar{Z}$  of  $\bar{Z}^{\star}$ , that is,  $\bar{Z} = \bar{U}\bar{U}^{\top}$ , with  $\bar{U} \in \mathbb{R}^{d \times r}$ , and seek to minimize the squared loss  $F(\bar{U})$ , defined as
$$
\min  _ {\bar {U} \in \mathbb {R} ^ {d \times r}} F (\bar {U}) := \frac {1}{m} \sum_ {i = 1} ^ {m} \underbrace {\frac {1}{4 n} \sum_ {j \in \mathcal {D} _ {i}} \left(y _ {j} - \left\langle A _ {j} , \bar {U} \bar {U} ^ {\top} \right\rangle\right) ^ {2}} _ {:= f _ {i} (\bar {U})}, \tag {2}
$$
where  $f_{i}(\bar{U})$  is the loss function of agent  $i$ . Ideally, the number  $r$  of columns of  $\bar{U}$  should be set to  $r^{\star}$ . However,  $r^{\star}$  might not be known in advance. In this study, we consider the so-called over-parameterized regime where  $r \geq r^{\star}$ .
The formulation (2) poses multiple challenges. Firstly,  $F$  is nonconvex, lacks global smoothness (i.e., global Lipschitz continuity of  $\nabla F$ ), and is not entirely known to the agents. Secondly, the over-parameterized regime may intuitively suggest a risk of overfitting. However, recent studies (e.g., [34, 20]) have compellingly revealed that when centralized Gradient Descent (GD) is applied to (2) with a small random initialization, it induces an implicit bias towards simpler solutions with favorable generalization properties. This bias—often referred to as the simplicity bias or the incremental learning behavior of GD/(Stochastic Gradient Descent)—assists in the exact or approximate recovery of the ground truth  $\bar{Z}^{\star}$ . Interestingly, this phenomenon serves as a hidden mechanism in various other (deep) learning tasks that mitigates overfitting in highly over-parameterized models (e.g., [13, 30, 31]). However, the direct implementation of GD in mesh networks is not feasible due to the lack of access to  $\nabla F$  by the agents or of a server collecting agents' gradients  $\nabla f_{i}$ .
The goal of this paper is to uncover the possible simplicity bias of a decentralized instance of GD, solving the matrix sensing problem (2) over mesh networks. To the best of our knowledge, this study is unique in the realm of decentralized optimization, establishing the first sample, convergence rate, and generalization guarantees of a decentralized gradient-based algorithm tailored for matrix sensing over mesh networks. We delve into the relevant existing literature in the subsequent section.
The literature offers numerous decentralized algorithms which could in principle be used to tackle the matrix sensing problem, directly or indirectly. However, the accompanying statistical and computational guarantees fall short, being either non-existent or inadequate, as we elaborate next.
- Off-the-shelf decentralized algorithms for nonconvex problems: The matrix sensing problem (2) naturally invites the application of decentralized algorithms specifically designed for nonconvex losses in the form  $F = (1/m)\sum_{i=1}^{m}f_{i}$  (summation of agent functions). Noteworthy examples of such algorithms include (i) decentralizations of the GD that merge local gradient updates with (push-sum) consensus algorithms [47, 2, 37], (ii) decentralized first-order methods employing gradient tracking strategies [8, 32, 42, 17], and (iii) decentralized algorithms grounded on primal-dual decomposition or penalization of lifted reformulations incorporating explicitly consensus constraints [15, 14, 50]. However, despite their initial appeal, when applied to (2), these algorithms either lack of any convergence guarantee—the requirement that  $F$  is globally smooth [47, 2, 37, 8, 32, 42, 15, 14, 50] and has a (uniformly) bounded gradient [47, 2, 37, 8, 42] is not met by the matrix sensing loss in (2)—or they converge at sublinear rate to some critical points of  $F$  (which may not be the global minimizers), whose generalization properties remain unexplored and obscure [17].
- Ad-hoc decentralized algorithms for some matrix recovery problems: This line of works comprises decentralized schemes designed specifically for the structured matrix-related optimization problem under consideration. Relevant examples are briefly highlighted next.
(i) Dictionary learning & matrix factorization problems [5, 51, 49]: In [5], convergence of a decentralized gradient tracking method for certain dictionary learning problems is established; a similar problem class is further investigated in [51], where generalization properties of a penalized consensus algorithm are studied, albeit without a convergence rate analysis. Despite their differences, these studies share a common premise of a full observation model, leading to a loss in the form of  $F(\bar{U}\bar{V}^{\top}) = \| Y - \bar{U}\bar{V}^{\top}\|^{2}$ . Here,  $Y = \bar{Z}^{\star} + N$  is the data matrix with  $N$  denoting noise. This full observation model contrasts with the matrix sensing model in (2), which is based on partial (noiseless) measurements. Lastly, [49] proposes a distributed Frank-Wolfe algorithm to address a low-rank matrix factorization problem, formulated as a trace (nuclear) norm convex minimization
problem (the nonconvex rank constraint is substituted by a nuclear norm constraint).
(ii) Distributed spectral methods [45, 19, 11, 12, 10, 44, 43]: Spectral methods have been established as effective strategies for obtaining reliable estimates of leading eigenvectors of a specified data matrix, as well as for providing a promising "warm start" for numerous iterative nonconvex matrix factorization algorithms [47]. Recent developments [19, 11, 12, 10, 44, 43] have successfully extended spectral methods—particularly principal component analysis—to decentralized contexts, achieving linear convergence rates, communications per iteration on the order of  $\mathcal{O}(dr)$ , and precise recovery up to a desired accuracy. A good tutorial on this subject can be found in [45]. These methods, in principle, can tackle the decentralized matrix sensing problem as formulated in this work through the estimation of the leading eigenspace of the surrogate matrix  $Y = \sum_{i=1}^{m} \sum_{j \in \mathcal{D}_i} y_j A_j$  where  $\sum_{j \in \mathcal{D}_i} y_j A_j$  is held by agent  $i$ . In fact, under suitable RIP on the linear mapping associated with  $Y$ , one has  $(1/N) \sum_{i=1}^{m} \sum_{j \in \mathcal{D}_i} y_j A_j \approx \bar{Z}^\star$  [34]. While such an approach can yield valuable insights about the ground truth  $\bar{Z}^\star$ , exact recovery of  $\bar{Z}^\star$  to arbitrary precision cannot be guaranteed [38]. This limitation starkly contrasts with the robust guarantees attainable by the gradient algorithm applied to the centralized matrix sensing problem with small random initialization (e.g., [34, 20]).
- Generic saddle-escaping decentralized algorithms: Under a sufficiently small RIP constant of the linear mapping associated with the signal model (1), the matrix sensing loss in (2) is shown to have no spurious local minima and all strict saddle points [1, 22]. Consequently, the task becomes escaping strict saddle points and computing second-order critical points. In the distributed optimization context, recent works studied the escape properties of several decentralized algorithms. Early works showed that certain decentralized schemes—the deterministic DGD [6, 18], the subgradient-flow [36], gradient-tracking algorithms [6], and primal-dual based methods [16, 24]—with random initialization, converge asymptotically towards a second-order critical point of a smooth function (subject to mild regularity conditions), with high probability. However, this near-certain convergence does not necessarily imply fast convergence. There exist non-pathological functions for which randomly initialized GD requires exponential time (in the ambient dimension) to escape saddle points [9]. It remains uncertain whether the inherent structure of the matrix sensing problem could yield superior convergence guarantees. Subsequent research has investigated the impact of decaying, additive noise perturbation on the agents' gradients of (stochastic) DGD in the Adapt-then-Combine (ATC) form [40, 39, 41]. While convergence to approximately second-order stationary points is assured within a polynomial number of iterations, the prerequisite that the loss has a globally Lipschitz gradient and Hessian matrix is not met by the matrix sensing loss in (2). This leaves decentralized saddle-escaping methods bereft of convergence rate and generalization guarantees when applied to (2).
1.2 Major contributions
We establish the first convergence rate and generalization guarantees of a decentralized gradient algorithm solving the matrix sensing problem via [2] over mesh networks. We borrow the following decentralized gradient descent [25], [46]: for each agent  $i = 1, \dots, m$ ,
$$
\bar {U} _ {i} ^ {t + 1} = \sum_ {j = 1} ^ {m} w _ {i j} \bar {U} _ {j} ^ {t + 1 / 2} \quad \text {a n d} \quad \bar {U} _ {i} ^ {t + 1 / 2} = \sum_ {j = 1} ^ {m} w _ {i j} \bar {U} _ {j} ^ {t} - \alpha \nabla f _ {i} \left(\sum_ {j = 1} ^ {m} w _ {i j} \bar {U} _ {j} ^ {t}\right), \tag {3}
$$
Here,  $\bar{U}_i^t$  is an estimate at iteration  $t$  of the optimization, common matrix  $\bar{U}$  in (2) held by agent  $i$ ;  $\alpha \in (0,1]$  is the stepsize; and  $w_{ij}$ 's are appropriately chosen nonnegative weights. We have  $w_{ii} > 0$ ,  $i = 1,\dots,m$ , and  $w_{ij} > 0$  if agents  $i$  and  $j$ ,  $i \neq j$ , can communicate; otherwise  $w_{ij} = 0$ . The algorithm employs two communication steps/iteration, aiming to enforce an agreement on both iterates  $\bar{U}_i^t$  and local gradients  $\nabla f_i$ . One could reduce the communication steps to one per iteration via a suitable variable change, resulting in the DGD-ATC form [46]. However, for the sake of clarity and ease of analysis, we opt to keep the form in (3), without any loss of generality.
- Guarantees: Our study presents a thorough statistical and convergence analysis of (3), yielding the following key insights. (i) Convergence to low-rank solutions: We demonstrate that, regardless the degree of overparametrization  $r$ , the iterates generated by (3) from a small random initialization converge towards low-rank solutions. We also provide an estimate of the worst-case iteration complexity. Improving results of the GD in a centralized setting (e.g., [34]), we specify an entire interval for algorithm termination, within which the generalization error is guaranteed to remain below the desired accuracy. This interval expands as the initialization becomes smaller. (ii) Two-phase convergence: Our analysis reveals a two-phase convergence behavior of (3). The initial
"spectral" phase sees the iterates mimic a theoretical centralized power method with full data access, while the subsequent "refinement" phase steers the trajectory towards the ground-truth solution. To the best of our knowledge, this provides the first evidence of a simplicity bias in a decentralized algorithm, aligning with the observed behavior of centralized GD [34]. (iii) Generalization error and communication complexity: The generalization error is shown to scale polynomially with the initialization size while the communication complexity scales logarithmically. Consequently, one can achieve arbitrarily small estimation errors with only a modest increase in communication and computation cost. (iv) RIP and network connectivity: Our findings hold under the conditions that the centralized measurement operator satisfies the standard RIP, and that network connectivity is sufficiently small. The former condition implies that our algorithm operates under the same sample complexity requested in the centralized setting. The latter is distinctive of the decentralized settings and shown to be unavoidable. (v) Almost performance invariance with the network size: We demonstrate that with an increase in the network size, the generalization error maintains its consistency whereas the communication cost grows logarithmically. Thus, the algorithm ensures effective error control with a marginal increase in communication overhead as the network expands.
- Convergence analysis: Although our analysis draws some insights from [34], the proof techniques employed diverge from those used therein. The decentralized nature of our setting introduces additional error terms, thereby making the analysis substantially more complex. Our methodology hinges on a newly introduced concept of RIP, termed in-network RIP. This concept harnesses the RIP of the measurement operator, much like the centralized GD, and intertwines it with the network's connectivity to derive favorable attributes of the new, overarching network-wide measurement operator. Furthermore, it reveals the interplay between sample complexity, network connectivity & topology, and communication complexity towards achieving statistical and computational guarantees over networks. Although we have defined the in-network RIP in the context of our specific algorithm dynamics, we posit that it possesses independent significance and could potentially pave the way for performance analysis of other distributed schemes.
In this section, we first list the notations used in the paper, and then provide details of our theoretical setup and necessary preliminary results.
For any positive integer  $m$ , we define  $[m] \triangleq \{1, \ldots, m\}$ ;  $1_m$  is the  $m$ -dimensional vector of all ones;  $I_d$  is the  $d \times d$  identity matrix;  $\otimes$  denotes the Kronecker product; and  $\mathrm{range}(M)$  (resp.  $\mathrm{rank}(M)$ ) denotes the range space (resp. rank) of the matrix  $M$ . When considering a matrix  $M \in \mathbb{R}^{md \times md}$ , partitioned into blocks of size  $d \times d$ , we will denote the block at the  $i$ -th row and  $j$ -th column as  $[M]_{ij}$ , for  $i, j \in [m]$ . Here,  $i$  and  $j$  indices increment by  $d$ , reflecting the size of the blocks.
We use  $\| \cdot \|$  to denote the Euclidean norm. When applied to matrices,  $\| \cdot \|$  is the operator norm induced by  $\| \cdot \|$ , and  $\| \cdot \|_F$  denotes the Frobenius norm of the argument matrix. We order the eigenvalues of any symmetric matrix  $M \in \mathbb{R}^{d \times d}$  in nonincreasing fashion, i.e.,  $\lambda_1(M) \geq \ldots \geq \lambda_d(M)$ . The singular values of (a rectangular) matrix  $M$  of rank  $r$  are denoted as  $\sigma_1(M) \geq \sigma_2(M) \geq \dots \sigma_r(M) > 0$ .
Truncated SVD: For any given matrix  $M \in \mathbb{R}^{d_1 \times d_2}$ , with rank  $r^\star > 0$ , we write the truncated SVD as  $M = V_M \Lambda_M Q_M^\top$ , where  $V_M \in \mathbb{R}^{d_1 \times r^\star}$  and  $Q_M \in \mathbb{R}^{d_2 \times r^\star}$  satisfy  $V_M^\top V_M = I_{d_1}$  and  $Q_M^\top Q_M = I_{d_2}$ , and  $\Lambda_M \in \mathbb{R}^{r^\star \times r^\star}$  is a diagonal matrix.
Augmented matrices: It is convenient to introduce the following "augmented" matrices suitable to rewrite the decentralized algorithm (3) in a concise block-stacked form:
$$
\mathcal {W} := W \otimes I _ {d}, \quad \mathcal {J} := (1 / m) 1 _ {m} 1 _ {m} ^ {\top} \otimes I _ {d}. \tag {4}
$$
where  $W \in \mathbb{R}^{m \times m}$  is the matrix of the gossip weights in (3), defined as  $[W]_{ij} = w_{ij}$ .
2.2 Basic definitions and assumptions
We develop our theoretical analysis under the following standard assumptions on the matrix sensing problem, algorithm parameters, and network connectivity.
- On the matrix sensing problem: Given the signal model (1), we decompose the ground-truth matrix as  $\bar{Z}^{\star} = \bar{X}\bar{X}^{\top}$ , for some  $\bar{X}\in \mathbb{R}^{d\times r^{\star}}$  (recall,  $r^\star$  is the rank of  $\bar{Z}^{\star}$ ).
Definition 1 (condition number). We define the condition number of  $\bar{X} \in \mathbb{R}^{d \times r^{\star}}$  as  $\kappa = \frac{\|\bar{X}\|}{\sigma_{r^{\star}}(\bar{X})}$ .
We associate to the signal model (1) the measurement linear operator  $\mathbb{R}^{d\times d}\ni \bar{Z}\mapsto \bar{\mathcal{A}} (\bar{Z})\in \mathbb{R}^N$  and its adjoint  $\mathbb{R}^N\ni w\mapsto \bar{\mathcal{A}}^* (w)\in \mathbb{R}^{d\times d}$ , defined as
$$
\bar {\mathcal {A}} (\bar {Z}) := \frac {1}{\sqrt {N}} \left(\left\langle A _ {j}, \bar {Z} \right\rangle\right) _ {j \in \mathcal {D} _ {i}, i \in [ m ]} \quad \text {a n d} \quad \bar {\mathcal {A}} ^ {*} (w) = \frac {1}{\sqrt {N}} \sum_ {i = 1} ^ {m} \sum_ {j \in \mathcal {D} _ {i}} w _ {j} A _ {j}. \tag {5}
$$
A standard assumption in the matrix sensing literature is requiring the RIP for the operator  $\bar{A}$ .
Definition 2 (RIP). The measurement operator  $\bar{\mathcal{A}}:\mathbb{R}^{d\times d}\to \mathbb{R}^N$  satisfies the  $(\delta ,r)$ -RIP condition if
$$
(1 - \delta) \| \bar {Z} \| _ {F} ^ {2} \leq \| \bar {\mathcal {A}} (\bar {Z}) \| ^ {2} \leq (1 + \delta) \| \bar {Z} \| _ {F} ^ {2}, \tag {6}
$$
for all matrices  $\bar{Z} \in \mathbb{R}^{d \times d}$  with  $\text{rank}(\bar{Z}) \leq r$ .
The RIP condition is the key to ensure the ground truth  $\bar{Z}^{\star}$  to be recoverable with partial observations. In fact, an important consequence of RIP is that  $\bar{\mathcal{A}}^{*}\bar{\mathcal{A}} (\bar{Z}) = (1 / N)\sum_{i = 1}^{m}\sum_{j\in \mathcal{D}_{j}}\langle A_{j},\bar{Z}\rangle A_{j}\approx \bar{Z}$ , for all  $\bar{Z}$  low-rank (see, e.g., [34]). Notice that when all entries of the matrices  $A_{j}$  are drawn i.i.d. with distribution  $\mathcal{N}(0,1)$  on the off-diagonal entries and distribution  $\mathcal{N}(0,1 / \sqrt{2})$  on the diagonal, the  $(\delta ,r)$ -RIP holds with high probability, if the number of observations  $N = \Omega (dr / \delta^2)$  (e.g., [3]).
- Network setup and gossip matrices: Agents are embedded in a communication network, modelled as an undirected graph  $\mathcal{G} = \{\mathcal{V},\mathcal{E}\}$ , where the vertices  $\mathcal{V} = [m]\triangleq \{1,\dots,m\}$  correspond to the agents and  $\mathcal{E}$  is the set of edges of the graph;  $(i,j)\in \mathcal{E}$  if and only if there is a communication link between agents  $i$  and  $j$ . We study the decentralized algorithm (3) using gossip weight matrices satisfying the following standard assumption in the literature of distributed optimization.
Assumption 1.  $W = [w_{ij}]_{ij=1}^{m}$  satisfies: (i)  $w_{ij} > 0$ , if  $(i,j) \in \mathcal{E}$ ; otherwise  $w_{ij} = 0$ ; furthermore,  $w_{ii} > 0$ , for all  $i \in [m]$ ; (ii)  $W = W^{\top}$  and  $W1 = 1$  (stochastic); (iii)  $W$  is positive semidefinite; and (iv) there holds  $\rho \triangleq \|W - 1_m 1_m^\top / m\|_2 < 1$ .
Assumption  $\mathbb{I}$  is standard in the literature of distributed algorithms and is satisfied by several weight matrices; see, e.g., [26]. Note that  $\rho < 1$  holds true by construction for connected graphs. Roughly speaking,  $\rho$  measures how fast the network mixes information; the smaller  $\rho$ , the faster the mixing.
2.3 Augmented mapping and in-network RIP
Fundamental to our analysis is a novel RIP-like property associated with an augmented linear mapping tied to the decentralized algorithm (3). This new property effectively captures the admissible "degree of distortion" on the signal information  $\bar{Z}^{\star}$ , taking into account both the partial observability of  $\bar{Z}^{\star}$  as postulated in (1) (through the measurement operator  $\bar{A}$  defined in (5)) and the intricacies of the in-network optimization process (regulated by the network operator  $\mathcal{W}$ , as defined in (4)).
We begin rewriting the decentralized algorithm (3) in a compact form. To do so, we define the following quantities: (i) the stacked block matrices  $U^{t} \in \mathbb{R}^{md \times r}$  and  $Z^{\star} \in \mathbb{R}^{md \times md}$ :
$$
U ^ {t} := \left(\bar {U} _ {i} ^ {t}\right) _ {i \in [ m ]} \quad \text {a n d} \quad Z ^ {\star} := 1 _ {m} 1 _ {m} ^ {\top} \otimes \bar {Z} ^ {\star}, \tag {7}
$$
respectively; (ii) the augmented mapping  $\mathcal{A}:\mathbb{R}^{md\times md}\to \mathbb{R}^N$  and its adjoint  $\mathcal{A}^*:\mathbb{R}^N\to \mathbb{R}^{md\times md}$
$$
[ \mathcal {A} (Z) ] _ {\ell} \triangleq \frac {1}{\sqrt {m n}} \left\langle m \left(w _ {\mathcal {V} (\ell)} w _ {\mathcal {V} (\ell)} ^ {\top}\right) \otimes A _ {\ell}, Z \right\rangle , \quad \mathcal {A} ^ {*} (q) \triangleq \sum_ {\ell = 1} ^ {N} \frac {q _ {\ell}}{\sqrt {m n}} \left(m w _ {\mathcal {V} (\ell)} w _ {\mathcal {V} (\ell)} ^ {\top}\right) \otimes A _ {\ell}, \tag {8}
$$
where  $\mathcal{V}(\ell): \cup_{i=1}^{m} \mathcal{D}_i \to \mathcal{V}$  returns the index  $i$  such that  $\ell \in \mathcal{D}_i$ ,  $\ell \in [N]$ . Note that these operators depend on both data measures (via  $\bar{\mathcal{A}}$ ) and the network (via  $\mathcal{W}$ ). Using the definitions in (4), (7), and (8), it is not difficult to check that (3) can be rewritten equivalently as: