Monarch矩阵

文章来源：Monarch矩阵-计算高效的稀疏型矩阵分解

最佳排版请看原博客：Monarch矩阵：计算高效的稀疏型矩阵分解 - 科学空间|Scientific Spaces

在矩阵压缩这个问题上，我们通常有两个策略可以选择，分别是低秩化和稀疏化。低秩化通过寻找矩阵的低秩近似来减少矩阵尺寸，而稀疏化则是通过减少矩阵中的非零元素来降低矩阵的复杂性。如果说SVD是奔着矩阵的低秩近似去的，那么相应地寻找矩阵稀疏近似的算法又是什么呢？

接下来我们要学习的是论文《Monarch: Expressive Structured Matrices for Efficient and Accurate Training》，它为上述问题给出了一个答案——“Monarch矩阵”，这是一簇能够分解为若干置换矩阵与稀疏矩阵乘积的矩阵，同时具备计算高效且表达能力强的特点，论文还讨论了如何求一般矩阵的Monarch近似，以及利用Monarch矩阵参数化LLM来提高LLM速度等内容。

值得指出的是，该论文的作者也正是著名的Flash Attention的作者Tri Dao，其工作几乎都在致力于改进LLM的性能，这篇Monarch也是个人主页上特意展示的几篇论文之一，单从这一点看就非常值得学习一番。

SVD回顾

首先我们来简单回顾一下SVD（奇异值分解）。对于矩阵$n\times m$ 小的矩阵$A$，SVD将它分解为

$$A = U\Sigma V \$$

其中$U,V$ 别是形状为$n\times n$、$m\times m$ 正交矩阵，$\Sigma$ 是$n\times m$ 对角矩阵，对角线元素非负且从大到小排列。当我们只保留$\Sigma$ 前$r$ 对角线元素时，就得到了$A$ 一个秩不超过$r$ 近似分解：

$$A \approx U_{[:,:r]}\Sigma_{[:r,:r]} V_{[:r,:]} \$$

这里下标就按照Python的切片来执行，所以$U_{[:,:r]}$ 形状为$n\times r$、 $\Sigma_{[:r,:r]}$ 形状为$r\times r$ 及$V_{[:r,:]}$ 形状为$r\times m$，这意味着$U_{[:,:r]}\Sigma_{[:r,:r]} V_{[:r,:]}$ 秩至多为$r$。

特别地，由SVD得到的如上低秩近似，正好是如下优化问题的精确解：

$$U_{[:,:r]}\Sigma_{[:r,:r]} V_{[:r,:]} = \mathop{\text{argmin}}_{rank(B)\leq r} \Vert A - B\Vert_F^2 \$$

其中$\Vert\cdot\Vert_F^2$ 矩阵的Frobenius范数的平方，即矩阵每个元素的平方和。也就是说，在Frobenius范数下，矩阵$A$ 最优$r$ 近似就是$U_{[:,:r]}\Sigma_{[:r,:r]} V_{[:r,:]}$，该结论被称为“Eckart-Young-Mirsky定理”。也正是因为这个结论，我们在文章开头才说“SVD是奔着矩阵的低秩近似去的”。

SVD可以展开讨论的内容非常多，甚至写成一本书也不为过，这里就不再继续深入了。最后说一下，SVD的计算复杂度是$\mathcal{O}(nm\cdot\min(m,n))$，因为我们至少要对$A^{\top} A$ $A A^{\top}$ 一做特征值分解。如果我们确定做SVD是为了寻找$r$ 近似，那么复杂度可以有所降低，这便是Truncated SVD。

Monarch矩阵

低秩分解应用非常广，但它未必总是符合我们的需求，比如可逆方阵的低秩近似必然不可逆，这意味着低秩近似不适合需要求逆的场景。此时另一个选择是稀疏近似，稀疏矩阵通常能够保证秩不退化。

注意稀疏和低秩并无必然联系，比如单位阵就是很稀疏的矩阵，但它可逆（满秩）。寻找矩阵的稀疏近似并不难，比如将绝对值最大的$k$ 元素外的所有元素都置零就是一个很朴素的稀疏近似，但问题是它通常不实用，所以难在寻找实用的稀疏近似。所谓“实用”，指的是保留足够表达能力或近似程度的同时，实现一定程度的稀疏化，并且这种稀疏化具有适当的结构，有助于矩阵运算（比如乘法、求逆）的提速。

