随机矩阵乘法：外积抽样与 importance weighting

把矩阵乘法写成外积之和

现在回到矩阵乘法。设

$$\mathbf{A}\in\mathbb{R}^{m\times n}, \qquad \mathbf{B}\in\mathbb{R}^{n\times p}$$

把 $\mathbf{A}$ 按列写成

$$\mathbf{A} = \begin{bmatrix} \mathbf{a}_1 & \mathbf{a}_2 & \cdots & \mathbf{a}_n \end{bmatrix}$$

把 $\mathbf{B}$ 按行写成

$$\mathbf{B} = \begin{bmatrix} \mathbf{b}_1^\top \\ \mathbf{b}_2^\top \\ \vdots \\ \mathbf{b}_n^\top \end{bmatrix}$$

那么矩阵乘法可以拆成

$$\boxed{ \mathbf{A}\mathbf{B} = \sum_{j=1}^n \mathbf{a}_j\mathbf{b}_j^\top }$$

每一项

$$\mathbf{a}_j\mathbf{b}_j^\top$$

都是一个 rank one matrix

记

$$\mathbf{M}_j = \mathbf{a}_j\mathbf{b}_j^\top$$

于是目标就是

$$\mathbf{A}\mathbf{B} = \sum_{j=1}^n \mathbf{M}_j$$

如果 $n$ 很大，完整求和可能很贵。随机矩阵乘法的想法是：不计算所有 $\mathbf{M}_j$，而是随机抽一些项，用它们构造一个对 $\mathbf{A}\mathbf{B}$ 的估计

把统计概念放回矩阵乘法

到这里，前面那些统计概念可以一一对应到矩阵乘法问题里

统计概念	在随机矩阵乘法中的对象
真实目标 $\theta$	$\mathbf{A}\mathbf{B}$
一次随机结果	抽到某个下标 $J=j$
抽样概率分布	抽到某个下标的可能性有多大，$\mathbb{P}(J=j)=p_j$
单次随机估计量	$\mathbf{X}=\frac{1}{p_J}\mathbf{M}_J$
多次抽样后的估计量	$\widehat{\mathbf{A}\mathbf{B}}=\frac{1}{s}\sum_{\ell=1}^s \frac{1}{p_{j_\ell}}\mathbf{M}_{j_\ell}$
抽样分布	重复运行整个随机算法时，$\widehat{\mathbf{A}\mathbf{B}}$ 的分布
无偏	$\mathbb{E}[\widehat{\mathbf{A}\mathbf{B}}]=\mathbf{A}\mathbf{B}$
有偏	$\mathbb{E}[\widehat{\mathbf{A}\mathbf{B}}]\neq\mathbf{A}\mathbf{B}$

这个表里最重要的是区分两层随机性

第一层是一次抽样抽到哪个下标 $J$

$$\mathbb{P}(J=j)=p_j$$

第二层是整个估计量 $\widehat{\mathbf{A}\mathbf{B}}$ 在重复运行算法时如何波动

$$\widehat{\mathbf{A}\mathbf{B}} \quad \text{本身也是随机的}$$

无偏性讨论的是第二层：估计量的长期平均是否等于真实目标，而不是问每次抽到的下标是不是均匀，也不是问单次结果是不是已经等于 $\mathbf{A}\mathbf{B}$

为什么抽到某项 M_j 以后要乘 1/p_j

设随机变量 $J$ 表示本次抽到的下标，我们一共有 $n$ 个 $\mathbf{M}_j$ 可选

$$\mathbb{P}(J=j)=p_j$$

所以

$$\sum_{j=1}^n p_j=1$$

注意，这里完全不要求 $p_j$ 相等。大的项可以被更频繁地抽到，小的项可以少抽一些

如果抽到 $j$ 之后直接输出

$$\mathbf{X}=\mathbf{M}_j$$

那么期望是

$$\mathbb{E}[\mathbf{X}] = \sum_{j=1}^n p_j\mathbf{M}_j$$

这里求和到 $n$，是因为 $\mathbf{X}$ 是 单次抽样 得到的随机矩阵。单次抽样时，可能结果有 $n$ 种：

