引言: 注意力机制的难以忍受的缓慢
基于 Transformer 的模型,如 BERT 和 GPT,已经彻底改变了自然语言处理 (NLP) ,在从情感分析到文本生成的各种任务上都取得了最先进 (state-of-the-art) 的成果。它们可以编写代码、总结文章,甚至进行令人惊讶的连贯对话。但这种强大的能力背后是高昂的代价: 计算成本。
这些模型的秘诀在于注意力机制,这是一种巧妙的技术,让模型能够衡量句子中不同词的重要性。问题在于,注意力的复杂度是二次方的,这意味着其计算成本会随着输入句子长度的平方增加。处理一个 100 词的句子是一回事,但处理一篇 1000 词的文档,成本就要高出 100 倍。
这种低效率是一个主要瓶颈。在高端 TITAN Xp GPU 上,使用 GPT-2 生成一个 30-token 的句子可能需要 370 毫秒——比用现代 CNN 分类一张图片慢了两个数量级。在像树莓派这样的低功耗设备上,这个时间会激增到令人痛苦的 43 秒。这一性能障碍阻碍了这些强大的模型在手机和其他边缘设备上的广泛部署。
如果我们能让这些模型更“斯巴达”呢?如果我们能剔除冗余,只保留必要的计算,同时不牺牲准确率呢?这正是 SpAtten 背后的核心思想 —— 由麻省理工学院研究人员提出的一项开创性算法-架构协同设计。SpAtten 通过利用人类语言中的固有冗余,大幅加速注意力机制。它能智能地、动态地剪除不重要的词和模型部分,从而减少计算量和内存访问。
图 1: 级联 token 和 head 剪枝在不影响准确率的情况下,逐层渐进地移除冗余。
在 BERT 模型的第三层,SpAtten 可能只处理核心概念 —— “film perfect” —— 剪掉了 88% 的原始工作量,仍然能正确完成情感分类。本文将深入探讨 SpAtten 如何通过三项关键创新实现这些效率提升:** 级联 Token 剪枝**、级联 Head 剪枝和渐进式量化。
背景: 为什么注意力机制如此缓慢?
要理解 SpAtten 的解决方案,我们首先需要弄清基于注意力的模型的瓶颈所在。
Transformer 架构剖析
NLP 模型通常分为两类:
- 判别式模型 (如 BERT) : 用于理解和分类文本,输出一个分类结果 (如“正面”情感) 或回归分数。这被称为摘要阶段,模型将输入信息提炼成最终预测。
- 生成式模型 (如 GPT-2) : 用于生成新文本,先对输入提示进行摘要阶段处理,再进入生成阶段,一次生成一个 token。
图 3: BERT 仅有摘要阶段,而 GPT-2 有摘要和生成两个阶段。
这两种模型都由相同的模块堆叠而成,每个模块包含注意力层和前馈层。注意力机制需要三个输入:** 查询 (Query, Q)** 、键 (Key, K) 和值 (Value, V) 。简单来说,Q 表示一个词在“寻找”什么,K 表示一个词“提供”什么,而 V 是其内容。注意力分数通过 \( Q \times K^T \) 计算,经过 Softmax 归一化为概率,再用于加权求和 V (\(\text{attention\_prob} \times V\)) 。这样,每个词都能关注其他所有词。
新架构的动机
那么成本来自哪里呢?我们看看 GPT-2 的性能分析。
图 2: 注意力机制占了 GPT-2 运行时间的一半以上;大部分延迟来自数据搬移,而非运算。
左图显示,尽管注意力机制的 FLOPs 占比很小,却占用了超过一半的总时间。右图揭示了注意力延迟的 73% 源于数据搬移——如拆分 head、重塑形状、转置等。这是一个内存受限问题。仅提升计算速度不足以解决问题;我们需要减少内存流量。这正是 SpAtten 的指导原则。
核心方法: SpAtten 的三管齐下策略
SpAtten 结合了三种互补的算法优化:
1. 级联 Token 剪枝
自然语言充满冗余。例如句子 “It is a very fun movie, I think.”,情感的核心是 “fun movie”,其余都是“填充”。
SpAtten 基于注意力概率计算累积 Token 重要性分数: 如果许多 token 都关注某个 token,那么它的重要性就高。
图 4: 级联剪枝利用累积的重要性分数和 top-k 选择逐步移除 token 和 head。
这些重要性分数会跨 head 和层累积 (图 5 示出了相关分数) 。那些分数持续较低的 token,如 “a” 或 “the”,会被剪掉。
图 5: 累积重要性分数低的 token 会被移除。
一旦某个 token 被剪掉,其 Q、K 和 V 向量会从所有后续层删去。这种级联效应在网络更深处同时减少注意力和前馈网络 (FFN) 的计算量。
2. 级联 Head 剪枝
多头注意力使模型可关注不同类型的关系,但很多 head 是冗余的。
SpAtten 通过输出向量的幅度来评分 head。输出幅度大的 head 影响力更大,会保留;影响小的 head 会被级联剪掉,减少特征维度,并与 token 剪枝形成互补。
3. 渐进式量化
SpAtten 还通过量化数值精度进一步减少内存访问,并且是自适应的。
图 7: 平坦分布会放大量化误差;尖峰分布则抑制误差。
量化注意力分数 \( \Delta s \) 产生的误差取决于概率分布形状。Softmax 的导数为:
[ \frac{\partial p_i}{\partial s_j} = \begin{cases} p_i (1 - p_i), & i = j \
- p_i p_j, & i \neq j \end{cases} ]
如果 \(p_i\) 接近 1 或 0,则误差较小;均匀分布则误差更大。
SpAtten 的渐进式量化流程:
图 6: 先获取高位比特 (MSB) ;仅当分布平坦时获取低位比特 (LSB) 。
- **先获取高位比特 (MSB) **(如前 4 位) 。
- 检查概率分布形状 —— 若最大概率 < 阈值,则分布平坦。
- 分布平坦时获取低位比特 (LSB) 并重新计算;尖峰分布则跳过。
这样在需要时保证精度,同时最小化平均内存流量。
SpAtten 硬件架构
为充分发挥这些算法优势,SpAtten 配备专用加速器。
图 8: 全流水线架构整合了剪枝与量化逻辑。
流水线步骤:
- Top-k 选择: 通过累积 token 分数选出最重要的 K。
- 数据获取: 从 DRAM 读取 K 向量。
- Q×K 乘法: 并行乘法器阵列计算得分。
- Softmax 与量化检查: 判断是否获取 LSB。
- 局部值剪枝: Top-k 引擎选出重要的 V。
- Attention×V 乘法: 输出注意力结果。
高并行 Top-k 引擎
标准的 top-k 排序太慢。SpAtten 使用基于快速选择的硬件引擎,平均复杂度为 \(O(n)\)。
图 9: 快速选择将输入划分为两个 FIFO;重复基准选择可迅速锁定 top-k 阈值。
该引擎避免成为瓶颈,实现实时剪枝。
专用计算模块
Query-Key 乘法模块拥有 512 个乘法器和可配置加法树,以支持不同向量尺寸。
图 11: 可重构加法树可在全吞吐量下支持多种 Query/Key 维度。
Softmax-量化模块专为 SpAtten 的渐进式策略定制。
图 12: Softmax 结果进入基于阈值的检查,判断是否获取低位比特。
实验与结果: 验证 SpAtten
SpAtten 在 30 个基准测试中评估,涵盖 BERT 与 GPT-2 任务,并与 CPU、GPU 及其他加速器对比。
相比通用硬件的显著加速
SpAtten 在不损失准确率的情况下,将 DRAM 访问减少 10 倍。
图 14: 相比 TITAN Xp 加速 162 倍,相比 Xeon 加速 347 倍,在移动平台上超过 1000 倍;能耗降低高达 4000 倍。
超越其他加速器
对比 A³ 和 MNNFast,SpAtten 在速度、能效和功能上全面领先。
特性 | MNNFast | A³ | SpAtten |
---|---|---|---|
级联 Head 剪枝 | ❌ | ❌ | ✅ |
级联 Token 剪枝 | ❌ | ❌ | ✅ |
渐进式量化 | ❌ | ❌ | ✅ |
减少 DRAM 访问 | ❌ | ❌ | ✅ |
加速对象 | 仅 BERT | 仅 BERT | BERT & GPT-2 |
吞吐量 (GOP/s) | 120 | 221 | 360 |
能效 (GOP/J) | 120 | 269 | 382 |
性能提升来源分析
图 20: 数据路径: +22.1×;级联剪枝: +3.4×;top-k 引擎: +3×;渐进式量化: +2.8× → 总计 >200×。
剪掉了什么?解读 SpAtten
分类任务中被剪的 token 通常是语法填充词 (如“it”、“is”) ,保留语义更强的词。相似性任务中保留主要名词/动词。GPT-2 生成任务中,剪枝会移除无关上下文,聚焦于相关源词。
图 23: 如“published”和“researcher”等 token 始终重要,体现稳定的语义焦点。
结论与启示
SpAtten 展示了算法-硬件协同设计在加速 NLP 中的巨大潜力:
- 利用稀疏性: 利用语言冗余剪掉 token 与 head。
- 级联决策: 将剪枝决策传递至所有后续层以最大化节省。
- 自适应性: 通过渐进式量化将精度与输入难度匹配。
- 专用硬件: 定制模块尤其是 top-k 引擎在实现性能收益中至关重要。
通过采用一种斯巴达式 —— 计算与内存均高度节俭 —— 的方法,SpAtten 为现代 NLP 模型在从数据中心 GPU 到智能手机等多种设备上的实际应用打开了可能性。