Monarch矩阵正是为此而生，假设$n=m^2$ 一个平方数，那么Monarch矩阵是全体$n$ 矩阵的一个子集，我们记为$\mathcal{M}^{(n)}$，它定义为如下形式的矩阵的集合：

$$M = PLPR \$$

其中$P$ $n\times n$ 置换矩阵（正交矩阵），$L,R$ 分块对角矩阵。下面我们来逐一介绍它们。

置换矩阵

置换矩阵$P$ 现的效果是将向量$[x_1,x_2,\cdots,x_n]$ 换成新的向量

$$[x_1, x_{1+m}, \cdots , x_{1+(m−1)m}, x_2, x_{2+m}, \cdots , x_{2+(m−1)m}, \cdots , x_m, x_{2m}, \cdots , x_n] \$$

当然这样写大家可能依然觉得迷糊，然事实上用代码实现非常简单：

1	`Px = x.reshape(m, m).transpose().reshape(n)`

如下图所示：

之前做CV的读者可能会觉得这个操作有点熟悉，它其实就是ShuffleNet中的“Shuffle”操作，这样对向量先reshape然后transpose最后再reshape回来的组合运算，起到一种“伪Shuffle”的效果，它也可以视为$m$ 制的“位反转排序”。很明显，这样的操作做两次，所得向量将复原为原始向量，所以我们有$P^2=I$，所以$P^{-1}=P^{\top}=P$。

分块对角

说完$P$，我们再来说$L,R$，它们也是$n\times n$ 小的矩阵，不过它们还是$m\times m$ 分块对角矩阵，每个块是$m\times m$ 小，如下图所示：

当$n$ 够大时，$L,R$ 零的数量占主导，所以$L,R$ 是稀疏矩阵，即Monarch矩阵是具备稀疏特性的矩阵分解形式。由于$P$ 固定的，所以$PLPR$ 的可变元素就来源于$L,R$ 非零元素，因此，矩阵$M$ 然是$n\times n$ 矩阵，但它实际自由参数不超过$2m^3=2n^{1.5}$ 。从$1.5$ 个数字我们就可以窥见Monarch矩阵的意图了，它希望将原本需要平方复杂度的运算，通过Monarch矩阵近似降低到1.5次方复杂度。

效率简析

那么Monarch矩阵能否达到这个目的呢？换句话说Monarch矩阵能否达到前面说的“实用”标准？表达能力方面我们后面再谈，我们先看计算高效方面。

比如“矩阵-向量”乘法，标准的复杂度是$\mathcal{O}(n^2)$，但如果是Monarch矩阵的话我们有$Mx = P(L(P(Rx)))$，由于乘$P$ 是简单的reshape和transpose，所以它几乎不占计算量，主要计算量来源于$L$ $R$ 一个向量相乘。由于$L,R$ 分块对角矩阵的特点，我们可以将向量为$m$ ，继而转化为$m$ $m\times m$ 矩阵与$m$ 向量相乘，总复杂度是$2m\times\mathcal{O}(m^2)=\mathcal{O}(2n^{1.5})$，比$\mathcal{O}(n^2)$ 低。

再比如求逆，我们考虑$M^{-1}x$，$n$ 矩阵求逆的标准复杂度是$\mathcal{O}(n^3)$，但对于Monarch矩阵我们有$M^{-1} x =R^{-1}PL^{-1}P x$，主要计算量来源于$L^{-1}$、$R^{-1}$ 及对应的“矩阵-向量”乘法，由于$L,R$ 是分块对角阵，我们只需要分别对每个对角线上的块矩阵求逆，也就是共有$2m$ $m\times m$ 矩阵求逆，复杂度是$2m\times\mathcal{O}(m^3)=\mathcal{O}(2n^2)$，同样低于标准的$\mathcal{O}(n^3)$。要单独写出$M^{-1}$ 是可以的，但需要利用到后面的恒等式(8)。

所以结论就是，由于$P$ 法几乎不占计算量以及$L,R$ 分块对角矩阵的特点，$n$ Monarch矩阵相关运算，基本上可以转化为$2m$ $m\times m$ 阵的独立运算，从而降低总的计算复杂度。所以至少计算高效这一点，Monarch矩阵是没有问题的，并且由于$L,R$ 非零元素本身已经方形结构，实现上也很方便，可以充分利用GPU进行计算，不会带来不必要的浪费。