$$\mathbf{M}_1,\mathbf{M}_2,\dots,\mathbf{M}_n$$

其中抽到 $\mathbf{M}_j$ 的概率是 $p_j$。所以按照期望定义，要对所有可能结果 $j=1,\dots,n$ 做概率加权求和

这一般不等于

$$\sum_{j=1}^n \mathbf{M}_j$$

因此它通常不是 $\mathbf{A}\mathbf{B}$ 的无偏估计

正确的做法是抽到 $j$ 之后输出

$$\mathbf{X} = \frac{1}{p_j}\mathbf{M}_j = \frac{1}{p_j}\mathbf{a}_j\mathbf{b}_j^\top$$

这时

$$\begin{align*} \mathbb{E}[\mathbf{X}] &= \sum_{j=1}^n p_j \left( \frac{1}{p_j}\mathbf{M}_j \right) \\ &= \sum_{j=1}^n \mathbf{M}_j \\ &= \mathbf{A}\mathbf{B} \end{align*}$$

所以

$$\boxed{ \mathbb{E}[\mathbf{X}]=\mathbf{A}\mathbf{B} }$$

这就是无偏估计了

如果抽 $s$ 次，并且每次独立地按同一个分布 $p_j$ 抽样，那么可以定义

$$\widehat{\mathbf{A}\mathbf{B}} = \frac{1}{s} \sum_{\ell=1}^s \frac{1}{p_{j_\ell}} \mathbf{a}_{j_\ell}\mathbf{b}_{j_\ell}^\top$$

这里要区分 $n$ 和 $s$：

• $n$ 是候选外积项的总数，也就是有多少个 $\mathbf{M}_j$ 可以被抽
• $s$ 是重复抽样的次数，也就是最终平均时一共抽了多少次

这里的 $\ell$ 表示第几次抽样：

$$\ell=1,2,\dots,s$$

而 $j_\ell$ 表示第 $\ell$ 次抽样时实际抽到的下标。例如第 $3$ 次抽样如果抽到了第 $7$ 个外积项，那么

$$j_3=7$$

对应的样本项就是

$$\frac{1}{p_{j_3}} \mathbf{a}_{j_3}\mathbf{b}_{j_3}^\top = \frac{1}{p_7} \mathbf{a}_7\mathbf{b}_7^\top$$

所以这个公式里的

$$\sum_{\ell=1}^s$$

是在把实际抽到的 $s$ 次随机估计结果加起来；而前面计算单次随机变量 $\mathbf{X}$ 的期望时，

$$\sum_{j=1}^n$$

是在对所有可能被抽到的 $n$ 个候选结果做概率加权求和

由期望的线性性，

$$\begin{align*} \mathbb{E} \left[ \widehat{\mathbf{A}\mathbf{B}} \right] &= \mathbb{E} \left[ \frac{1}{s} \sum_{\ell=1}^s \frac{1}{p_{j_\ell}} \mathbf{a}_{j_\ell}\mathbf{b}_{j_\ell}^\top \right] \\ &= \frac{1}{s} \sum_{\ell=1}^s \mathbb{E} \left[ \frac{1}{p_{j_\ell}} \mathbf{a}_{j_\ell}\mathbf{b}_{j_\ell}^\top \right] \end{align*}$$

而每一次抽样都使用同一个分布 $p_j$，所以对任意一次 $\ell$ 都有

$$\begin{align*} \mathbb{E} \left[ \frac{1}{p_{j_\ell}} \mathbf{a}_{j_\ell}\mathbf{b}_{j_\ell}^\top \right] &= \sum_{j=1}^n p_j \left( \frac{1}{p_j} \mathbf{a}_j\mathbf{b}_j^\top \right) \\ &= \sum_{j=1}^n \mathbf{a}_j\mathbf{b}_j^\top \\ &= \mathbf{A}\mathbf{B} \end{align*}$$

代回去得到

$$\begin{align*} \mathbb{E} \left[ \widehat{\mathbf{A}\mathbf{B}} \right] &= \frac{1}{s} \sum_{\ell=1}^s \mathbf{A}\mathbf{B} \\ &= \frac{1}{s} s\mathbf{A}\mathbf{B} \\ &= \mathbf{A}\mathbf{B} \end{align*}$$

