土法炼钢兴趣小组的算法知识备份

【Transformer 与注意力机制】43|稀疏与局部注意力:Longformer、BigBird、Sparse Transformer

文章导航

分类入口
transformer
标签入口
#transformer#sparse-attention#longformer#bigbird#long-context

目录

42|FlashAttention 让标准 attention 更快、更省显存,但它没有改变一个事实:每个 token 仍然要和每个 token 建立关系。稀疏注意力问的是另一个问题:如果 full attention 太贵,能不能只保留一部分重要连接?

这个问题可以用图来理解。full attention 是完全图,每个位置都连向所有位置。稀疏 attention 则设计一张更便宜的图:局部相邻 token 互相连接,少数全局 token 负责跨段通信,有时再加随机边或固定 stride 边,让远距离信息不至于完全断开。

本篇能让你学会三件事:

  1. 稀疏 attention 和 FlashAttention 的根本区别;
  2. Longformer、BigBird、Sparse Transformer 各自怎么设计连接;
  3. 为什么稀疏 attention 很有用,却没有完全取代 full attention。

一、从完全图到稀疏图

标准 self-attention 让每个 query 看所有 key。这个设计表达能力强,也非常通用;代价是序列长度为 \(n\) 时有 \(n^2\) 条关系。

稀疏 attention 的核心就是减少边。最朴素的做法是局部窗口:每个 token 只看附近若干 token。语言和代码都有强局部性,很多依赖确实发生在附近。这样复杂度可以从二次降低到接近线性。

但只看局部会丢失远距离信息。长文档中的题目、摘要、段落标题、特殊分类 token,可能需要和全篇交互。于是很多稀疏方案会加入全局 token 或随机连接,让信息可以跨窗口传播。


二、Longformer:local + global

Longformer 面向长文档任务,把 attention pattern 改成局部窗口加少量全局 attention。普通 token 只看附近窗口;被指定为 global 的 token 可以看所有 token,也被所有 token 看见。

这很适合文档分类、长文问答等任务。比如 [CLS] 或问题 token 可以设为全局位置,用来汇聚整篇信息;正文 token 保持局部窗口,降低成本。

Longformer 的启发是:不是所有 token 都需要同等全局视野。任务相关的少数位置可以承担中转站角色,大部分位置只处理局部上下文。


三、BigBird:local + global + random

BigBird 在局部和全局之外加入随机连接。随机边的作用是改善图连通性,让远处 token 之间可以通过较短路径互相影响。

BigBird 论文还讨论了这种稀疏模式在理论上保留某些 full attention 能力的条件。对工程读者来说,更重要的直觉是:长序列不一定需要完整连接,但稀疏图必须避免被切成互不通信的小块。

local 负责近邻结构,global 负责任务级汇聚,random 负责远距离桥接。三者组合比单纯局部窗口更稳健。


四、Sparse Transformer:结构化稀疏

Sparse Transformer 更早地探索了结构化稀疏模式,例如固定 stride、block pattern 等。它试图让每个位置只 attend 到一组规则选择的位置,降低计算成本,同时保留跨距离依赖。

结构化稀疏的优势是模式规则,便于实现和分析。缺点是归纳偏置更强:如果真实任务依赖的关系不符合预设模式,模型可能很难补救。

这也是稀疏 attention 的核心 trade-off:你用连接规则换效率,但也把“什么关系重要”提前写进了模型。


五、工程难点

稀疏 attention 理论上少算,但真实硬件不一定自动变快。GPU 喜欢规则、密集、大块矩阵计算。稀疏模式如果太碎,kernel 调度和内存访问可能抵消理论收益。

这就是为什么 FlashAttention 这类 full attention 优化非常有竞争力。它不改变连接结构,仍能利用密集矩阵和硬件特性。稀疏 attention 则需要同时赢复杂度和硬件实现。

此外,不同任务需要不同稀疏模式。长文档、代码、视频、基因序列的结构不一样,一套 pattern 未必通吃。


六、为什么没有完全取代 full attention

full attention 的优势是简单、通用、稳定。模型不用预先假设哪些 token 应该互相连接,而是在训练中自己学。对于中短序列,full attention 的成本还能接受,稀疏化的收益未必大于复杂度。

稀疏 attention 的价值主要在长序列场景。它是把上下文窗口推长的重要工具,但不是标准 attention 的全面替代。很多系统会把它和 FlashAttention、滑动窗口、检索、压缩记忆一起使用。


七、关键概念回顾


八、常见误解

8.1 “稀疏 attention 一定更快”

理论计算少,不代表实际 GPU 吞吐一定更高。稀疏模式需要硬件友好实现。

8.2 “稀疏 attention 一定损失能力”

也不一定。对长文档等结构明显的任务,合适的稀疏 pattern 可以用更低成本取得好效果。

8.3 “FlashAttention 和稀疏 attention 是同一种东西”

不是。FlashAttention 是 exact full attention 的 I/O 优化;稀疏 attention 改变连接关系。


九、下一步

稀疏 attention 省的是 token 之间的连接。下一篇看另一种稀疏:参数很多,但每个 token 只激活一小部分专家,这就是 MoE。


十、参考文献

  1. Beltagy, I. et al. “Longformer: The Long-Document Transformer.” arXiv:2004.05150, 2020.
  2. Zaheer, M. et al. “Big Bird: Transformers for Longer Sequences.” NeurIPS 2020.
  3. Child, R. et al. “Generating Long Sequences with Sparse Transformers.” arXiv:1904.10509, 2019.
  4. Dao, T. et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022.

← 上一篇:42|FlashAttention | 下一篇:44|MoE

同主题继续阅读

把当前热点继续串成多页阅读,而不是停在单篇消费。


By .