Monarch分解

确认Monarch矩阵的有效性后，接下来应用方面的一个关键问题就是：给定任意的$n=m^2$ 矩阵$A$，如何求它的Monarch近似呢？跟SVD类似，我们定义如下优化问题

$$\mathop{\text{argmin}}_{M\in\mathcal{M}^{(n)}} \Vert A - M\Vert_F^2 \$$

非常幸运的是，这个问题有一个复杂度不超过$\mathcal{O}(n^{2.5})$ 求解算法，这比SVD的$\mathcal{O}(n^3)$ 要更高效一些。

高维数组

理解这个算法的关键一步，是将Monarch相关的矩阵、向量都转化为更高维数组的形式。具体来说，Monarch矩阵$M$ 来是一个二维数组，每个元素记为$M_{i,j}$，表示该元素位于第$i$ 、第$j$ ，现在我们要按照分块矩阵的特点，将它等价地表示为四维数组，每个元素记为$M_{i,j,k,l}$，表示第$i$ 行、第$j$ 行、第$k$ 列、第$l$ 列的元素，如下图所示：

虽然说起来挺费劲的，但事实上代码就一行

1	`M.reshape(m, m, m, m)`

同理，$n$ （列）向量$x$ 被转为$m\times m$ 二维数据，代码也是一行x.reshape(m, m)。剩下的$L,R$ 然是表示为$m\times m\times m$ 三维数组，如$L_{i,j,k}$ 示第$i$ 、第$j$ 行、第$k$ 列的元素，这本来也是$L,R$ 高效的储存方式，但为了统一处理，我们也可以用Kronecker delta符号将它们升到四维，比如$L_{i,j,k,l} = \delta_{i,k}L_{i,j,l}$、$R_{i,j,k,l} = \delta_{i,k}R_{i,j,l}$。

新恒等式

接下来，我们将推出$M$ $L,R$ 一个新关系式。首先，可以证明在二维表示中，矩阵$P$ 向量$x$ 乘法变得更简单了，结果就是$x$ 转置，即$(Px){i,j} = x{j,i}$，所以我们有$(PR){i,j,k,l} = R{j,i,k,l} = \delta_{j,k}R_{j,i,l}$；接着，两个矩阵的乘法，在四维表示之下求和指标也有两个，所以

$$(L P R){\alpha,\beta,k,l} = \sum{i,j} L_{\alpha,\beta,i,j}(PR){i,j,k,l} = \sum{i,j} \delta_{\alpha, i} L_{\alpha,\beta,j}\delta_{j,k}R_{j,i,l} = L_{\alpha,\beta,k}R_{k,\alpha,l} \$$

最后就是$(P L P R){\alpha,\beta,k,l}=L{\beta,\alpha,k}R_{k,\beta,l}$，将$\alpha,\beta$ 回$i,j$ 到$(P L P R){i,j,k,l}=L{j,i,k}R_{k,j,l}$，又因为$M=PLPR$，所以有

$$M_{i,j,k,l} = L_{j,i,k}R_{k,j,l} \tag{8}$$

从这个等式可以看出，当我们固定一对$(j,k)$ ，左边是一个子矩阵，右边是两个向量的外积，这意味着如果我们要给矩阵$A$ Monarch近似，只需要将$A$ 照同样方式转为四维数组，并固定一对$(j,k)$，那么问题就变成了找对应子矩阵的“秩-1近似”！换句话说，有了这个恒等式之后，给矩阵$A$ Monarch近似可以转化为给$m^2$ 子矩阵找“秩-1近似”，这可以用SVD完成，每个复杂度不超过$\mathcal{O}(m^3)$，所以总复杂度不超过$m^2\times\mathcal{O}(m^3) = \mathcal{O}(n^{2.5})$。

参考实现

笔者简单用Numpy写的参考实现如下：

import numpy as np

def monarch_factorize(A):
    M = A.reshape(m, m, m, m).transpose(1, 2, 0, 3)
    U, S, V = np.linalg.svd(M)
    L = (U[:, :, :, 0] * S[:, :, :1]**0.5).transpose(0, 2, 1)
    R = (V[:, :, 0] * S[..., :1]**0.5).transpose(1, 0, 2)
    return L, R