所以 $\widehat{\mathbf{A}\mathbf{B}}$ 仍然是 $\mathbf{A}\mathbf{B}$ 的无偏估计

这里的 $1/s$ 是在平均这 $s$ 次随机估计结果，$1/p_{j_\ell}$ 是在修正第 $j_\ell$ 项被抽中的概率

这两个因子的含义不同

$$\boxed{ \frac{1}{s} \text{ 平均样本记录，} \qquad \frac{1}{p_j} \text{ 修正抽样概率} }$$

这就是随机矩阵乘法的意义：我们不必显式计算所有外积项来得到完整的 $\mathbf{A}\mathbf{B}$，而是多次重复这个随机过程，并对结果取平均。抽样次数 $s$ 越大，估计量的方差越小，$\widehat{\mathbf{A}\mathbf{B}}$ 的结果通常越接近真实的 $\mathbf{A}\mathbf{B}$

有偏估计会是什么样

有偏的核心形式是

$$\mathbb{E} \left[ \widehat{\mathbf{A}\mathbf{B}} \right] \neq \mathbf{A}\mathbf{B}$$

最简单的有偏情形就是抽到以后不做缩放

$$\mathbf{X}=\mathbf{M}_J$$

于是

$$\mathbb{E}[\mathbf{X}] = \sum_{j=1}^n p_j\mathbf{M}_j$$

这估计的是按 $p_j$ 加权后的平均，而不是所有外积项的总和

如果采用均匀抽样

$$p_j=\frac{1}{n}$$

但仍然直接输出 $\mathbf{M}_J$，那么

$$\mathbb{E}[\mathbf{X}] = \frac{1}{n} \sum_{j=1}^n \mathbf{M}_j = \frac{1}{n}\mathbf{A}\mathbf{B}$$

它少了一个因子 $n$

所以均匀抽样下，单次无偏估计应该是

$$\mathbf{X}=n\mathbf{M}_J$$

因为

$$\mathbb{E}[\mathbf{X}] = \sum_{j=1}^n \frac{1}{n} n\mathbf{M}_j = \mathbf{A}\mathbf{B}$$

这也说明了一个重要事实：无偏性不是由“是否均匀抽样”决定的，而是由目标、抽样概率和估计量构造方式共同决定的

方差与抽样概率的选择

无偏只保证平均方向正确，但不保证每次结果稳定

如果某个 $\mathbf{M}_j$ 很大，而 $p_j$ 又很小，那么一旦抽到它，估计量里会出现

$$\frac{1}{p_j}\mathbf{M}_j$$

这个项可能非常大，从而带来很大的方差

为了量化矩阵估计的波动，可以使用 Frobenius norm 下的均方误差。令

$$\mathbf{C}=\mathbf{A}\mathbf{B}$$

单次估计量是

$$\mathbf{X}=\frac{1}{p_J}\mathbf{M}_J$$

因为 $\mathbb{E}[\mathbf{X}]=\mathbf{C}$，所以

$$\mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 = \mathbb{E}\|\mathbf{X}\|_F^2-\|\mathbf{C}\|_F^2$$

这一点可以直接展开 Frobenius inner product 来看。由于

$$\|\mathbf{X}-\mathbf{C}\|_F^2 = \langle \mathbf{X}-\mathbf{C},\mathbf{X}-\mathbf{C}\rangle_F$$

所以

$$\begin{align*} \|\mathbf{X}-\mathbf{C}\|_F^2 &= \|\mathbf{X}\|_F^2 -2\langle \mathbf{X},\mathbf{C}\rangle_F +\|\mathbf{C}\|_F^2 \end{align*}$$

两边取期望：

$$\begin{align*} \mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 &= \mathbb{E}\|\mathbf{X}\|_F^2 -2\mathbb{E}\langle \mathbf{X},\mathbf{C}\rangle_F +\|\mathbf{C}\|_F^2 \end{align*}$$

因为 $\mathbf{C}$ 是固定目标，不是随机变量，所以

