矩阵乘法加速计算优化最新进展

矩阵乘法作为众多GPU算子的基础操作，是高性能计算的重要问题之一，也是 AI 等应用的基石。它的算法机制本身相当简单，但为了达到更快的速度，人们多年来不懈努力，优化程度却一直有限。

现在，加速矩阵乘法过程的任务成为数学和计算机科学的交叉点。研究人员至今仍在继续改进该过程，尽管近几十年来进展相当有限。名古屋大学计算机科学家 François Le Gall 表示，自 1987 年以来，矩阵乘法的数值改进「一直很小，而且极其难以实现」。

一、相关论文进展

清华大学的段然（Ran Duan）、周任飞（Renfei Zhou）和加州大学伯克利分校的 Hongxun Wu 在解决这个长期存在的问题上迈出了重要一步，撰写的论文足足有 87 页。对于三位研究者的成果， Le Gall 表示尽管改进本身相对较小，但是「从概念上讲比以往的改进都大。」该论文被计算机科学领域的顶会 FOCS 2023 接收。

论文 v1 发布在 2022 年 10 月，v5 在 2023 年 11 月。论文地址：https://arxiv.org/abs/2210.10173

段然为清华大学交叉信息研究院副教授，主要研究方向为图论算法、数据结构、计算理论。Hongxun Wu 为加州大学伯克利分校二年级博士生，也是清华姚班出身。周任飞为清华姚班 2020 级的大四本科生，主修理论计算机科学（TCS）。他主要研究（简洁）数据结构和快速矩阵乘法，并对 TCS 的其他领域具有广泛兴趣，比如流算法、博弈论和在线算法等。此前，周任飞曾在理论计算机科学顶级会议 FOCS/SODA 上发表多篇论文。三位研究者的论文揭示了以前未知且未开发的潜在改进来源，并且已经取得了成果。

Faster Matrix Multiplication via Asymmetric Hashing - Ran Duan, Hongxun Wu, Renfei Zhou - 2022/10/18 - arXiv

2024 年 1 月在工业和应用数学学会上发表的第二篇论文（周任飞同样参与撰写）以此为基础，展示了如何进一步增强矩阵乘法。哈佛大学理论计算机科学家 William Kuszmaul 对此表示，这是一项重大的技术突破，是十多年来我们所看到的矩阵乘法的最大改进。论文地址：https://epubs.siam.org/doi/10.1137/1.9781611977912.134

New Bounds for Matrix Multiplication: from Alpha to Omega - Virginia Vassilevska Williams, Yinzhan Xu, Zixuan Xu, Renfei Zhou - 2023/07/16 - arXiv

二、矩阵乘法

高中教科书告诉你，计算规则是，第一个矩阵第一行的每个数字（2和1），各自乘以第二个矩阵第一列对应位置的数字（1和1），然后将乘积相加（ 2 x 1 + 1 x 1），得到结果矩阵左上角的那个值3。

也就是说，结果矩阵第m行与第n列交叉位置的那个值，等于第一个矩阵第m行与第二个矩阵第n列，对应位置的每个值的乘积之和。

矩阵乘法到底是什么东西。关键就是一句话，矩阵的本质就是线性方程式，两者是一一对应关系。如果从线性方程式的角度，理解矩阵乘法就毫无难度。

三、矩阵乘法要改进什么问题

矩阵乘法可能看起来是一个晦涩的问题，但它是一种基本的计算操作。它被融入了人们每天使用的大部分算法中，用于各种任务，从显示更清晰的计算机图形到解决网络理论中的物流问题。就像在计算的其他领域一样，速度至关重要。即使是微小的改进最终也可能大大减少所需要的时间、计算能力和金钱。

传统的两个 n×n 矩阵相乘的方法 —— 即将第一个矩阵中每一行的数字与第二个矩阵中每一列的数字相乘 —— 需要进行 n³ 次独立的乘法操作。对于 2 乘 2 的矩阵而言，这意味着需要进行 2³，也就是 8 次乘法操作。

1969 年，数学家 Volker Strassen 发现了一种更精巧的方法，只需 7 个乘法步骤和 18 个加法步骤，就能完成 2×2 矩阵的乘法运算。两年后，计算机科学家 Shmuel Winograd 证明，对于 2×2 矩阵来说，7 步乘法确实是绝对最小值。

Strassen 利用同样的想法证明，所有较大的 n×n 矩阵也可以用少于 n³步的方法进行乘法运算。这一策略中的一个关键因素涉及一个称为分解的程序：将一个大矩阵分解成一个个更小的子矩阵，这些子矩阵最终可能小到 2×2 甚至 1×1（只是单个数字）。

对于将巨型数组分解成小块的理由相当简单，麻省理工学院的计算机科学家 Virginia Vassilevska Williams 说：「对于一个大矩阵（比如 100×100 的矩阵），人类很难想到最佳的算法。」即使是 3 乘 3 的矩阵也还没有完全解决。「然而，人们可以使用已经为小矩阵开发的快速算法来获得更大矩阵的快速算法。」

研究人员确定，速度的关键在于减少乘法步骤的数量，尽可能将指数从 n³（传统方法）降低。可能的最低值 n² 基本上就是写出答案所需的时间。计算机科学家把这个指数称为 Ω，即 ω。 是当 n 越来越大时，成功将两个 n×n 矩阵相乘所需的最少步骤。同为 2024 年 1 月论文合著者的周任飞说：「这项工作的重点，是看你能接近 2 多少，并且是否可以在理论上实现。」

