用 FlexAttention 尝试一种新的注意力模式。
理论上,注意力机制就是你所需要的一切。然而在实际操作中,我们还需要优化像 FlashAttention 这样的注意力机制的实现。
尽管这些融合的注意力机制大大提高了性能,且支持长上下文,但这种效率的提升也伴随着灵活性的丧失。对于机器学习研究人员来说,这就像是一种「软件彩票」—— 如果你的注意力变体不适合现有的优化内核,你将面临运行缓慢和 CUDA 内存不足的困境。
一些注意力变体包括因果注意力、相对位置嵌入、Alibi、滑动窗口注意力、PrefixLM、文档掩码、不规则张量、PagedAttention 等。更糟糕的是,人们通常希望将这些变体组合在一起!比如滑动窗口注意力 + 文档掩码 + 因果注意力 + 上下文并行,又比如 PagedAttention + 滑动窗口的组合。
下图左侧代表了当今的现状 —— 一些掩码 + 偏置 + 设置的组合已经有现成的内核实现。然而,各种选项的添加会导致设置呈指数级增长。更糟糕的是,这种方式不会支持新的注意力变体。
![](http://n.sinaimg.cn/spider20240810/465/w1080h985/20240810/8755-d8e8b09726aca55f746b1c01485d0062.png)
为了彻底地解决这个超立方体问题,PyTorch 团队引入了 FlexAttention,一个新的 PyTorch API。
FlexAttention 是一个灵活的 API,允许用户使用几行惯用的 PyTorch 代码就能实现多个注意力变体。
团队人员通过 torch.compile 将其降低到一个融合的 FlashAttention 内核中 ,生成了一个不会占用额外内存且性能可与手写内核相媲美的 FlashAttention 内核。
利用 PyTorch 的自动求导机制自动生成反向传播。
最后,PyTorch 团队还可以利用注意力掩码中的稀疏性,从而显著改善标准注意力实现。
![](http://n.sinaimg.cn/spider20240810/154/w1080h674/20240810/c279-c51ac6cac52c5f0d11f0b8f9e7eeccbe.png)
FlashAttention 1-3 版本的参与者 Tri Dao 对这项研究进行了转发并评论:这项研究使得很多技术都融合在一起了。
![](http://n.sinaimg.cn/spider20240810/532/w1080h252/20240810/7dcb-d5e3dc1837e0022007193bc0b3455414.png)
FlexAttention
经典的注意力方程式如下:
![](http://n.sinaimg.cn/spider20240810/250/w894h156/20240810/6805-a706c86f0afa530f56a4f395c86ea50a.png)
代码形式:
![](http://n.sinaimg.cn/spider20240810/422/w1080h142/20240810/730d-4162bb9a4600892de39709dea7234771.png)
FlexAttention 形式如下,其通过接受用户定义的函数 score_mod 来解决上述问题。
![](http://n.sinaimg.cn/spider20240810/411/w1080h131/20240810/80aa-c8c2280a29e864d8d7a686c3c9b397ab.png)
代码形式:
![](http://n.sinaimg.cn/spider20240810/443/w1080h163/20240810/dd04-0f943c44f7db7f942dfe1294b91770f0.png)
此函数允许用户在 softmax 之前修改注意力分数。研究人员发现,该函数最终足以满足大多数用户对注意力变体的需求。
具体而言,score_mod 如下:
![](http://n.sinaimg.cn/spider20240810/374/w1080h94/20240810/118f-e3f1e19503933d48cb7e980491d31f35.png)
要应用此函数,可以将其实现为:
for b in range (batch_size):
for h in range (num_heads):
for q_idx in range (sequence_length):
for kv_idx in range (sequence_length):
modified_scores [b, h, q_idx, kv_idx] = score_mod (scores [b, h, q_idx, kv_idx], b, h, q_idx, kv_idx)
最终的 API 具有令人惊讶的表达能力。
Score Mod 示例
全注意力
在这种情况下,score_mod 无操作,它接受分数作为输入,然后原样返回它们。
![](http://n.sinaimg.cn/spider20240810/46/w742h104/20240810/6e90-fb93c845a6bc9c01bbe6ca3a13f4fd5f.png)
然后端到端的使用。
![](http://n.sinaimg.cn/spider20240810/248/w927h121/20240810/8193-a9b81ec3cc54133f2db5e4c9400e3183.png)
相对位置编码
一种常见的注意力变体是相对位置编码。相对位置编码不是对查询和键中的绝对距离进行编码,而是根据查询和键之间的距离调整分数。
![](http://n.sinaimg.cn/spider20240810/83/w790h93/20240810/3a3c-4c29571d33368934cf21b8c54af2b7a3.png)
需要注意的是,与典型实现不同,这不需要具体化 SxS 张量。相反,FlexAttention 会在内核中动态计算偏差值,从而显著提高内存和性能。
![](http://n.sinaimg.cn/spider20240810/15/w1080h535/20240810/b7e7-e5b18c20c41de2678dcd777e7e83b763.png)
Soft-capping
Soft-capping 是 Gemma 2 和 Grok-1 使用的一种技术,在 FlexAttention 中,它的形式是这样的:
![](http://n.sinaimg.cn/spider20240810/394/w960h234/20240810/f260-a04d0606ecf947b02ecaf6a2ee7f8cef.png)
Causal Mask
尽管双向注意力很简单,但在论文《Attention is All You Need》,以及其他的 LLM 中,它们的设置都是仅解码器的注意力,其中每个 token 只能关注它之前的 token。如果用户使用 score_mod API ,可以将其表示为:
![](http://n.sinaimg.cn/spider20240810/278/w973h105/20240810/48d9-b23f13835d6474b390375625aec5939d.png)
Sliding Window + Causal
![图源:https://arxiv.org/abs/2310.06825](http://n.sinaimg.cn/spider20240810/702/w1080h422/20240810/9d7d-a9aedddfa69e989d3c2fa95775b51d17.png)
Mistral 一直在推广滑动窗口注意力(也称为局部注意力),它允许查询 token 仅关注最近的 1024 个 token,通常与因果注意力一起使用。
![](http://n.sinaimg.cn/spider20240810/605/w940h465/20240810/bb6b-765ffc1ffaed454f5e65058842b82b43.png)
研究者对带有滑动窗口掩码的 F.scaled_dot_product_attention 以及带有因果掩码的 FA2 进行基准测试。结果表明,FlexAttention 不仅明显快于 F.scaled_dot_product_attention,也明显快于带有因果掩码的 FA2。
性能
总体而言,FlexAttention 的性能几乎与手写的 Triton 内核一样好。然而,由于 FlexAttention 具有通用性,因此会遭受轻微的性能损失。例如,用户必须承受一些额外的延迟。
FlexAttention 在前向传播中实现了 FlashAttention2 性能的 90%,在反向传播中实现了 85%。FlexAttention 目前正在使用一种确定性算法,该算法比 FAv2 重新计算了更多的中间体,研究者计划改进 FlexAttention 的反向算法,来缩小这一差距!
![参考链接:https://pytorch.org/blog/flexattention/](http://n.sinaimg.cn/spider20240810/52/w1080h572/20240810/f2dd-0982815f96cf5b1fb8af212bd2e40461.png)
![© THE END](http://n.sinaimg.cn/spider20240810/377/w1080h897/20240810/d06b-fea60c3d0340f75eb65db70a2f01bde8.jpg)
![新浪科技公众号 新浪科技公众号](http://n.sinaimg.cn/tech/content/tech_qr2x.png)
“掌”握科技鲜闻 (微信搜索techsina或扫描左侧二维码关注)
![](http://n.sinaimg.cn/tech/content/tech_weixin2.png)