Optimal Oblivious Subspace Embeddings with Near-Optimal Sparsity

Our subspace embedding guarantees are achieved by a family of OSEs which have a fixed number of non-zero entries in each column, a key property that was also required of sparse OSE distributions called OSNAP described by Nelson and Nguyen [26]. As we explain later, our analysis techniques apply to other natural families of sparse embedding distributions, including those with i.i.d. entries [1], however the OSNAP-style construction is crucial for achieving the near-optimal sparsity $s=\tilde{O}(1/\epsilon )$.

In our construction of the $m\times n$ OSE matrix $\mathrm{\Pi }$, we start by defining an unscaled version of the matrix, called $S$, which has entries in $\{-1,0,1\}$. We then scale $S$ to appropriately normalize the entry-wise variances, obtaining $\mathrm{\Pi }$. Concretely, we wish to obtain an $m\times n$ sparse random matrix $S$ which has exactly $s$ non-zero $\pm 1$ entries in each column. Assume $s$ exactly divides $m$. Then we can divide each column of $S$ into $s$ subcolumns and randomly populate one entry in each subcolumn by a Rademacher random variable (see Figure 1). We call this family of distributions (unscaled) OSNAP, carrying over Nelson and Nguyen’s terminology (technically, their definition is somewhat broader than ours).

Each non-zero entry in the matrix $S$ can be identified by a tuple $(l,\gamma )\in [n]\times [s]$ where $l$ identifies the column of the non-zero entry and $\gamma$ is the index of the entry in that column. Thus the $(l,\gamma )\text{<sup class="ltx\_sup">th</sup>}$ non-zero entry in $S$ is located in column $l$ and row ${\mu }_{(l,\gamma )}$, where ${\mu }_{(l,\gamma )}$ is a uniformly chosen integer from the interval $[(m/s)(\gamma -1)+1:(m/s)\gamma ]$. For example, the $(1,1)\text{<sup class="ltx\_sup">th</sup>}$ non-zero entry in $S$ is located in column $1$ and some row in the interval $[1:m/s]$. An $m\times n$ matrix with a non-zero entry in column $l$ and row ${\mu }_{(l,\gamma )}$ is given by $e_{{\mu }_{(l,\gamma )}}{e}_l^T$, where $e_{{\mu }_{(l,\gamma )}}$ and $e_l$ represent standard basis vectors in ${ℝ}^m$ and ${ℝ}^n$ respectively, and for $S$ we wish to place a random sign ${\xi }_{(l,\gamma )}$ at this position. This motivates our formal definition for OSNAP,

Thus, each column of the OSNAP matrix $\mathrm{\Pi }$ has $s=pm$ many non-zero entries, and the sparsity level can be varied by setting the parameter $p\in [0,1]$ appropriately. With the distribution formally defined, we now provide the full statement of our subspace embedding guarantee for OSNAP,

This is a direct improvement over Theorem 1.2 of [6], which requires sparsity $s\ge c{\log }^4(d/\delta )/{\epsilon }^6$ where $c$ is an absolute constant, with the same condition on $m$. The primary gain lies in the polynomial dependence on $1/\epsilon$, but we note that our result also achieves a better logarithmic dependence on $d$, which means that an improvement is obtained even for $\epsilon =\mathrm{\Theta }(1)$.

Our techniques can be used to obtain a similar result for a simple OSE model with i.i.d. sparse Rademacher entries [1], which was also considered by [6]. However, in this case, we need an additional requirement of $s=pm\ge c\log (\frac{d}{\epsilon \delta })/{\epsilon }^2$ for the sparsity (see Section 9 of the technical report for details; this is again a direct improvement over a result of [6]).

Our proof techniques for Theorem 7 are based on the moment method, and thus, they naturally imply the following slightly stronger moment-based characterization of an oblivious subspace embedding, which was proposed by [13] as an extension of the corresponding moment-based characterization of a Johnson-Lindenstrauss embedding [22].

Note that a simple application of Markov’s inequality recovers the guarantee in Definition 1 from the $(\epsilon ,\delta ,d,\mathrm{\ell })$-OSE moments property with any $\mathrm{\ell }\ge 1$. Moreover, [13] showed that this moment-based OSE characterization implies several other desirable guarantees of embedding matrices in the context of approximate matrix multiplication, generalized regression and low-rank approximation.

As an immediate consequence of our analysis, we obtain the following OSE moment guarantee for the OSNAP distribution.