$$\begin{align*} \mathbb{E}\langle \mathbf{X},\mathbf{C}\rangle_F &= \langle \mathbb{E}[\mathbf{X}],\mathbf{C}\rangle_F \\ &= \langle \mathbf{C},\mathbf{C}\rangle_F \\ &= \|\mathbf{C}\|_F^2 \end{align*}$$

因此

$$\begin{align*} \mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 &= \mathbb{E}\|\mathbf{X}\|_F^2 -2\|\mathbf{C}\|_F^2 +\|\mathbf{C}\|_F^2 \\ &= \mathbb{E}\|\mathbf{X}\|_F^2-\|\mathbf{C}\|_F^2 \end{align*}$$

于是

$$\begin{align*} \mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 &= \mathbb{E}\|\mathbf{X}\|_F^2-\|\mathbf{C}\|_F^2 \\ &= \sum_{j=1}^n p_j \left\| \frac{1}{p_j}\mathbf{M}_j \right\|_F^2 -\|\mathbf{C}\|_F^2 \\ &= \sum_{j=1}^n \frac{\|\mathbf{M}_j\|_F^2}{p_j} -\|\mathbf{A}\mathbf{B}\|_F^2 \end{align*}$$

而 rank one matrix 满足

$$\|\mathbf{M}_j\|_F = \|\mathbf{a}_j\mathbf{b}_j^\top\|_F = \|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$$

因此要控制的主要部分是

$$\sum_{j=1}^n \frac{ \|\mathbf{a}_j\|_2^2 \|\mathbf{b}_j\|_2^2 }{p_j}$$

和前面推导的平均估计量方差下降规律一样，抽 $s$ 次再平均会把波动缩小到单次估计的约 $1/s$

$$\mathbb{E} \left\| \widehat{\mathbf{A}\mathbf{B}}-\mathbf{A}\mathbf{B} \right\|_F^2 = \frac{1}{s} \left( \sum_{j=1}^n \frac{\|\mathbf{M}_j\|_F^2}{p_j} -\|\mathbf{A}\mathbf{B}\|_F^2 \right)$$

现在问题变成：在约束

$$\sum_{j=1}^n p_j=1$$

下，怎样选择 $p_j$ 让方差小？

设

$$c_j=\|\mathbf{M}_j\|_F = \|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$$

那么均方误差里的关键部分可以写成

$$\sum_{j=1}^n \frac{\|\mathbf{M}_j\|_F^2}{p_j} - \|\mathbf{A}\mathbf{B}\|_F^2 = \sum_{j=1}^n \frac{c_j^2}{p_j} - \|\mathbf{A}\mathbf{B}\|_F^2$$

这里

$$\|\mathbf{A}\mathbf{B}\|_F^2$$

是目标矩阵本身的大小，和抽样概率 $p_j$ 无关。我们现在能选择的只有概率分布

$$p_1,p_2,\dots,p_n$$

所以如果想通过选择 $p_j$ 来降低随机估计的波动，就只需要关注依赖 $p_j$ 的那一项，也就是最小化

$$\sum_{j=1}^n \frac{c_j^2}{p_j}$$

也可以这样理解：真正的误差项是

$$\sum_{j=1}^n \frac{c_j^2}{p_j} - \|\mathbf{A}\mathbf{B}\|_F^2$$

我们当然希望前一项尽量接近后面的

$$\|\mathbf{A}\mathbf{B}\|_F^2$$

但 $\|\mathbf{A}\mathbf{B}\|_F^2$ 是固定常数，无法通过选择 $p_j$ 改变，所以最小化整个误差项等价于最小化

$$\sum_{j=1}^n \frac{c_j^2}{p_j}$$

而且这项只能从上方接近 $\|\mathbf{A}\mathbf{B}\|_F^2$。因为

$$\mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 \geq 0$$

所以由

$$\mathbb{E}\|\mathbf{X}-\mathbf{C}\|_F^2 = \sum_{j=1}^n \frac{c_j^2}{p_j} - \|\mathbf{A}\mathbf{B}\|_F^2$$

可得

$$\sum_{j=1}^n \frac{c_j^2}{p_j} \geq \|\mathbf{A}\mathbf{B}\|_F^2$$

