DeepMind前不久发在Nature上的论文Discovering faster matrix multiplication algorithms with reinforcement learning引发热议。
这篇论文在德国数学家Volken Strassen「用加法换乘法」思路和算法的根本上，构建了一个基于AlphaZero的强化学习模型，更高效地探索进一步提高矩阵乘法速率的通用方法。
最近，Youtube播主Yannic Kilcher发布了一个长达近1小时的低廉甜头视频，由浅入深地沿着论文的脉络，对这个登上Nature封面的事情进行理解读。
基本思路：用加法换乘法
众所周知，矩阵乘法的传统算法是：两个矩阵行列交流相乘，然后求和，作为新矩阵的对应元素。个中涉及到大量的加法和乘法运算。
对付打算机来说，运算加法的速率要远远快于乘法，以是提升运算速率的关键，便是只管即便减少乘法运算的次数，纵然为此增加加法运算次数，对付打算加速的效果也是非常明显的。
遵照这个「用加法换乘法」的基本思路，德国数学家Volken Strassen于1969年创造了更高效、占用打算资源更少的矩阵乘法算法。
实际上，这个思路在一些最根本的数学公式中就已经有充分表示。比如平方差公式：
a^2-b^2 =（a+b)(a-b)
等号左侧打算两次乘法、一次加法，等号右侧打算一次乘法、两次加法。实际上，如果按照多项式乘法对等号右侧展开，实际上发生了正负ab的消去，将乘法运算的次数从4次降落为2次。
Strassen的算法是，利用原矩阵布局一些加乘结合的中间量，每个中间量只包含一次乘法打算，将原矩阵乘法转换为这些中间量的加法运算，将一些符号相反的乘法消去，实现降落乘法运算次数的目的。
在22矩阵的乘法中，Strassen的算法将乘法运算次数由8次降为7次。
矩阵乘法的张量表示和低秩分解那么下一个问题便是，如何找到一种算法，构建能够消去乘法运算的中间量，同时更方便地利用强化学习技能？
DeepMind给出的答案是：将矩阵乘法转换为「低秩分解」问题。
同样以22矩阵为例，利用三维张量来表示 AB=C 的矩阵乘法运算过程，个中旁边维度（列）为A，高下维度（行）为B，前后维度（深）为C。
用{0,1}对这个表示张量进行添补。C中取到值的部分，添补为1，别的添补为0。如下图所示。
比如，c1=a1b1+a2b3，在「最深一层」所表示的c1上，可以看到左上方（第1行第1列）的a1b1，和第3行第2列的a2b3被表示为紫色1，别的为白色0。
在张量表示后，可以通过对矩阵的「低秩分解」，设张量Tn为两个 n×n 矩阵相乘的表示张量。将Tn分解为r个秩一项(rank-one term)的外积。
两个n维向量的外积可以得到一个n×n的矩阵，三个n维向量的外积可以得到一个 n×n×n 的张量。
仍以Strassen的算法为例，低秩分解后的结果，即上式中的U、V、W对应为3个7秩矩阵。这里的分解矩阵的秩决定原矩阵乘法中乘法运算的次数。
实际上，用这个方法可以将n×n矩阵乘法的打算繁芜度降落至 O(Nlogn(R)) 。
由此可以设计一种规则，逐一对应地得到图（b）中的矩阵乘法算法，即论文中的「算法1」：
建模：基于强化学习的AlphaTensor
DeepMind利用强化学习演习了一个AlphaTensor智能体来玩一个单人游戏（Tensor Game），开始时没有任何关于现有矩阵乘法算法的知识。
这个强化学习模型正是基于此前的AI围棋大师AlphaZero。
那么这个游戏要如何设计，才能将其与矩阵乘法的简化建立联系，从而办理实际问题呢？
运用AlphaZero时，作者有一些分外的网络架构技巧。
他们利用了线性代数的某些属性，比如，纵然我们改变了线性运算的某些根本，问题也是同样的。因此，纵然我们改变了矩阵的根本，它在实质上仍旧代表同样的转换。
然而，对付这个算法来说，却不是这样的。
有了不同的数字，算法看起来就不同了，由于它是一种对彼此的转换。在这里，作者就很好地利用了线性代数的基本属性，创建出了更多的演习数据。
其余，分解3D张量看起来很难，但创造一个3D张量，就很随意马虎。
我们只需对添加的3个向量进行采样，把它们加在一起，就有了一个三维张量。经由精确的分解，它们还可以创建合成演习数据。
这些技巧都非常聪明，供应了更多的数据给系统。系统经由演习，可以准确地供应这些分解。
让我们剖析一下神经网络架构，它是一个基于Transformer的网络。
实质上，它是一个强化学习算法。
首先要输入当前的张量以及张量的历史，接着是躯干(Torso)，然后是嵌入(Embedding)，末了是Policy Head和Value Head。
在上图所指的位置，我们要选择三个向量u,v,w，进行相应打算。
一旦我们有三个向量的动作，我们就可以从原始张量中减去它。然后的目标是，找到从原始张量中减去的下一个动作。所有张量的Entry都是0的时候，游戏恰好结束。
这显然是一个离散问题。如果张量的阶数高于2，就属于NP hard。
这个任务实际上很艰巨，我们利用的是3个向量，每个向量都有对应的Entry，因此这是一个巨大的动作空间，比国际象棋或围棋之类的空间都大得多，因此也困难得多。
这是一个更风雅的架构图。他们把末了一个韶光步中涌现的张量的历史，用各种办法把投影到这个网格层上，然后线性层Grid 2将其转换为某种C维向量（这里韶光维度就减少了）。
在这里，我们输出一个策略，这个策略是我们动作空间上的一个分布，还有一个输出到Value Head。
Value Head是从Policy Head中获取嵌入，然后通过一些神经网络推动。
要点便是，将网络与蒙特卡洛树搜索匹配。
总结一下：为理解决这些游戏，开始，我们的矩阵是满的，棋盘处于初始状态，然后就要考虑不同的动作，每一步动作都会包含更多的动作，包括你的对手可能考虑到的动作。
这实在便是一个树搜索算法。现在Alpha Zero style的蒙特卡洛树搜索，便是通过神经网络的策略和代价函数，勾引我们完成这个树搜索。
它在用蓝线圈出的节点，就会向你提出建议，让你得到更成功的张量分解，也便是说，让你有更高的机率得胜。并且，它会直接打消掉你不该考试测验的步骤，缩小你的考虑范围。
你只须要搜索，然后通过迭代演习，在某个节点，得到Zero Tensor，就意味着你胜利了。
没有完成游戏的话，褒奖就非常低，反馈到演习神经网络之后，会做出更好的预测。
实际上，褒奖不止是0或1， 为了鼓励模型创造最短路径， 作者还设定了一个-1的褒奖。
这就比只给0或1的褒奖好得多，由于它鼓励了低阶的分解，还供应了更密集的褒奖旗子暗记。
由于问题很难，胜利具有很高的有时性，褒奖是稀少的。而如果走每一步都会得到褒奖，也有可能是-1的褒奖，就会敦促模型采纳更少的步骤。
更主要的是，在这个合成演示中，他们会匹配一个监督褒奖。
由于作者不仅可以天生数据，他们实际上是知道精确的步骤的，以是他们可以以监督的办法演习神经网络——由于是我们提出的问题，以是我们已经知道你该采纳哪些步骤了。
再回顾一下全体算法。
针对原始游戏，作者改变了basis，将数据增强，然后进行蒙特卡洛树搜索。几个树搜索之后，游戏结束，根据结果的输赢，会得到相应的褒奖，然后来演习。
把它放在游戏缓冲区，就可以更好地预测要实行的操作。
Policy Head会辅导你走哪条路，在某个节点，你可以问Value Head：现在的状态值是多少？把所有内容汇总到顶部，选择最有希望的步骤。这便是MCTS Alpha Zero style的简介。
作者的另一个巧思是：除了-1的褒奖，还在终端供应额外的褒奖。如果算法在英伟达V100或TPUv2上运行得很快，还会得到额外的褒奖。
AlphaTensor当然不知道V100是什么，但通过强化学习的力量，我们就可以找到在特定硬件上速率非常快的算法。
这样，我们就可以让算法提出定制的办理方案。
不仅是矩阵乘法，编译器也是这种事理。我们可以用这种方法，为特定的硬件优化速率、内存等。显然，它的运用领域已经远远超出了矩阵乘法。
对付数学的变革作者还创造，对付两个四乘四矩阵相乘的得到的T4，AlphaTensor创造了超过14,000个非等价分解。
每种大小的矩阵乘法算法多达数千种，表明矩阵乘法算法的空间比以前想象的要丰富。
对付关心繁芜性理论的数学家来说，这是一个巨大的创造。
参考资料：
https://www.reddit.com/r/MachineLearning/comments/xycz6y/d_alphatensor_explained_video_walkthrough/
https://youtu.be/3N3Bl5AA5QU