四、激光法

1986 年，Strassen 取得了另一项重大突破，他推出了矩阵乘法的激光法。Strassen 用它确定了 ω 的上限值为 2.48。虽然该方法只是大型矩阵乘法的一个步骤，但却是最重要的步骤之一，因为研究人员一直在不断改进它。

一年后，Winograd 和 Don Coppersmith 推出了一种新算法，对激光法进行了完美的补充。这套工具的组合在后来几乎所有加速矩阵乘法的研究中都得到了应用。

下面是一个简化的方法，让我们来看看这些不同的元素是如何结合在一起的。让我们从两个大型矩阵 A 和 B 开始，将它们相乘。首先，你要把它们分解成许多较小的子矩阵，有时也叫块。接下来，你就可以使用 Coppersmith 和 Winograd 的算法，将其作为处理并最终组装这些块的指导手册。Vassilevska Williams 说：「它告诉我在乘积矩阵 C 中要乘什么、加什么，以及哪些元素在哪里。」「它只是一个从 A 和 B 建立 C 的『配方』」。

然而，这里有一个问题：有时你会得到具有共同元素的块。保留这些共同元素会相当于将这些元素计算两次，因此在某个时候，需要消除这些重叠部分。研究人员通过「消灭」它们所在的块来解决这个问题 —— 将它们的分量设置为零以将它们从计算中移除。

这就是 Strassen 的激光法最终发挥作用的地方。Le Gall 说，「激光法通常非常有效，并且通常能找到消除重叠的子块的好方法」。在激光消除了所有重叠之后，你就可以构建最终的乘积矩阵 C。

将这些各种技术结合起来，就得到了一种用尽量少的乘法总数来乘两个矩阵的算法，至少在理论上是这样。激光法并不是为了实际应用；它只是一种思考矩阵相乘的理想方式。周任飞表示，「我们从未在计算机上运行这种方法，我们进行对它的分析。」

正是这种分析促成了 ω 十多年来的最大改进。

五、被发现的「隐藏损失」

在段然、周任飞和 Hongxun Wu 的第一篇论文《Faster Matrix Multiplication via Asymmetric Hashing》中，他们表明，施特拉森算法的进程可以大大加快。这一切要得益于他们称之为「隐藏损失」（hidden loss）的概念。周任飞表示，该概念深深地隐藏在以前的分析中，是无意中消除了太多块的结果。

激光法的工作原理是将重叠的块标记为垃圾，并安排处理，而其他块被认为有价值并将被保存。不过，选择过程有些随机。事实上，被标记为垃圾的块可能最终还是有用的。

这并不完全令人惊讶，但通过检查许多随机选择，段然团队确定激光法系统性地低估了块的价值，因此应该保存更多的块，减少扔掉的块。而且，正如通常的情况一样，更少的浪费可以转化为更高的效率。

对于段然团队的做法，Le Gall 认为，「能够保留更多块而不重叠，这种做法实现了更快的矩阵乘法算法。」

在证明了这种损失的存在后，段然团队修改了激光法标记块的方式，从而大大减少了浪费。他们将 ω 的新上限设定在了 2.371866 左右，这要比 Josh Alman 和 Vassilevska Williams 在 2020 年设定的上限 2.3728596 有所改进。

这看起来是一个不大的变化，将上限降低了大约 0.001，但这是自 2010 年以来科学家们看到的最大进步。相比之下，Vassilevska Williams 和 Alman 2020 年的结果只比之前的结果提高了 0.00001。

当然，对研究人员来说，最令人兴奋的不仅仅是新纪录本身，该记录并没有持续多久。事实上，这篇论文揭示了一种新的改进途径，而在此之前，这种途径完全没有被注意到。

Le Gall 称，近四十年来，每个人都依赖相同的激光法。随着段然等三位研究者的论文出现，我们可以做得更好。

因此，周任飞参与撰写的 2024 年 1 月的论文改善了这种新方法，进一步减少了隐藏损失。他们又进一步提高了 ω 的上限，使它降低到了 2.371552。

研究者还使用同样的方法来改进矩形（n×m）矩阵的乘法过程，该乘法过程在图论、机器学习和其他领域均有广泛应用。

沿着这些方向取得一些进一步的进展几乎是肯定的，但这是有限度的。2015 年，Le Gall 和两位合作者证明，目前的方法，也就是激光法，再加上 Coppersmith 和 Winograd 的方法，无法得到低于 2.3078 的 ω。

Le Gall 说：「要想进一步改进，就必须在 Coppersmith and Winograd 的原始方法基础上加以改进，而这种方法自 1987 年以来就没有真正改变过。」但到目前为止，还没有人提出更好的方法。也许根本就没有。

周任飞说：「改进 ω 实际上是理解这个问题的一部分。如果我们能很好地理解这个问题，就能设计出更好的算法。不过，人们对这个古老问题的理解还处于非常初级的阶段。」

关于文中介绍的Strassen算法和Coppersmith–Winograd算法的更多细节，可学习《通用矩阵乘（GEMM）优化算法 - 黎明灰烬》，博主除了介绍了这类基于算法分析的优化方法外，还介绍了基于软件优化的方法，根据计算机存储系统的层次结构特性，选择性地调整计算顺序，主要有循环拆分向量化、内存重排等。