所以我们能做的是把它压到尽可能小，也就是让它尽量贴近这个下界

直观上，如果某个外积项很大，也就是 $c_j$ 很大，但给它的抽样概率 $p_j$ 很小，那么

$$\frac{c_j^2}{p_j}$$

会变得很大，从而让方差变大。因此大的外积项应该被更频繁地抽到

现在这是一个带约束的最优化问题：

$$\min_{p_1,\dots,p_n} \sum_{j=1}^n \frac{c_j^2}{p_j} \qquad \text{s.t.} \qquad \sum_{j=1}^n p_j=1$$

约束来自概率分布本身：所有抽样概率必须加起来等于 $1$

Lagrange multiplier 就是处理这类带约束最优化问题的方法。它的做法是引入一个新变量 $\lambda$，把约束也放进目标函数里。直观上，$\lambda$ 用来记录“为了满足约束，目标函数的最优方向需要怎样被修正”

因此构造

$$L(p_1,\dots,p_n,\lambda) = \sum_{j=1}^n \frac{c_j^2}{p_j} + \lambda \left( \sum_{j=1}^n p_j-1 \right)$$

接下来要求最优点。对于一元函数，最小值点通常满足

$$f'(x)=0$$

这里变量有很多个：

$$p_1,p_2,\dots,p_n,\lambda$$

所以对应地，要对每个变量求偏导，并令偏导为 $0$。这就是一阶最优条件

对 $p_j$ 求偏导，是为了找到让目标函数最小的概率分配：

$$\frac{\partial L}{\partial p_j} = - \frac{c_j^2}{p_j^2} + \lambda =0$$

所以

$$p_j = \frac{c_j}{\sqrt{\lambda}}$$

对 $\lambda$ 求偏导，则会把约束条件带回来：

$$\frac{\partial L}{\partial \lambda} = \sum_{j=1}^n p_j-1 = 0$$

也就是

$$\sum_{j=1}^n p_j=1$$

现在把

$$p_j = \frac{c_j}{\sqrt{\lambda}}$$

代入约束：

$$\begin{align*} \sum_{j=1}^n p_j &= \sum_{j=1}^n \frac{c_j}{\sqrt{\lambda}} \\ &= \frac{1}{\sqrt{\lambda}} \sum_{j=1}^n c_j \\ &= 1 \end{align*}$$

因此

$$\sqrt{\lambda} = \sum_{j=1}^n c_j$$

再代回 $p_j=\frac{c_j}{\sqrt{\lambda}}$，得到

$$\boxed{ p_j = \frac{c_j}{\sum_{k=1}^n c_k} = \frac{\|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2} {\sum_{k=1}^n \|\mathbf{a}_k\|_2\|\mathbf{b}_k\|_2} }$$

这里分母里的 $k$ 只是求和下标，用来避免和分子里的 $j$ 混淆。它表示把所有 $c_1,c_2,\dots,c_n$ 加起来：

$$\sum_{k=1}^n c_k = c_1+c_2+\cdots+c_n$$

所以这个步骤本质上是在做归一化：先得到 $p_j\propto c_j$，再用 $\sum_j p_j=1$ 把它变成真正的概率分布

如果某个 $c_j=0$，也就是对应的外积项本来就是零矩阵，那么它对总和没有贡献，实际算法里可以不抽它。上面的概率公式主要针对 $c_j>0$ 的项，此时需要保证 $p_j>0$，否则 $1/p_j$ 没有意义

直观上就是

$$\boxed{ \text{越可能贡献大的外积项，越应该更频繁地抽到} }$$

这样不会改变无偏性，但会降低估计量的方差

实际计算时还要注意一个成本问题：这个最优概率需要先计算所有

$$\|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$$

并把它们加起来。看起来好像也要扫过所有 $j=1,\dots,n$，但这通常仍然比完整计算 $\mathbf{A}\mathbf{B}$ 便宜很多

如果

$$\mathbf{a}_j\in\mathbb{R}^m, \qquad \mathbf{b}_j\in\mathbb{R}^p$$

那么形成一个外积

$$\mathbf{a}_j\mathbf{b}_j^\top$$

需要大约 $O(mp)$ 的工作量，完整计算所有 $n$ 个外积大约是