def convert_3D_to_2D(LR):
    X = np.zeros((m, m, m, m))
    for i in range(m):
        X[i, i] += LR[i]
    return X.transpose(0, 2, 1, 3).reshape(n, n)

m = 8
n = m**2
A = np.where(np.random.rand(n, n) > 0.8, np.random.randn(n, n), 0)

L, R = monarch_factorize(A)
L = convert_3D_to_2D(L)
R = convert_3D_to_2D(R)
PL = L.reshape(m, m, n).transpose(1, 0, 2).reshape(n, n)
PR = R.reshape(m, m, n).transpose(1, 0, 2).reshape(n, n)

U, S, V = np.linalg.svd(A)

print('Monarch error:', np.square(A - PL.dot(PR)).mean())
print('Low-Rank error:', np.square(A - (U[:, :m] * S[:m]).dot(V[:m])).mean())

笔者简单对比了一下SVD求出的秩-$m$ 似（此时低秩近似跟Monarch近似参数量相当），发现如果是完全稠密的矩阵，那么秩-$m$ 似的平方误差往往优于Monarch近似（但不多），这也是意料之中，因为从Monarch近似的算法就可以看出它本质上也是个定制版的SVD。不过，如果待逼近矩阵是稀疏矩阵时，那么Monarch近似的误差往往更优，并且越稀疏越优。

Monarch推广

到目前为止，我们约定所讨论的矩阵都是$n$ 方阵，并且$n=m^2$ 一个平方数。如果说方阵这个条件尚能接受，那么$n=m^2$ 个条件终究还是太多限制了，因此有必要至少将Monarch矩阵的概念推广到非平方数$n$。

非平方阶

