Flash attention 公式解读

论文链接

https://zhuanlan.zhihu.com/p/663932651

这个论文解释的挺好，不过不足之处在于对于公式为什么这样推倒没有做解释，对于小白来说可能未必很友好（当然我也是小白，所以开始读的时候有点费劲），所以此处增加几处注释。

以下章节为注释点。

为什么可以从3个循环，变成两个循环（原文2.4）

朴素的softmax方式我们已经知道怎么做了，是一个比较容易理解的三层循环。

我们写个简单的例子

2.4章节里所谓的3个循环，变成两个循环，其实就是把第二个循环里的逻辑和第一个循环里的逻辑进行了融合。

我们可以看到第二个循环里，标红的mN，这个是第二次循环必须存在的理由。

所以出现了如下公式

该公式上来先看第一行，也即红框部分。我们试着写在上面的图里。

我们可以看到，最终红框里两个结果是一模一样的，证明按照变成循环的方式，是可行的，是完全可以把第二层循环融入到第一层循环的。

但是这里出现了新问题，我们这里四个元素，循环四次，按照新的方式，也是每一次都要累加啊，尤其是最后一次循环，累加了三个元素，有点拆东墙补西墙的意思，如果第二次循环逻辑融入到第一次循环里，如下图所示

甲逻辑变成乙逻辑，我们肯定一定有优化，但是甲逻辑变成丙逻辑，我们就有点不确定了。

所以，别急。

红框下面的公式推导，成功的把红框里的累加，变成了一个非累加操作。说白了就是借助最外层的i循环，进行递推。

所以可以清楚的看到

3 层循环变成了2层(图中代码缩进有点问题).

学到的一个重要的形式

如上章节我们知道了三层怎么变成两层的，但是还有一个重要的形式需要记住。

如果我们再碰到这种形式，就能马上联想到上面的第二层循环优化，看看这个xxx逻辑能不能和上面的循环进行融合(方式如上一个章节)，后文的attention阶段，把2层循环变成1层循环，用的就是这个解决方式。

attention： 2层循环优化成了1层循环（原文2.5）

正文来啦。

这里写出了2层循环逻辑的attention逻辑。

我们注意到第二层循环用到了mN这个变量。我们想想办法，能不能照葫芦画瓢，把最后一层循环给优化掉呢？

第一步，对公式优化下。

上面的红框优化成了下面的红框，变成了我们刚学到的那个重要的形式。

第二步，

开始改写下面的红框，变成累加形式。

我们推测，这种累加形式，应该和上面一样，可以变成非累加形式，这里推导过程如下。

最终，优化成了一层。

经验

学习东西，最好学完后自己给自己讲一遍，问题上什么，怎么解决的，为什么这么解决，有哪些隐藏的细节自己没有发现，书要越读越厚。