人工智能揭示矩阵乘法的新可能性

人工智能揭示矩阵乘法的新可能性

---

数学家酷爱漂亮的谜题。当你尝试找到最有效的方法时，即使像乘法矩阵（二维数字表）这样抽象的东西也会感觉像玩一场游戏。这有点像尝试用尽可能少的步骤解开魔方——具有挑战性，但也很诱人。除了魔方，每一步可能的步数为 18；对于矩阵乘法，即使在相对简单的情况下，每一步都可以呈现超过 10^12 个选项。

在过去的 50 年里，研究人员以多种方式解决了这个问题，所有这些都是基于人类直觉辅助的计算机搜索。2022 年 10 月，人工智能公司 DeepMind 的一个团队展示了如何从一个新的方向解决这个问题，在《Nature》杂志的一篇论文中报告说，他们已经成功地训练了一个神经网络来发现新的快速矩阵乘法算法。就仿佛 AI 找到了解决极其复杂的魔方的未知策略。

论文链接：https://www.nature.com/articles/s41586-022-05172-4

「这是一个非常巧妙的结果。」哥伦比亚大学计算机科学家 Josh Alman 说，但他和其他矩阵乘法专家也强调，这种人工智能辅助将补充而不是取代现有方法——至少在短期内是这样。「这就像对可能成为突破的事物的概念验证。」Alman 说。结果只会帮助研究人员完成他们的任务。

仿佛是为了证明这一点，《自然》杂志的论文发表三天后，两位奥地利研究人员展示了新旧方法如何相互补充。他们使用传统的计算机辅助搜索来进一步改进神经网络发现的一种算法。

论文链接：https://arxiv.org/pdf/2210.04045.pdf

结果表明，就像解决魔方的过程一样，通往更好算法的道路将充满曲折。

乘法矩阵

矩阵乘法是所有数学中最基本和最普遍的运算之一。要将一对 n×n 矩阵相乘，每个矩阵都有 n^2 个元素，你可以将这些元素以特定组合相乘并相加以生成乘积，即第三个 n×n 矩阵。将两个 n×n 矩阵相乘的标准方法需要 n^3 次乘法运算，因此，例如，一个 2×2 矩阵需要八次乘法。

对于具有数千行和列的较大矩阵，此过程很快就会变得麻烦。但在 1969 年，数学家 Volker Strassen 发现了一种使用七个而不是八个乘法步骤将一对 2×2 矩阵相乘的过程，代价是引入了更多的加法步骤。

如果你只想乘以一对 2×2 矩阵，则 Strassen 算法不必要地复杂化。但他意识到它也适用于更大的矩阵。那是因为矩阵的元素本身可以是矩阵。例如，可以将具有 20,000 行和 20,000 列的矩阵重新设想为一个 2×2 矩阵，其四个元素各为 10,000×10,000 矩阵。然后可以将这些矩阵中的每一个细分为四个 5,000×5,000 块，依此类推。Strassen 可以应用他的方法在此嵌套层次结构的每一层上乘以 2×2 矩阵。随着矩阵大小的增加，减少乘法所节省的成本也会增加。

Strassen 的发现促使人们开始寻找高效的矩阵乘法算法，此后启发了两个不同的子领域。一个侧重于一个原则问题：如果你想象将两个 n×n 矩阵相乘并让 n 趋于无穷大，那么最快的算法中的乘法步骤数如何与 n 成比例？

最佳缩放比例的当前记录 n^2.3728596 属于麻省理工学院计算机科学家 Alman 和 Virginia Vassilevska Williams。（最近发布的预印本报告了使用新技术的微小改进。）但这些算法纯粹是理论上的兴趣，仅在荒谬的大矩阵上胜过 Strassen 等方法。

论文链接：https://arxiv.org/abs/2210.10173