$$O(mnp)$$

而计算一个范数乘积

$$\|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$$

只需要扫过 $\mathbf{a}_j$ 和 $\mathbf{b}_j$ 的元素，代价大约是 $O(m+p)$。计算所有 $j$ 的范数乘积大约是

$$O(n(m+p))$$

这通常只是扫描输入数据的成本，而不是形成并累加完整输出矩阵的成本。因此在 $m,p$ 较大时，它往往远小于完整矩阵乘法

当然，如果连扫描所有列和行都很贵，或者数据是 streaming / distributed 的，实际算法也可能使用 uniform sampling、近似的 sampling probabilities、近似 leverage scores 等更便宜的方案。上面的 $p_j\propto\|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$ 是方差最优意义下的理想选择，不一定总是实现成本最低的选择

1/p_j 的另一种理解

上面的 $1/p_j$ 可以理解成一种补偿因子。结合前面的说法，这一整套做法可以看成 Monte Carlo estimation + importance weighting：

• Monte Carlo estimation：用随机抽样的平均来估计完整总和
• importance weighting：用 $1/p_j$ 修正抽样概率，保证估计量无偏

在矩阵乘法里，我们的目标是完整的总和

$$\sum_{j=1}^n \mathbf{M}_j$$

也就是说，每一项 $\mathbf{M}_j$ 在目标里都应该完整贡献一次

但实际随机抽样时，第 $j$ 项只以概率 $p_j$ 被抽到。如果抽到以后直接输出 $\mathbf{M}_j$，那么长期平均只会得到

$$\sum_{j=1}^n p_j\mathbf{M}_j$$

这不是我们想要的完整总和

所以抽到第 $j$ 项时，要把它放大

$$\frac{1}{p_j}$$

倍，用来补偿“它只以概率 $p_j$ 出现”这件事。于是

$$\begin{align*} \mathbb{E} \left[ \frac{1}{p_J}\mathbf{M}_J \right] &= \sum_{j=1}^n p_j \left( \frac{1}{p_j}\mathbf{M}_j \right) \\ &= \sum_{j=1}^n \mathbf{M}_j \\ &= \mathbf{A}\mathbf{B} \end{align*}$$

所以可以把 $1/p_j$ 认为是：

抽得越少，抽中以后就要补得越多。

这类思想通常叫 importance weighting。更一般地，如果目标权重是 $\pi_j$，而实际抽样概率是 $q_j$，修正权重就是

$$\frac{\pi_j}{q_j}$$

随机矩阵乘法这一节里，目标是普通总和，可以理解成每一项的目标权重都是 $1$，实际抽样概率是 $p_j$，所以修正权重就是 $1/p_j$

总结

随机矩阵乘法的核心可以压缩成：

把 $\mathbf{A}\mathbf{B}$ 写成 rank one terms 的总和，然后随机抽样其中一部分，用概率修正保证无偏。

具体来说，

$$\mathbf{A}\mathbf{B} = \sum_{j=1}^n \mathbf{a}_j\mathbf{b}_j^\top$$

抽到第 $j$ 项的概率是 $p_j$，单次估计量取

$$\mathbf{X} = \frac{1}{p_j} \mathbf{a}_j\mathbf{b}_j^\top$$

于是

$$\mathbb{E}[\mathbf{X}] = \mathbf{A}\mathbf{B}$$

多次抽样后取平均

$$\widehat{\mathbf{A}\mathbf{B}} = \frac{1}{s} \sum_{\ell=1}^s \frac{1}{p_{j_\ell}} \mathbf{a}_{j_\ell}\mathbf{b}_{j_\ell}^\top$$

仍然无偏，并且方差随样本数增加而下降

$$\text{variance} \sim \frac{1}{s}$$

但无偏不代表稳定。要让估计稳定，需要合理选择抽样概率

$$p_j \propto \|\mathbf{a}_j\|_2\|\mathbf{b}_j\|_2$$

所以这节真正想建立的直觉是：随机算法不是说随便丢掉信息，我们要用概率分布来决定需要哪些信息，用权重修正保证平均正确，再用方差分析决定怎样抽得更稳