
Monarch矩阵(Monarch Matrices)是近年来在深度学习和计算数学交叉领域提出的一种非常优雅且高效的矩阵参数化方法。它的出现主要是为了解决现代神经网络在规模不断扩张时所面临的算力和内存瓶颈。
稠密与稀疏的进退两难
在深度学习中,尤其是在Transformer等大语言模型中,线性层(即矩阵乘法)占据了绝大部分的计算量和参数量。传统的全连接层使用的是稠密矩阵,它的表达能力极强,能够捕捉复杂的数据特征,但其代价是高昂的:计算复杂度和内存消耗都随着输入维度的增加呈 O ( N 2 ) O(N^2) O(N2) 的二次方增长。当模型维度达到上万甚至十万级别时,这种复杂度便成为了不可承受之重,使得模型在扩展上下文长度或增加隐藏层维度时,很容易撞到硬件资源的"天花板"。
为了降低计算量,通常会考虑稀疏矩阵和结构化矩阵两个主要方向,但在实际应用中都存在致命痛点。非结构化的稀疏矩阵虽然在理论上减少了浮点运算的次数,但由于其内存访问模式极度不规则,在高度依赖并行计算和连续内存读取的现代硬件(如GPU或TPU)上不仅跑不快,反而可能比稠密矩阵还要慢。另一方面,传统的结构化矩阵(如低秩矩阵或循环矩阵)虽然计算复杂度低,但表达能力往往不够,强行替换会导致模型预测准确率严重下降。
Monarch矩阵正是在这种"稠密矩阵太慢,稀疏矩阵对硬件不友好,传统结构化矩阵表达能力差"的困境中诞生的。它的核心目标是在计算效率、硬件友好性和模型表达能力这三者之间找到一个完美的平衡点。
化繁为简的蝴蝶网络
Monarch矩阵的设计灵感直接来源于著名的快速傅里叶变换(FFT)算法,特别是其中的Cooley-Tukey算法。该算法通过一种被称为"蝴蝶操作"的结构,将一个极其复杂的全局计算分解为多个局部的、小规模的计算。
在数学构造上,Monarch矩阵并不试图用单一的复杂矩阵去逼近稠密矩阵,而是将其定义为多个块对角矩阵(Block-Diagonal Matrices)与排列操作(Permutations)的乘积。块对角矩阵的一个巨大优势是对现代GPU极其友好,而中间穿插的排列操作则起到了至关重要的全局路由作用。原本在不同对角块中进行局部处理的信息,通过排列得以进行全局的交换和混合。这种机制确保了即使网络权重是高度结构化和局部化的,输入数据依然能够像在稠密矩阵中一样,实现全方位的特征融合。
从计算机视觉的启发到硬件感知设计
要真正理解Monarch矩阵的历史定位,我们需要追溯它的提出者以及它与早期计算机视觉算法的奇妙联系。Monarch矩阵由斯坦福大学的Tri Dao等人在ICML 2022 上正式提出 [1]。值得一提的是,Tri Dao也是后来的 FlashAttention [2] 的作者。无论是FlashAttention还是Monarch矩阵,它们都贯彻了一个极具前瞻性的核心理念,硬件感知算法设计(Hardware-aware Algorithm Design)。研究团队致力于在算法的数学表达和现代GPU底层的Tensor Core运算逻辑之间寻找最优解。
这种"局部计算结合全局路由"的思想并非凭空出现。如果你熟悉计算机视觉,会发现Monarch矩阵在底层操作上与2017年提出的轻量级卷积网络ShuffleNet [3] 有着极其深刻的渊源和惊人的相似性。
在ShuffleNet中,为了减少计算量,研究者使用了分组卷积(Group Convolution),让不同通道组之间各自独立计算。但在数学上,一个分组卷积就等价于乘以一个块对角矩阵。为了打破分组卷积带来的"信息隔离",ShuffleNet引入了通道混洗(Channel Shuffle)操作,将各个组的通道打散重排。而这恰恰完全等同于Monarch矩阵中间穿插的排列矩阵操作。可以说,ShuffleNet的一个基本模块(分组卷积 → \rightarrow → 通道混洗 → \rightarrow →分组卷积)在结构上完美等价于Monarch矩阵的二因子分解形式。
既然如此相似,Monarch矩阵的突破性究竟在哪里?答案在于它完成了从"工程启发"到"严密数学"的飞跃。ShuffleNet是工程启发式架构,且局限于卷积通道维度;而Monarch矩阵则建立了一套严密的数学框架,证明了这种结构可以作为一种通用的线性层参数化方法。它不仅证明了这种分解的通用性(可以任意精度逼近任何稠密变换),更将其优化目标直接锚定在了现代AI加速器的内存层级上。可以说,ShuffleNet是在这个结构上建造了一栋实用的木屋,而Monarch矩阵则是为这个结构打下了坚固的数学地基,并建起了能容纳Transformer大模型的摩天大楼。
核心数学原理与推导
正是基于上述的跨领域启发,需要将直觉转化为具体的数学形式。重新审视标准神经网络中的线性层计算模型,假设我们有一个输入向量 x ∈ R N x \in \mathbb{R}^N x∈RN,标准的稠密线性变换通过一个权重矩阵 W ∈ R N × N W \in \mathbb{R}^{N \times N} W∈RN×N 将其映射为输出向量 y ∈ R N y \in \mathbb{R}^N y∈RN。
块对角矩阵的局限性
为了打破 O ( N 2 ) O(N^2) O(N2) 的复杂度瓶颈,最直观的数学思路是强制矩阵 W W W 为块对角矩阵。假设我们将 N N N 维空间均匀划分为 p p p 个子空间,每个子空间的维度为 q q q,使得 N = p × q N = p \times q N=p×q。一个典型的块对角矩阵 B 1 B_1 B1 可以看作是由 p p p 个大小为 q × q q \times q q×q 的小矩阵沿对角线排列而成。
当我们将 B 1 B_1 B1 作用于向量 x x x 时,实际上是将 x x x 切分成了 p p p 个独立的 q q q 维小向量,并分别与这 p p p 个小矩阵相乘。这种操作的计算复杂度仅为 p × q 2 = N q p \times q^2 = Nq p×q2=Nq。如果取 p ≈ q ≈ N p \approx q \approx \sqrt{N} p≈q≈N ,复杂度就从 O ( N 2 ) O(N^2) O(N2) 断崖式下降到了 O ( N 3 2 ) O(N^{\frac{3}{2}}) O(N23)。然而,正如前面提到的分组卷积的缺陷一样,纯粹的块对角矩阵切断了不同子空间之间的信息交流,使得输入向量的各个部分沦为"信息孤岛",无法实现深度学习模型高度依赖的全局特征融合。
Monarch矩阵的严密定义
Monarch矩阵的破局之道在于引入了排列矩阵(Permutation Matrix) P P P。在数学上,一个标准的二因子Monarch矩阵 M ∈ R N × N M \in \mathbb{R}^{N \times N} M∈RN×N 被定义为以下连乘形式:
M = P B 2 P T B 1 M = P B_2 P^T B_1 M=PB2PTB1
在这个公式中, B 1 B_1 B1 是由 p p p 个 q × q q \times q q×q 块组成的块对角矩阵; B 2 B_2 B2 是由 q q q 个 p × p p \times p p×p 块组成的块对角矩阵;而 P P P 则是块转置排列矩阵(Block-transposition Permutation)。
张量视角的维度变换
为了更清晰地推导这个过程,我们可以将一维向量 x x x 的变换过程在二维张量空间中展开。假设我们将长度为 N N N 的输入向量 x x x 重新排列(Reshape)为一个大小为 p × q p \times q p×q 的矩阵 X X X。
- 第一步,应用 B 1 B_1 B1。等价于对二维矩阵 X X X 的每一行(长度为 q q q)独立地进行线性变换。此时信息在局部的行维度内部得到了充分混合,但行与行之间依然是隔离的。
- 第二步,应用排列算子 P T P^T PT。在二维视角下,这本质上就是对矩阵 X X X 进行转置(Transpose),将其形状从 p × q p \times q p×q 变为 q × p q \times p q×p。这一步没有任何浮点运算,但发生了一次信息交互:原本在不同行中被隔离的元素,现在被重新分配到了同一行中。局部信息由此跨越了全局。
- 第三步,应用 B 2 B_2 B2。 B 2 B_2 B2 再次对转置后矩阵的每一行(此时长度为 p p p)独立进行线性变换。由于经过了转置,这一次的局部混合实际上是对第一步中的跨组信息进行二次融合。最后,通过最外层的 P P P 将张量重新展平为一维向量,便得到了具备全局视野的最终输出 y y y。
复杂度计算与硬件亲和性
从复杂度角度推导,第一步为 O ( N q ) O(Nq) O(Nq),第三步为 O ( N p ) O(Np) O(Np),转置操作复杂度为 O ( 1 ) O(1) O(1) 的浮点运算。整个Monarch矩阵向量乘法的总复杂度严格限制在 O ( N ( p + q ) ) O(N(p+q)) O(N(p+q))。在 p = q = N p = q = \sqrt{N} p=q=N 的最优分割下,复杂度定格在 O ( N 3 2 ) O(N^{\frac{3}{2}}) O(N23)。如果将这种分解思想递归应用(如持续对分),甚至可以逼近 O ( N log N ) O(N \log N) O(NlogN)。
这种优雅的数学分解在实际中取得巨大成功,更深层的原因在于它极度契合现代张量处理硬件的底层执行逻辑。在传统的非结构化稀疏矩阵中,不连续的内存读取严重破坏了缓存命中率。Monarch矩阵的构成要素直接对应着高效的底层算子:块对角矩阵乘法等价于批量矩阵乘法,能够让计算单元保持极高吞吐量;而排列矩阵直接对应于张量的张量重塑和转置操作,这在高度优化的底层CUDA实现中甚至可以直接融合在寄存器级别完成,几乎不带来额外的显存读写开销。
实际应用:为大模型"瘦身"与"提速"
凭借着理论上的优雅和硬件上的高效,Monarch矩阵在深度学习领域展现出了广阔的应用前景,尤其是在大语言模型和高效架构的设计中。
- 直接替代传统网络中的稠密层。在许多模型训练实验中,研究人员使用Monarch矩阵替换了多层感知机(MLP)和自注意力机制中的投影层。结果表明,这种替换不仅大幅减少了模型的参数量,缩短了端到端的训练时间,而且最终模型的困惑度(Perplexity)和下游任务表现与稠密基线模型几乎一致。
- 全新的高效网络架构,例如Monarch Mixer (M2) [4]。随着大模型在长文本处理上面临二次方复杂度的挑战,Monarch Mixer试图完全抛弃传统的注意力机制。它利用Monarch矩阵在时间和空间维度上沿着序列和通道进行混合,实现了真正的次二次方甚至线性复杂度,在极长的序列任务中不仅训练速度快,而且推理时的显存占用极低。
- 在稀疏训练和模型微调中也大放异彩。通过在结构化的Monarch空间中寻找最优解,大模型可以在资源受限的边缘设备或单张消费级显卡上完成原本需要集群才能进行的训练任务。
参考文献
- Monarch矩阵基础论文:Monarch: Expressive structured matrices for efficient and accurate training
- FlashAttention相关硬件感知工作:Flashattention: Fast and memory-efficient exact attention with io-awareness
- 计算机视觉启发渊源:Shufflenet: An extremely efficient convolutional neural network for mobile devices
- Monarch Mixer (M2) 高效架构:Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture
- 进一步扩展阅读:https://spaces.ac.cn/archives/10249