第二个子领域的思考规模较小。在 Strassen 的工作之后不久，以色列裔美国计算机科学家 Shmuel Winograd 表明 Strassen 已经达到了理论极限：不可能用少于七个乘法步骤来乘以 2×2 矩阵。但对于所有其他矩阵大小，所需乘法的最小数量仍然是一个悬而未决的问题。小矩阵的快速算法可能会产生巨大的影响，因为当乘以合理大小的矩阵时，这种算法的重复迭代可能会击败 Strassen 的算法。

论文链接：https://www.sciencedirect.com/science/article/pii/0024379571900097

不幸的是，可能性的数量是巨大的。即使对于 3×3 矩阵，「可能的算法数量也超过了宇宙中的原子数量，」DeepMind 研究员兼新工作的负责人之一 Alhussein Fawzi 说。

面对这些令人眼花缭乱的选项，研究人员通过将矩阵乘法转化为一个看起来完全不同的数学问题——一个计算机更容易处理的问题——取得了进展。可以将两个矩阵相乘的抽象任务表示为一种特定类型的数学对象：称为张量的三维数字数组。然后，研究人员可以将这个张量分解为基本分量的总和，称为「rank-1」张量；这些中的每一个都代表相应矩阵乘法算法中的不同步骤。这意味着找到一个有效的乘法算法相当于最小化张量分解中的项数——项越少，涉及的步骤就越少。

通过这种方式，研究人员发现了新的算法，可以使用比标准 n^3 更少的乘法步骤来乘以许多小矩阵大小的 n×n 矩阵。但是，不仅优于标准而且优于 Strassen 小矩阵算法的算法仍然遥不可及——直到现在。

Game On

DeepMind 团队通过将张量分解变成单人游戏来解决这个问题。他们从 AlphaGo 的深度学习算法入手——AlphaGo 是另一个 DeepMind AI，它在 2016 年学会了玩棋盘游戏 Go，足以击败顶尖的人类棋手。

所有的深度学习算法都是围绕神经网络构建的：人工神经元网络被分类成层，连接强度可以变化，代表每个神经元对下一层神经元的影响程度。这些连接的强度在训练过程的多次迭代中得到调整，在此期间神经网络学习将它接收到的每个输入转换为有助于算法实现其总体目标的输出。

在 DeepMind 的新算法（称为 AlphaTensor）中，输入代表通往有效矩阵乘法方案的步骤。神经网络的第一个输入是原始矩阵乘法张量，其输出是 AlphaTensor 为其第一步选择的 rank-1 张量。该算法从初始输入中减去这个 rank-1 张量，产生一个更新的张量，该张量作为新输入反馈到网络中。重复该过程，直到最终起始张量中的每个元素都减少为零，这意味着没有更多的 rank-1 张量可以取出。

在这一点上，神经网络发现了一个有效的张量分解，因为它在数学上保证了所有 rank-1 张量的总和恰好等于起始张量。到达那里所采取的步骤可以转换回相应的矩阵乘法算法的步骤。

游戏是这样的：AlphaTensor 反复将张量分解为一组 rank-1 分量。每次，如果 AlphaTensor 找到减少步数的方法，它就会获得奖励。但胜利的捷径一点也不直观，就像你有时必须在魔方上拼凑出一张完美有序的脸，然后才能解决整个问题。

该团队现在有了一种算法，理论上可以解决他们的问题。他们只需要先训练它。

新路径

与所有神经网络一样，AlphaTensor 需要大量数据进行训练，但张量分解是一个众所周知的难题。研究人员可以为网络提供有效分解的例子很少。相反，他们通过在更简单的逆问题上进行训练来帮助算法开始：将一堆随机生成的 rank-1 张量相加。

布朗大学计算机科学家 Michael Littman 说：「他们正在使用简单的问题为困难的问题生成更多数据。」将这种向后训练过程与强化学习相结合，其中 AlphaTensor 在寻找有效分解时会产生自己的训练数据，其效果比单独使用任何一种训练方法都要好得多。

DeepMind 团队训练 AlphaTensor 来分解代表矩阵乘法的张量，最高可达 12×12。它寻求用于将普通实数矩阵相乘的快速算法，以及特定于更受约束的设置（称为模 2 运算）的算法。（这是仅基于两个数字的数学，因此矩阵元素只能是 0 或 1，并且 1 + 1 = 0。）研究人员通常从这个更受限制但仍然广阔的空间开始，希望这里发现的算法可以适用于实数矩阵。

