This html page is created by scribe. If you find any mistakes, please feel free to drop a mail here.
Efficient AI · Lecture notes
JL Lemma, Subspace Embeddings, and CountSketch
Random projections that preserve distances, how they speed up
least-squares regression, how a net argument turns them into
$\ell_2$-subspace embeddings, how sparse such maps can be, and the
CountSketch matrix, which uses just one non-zero per column.
1. The Johnson–Lindenstrauss lemma (existence)
Lemma — JL lemma
Given $n$ points $\{x_1,\dots,x_n\}\subset\mathbb R^d$ and
$0<\varepsilon<1$, there exists a linear map $f(x)=Ax$ with
$A\in\mathbb R^{r\times d}$ such that for all $i\ne j$
\[
(1-\varepsilon)
\;\le\;
\frac{\|f(x_i)-f(x_j)\|_2^2}{\|x_i-x_j\|_2^2}
\;\le\;
(1+\varepsilon),
\qquad\text{provided}\quad
r\;\ge\;\frac{c\log n}{\varepsilon^2}
\]
for an absolute constant $c$.
Here $A$ is a random map that reduces the dimension from $d$ to
$r$. Note that $r$ depends only on the number of points $n$, not on
the original dimension $d$.
Proof
Construction
Pick each entry independently:
\[
A_{kl}=\frac{1}{\sqrt r}\,g_{kl},
\qquad
g_{kl}\overset{\text{i.i.d.}}{\sim}\mathcal N(0,1).
\]
One pair
Fix a pair and let $u=u_{ij}=x_i-x_j$. For each coordinate
$k\in[r]$ of $Au$,
\[
(Au)_k
=
\frac{1}{\sqrt r}\sum_l g_{kl}\,u_l
=
\frac{\|u\|_2}{\sqrt r}\;
\underbrace{\mathcal N(0,1)}_{y_k},
\]
because a sum of independent Gaussians is Gaussian:
$\sum_l g_{kl}u_l\sim\mathcal N(0,\|u\|_2^2)$. The $y_k$ are
independent, since different rows of $A$ are independent. Therefore
\[
\|Au\|_2^2
=
\sum_k (Au)_k^2
=
\frac{\|u\|_2^2}{r}\sum_{k=1}^{r}y_k^2 .
\]
So $\|Au\|_2^2/\|u\|_2^2$ is a $\chi^2_r$ random variable divided by
$r$, and its mean is exactly $1$.
Concentration
For $0<\varepsilon<1$,
\[
\Pr\!\left(
\left|\frac{\sum_k y_k^2-r}{r}\right|>\varepsilon
\right)
\;\le\;
2\,e^{-r\varepsilon^2/4}.
\]
Equivalently,
\[
\Pr\!\left(
\left|\frac{\|Au\|_2^2}{\|u\|_2^2}-1\right|>\varepsilon
\right)
\;\le\;
2\,e^{-r\varepsilon^2/4}
\;<\;\delta .
\]
This gives an $(\varepsilon,\delta)$ guarantee for one fixed pair. For
this, it suffices that
\[
-\frac{r\varepsilon^2}{4}<\log\frac{\delta}{2}
\quad\Longleftrightarrow\quad
r>\frac{4\log(2/\delta)}{\varepsilon^2}.
\]
All pairs (union bound)
There are $\binom n2$ pairs, so
\[
\Pr\!\left(
\exists\,i,j:\
\left|\frac{\|Au_{ij}\|_2^2}{\|u_{ij}\|_2^2}-1\right|>\varepsilon
\right)
\;\le\;
2\binom n2\,e^{-r\varepsilon^2/4}
\;<\;\delta,
\]
which holds whenever
\[
r\;>\;\frac{4}{\varepsilon^2}\log\!\left(\frac{n(n-1)}{\delta}\right).
\]
For constant $\delta$ this is $r=O(\log n/\varepsilon^2)$. Since the
failure probability is below $1$, a good $A$ exists.
2. Application to linear regression
Let $A\in\mathbb R^{n\times d}$ with $n\gg d$, and $b\in\mathbb R^n$.
The least-squares solution is
\[
x^*=\arg\min_x\|Ax-b\|_2 .
\]
Goal.
Find $x^+$ such that
$\|Ax^+-b\|_2\le(1+\varepsilon)\,\|Ax^*-b\|_2$.
Sketch and solve: let
$x^+=\arg\min_x\|SAx-Sb\|_2$, the solution of the smaller problem,
where $S$ has $r$ rows ($r$ is the reduced dimension). Then
\[
\|SAx^+-Sb\|_2
\;\le\;
\|SAx^*-Sb\|_2
\;=\;
(1\pm\varepsilon)\,\|Ax^*-b\|_2
\qquad\text{w.p. }\ge 1-\delta .
\]
3. JL transforms
Definition — $\mathrm{JLT}(\varepsilon,\delta,f)$
A random matrix $S\in\mathbb R^{r\times n}$, drawn from a distribution
$\Pi$, is a
Johnson–Lindenstrauss transform with
parameters $\varepsilon,\delta,f$, where $f$ is the number of points,
if, for any subset $V\subset\mathbb R^n$ with $|V|=f$, with
probability at least $1-\delta$,
\[
\forall\, v_i,v_j\in V:\qquad
\langle Sv_i,Sv_j\rangle
=
\langle v_i,v_j\rangle
\pm
\varepsilon\,\|v_i\|_2\,\|v_j\|_2 .
\]
That is,
inner products are preserved up to $\varepsilon$.
From $n$ points to a whole subspace
The JL lemma preserves $n$ fixed points. If the data is a matrix
$A\in\mathbb R^{n\times d}$ with column space $\mathcal C(A)$, we
instead want to preserve every vector in $\mathcal C(A)$, which is
infinitely many points:
\[
\|SAx\|_2^2=(1\pm\varepsilon)\,\|Ax\|_2^2
\qquad\forall\,x\in\mathbb R^d .
\]
4. $\ell_2$-subspace embeddings
Definition — $\ell_2$-subspace embedding
Let $A\in\mathbb R^{n\times d}$. A matrix $S\in\mathbb R^{r\times n}$ is
a $(1\pm\varepsilon)$
$\ell_2$-subspace embedding for
the column space $\mathcal C(A)$ if
\[
\forall\, x\in\mathbb R^d:\qquad
\|SAx\|_2^2=(1\pm\varepsilon)\,\|Ax\|_2^2 .
\]
Whether the definition uses $\|\cdot\|_2$ or $\|\cdot\|_2^2$ only
changes $\varepsilon$ by a constant factor, since
$\sqrt{1\pm\varepsilon}=1\pm O(\varepsilon)$.
Reducing to a finite set: the net argument
Definition — $\gamma$-net
Let $\mathcal C$ be the set of unit vectors in $\mathcal C(A)$. A set
$N$ is a
$\gamma$-net of $\mathcal C$ if
\[
\forall\, y\in\mathcal C\ \ \exists\, x\in N:\qquad
\|x-y\|_2<\gamma .
\]
Fact — size of a net
There is a $\gamma$-net with
\[
|N|\;\le\;\left(1+\frac{2}{\gamma}\right)^d .
\]
Let $N\subset\mathcal C(A)$ be such a net. If $S$ preserves the inner
products of all pairs in the net,
\[
\forall\, w_i,w_j\in N:\qquad
\langle Sw_i,Sw_j\rangle
=
\langle w_i,w_j\rangle
\pm
\varepsilon\,\|w_i\|_2\,\|w_j\|_2 ,
\]
then the guarantee extends from the finite net to the whole subspace:
\[
\|SAx\|_2^2=(1\pm O(\varepsilon))\,\|Ax\|_2^2
\qquad\forall\,x\in\mathbb R^d .
\]
Choosing $\gamma=\tfrac12$.
Then $|N|\le 5^d$, and applying a
$\mathrm{JLT}(\varepsilon,\delta,f=5^d)$ to $N$ gives
$\|SAx\|_2^2=(1\pm O(\varepsilon))\|Ax\|_2^2$ for all
$x\in\mathbb R^d$.
5. JLTs give $\ell_2$-subspace embeddings
Theorem
A $\mathrm{JLT}(\varepsilon,\delta,f=5^d)$ is an $\ell_2$-subspace
embedding for the column space of any fixed $n\times d$ matrix $A$,
i.e. for a $d$-dimensional subspace of $\mathbb R^n$.
Constructions of $S\in\mathbb R^{r\times n}$
Gaussian
\[
S=\tfrac{1}{\sqrt r}\,R
\]
$R_{ij}\overset{\text{i.i.d.}}{\sim}\mathcal N(0,1)$
Random signs
\[
S=\tfrac{1}{\sqrt r}\,R
\]
$R_{ij}=\pm1$, each w.p. $\tfrac12$
Sparse · Achlioptas
\[
S=\sqrt{\tfrac{3}{r}}\,R
\]
$R_{ij}=0$ w.p. $\tfrac23$,
$+1$ w.p. $\tfrac16$, $-1$ w.p. $\tfrac16$
In each case $S$ is $r\times n$ with
\[
r=\Theta\!\left(\frac{\log(f/\delta)}{\varepsilon^2}\right),
\qquad f=5^d,
\qquad\text{so}\qquad
r=\Theta\!\left(\frac{d+\log(1/\delta)}{\varepsilon^2}\right).
\]
Sparsity (non-zeros per column)
Computing $SA$ costs time proportional to the number of non-zeros in
$S$ that touch $A$, so we want each column of $S$ to be as sparse as
possible. How sparse can a JLT be?
Upper bound — Kane–Nelson
There is a construction with
\[
s=O\!\left(\frac1\varepsilon\log\frac{f}{\delta}\right)
\]
non-zero entries per column.
Lower bound — Nelson–Nguyen
Any JLT needs at least this many non-zeros per column:
\[
s=\Omega\!\left(
\frac{\frac1\varepsilon\log\!\left(\frac f\delta\right)}
{\log\!\left(\frac1\varepsilon\right)}
\right),
\qquad f=5^d.
\]
Consequence.
A JLT-based subspace embedding needs about $d/\varepsilon$ non-zeros
per column (up to the $\log(1/\varepsilon)$ factor). CountSketch gets
around this lower bound: it is not a JLT for arbitrary sets of $5^d$
points, only for points lying in one $d$-dimensional subspace.
6. The CountSketch matrix
Write $[k]:=\{0,1,\dots,k-1\}$. Choose two random functions:
\[
h:[n]\to[r],
\qquad\qquad
g:[n]\to\{-1,+1\}.
\]
The matrix $S\in\mathbb R^{r\times n}$ has exactly one non-zero
entry per column. Column $i$ has the entry $g(i)=\pm1$ in row
$h(i)$, and zeros everywhere else:
\[
S_{h(i),\,i}=g(i),
\qquad
S_{j,\,i}=0\ \text{ for } j\ne h(i).
\]
Because each column has one non-zero, $SA$ can be computed in
$O(\mathrm{nnz}(A))$ time. Each row of $A$ is multiplied by $\pm1$ and
added into one of the $r$ output rows.
Theorem — CountSketch is a subspace embedding
CountSketch with
\[
r=O\!\left(\frac{d^2}{\varepsilon^2}\log\!\left(\frac{d}{\varepsilon\delta}\right)\right)
\]
rows is a $(1\pm\varepsilon)$ $\ell_2$-subspace embedding for any fixed
$A\in\mathbb R^{n\times d}$ with probability at least $1-\delta$.
Theorem — Limited independence suffices
If $h$ is
2-wise independent and $g$ is
4-wise independent, then
\[
r=O\!\left(\frac{d^2}{\delta\,\varepsilon^2}\right)
\]
rows suffice.
Summary
| Sketch |
Non-zeros per column |
Rows $r$ for a subspace embedding |
| Gaussian / random signs |
$r$ (dense) |
$\Theta\!\left(\frac{d+\log(1/\delta)}{\varepsilon^2}\right)$ |
| Achlioptas |
$\approx r/3$ in expectation |
$\Theta\!\left(\frac{d+\log(1/\delta)}{\varepsilon^2}\right)$ |
| Sparsest JLT |
$\approx d/\varepsilon$ (tight up to $\log(1/\varepsilon)$) |
$\Theta\!\left(\frac{d+\log(1/\delta)}{\varepsilon^2}\right)$ |
| CountSketch |
$1$ |
$O\!\left(\frac{d^2}{\delta\varepsilon^2}\right)$ |
Takeaway.
CountSketch gives up a polynomial factor in the number of rows
($d^2$ instead of $d$) in return for the cheapest possible
application cost: one non-zero per column, so $SA$ takes
$O(\mathrm{nnz}(A))$ time.
Reference:
Woodruff, Sketching as a Tool for Numerical Linear Algebra, Sections 1 and 2.1, and Lemma 18.