【论文阅读】原生稀疏注意力NSA:Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention

上周才刚看完DeepSeek的R1,这周它又出新文章了,他们的产出速度确实快。

1. 简介

这篇文章相较于DeepSeek-R1这篇文章,聚焦范围更加小,R1的重点在于提出了一个文本生成的训练策略和蒸馏策略,这篇文章则是提出了一个注意力机制NSA,主要解决的是长序列做注意力时带来的效率问题。通篇文章看下来,它的实际意义可能比较局限,因此本文仅关注其主要内容,对于具体细节和实验结果并不进一步细究。

论文标题:Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
论文链接:https://arxiv.org/abs/2502.11089

2. 动机和相关工作

动机:目前大模型的底层核心架构还是Transformer,Transformer的一个特点是在训练时,需要计算每个Token和其它Token之间的注意力,计算效率为 O ( n 2 ) O(n^2) O(n2)。当输入的token很长时,就需要花费很长时间进行计算,非常消耗资源。

目前,已有如下几个思路进行解决长序列注意力的效率问题:

  • 固定稀疏模式(Fixed Sparse Pattern):大致思路是取一个局部窗口,只在这个窗口中进行注意力的计算。
  • 动态token裁剪(Dynamic Token Pruning):有点类似于模型剪枝,大致思路是仅仅动态保留一部分最重要的token计算注意力。
  • 序列感知选择(Query-Aware Selection):大致思路是将几个token打包成块,在块之间计算注意力。

上述方法虽然比直接做全部Token注意力(Full Attention)计算效率要高,但模型的性能同样容易受到较大影响。

NSA就说,我既要加速计算,又要让模型性能不受影响,甚至更高。

下面这张图是NSA的直观效果,速度比全做注意力提升了数倍,性能也略有提升。

在这里插入图片描述

3. 主要方法

NSA主要采用以下几个方法:

  1. 粗粒度压缩(Token Compression)
    具体做法是:将Token按照位置划分为若干连续块(block),然后将每个块中的内容压缩成单个表示。

  2. 细粒度选择(Token Selection)
    考虑到直接将Token进行压缩,会让一些细节信息,比如某些重要token,会被块直接给平均表示掉。因此,这里再将一些重要的细粒度Token筛选出来,具体做法是:再次将Token执行分块(可以和上一步分块一样,也可以不一样),利用注意力分数对其进行重要性评估,最后选取Top-K个最重要的块,将里面的Token取出来和其它的块进行注意力计算。

  3. 局部滑动窗口(Local Sliding Window)
    分块操作带来的另一个弊病时局部信息丢失。比如,两个原本连续的Token,因划分成不同的block,导致这两个连续信息会丢失。补救方式是运用一个滑动窗口,使其处理局部的上下文信息。

  4. 门控网络融合
    上述三点实际上视作不同的分支,每个分支各有侧重点。最终,总的注意力为三路分支的加权融合。这里权重通过门控网络来学习得到。由于所有的策略参数都是可微,可学习的,作为整个模型的一部分,因此这个稀疏注意力框架也叫原生(Native)。

此外,在主要方法一节,文章还单独提到了一些硬件底层的创新。下图展示了NSA的内核设计,有点类似于FlashAttention,通过算法在静态随机存取存储器(SRAM)和高带宽存储器(HBM)的优化,进一步提升了计算效率。

在这里插入图片描述

总结

这篇文章主要只是展现了一项新的注意力计算方法。这个方法是一些基础策略的组合,虽然在文章中测试的效果不错,但未必值得借鉴。

其一是它额外引入了很多策略性的超参数,无疑增加了调参的成本;
其二是它引入了额外的模型参数,这可能会影响到模型整体的稳定性,虽然在文中实验性能不降反升,但换个场景未必有效;
其三是它的效率提升很大程度上来自于硬件底层的优化,这种方式基本于GPU设备强绑定,如果设备不同,可能需要额外的工程优化。

因此,我觉得这篇文章的影响力不会像R1这么大,目前受到关注的原因主要还是媒体对DeepSeek的炒作。

参考资料

[1] (3万字长文)深入解析 DeepSeek 最新论文提出的 NSA 注意力机制 https://blog.csdn.net/qq_38334677/article/details/145729713
[2] 刚刚,DeepSeek放出重磅论文!梁文锋亲自参与!
https://mp.weixin.qq.com/s/mlC5CFJ06tSjVJdATuM69Q
[3] 刚刚!DeepSeek梁文锋亲自挂名,公开新注意力架构NSA
https://mp.weixin.qq.com/s/SlhNdGNFpGlicIkMUE21og

### NSA 原生稀疏注意力机制的工作原理与实现 NSA (Native Sparse Attention) 是一种优化的注意力计算方法,旨在通过减少全连接矩阵运算中的冗余来提高效率。其核心思想在于利用稀疏化技术降低计算复杂度并保持模型性能。 #### 工作原理概述 传统的自注意力机制涉及全局范围内的 token 对之间相互作用的建模,这通常会带来 \(O(n^2)\) 的时间复杂度和内存消耗问题[^1]。而 NSAM(NSA Native Sparse Attention Mechanism)则通过对注意力权重矩阵进行结构化的稀疏处理,在保留重要上下文关系的同时显著减少了不必要的计算开销。 具体而言,该机制采用了一种局部窗口策略以及跨距离依赖捕获方式: - **局部窗口注意力建模**: 将序列划分为若干固定大小的小片段,并仅允许这些片段内部 tokens 彼此交互。这种设计可以有效控制每一步所需访问的数据量从而达到降本增效的目的。 ```python def local_window_attention(query, key, value, window_size=4): """ 实现基于滑动窗口的本地注意力机制 参数: query/key/value: 输入张量形状为 [batch_size, seq_len, dim] window_size: 定义每个token能关注到的最大跨度长度 返回: 输出经过加权求和后的value向量 """ batch_size, seq_len, _ = query.shape pad_length = ((seq_len // window_size)+1)*window_size - seq_len # 计算填充长度使能够整除窗宽 padded_query = torch.nn.functional.pad(query,(0,pad_length)) # 零填充至可被窗宽整除 ... ``` 上述代码展示了如何构建一个简单的局部窗口注意力函数。 - **跨距离依赖捕捉**: 此外还引入了一些特殊位置索引来建立远程联系,比如每隔一定步数选取特定节点形成全局视角下的稀疏链接网络。这种方法既维持了必要的远端信息交流又不会过度增加资源负担。 总体来说,NSA 方法不仅大幅削减了传统 Transformer 架构里密集型操作所带来的高昂代价,而且还能很好地适应多种自然语言处理任务需求。 #### 关键优势分析 相比其他类型的稀疏化方案,NSA 更注重平衡速度提升与表达能力之间的取舍。它具备如下几个主要优点: - 显著降低了时间和空间上的要求; - 能够灵活调整参数配置以匹配不同规模数据集的要求; - 经过实验验证表明即使是在大规模预训练场景下也能取得接近甚至优于常规版本的表现水平。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

zstar-_

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值