训练后，AlphaTensor 在几分钟内重新发现了 Strassen 的算法。然后，它针对每种矩阵大小发现了多达数千种新的快速算法。这些与最著名的算法不同，但乘法步骤数相同。

在少数情况下，AlphaTensor 甚至打破了现有记录。它最令人惊讶的发现发生在模 2 运算中，它发现了一种新算法，可以在 47 个乘法步骤中将 4×4 矩阵相乘，比 Strassen 算法两次迭代所需的 49 个步骤有所改进。它还打破了最著名的 5×5 模 2 矩阵算法，将所需的乘法次数从之前的 98 次记录减少到 96 次。（但这个新记录仍然落后于使用 5×5 矩阵击败 Strassen 算法所需的 91 步。）

这一引人注目的新结果引起了很多兴奋，一些研究人员对基于 AI 的现状改进大加赞赏。但并非矩阵乘法领域中的每个人都对此印象深刻。「我觉得它有点被夸大了，」Vassilevska Williams 说。「这是另一种工具。这不像是，[哦，计算机打败了人类，] 你知道吗？」

研究人员还强调，破纪录的 4×4 算法的直接应用将受到限制：它不仅只在模 2 算法中有效，而且在现实生活中，除了速度之外还有其他重要的考虑因素。

Fawzi 也认为，这仅仅是个开始。「有很大的改进和研究空间，这是一件好事，」他说。

最后的转折

相对于成熟的计算机搜索方法，AlphaTensor 的最大优势也是它最大的弱点：它不受人类直觉的约束，无法判断好的算法是什么样子的，因此它无法解释自己的选择。这使得研究人员很难从其成就中学习。

但这可能并没有看上去那么大的缺点。在 AlphaTensor 结果公布几天后，奥地利林茨大学（JKU）的数学家 Manuel Kauers 和他的研究生 Jakob Moosbauer 报告说又向前迈进了一步。

Manuel Kauers 调整了 DeepMind 的方法以产生进一步的改进。——Jakob Moosbauer

当 DeepMind 论文发表时，Kauers 和 Moosbauer 正在使用传统的计算机辅助搜索来寻找新的乘法算法。但与大多数以新的指导原则重新开始的此类搜索不同，他们的方法通过反复调整现有算法来工作，希望从中节省更多的乘法。以 AlphaTensor 的 5×5 模 2 矩阵算法为起点，他们惊奇地发现，他们的方法在短短几秒钟的计算之后，就将乘法步骤从 96 步减少到了 95 步。

AlphaTensor 还间接帮助他们进行了另一项改进。此前，Kauers 和 Moosbauer 并没有费心去探索 4×4 矩阵的空间，他们认为不可能击败 Strassen 算法的两次迭代。AlphaTensor 的结果促使他们重新考虑，在从头开始计算一周后，他们的方法出现了另一种 47 步算法，与 AlphaTensor 发现的算法无关。「如果有人告诉我们 4×4 有什么值得发现的东西，我们本可以早点做到这一点，」Kauers 说。「但是，好吧，这就是它的工作原理。」

Littman 预计会有更多这样的惊喜，他将这种情况比作跑步者第一次在四分钟内跑完一英里，这一壮举曾被广泛认为是不可能的。「人们就像，[哦，等等，我们可以做到这一点，] 现在很多人都可以做到，」他说。

展望未来，Fawzi 希望推广 AlphaTensor 以解决更广泛的数学和计算任务，就像它的祖先 AlphaGo 最终扩展到其他游戏一样。

Kauers 认为这是将机器学习应用于发现新算法的真正试金石。他指出，寻求快速矩阵乘法算法是一个组合问题，无论有无人工协助，计算机搜索都非常适合。但并不是所有的数学问题都那么容易确定。他说，如果机器学习能够发现一种全新的算法理念，「这将改变游戏规则。」