为此，我们先引入一些记号。假设$b$ $n$ 一个因数，$\mathcal{BD}^{(b,n)}$ 示全体$\frac{n}{b}\times \frac{n}{b}$ 小的分块对角矩阵，其中每个块大小是$b\times b$ 子矩阵，很明显它是前面$L,R$ 一般化，按照这个记号我们可以写出$L,R\in\mathcal{BD}^{(\sqrt{n},n)}$。此外，我们还要一般化置换矩阵$P$，前面我们说了$P$ 实现是Px = x.reshape(m, m).transpose().reshape(n)，现在我们一般化为Px = x.reshape(n // b, b).transpose().reshape(n)，记为$P_{(\frac{n}{b},b)}$。

有了这些记号，我们可以定义一般的Monarch矩阵（原论文的附录）：

$$\mathcal{M}^{(b,n)} = \Bigg{M = P_{(b,\frac{n}{b})} L P_{(\frac{n}{b},b)} R,\Bigg|, L\in\mathcal{BD}^{(\frac{n}{b},n)}, R\in\mathcal{BD}^{(b,n)} \Bigg} \$$

示意图如下：

前面所定义的Monarch矩阵，在这里可以简单记为$\mathcal{M}^{(n)} = \mathcal{M}^{(\sqrt{n},n)}$。不难计算，$L$ 非零元素至多有$\frac{n^2}{b}$ ，$R$ 非零元素至多有$nb$ ，加起来是$\frac{n^2}{b} + nb$，它在$b=\sqrt{n}$ 得最小值，所以$b=\sqrt{n}$ 于最稀疏的一个例子。

只要形式

可能读者会困惑，为什么要区分$L\in\mathcal{BD}^{(\frac{n}{b},n)}, R\in\mathcal{BD}^{(b,n)}$，统一用一个不行吗？事实上，这样设计是为了保持高维表示下式(8)依然成立，从而可以推出类似的分解算法（请读者补充一下），以及可以从理论上保证它的表达能力。

如果我们不在意这些理论细节，只希望构造一个带有稀疏特性的矩阵参数化方式，那么就可以更灵活地对Monarch矩阵进行推广了，比如

$$M = \left(\prod_{i=1}^k P_i B_i\right)P_0 \$$

其中$B_1,B_2,\cdots,B_k \in \mathcal{BD}^{(b,n)}$，$P_0,P_1,\cdots,P_k$ 是置换矩阵，最后多乘一个$P_0$ 出于对称性的考虑，并不是必须的，如果你觉得有必要，还可以给每个$B_i$ 择不同的$b$，即$B_i\in \mathcal{BD}^{(b_i,n)}$。

甚至，你可以结合低秩分解的形式，推广到非方阵的块矩阵，如下图：

基于这个类比，我们还可以进一步将Monarch矩阵的概念推广到非方阵。总之，如果只是需要一种类似Monarch矩阵的稀疏化结构矩阵，而不在意理论细节，那么结果就仅限于我们的想象力了。

应用例子

目前看来，Monarch矩阵最大的特点就是对矩阵乘法比较友好，所以最大的用途无非就是替换全连接层的参数矩阵，从而提高全连接层的效率，这也是原论文实验部分的主要内容。

我们又可以将其分为“事前处理”和“事后处理”两类：“事前处理”就是在训练模型之前就将全连接层的参数矩阵改为Monarch矩阵，这样训练和推理都能提速，训练出来的模型也最贴合Monarch矩阵；“事后处理”就是已经有一个训练好的模型，此时我们可以用Monarch分解给全连接层的参数矩阵找一个Monarch近似，然后替换掉原来的矩阵，必要时再简单微调一下，以此提高原始模型的微调效率或推理效率。

除了替换全连接层外，《Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture》还讨论了更极端的做法——作为一个Token-Mixer模块直接替换Attention层。不过就笔者看来，Monarch-Mixer并不算太优雅，因为它跟MLP-Mixer一样，都是用一个可学的矩阵替换掉Attention矩阵，只不过在Monarch-Mixer这里换成了Monarch矩阵。这样的模式学到的是静态的注意力，个人对其普适性是存疑的。

最后，对如今的LLM来说，Monarch矩阵还可以用来构建参数高效的微调方案（Parameter-Efficient Fine-Tuning，PEFT）。我们知道，LoRA是从低秩分解出发设计的，既然低秩和稀疏是两条平行的路线，那么作为稀疏的代表作Monarch矩阵不应该也可以用来构建一种PEFT方案？Google了一下，还真有这样做的，论文名是《MoRe Fine-Tuning with 10x Fewer Parameters》，还很新鲜，是ICML2024的Workshop之一。

蝶之帝王

最后再简单说说Monarch矩阵的拟合能力。“Monarch”意为“帝王”、“君主”，取自“Monarch Butterfly（帝王蝴蝶）”一词，之所以这样命名，是因为它对标的是更早的“utterfly矩阵”。

什么是Butterfly矩阵？这个说起来还真有点费劲。Butterfly矩阵是一系列（$\log_2 n$ ）Butterfly因子矩阵的乘积，而Butterfly因子矩阵则是一个分块对角矩阵矩阵，其对角线上的矩阵叫做Butterfly因子（没有“矩阵”两个字），Butterfly因子则是一个$2\times 2$ 的分块矩阵，它的每个块矩阵则是一个对角阵（套娃结束）。如下图所示：

准确的Butterfly矩阵定义大家自行看论文就好，这里不详细展开。Butterfly这个名字来源于作者觉得每个Butterfly因子的形状像Butterfly（蝴蝶），当然像不像大家见仁见智，反正作者觉得像。从字面上来看，“Monarch Butterfly”比“Butterfly”更高级（毕竟是“帝王”），这暗示着Monarch矩阵比Butterfly矩阵更强。确实如此，Monarch论文附录证明了，不管$b$ 什么，$\mathcal{M}^{(b,n)}$ 能覆盖所有的$n$ Butterfly矩阵，并且$n > 512$ $\mathcal{M}^{(b,n)}$ 格大于全体$n$ Butterfly矩阵集合，换言之Butterfly矩阵能做到的Monarch矩阵也能做到，反之未必。

我们也可以从“矩阵-向量”乘法复杂度来直观感知Monarch矩阵表达能力。我们知道，一个$n\times n$ 阵乘以$n$ 向量的标准复杂度是$\mathcal{O}(n^2)$，但对于某些结构化矩阵可以更低，比如傅立叶变换可以做到$\mathcal{O}(n\log n)$，Butterfly矩阵也是$\mathcal{O}(n\log n)$，Monarch矩阵则是$\mathcal{O}(n^{1.5})$，所以Monarch矩阵“应该”是不弱于Butterfly矩阵的。当然，Butterfly矩阵也有它的好处，比如它的逆和行列式都比较好算，这对于Flow模型等需要求逆和行列式的场景更为方便。