迈向高效智能--Transformer的效率革命
近年来,随着 ChatGPT 的横空出世,大语言模型成功实现“破圈”,深入融入人们的工作与生活,显著提升了各行各业的智能化水平。作为 ChatGPT 背后的核心架构,Transformer 技术也随之成为人工智能研究的焦点。Transformer 虽以强大的建模能力在自然语言处理、计算机视觉等多个领域取得了突破性的 SOTA 成果,但其固有的 计算复杂度,尤其是在处理长文本或大规模图数据时,极大限制了其扩展性和实际应用效率。为应对这一挑战,近年来研究社区掀起了一场 Transformer 的效率革命,提出了一系列旨在提升计算效率的 Transformer 架构优化方法。本文围绕“迈向高效智能——Transformer 的效率革命”这一主题,系统梳理并解析近年来的几项代表性成果,包括 Google的Performer[1]、Meta的Linformer[2]、Stanford的FlashAttention[3]、DeepSeek的NSA[4]以及Kimi的Moba[5]。希望通过本文的介绍,读者能够全面理解这些高效架构的技术路线及其背后的设计动因,并为后续在长文本建模、低延迟部署与高效推理等实际场景中提供理论启发与工程指导。
Preliminary
首先介绍一下 Transformer 架构中最为核心的自注意力模块,其数学形式如下:
其中 是自注意力模块的输入特征矩阵,是三个线性变换矩阵,分别对应 query, key 和 value。从数学表达式不难看出,自注意力的主要计算及存储开销来源于 和 这两个矩阵的乘积导致的 复杂度。
Performer
既然Transformer 的高昂复杂度主要源于自注意力机制中的 矩阵乘积,那么一个自然的想法是:能否在不显式计算该注意力矩阵的前提下,降低计算复杂度?注意到,若忽略自注意力计算中的 函数,其表达式可近似为矩阵乘积形式:
根据矩阵乘法的结合律,可以按如下顺序计算:
因此,整体复杂度为:时间复杂度:,空间复杂度:。问题在于:如何处理 函数的非线性操作?Performer [1] 提出了一个巧妙的解决方案 —— 从核函数(Kernel Function)的角度出发,对自注意力进行近似。具体地,其将注意力矩阵 定义为:
其中 是注意力矩阵第 行第 列元素;表示预定义的核函数;是该核函数对应的向量变换函数。
在原文中,作者使用随机特征映射作为 的实现方式;而在后续的一些工作中(如 Roformer [6]),则使用 的形式。理论上已经证明,该核函数方法可以有效地近似 函数。因此,Performer 的自注意力机制最终被表示为:
其中:,表示对注意力矩阵的归一化操作。
Linformer
与 Performer 通过核函数近似突破 限制、利用矩阵乘法结合律实现线性时间复杂度的方式不同,Meta 团队提出的 Linformer[2] 则从另一角度出发:压缩与矩阵的维度,以实现高效的 Transformer 架构。继续回顾 Transformer 自注意力计算的过程,可以发现其复杂度的核心原因在于两个矩阵的大小都是的,如果我们可以将矩阵的大小压缩至,其中,那么整体的复杂度就可以降低到,从而实现线性复杂度的高效 Transformer。为了实现矩阵的压缩,Linformer 引入了两个可学的线性投影矩阵,在计算得到矩阵后分别对其进行线性投影,得到,从而 Linformer 的自注意力机制最终被表示为:
Flash-Attention
上述的 Performer 和 Linformer 分别通过核函数近似与键值压缩在接近线性的复杂度内实现了对原始自注意力的近似,斯坦福大学提出的 FlashAttention[3] 则是从硬件级别的 IO 优化出发,显著提升了自注意力计算的效率。具体来说,GPU 的存储体系通常分为两个层级:高速但容量有限的 SRAM,以及容量大但访问速度较慢的 HBM(高带宽内存)。一次完整的计算过程包括三个步骤:
从 HBM 加载数据至 SRAM, 在 SRAM 中进行计算, 将结果写回 HBM。
在标准的自注意力实现中,如 的乘积、操作以及注意力矩阵与的乘积,分别发生在多个独立步骤中,每一步都涉及对规模数据的频繁读写,严重影响了整体计算效率。FlashAttention 的核心思想是将上述多个操作融合为一个单一的计算核(kernel),在 SRAM 中一次性完成计算、Softmax 运算和乘积,避免了重复的数据搬运,同时也避免了将庞大的注意力矩阵显式写回 HBM。这种方法不仅降低了 IO 负担,也实现了更高效、更节能的自注意力计算。
下面我将给出一份简化的伪代码对其计算流程加以说明:
# 输入:# q: 查询向量 [d]# K, V: 键值序列,长度为 N,每个向量维度为 [d]# 输出:# attention_output: [d_v]max_score = -inf # 用于数值稳定的 log-sum-expsum_exp = 0# Softmax 分母累加器weighted_sum = [0] * d_v # 输出向量初始化for i in range(N): # 对每个 token i k_i = load(K[i]) # 每次加载一个 key v_i = load(V[i]) # 每次加载一个 value score = dot(q, k_i) # 计算 score_i = q · k_iif score > max_score:# 重新缩放已有 softmax 权重(避免溢出) scale = exp(max_score - score) sum_exp = sum_exp * scale + 1 weighted_sum = [x * scale for x in weighted_sum] weighted_sum = [x + y for x, y in zip(weighted_sum, v_i)] max_score = scoreelse:# 正常累加当前 softmax 权重 exp_score = exp(score - max_score) sum_exp += exp_score weighted_sum = [x + exp_score * y for x, y in zip(weighted_sum, v_i)]# 输出归一化的 attention 向量output = [x / sum_exp for x in weighted_sum]
更为详细的伪代码实现可参考 FlashAttention 原论文。通过分析其伪代码可以看出,FlashAttention 的核心之一在于如何以流式(Streaming)方式计算 ,该技巧也被称为在线 SoftMax(Online SoftMax)。下面我们通过一个简单的例子,直观展示在线 SoftMax 的具体计算过程:
Calculate
Init
Received 1,
Received 2,
Received 3,
NSA && MoBA
NSA[4] 和 MoBA[5] 是近期由 DeepSeek 和 Kimi 团队提出的两项旨在提升 Transformer 效率、解决大语言模型长文本建模难题的工作。这两篇方法在核心思想上高度一致:通过自适应地选择参与注意力计算的 key 和 value,从而显著降低计算和存储开销。
NSA 首先将长文本序列划分为多个段落,并对每个段落内部的 表征进行压缩,生成该段落的汇总表示 。对于任意查询向量 ,NSA 先通过粗粒度的注意力机制与各段的 交互,计算出段落级注意力分数,并基于这些分数选择若干相关段落。随后,仅与被选中段落内的原始词元进行细粒度的注意力计算,以提升效率并保留细节信息。此外,NSA 还引入基于滑动窗口的局部注意力掩码,以增强对邻近词元的建模能力。最终,NSA 将局部、段落级粗粒度、全局细粒度三种注意力结果通过一个门控机制进行融合,生成最终的输出。
MoBA 的整体框架与 NSA 类似,但在实现细节上更为简洁。它使用非可学习的池化操作替代 NSA 中用于段落表示的可学习压缩模块;同时也未采用滑动窗口掩码来强化局部建模,而是依赖梯度下降过程让模型自然学习这种结构性先验。这种简化设计使得 MoBA 更容易与 FlashAttention 等高效计算核进行集成,从而进一步提升推理效率;相比之下,NSA 则需专门设计新的计算内核来实现其高效计算。
Conclusion
本文系统梳理并深入解析了高效 Transformer 领域的几项代表性工作,包括 Performer、Linformer、FlashAttention、NSA 与 MoBA,旨在帮助读者把握其核心思想与关键技术路径,为理解这一领域的发展脉络提供清晰参考。尽管大语言模型近年来取得了飞跃式进展,但长文本的高效建模与理解依然是当前亟待攻克的重要挑战之一。未来,Transformer 架构的效率革命仍将持续演进,如何在模型架构层面实现更高效、更智能的注意力机制,将持续成为人工智能研究的热点前沿。
[1] Choromanski, Krzysztof Marcin, et al. "Rethinking Attention with Performers." International Conference on Learning Representations.
[2] Wang, Sinong, et al. "Linformer: Self-attention with linear complexity." arXiv preprint arXiv:2006.04768(2020).
[3] Dao, Tri, et al. "Flashattention: Fast and memory-efficient exact attention with io-awareness." Advances in neural information processing systems35 (2022): 16344-16359.
[4] Yuan, Jingyang, et al. "Native sparse attention: Hardware-aligned and natively trainable sparse attention." arXiv preprint arXiv:2502.11089(2025).
[5] Lu, Enzhe, et al. "MoBA: Mixture of Block Attention for Long-Context LLMs." arXiv preprint arXiv:2502.13189(2025).
[6] Su, Jianlin, et al. "Roformer: Enhanced transformer with rotary position embedding." Neurocomputing568 (2024): 127063.
推荐站内搜索:最好用的开发软件、免费开源系统、渗透测试工具云盘下载、最新渗透测试资料、最新黑客工具下载……
还没有评论,来说两句吧...