In a related problem, we seek to embed a subspace given by a fixed $U\in {ℝ}^{n\times d}$, with information about the squared row norms of $U$ being used to define the distribution of non-zero entries in $\mathrm{\Pi }$. Such distributions for $\mathrm{\Pi }$ are called non-oblivious (a.k.a. data-aware) subspace embeddings. Previous work [6] has dealt with one such family of distributions termed LESS embeddings [17, 16, 15], showing that they require $\tilde{O}(1/{\epsilon }^4)$ non-zero entries per row of $\mathrm{\Pi }$ to obtain an $\epsilon$-embedding guarantee. Since the embedding matrix is very wide, this leads to a much sparser embedding (sparser than any OSE) that can be applied in time sublinear in the input size, leading to fast subspace embedding algorithms.

In this work, we show that our new techniques also extend to LESS embeddings and enable us to prove sharper sparsity estimates than [6]. To fully leverage our approach, we define a new type of sparse embedding (LESS-IC), which can be viewed as a cross between CountSketch and LESS. Here, IC stands for independent columns. At a high level, the CountSketch part ensures that we can use our decoupling method to achieve optimal dependence on $1/\epsilon$, while the LESS part enables adaptivity to a fixed subspace.

Specifically, a LESS-IC embedding matrix $\mathrm{\Pi }$ has a fixed number of non-zero entries in each column, chosen so that it is proportional to the leverage score (i.e. the squared row norm) of the corresponding row of $U$. This is achieved by modifying the OSNAP distribution such that the number of subcolumns is no longer the same in each column. For columns corresponding to very small leverage scores, we only have one “subcolumn”. Thus, each column has at least one non-zero entry. This means that the cost of applying LESS-IC to an $n\times d$ matrix $A$ can no longer be sublinear (like it can in the existing LESS embedding constructions), but rather has a fixed linear term of $O(\mathrm{nnz}(A))$, plus an additional sublinear term. Given that the preprocessing step of approximating the leverage scores has to take at least $\mathrm{nnz}(A)$ time, the linear term in the cost of applying LESS-IC is negligible.

To generate an embedding matrix with the LESS-IC distribution, it suffices to have a good enough approximation for the leverage scores of the matrix $U$, in the following sense.

To see how approximate leverage scores determine the distribution of entries in the LESS-IC distribution, let us first consider a simpler distribution, LESS-IE from [6], based on a similar construction first proposed by [17]. Here, we once again start by defining an unscaled matrix $S$, which is then normalized to obtain the subspace embedding matrix $\mathrm{\Pi }$.

In the LESS-IE model, we have ${\beta }_{1}pm{z}_j$ many non-zero entries in column $j$ in expectation. However, to achieve $1/\epsilon$ dependency of the sparsity, we need to have exactly ${\beta }_{1}pm{z}_j$ many non-zero entries in the column in the LESS-IC model to fully take advantage of the error cancellation that occurs in our decoupling argument (See Section 7.2 and Section 9.1 of the technical report). Though these sections deal with oblivious subspace embeddings, the same arguments still apply in the LESS case). This is done by modifying the OSNAP construction so that the size (and consequently, the number) of subcolumns is different across columns.

Notice that to have ${\beta }_{1}pm{z}_j$ many non-zero entries in column $j$, we would need ${\beta }_{1}pm{z}_j$ many subcolumns in column $j$ each with one non-zero entry in a random position. This means that the size of each subcolumn needs to be $m/({\beta }_{1}pm{z}_j)=1/({\beta }_{1}p{z}_j)$. However, since $1/({\beta }_{1}p{z}_j)$ may not be an integer, we consider subcolumns of size $b_j:=\max \{\lfloor 1/({\beta }_{1}p{z}_j)\rfloor ,1\}$.

In column $j$, we stack subcolumns of size $b_j$ until we fill up all the rows up to $m$. Let $s_j$ be the smallest number of subcolumns to do this. Then, it may happen that the row indices of the bottom-most subcolumn exceed $m$. For example, consider the distribution on the first column of $\mathrm{\Pi }$ when $m=70$, and $b_1=15$. In this case $s_1=5$, so we can stack four subcolumns of size 15 and the $5\text{<sup class="ltx\_sup">th</sup>}$ subcolumn only spans row indices $[61:70]$. In each subcolumn, we randomly choose a row to place a non-zero entry, which would be a Rademacher random variable. (See Figure 2). The non-zero entries are appropriately scaled so that all entries of the matrix have the same variance (See Section 8 of the technical report for the full definition).

For the LESS-IC distribution, we show the following subspace embedding guarantee. The structure of the proof is similar to the case of OSNAP, and only the specific expressions change due to the different distribution.

Here, we briefly outline how our LESS-IC embedding yields a fast subspace embedding construction to recover the time complexity claimed in Theorem 4. This follows analogously to the construction from Theorem 1.6 of [6], and our improvement in the dependence on $1/\epsilon$ compared to their result (from $1/{\epsilon }^6$ to $1/\epsilon$) stems from the improved sparsity of our LESS-IC embedding.

The key preprocessing step for applying the LESS-IC embedding is approximating the leverage scores of the matrix $A$. Using Lemma 5.1 in [6] (adapted from Lemma 7.2 in [8]), we can construct coarse approximations of all leverage scores so that ${\beta }_1=O(n^{\gamma })$ and ${\beta }_2=O(1)$ in time $O({\gamma }^{-1}(\mathrm{nnz}(A)+d^2)+d^{\omega })$. Applying LESS-IC (Theorem 15) with these leverage scores and parameters ${\beta }_1,{\beta }_2$, computing $\mathrm{\Pi }A$ takes $O(\mathrm{nnz}(A)+n^{\gamma }{d}^2{\log }^3(d/\epsilon \delta )/\epsilon )$, where $\mathrm{nnz}(A)$ comes from the fact that every column of $\mathrm{\Pi }$ has at least one non-zero, while the second term accounts for the additional $O(\beta d{\log }^3(d/\epsilon \delta )/\epsilon )$ non-zeros.

Thus, if $d\ge n^c$ for, say, $c=0.1$, then we conclude the claim by appropriately scaling $\gamma$ by a constant factor. Now, suppose otherwise. First, note that without loss of generality we can assume that (through scaling the time complexity by a constant factor), $\mathrm{nnz}(A)\ge n$ (by removing empty rows) and $\epsilon \ge \sqrt{d/n}$ (because otherwise $m\ge n$ and we could use $\tilde{A}=A$). Thus, under our assumption that , we have $n^{\gamma }{d}^2/\epsilon \le n^{0.5+\gamma +2c}\le n^{0.8}\ll \mathrm{nnz}(A)$, and the time complexity is dominated by the $O({\gamma }^{-1}\mathrm{nnz}(A))$ term.

Finally, we note that Corollary 5 follows simply by constructing a subspace embedding $\mathrm{\Pi }$ via Theorem 4 with respect to matrix $[A\mid b]$, and computing $\tilde{A}=\mathrm{\Pi }A$, $\tilde{b}=\mathrm{\Pi }b$. The proof of the claim is identical to the proof of Theorem 1.8 in [6]. Our improvement comes directly from the faster runtime of our subspace embedding construction.

Section 4 provides a high level overview of the ideas used in the proofs of our main results, Theorem 7 and Theorem 15. Section 5 provides a sketch of the proof of Theorem 7, listing the main technical steps, leaving the full proof with all technical details to Section 7 of the technical report. The proof of Theorem 15 follows similarly and is covered in Section 8 of the technical report. The subspace embedding guarantee for a sparse matrix with independent entries is proved in Section 9 of the technical report.

The following notation and terminology will be used in the paper. The notation $[n]$ is used for the set $\{1,2,\mathrm{\dots },n\}$ and the notation $\mathrm{P}([n])$ denotes the set of all partitions of $[n]$. Also, for two integers $a$ and $b$ with $a\le b$, we use the notation $[a:b]$ for the set $\{k\in \mathbb{Z}:a\le k\le b\}$. For $x\in ℝ$, we use the notation $\lfloor x\rfloor$ to denote the greatest integer less than or equal to $x$ and $\lceil x\rceil$ to denote the least integer greater than or equal to $x$. In ${ℝ}^n$ (or ${ℝ}^m$ or ${ℝ}^d$), the $l$th coordinate vector is denoted by $e_l$. All matrices considered in this paper are real valued and the space of $m\times n$ matrices with real valued entries is denoted by $M_{m\times n}(ℝ)$. Also, for a matrix $X\in M_{d\times d}(ℝ)$, the notation $Tr(X)$ denotes the trace of the matrix $X$, and $tr(X)=\frac{1}{d}Tr(X)$ denotes the normalized trace. We write the operator norm of a matrix $X$ as $\parallel X\parallel$, and it is also denoted by ${\parallel X\parallel }_{op}$ in some places where other norms appear for clarity. The spectrum of a matrix $X$ is denoted by $\mathrm{spec}(X)$. The standard probability measure is denoted by $\mathbb{P}$, and the symbol $𝔼$ means taking the expectation with respect to this standard probability measure. To simplify the notation, we follow the convention from [4] and use the notation $𝔼{[X]}^{\alpha }$ for ${(𝔼(X))}^{\alpha }$, i.e., when a functional is followed by square brackets, it is applied before any other operations. The covariance of two random variables $X$ and $Y$ is denoted by $\mathrm{Cov}(X,Y)$. The standard $L_q$ norm of a random variable $\xi$ is denoted by ${\parallel \xi \parallel }_q$, for $1\le q\le \mathrm{\infty }$. Throughout the paper, the symbols $c_1,c_2,\mathrm{\dots }$, and $Const,Cons{t}^{\prime },\mathrm{\dots }$ denote absolute constants.