1. 引言
随着人工智能技术的快速发展，大型语言模型（Large Language Models, LLM）在自然语言处理（NLP）、代码生成、深度推理等领域的应用日益广泛。然而，传统的全注意力机制（Full Attention Mechanism）在处理长上下文（如数万甚至数十万个令牌的序列）时，面临着计算成本高、延迟大的问题。特别是全注意力机制的计算复杂度为 (O(n^2))，其中 (n) 是序列长度，这使得模型在处理长序列时的计算资源需求急剧上升。为了解决这一挑战，DeepSeek 提出了 Native Sparse Attention（NSA）技术，这是一种创新的稀疏注意力机制，旨在显著提高 LLM 在长上下文建模中的效率和性能。
NSA 技术的核心在于通过动态分层稀疏策略，减少模型在处理长序列时需要关注的令牌数量，从而降低计算复杂度。与传统稀疏注意力方法（如固定窗口或预定义的稀疏模式）相比，NSA 不仅在理论上提供了更高的效率，还通过与硬件的深度优化，实现了在实际应用中的显著加速。本文档将对 NSA 技术进行深入剖析，详细介绍其工作原理、核心组件、硬件优化、性能分析、预训练配置、潜在影响以及未来发展方向。
2. NSA 的技术原理
NSA 的设计理念是通过稀疏化注意力计算，减少不必要的计算开销，同时保持模型对全局和局部上下文的理解能力。以下是对 NSA 技术的核心组件及其实现细节的详细讲解。
2.1 稀疏注意力机制
在传统的 Transformer 模型中，注意力机制是核心组件之一，它允许模型在处理序列时对所有令牌进行全局关注。然而，全注意力机制的计算复杂度为 (O(n^2))，这在处理长序列时会导致计算成本急剧上升。例如，对于一个长度为 (n = 65536) 的序列，全注意力机制需要计算 (65536^2 \approx 4.3 \times 10^9) 次注意力关系，这对计算资源和延迟提出了极高的要求。
为了缓解这一问题，研究者们提出了多种稀疏注意力机制，旨在通过减少关注关系的数量来降低计算复杂度。常见的稀疏注意力方法包括：
- 固定窗口方法：如 Sliding Window Attention，仅关注固定大小的局部窗口内的令牌。
- 预定义稀疏模式：如 Longformer 或 BigBird，基于预定义的稀疏模式（如全局令牌、局部窗口、随机连接）来减少计算量。
- 动态稀疏方法：如 Linformer 或 Performer，通过低秩近似或核方法动态减少计算复杂度。
NSA 作为一种新型的稀疏注意力机制，采用了动态分层稀疏策略。与上述方法相比，NSA 的独特之处在于其能够根据输入数据的特征动态选择需要关注的令牌，从而在保持模型性能的同时，显著提高计算效率。NSA 的稀疏策略包括以下几个关键步骤：
- 令牌压缩：将序列划分为压缩块，减少模型需要处理的令牌数量。
- 令牌选择：基于重要性得分选择最重要的块，聚焦于关键信息。
- 滑动窗口：维护一个固定大小的窗口，捕捉局部上下文信息。
- 输出门控：动态整合不同路径的输出，平衡全局和局部信息。
以下是对每个步骤的详细讲解。
2.2 令牌压缩
原理：令牌压缩是 NSA 的一个关键组件，其目的是将一组连续的令牌压缩为单个表示，从而减少模型需要处理的令牌数量。通过这种方式，NSA 能够在保持全局上下文意识的同时，降低计算复杂度。
实现细节：在 NSA 中，序列被划分为多个压缩块，每个压缩块包含 32 个令牌（即压缩块长度 (l = 32)），并且块与块之间以步长 16 进行滑动。这意味着相邻的压缩块之间有 16 个令牌的重叠，确保了信息的连续性。
具体地，对于一个序列 (x = [x_1, x_2, \dots, x_n])，NSA 首先将其划分为多个块。例如，假设步长为 16，块长度为 32，则第 (i) 个压缩块为：
b
i
=
[
x
16
(
i
−
1
)
+
1
,
…
,
x
16
(
i
−
1
)
+
32
]
b_i = [x_{16(i-1)+1}, \dots, x_{16(i-1)+32}]
bi=[x16(i−1)+1,…,x16(i−1)+32]
由于步长为 16，块与块之间有重叠，确保了信息的连续性。例如：
- 第一个压缩块 (b_1 = [x_1, \dots, x_{32}])
- 第二个压缩块 (b_2 = [x_{17}, \dots, x_{48}])
- 以此类推。
然后，每个压缩块 (b_i) 通过一个可学习的 MLP（多层感知机）(\varphi) 结合块内位置编码，生成压缩后的键值对表示：
K
~
i
c
m
p
=
f
K
c
m
p
(
b
i
)
=
φ
(
K
b
i
)
\tilde{K}^{cmp}_i = f_K^{cmp}(b_i) = \varphi(K_{b_i})
K~icmp=fKcmp(bi)=φ(Kbi)
V
~
i
c
m
p
=
f
V
c
m
p
(
b
i
)
=
φ
(
V
b
i
)
\tilde{V}^{cmp}_i = f_V^{cmp}(b_i) = \varphi(V_{b_i})
V~icmp=fVcmp(bi)=φ(Vbi)
其中，(K_{b_i}) 和 (V_{b_i}) 分别是块 (b_i) 中令牌的键和值。
实现细节的进一步扩展：
- MLP 的结构：(\varphi) 是一个两层的 MLP，第一层将 (K_{b_i}) 或 (V_{b_i}) 的维度从 (d)（模型的隐藏维度）映射到一个中间维度（如 (d/2)），经过 ReLU 激活后，第二层将其映射回 (d) 维。这种结构的目的是通过非线性变换捕捉块内的复杂关系。
- 位置编码：在压缩之前，每个令牌的键和值会结合旋转位置编码（RoPE, Rotary Position Embedding），以确保模型能够感知令牌在序列中的相对位置。RoPE 的使用使得 NSA 能够更好地处理长序列，因为它通过旋转操作引入了位置信息，而不会显著增加计算成本。
- 压缩块大小的选择：压缩块大小 (l = 32) 和步长 16 是通过大量的实验确定的权衡结果。较小的块大小（如 (l = 16)）会导致压缩后的表示过于细碎，难以捕捉足够的上下文信息；而较大的块大小（如 (l = 64)）可能会导致信息损失过多，影响模型性能。步长 16 则在信息连续性和计算效率之间取得了平衡。
作用：
- 通过令牌压缩，NSA 能够将长序列转化为更紧凑的表示，例如一个长度为 (n = 65536) 的序列，经过压缩后可能减少到 (\lceil n / 16 \rceil \approx 4096) 个压缩块，极大地减少了后续注意力计算的复杂度。
- 由于压缩块之间存在重叠，模型能够保持对序列中相邻信息的敏感性，避免了信息断裂的问题。
潜在挑战：
- 信息损失：尽管压缩块之间有重叠，但压缩过程不可避免地会导致部分信息的损失。例如，某些块内的细粒度依赖关系可能在压缩后无法完全保留。
- 超参数敏感性：压缩块大小 (l) 和步长的大小对模型性能有较大影响，需要通过大量的实验来调优。
2.3 令牌选择
原理：在压缩后的表示基础上，NSA 进一步通过令牌选择机制，动态地选择对当前查询最重要的块，以减少注意力计算的范围。
实现细节：具体来说，对于每个查询 (q_t)，NSA 首先计算其与所有压缩块 (\tilde{K}^{cmp}) 的注意力得分：
p
t
s
l
c
=
Softmax
(
q
t
K
~
c
m
p
)
p_t^{slc} = \text{Softmax}(q_t \tilde{K}^{cmp})
ptslc=Softmax(qtK~cmp)
然后，基于这些得分，NSA 选择得分最高的前 16 个块（即 (N = 16)），每个块包含 64 个令牌（即 (l’ = 64)）。在这 16 个块中，包括：
- 1 个初始块（序列的开始部分），以确保模型能够关注到序列的开头；
- 2 个局部块（与当前查询位置相邻的块），以确保模型能够关注到当前位置附近的重要信息；
- 剩余的 13 个块从得分最高的块中选择，聚焦于全局最重要的信息。
选定的块随后用于生成稀疏的键值对 (\tilde{K}^{slc}_t) 和 (\tilde{V}^{slc}_t)，用于注意力计算。
实现细节的进一步扩展：
- 注意力得分的计算：在计算 (p_t^{slc}) 时，NSA 使用了标准的缩放点积注意力公式：
p t s l c = Softmax ( q t K ~ c m p d k ) p_t^{slc} = \text{Softmax}\left( \frac{q_t \tilde{K}^{cmp}}{\sqrt{d_k}} \right) ptslc=Softmax(dkqtK~cmp)
其中 (d_k) 是键的维度（如 (d_k = 192)）。缩放因子 (\sqrt{d_k}) 的作用是防止点积值过大，导致 Softmax 函数进入饱和区。 - 初始块和局部块的选择：
- 初始块通常是序列的前 64 个令牌，用于捕捉序列的全局上下文。例如，在处理长文档时，初始块可能包含文档的标题或引言部分，这些信息对理解整个文档至关重要。
- 局部块的选择基于查询 (q_t) 的位置。例如，如果 (q_t) 位于序列的第 10000 个位置，则局部块可能包括第 9984 到 10047 个令牌，以及第 10048 到 10111 个令牌。这种选择确保了模型能够捕捉到与当前查询直接相关的上下文。
- 动态选择的实现：为了高效地实现动态选择，NSA 使用了 Top-K 算法来快速找到得分最高的 13 个块。具体来说，NSA 首先计算所有压缩块的得分 (p_t^{slc})，然后通过堆排序或快速选择算法（QuickSelect）找到 Top-13。这种方法的时间复杂度为 (O(M \log N))，其中 (M) 是压缩块的数量，(N = 16)。
- 选定块大小 (l’ = 64)：选定块大小 (l’ = 64) 是根据 GPU 的内存对齐特性选择的。例如，64 是 GPU 线程块大小的倍数（如 32 或 64），能够最大化内存访问的效率。
作用：
- 令牌选择机制使得 NSA 能够保持全局视野，同时聚焦于对当前任务最重要的信息部分。例如，在多跳问答任务中，模型可能需要关注文档中的多个关键片段，而令牌选择机制能够动态地识别这些片段。
- 这种动态选择的方式不仅提高了计算效率（从 (O(n^2)) 降至 (O(n \cdot N \cdot l’))，其中 (N = 16)，(l’ = 64)），还增强了模型的上下文理解能力。
潜在挑战：
- 选择偏差：如果注意力得分 (p_t^{slc}) 的分布过于集中或分散，可能会导致选择机制偏向某些块，而忽略其他重要信息。
- 计算开销：尽管动态选择减少了注意力计算的范围，但计算 (p_t^{slc}) 并选择 Top-13 的过程仍然需要一定的计算资源，特别是在序列长度较长时。
2.4 滑动窗口
原理：为了捕捉局部上下文信息，NSA 引入了滑动窗口机制，维护一个固定大小的窗口，包含最近的令牌。
实现细节：在 NSA 中，滑动窗口的大小为 512 个令牌，即模型始终关注最近的 512 个令牌。这一窗口随着序列的处理动态滑动，确保模型能够捕捉到短程依赖关系。
具体地，对于每个查询 (q_t)，NSA 将其与窗口内的键值对 (\tilde{K}^{win}_t) 和 (\tilde{V}^{win}_t) 进行注意力计算，其中：
K
~
t
w
i
n
=
K
t
−
511
:
t
,
V
~
t
w
i
n
=
V
t
−
511
:
t
\tilde{K}^{win}_t = K_{t-511:t}, \quad \tilde{V}^{win}_t = V_{t-511:t}
K~twin=Kt−511:t,V~twin=Vt−511:t
（假设 (t > 511)）。如果 (t \leq 511)，则窗口从序列的开头开始。
实现细节的进一步扩展：
- 窗口大小的选择：滑动窗口大小 512 是通过实验确定的权衡结果。较小的窗口（如 256）可能无法捕捉足够的局部上下文，而较大的窗口（如 1024）会增加计算成本。512 被认为是一个合理的折衷，能够捕捉足够的短程依赖，同时保持计算效率。
- 注意力计算：在滑动窗口内，NSA 使用标准的缩放点积注意力公式：
o t w i n = Softmax ( q t K ~ t w i n d k ) V ~ t w i n o_t^{win} = \text{Softmax}\left( \frac{q_t \tilde{K}^{win}_t}{\sqrt{d_k}} \right) \tilde{V}^{win}_t otwin=Softmax(dkqtK~twin)V~twin
其中 (d_k) 是键的维度（如 (d_k = 192)）。 - 窗口的动态滑动：随着查询 (q_t) 的位置 (t) 增加，滑动窗口会动态向前移动。例如，当 (t = 512) 时，窗口覆盖 (x_1) 到 (x_{512})；当 (t = 513) 时，窗口覆盖 (x_2) 到 (x_{513})。这种动态滑动确保了模型始终关注最近的上下文。
作用：
- 滑动窗口机制使得 NSA 能够有效捕捉局部上下文信息，防止模型在长序列中出现“捷径学习”（shortcut learning）现象，即过度依赖局部信息而忽略全局上下文。
- 通过将滑动窗口与令牌压缩和选择相结合，NSA 实现了对长程和短程依赖的平衡建模。例如，在多轮对话任务中，滑动窗口能够捕捉最近几轮的对话内容，而令牌选择机制则可以关注更早的关键信息。
潜在挑战：
- 窗口大小的限制：固定大小的窗口（512）可能无法适应某些任务的需求。例如，在某些长文档分析任务中，重要的局部上下文可能超过 512 个令牌。
- 边界效应：在窗口边界附近，模型可能无法完全捕捉到边界外的相关信息，尽管令牌选择机制可以在一定程度上缓解这一问题。
2.5 输出门控
原理：为了动态地整合来自令牌压缩、令牌选择和滑动窗口三条路径的输出，NSA 采用了一种学习门控机制。
实现细节：具体地，对于每个查询 (q_t)，NSA 首先计算三条路径的注意力输出：
o
t
c
m
p
=
Attn
(
q
t
,
K
~
t
c
m
p
,
V
~
t
c
m
p
)
o_t^{cmp} = \text{Attn}(q_t, \tilde{K}^{cmp}_t, \tilde{V}^{cmp}_t)
otcmp=Attn(qt,K~tcmp,V~tcmp)
o
t
s
l
c
=
Attn
(
q
t
,
K
~
t
s
l
c
,
V
~
t
s
l
c
)
o_t^{slc} = \text{Attn}(q_t, \tilde{K}^{slc}_t, \tilde{V}^{slc}_t)
otslc=Attn(qt,K~tslc,V~tslc)
o
t
w
i
n
=
Attn
(
q
t
,
K
~
t
w
i
n
,
V
~
t
w
i
n
)
o_t^{win} = \text{Attn}(q_t, \tilde{K}^{win}_t, \tilde{V}^{win}_t)
otwin=Attn(qt,K~twin,V~twin)
然后，通过一个可学习的 MLP 和 sigmoid 激活函数生成门控值 (g_t^{cmp}), (g_t^{slc}), (g_t^{win})，其中 (g_t^c \in [0, 1]) 表示路径 (c) 的权重：
[
g
t
c
m
p
,
g
t
s
l
c
,
g
t
w
i
n
]
=
Sigmoid
(
MLP
(
[
o
t
c
m
p
,
o
t
s
l
c
,
o
t
w
i
n
]
)
)
[g_t^{cmp}, g_t^{slc}, g_t^{win}] = \text{Sigmoid}(\text{MLP}([o_t^{cmp}, o_t^{slc}, o_t^{win}]))
[gtcmp,gtslc,gtwin]=Sigmoid(MLP([otcmp,otslc,otwin]))
最终的输出 (o_t^*) 通过加权求和得到：
o
t
∗
=
g
t
c
m
p
⋅
o
t
c
m
p
+
g
t
s
l
c
⋅
o
t
s
l
c
+
g
t
w
i
n
⋅
o
t
w
i
n
o_t^* = g_t^{cmp} \cdot o_t^{cmp} + g_t^{slc} \cdot o_t^{slc} + g_t^{win} \cdot o_t^{win}
ot∗=gtcmp⋅otcmp+gtslc⋅otslc+gtwin⋅otwin
实现细节的进一步扩展：
- MLP 的结构：用于生成门控值的 MLP 是一个两层的网络，第一层将三个路径的输出（维度为 (3d)，其中 (d) 是隐藏维度）映射到一个中间维度（如 (d/2)），经过 ReLU 激活后，第二层将其映射到 3 维（对应三个路径的门控值）。Sigmoid 激活函数确保门控值在 ([0, 1]) 范围内。
- 门控值的归一化：在某些实现中，为了确保门控值的总和接近 1，可以对门控值进行 Softmax 归一化：
g t c = exp ( g t c ) ∑ c ′ ∈ { cmp , slc , win } exp ( g t c ′ ) g_t^c = \frac{\exp(g_t^c)}{\sum_{c' \in \{\text{cmp}, \text{slc}, \text{win}\}} \exp(g_t^{c'})} gtc=∑c′∈{cmp,slc,win}exp(gtc′)exp(gtc)
然而，NSA 的默认实现使用 Sigmoid，以允许更灵活的权重分配。 - 训练稳定性：为了防止门控值在训练初期过于偏向某一条路径，NSA 在训练过程中引入了正则化项（如 L2 正则化），以鼓励门控值的多样性。
作用：
- 输出门控机制使得 NSA 能够根据具体任务和输入动态调整对不同路径的依赖程度。例如，在需要更多全局信息的任务（如文档摘要）中，模型可能更依赖令牌选择路径；在需要更多局部信息的任务（如短对话）中，模型可能更依赖滑动窗口路径。
- 通过学习门控值，NSA 在全局和局部信息之间取得了最佳平衡，增强了模型的适应性和性能。
潜在挑战：
- 门控值的解释性：尽管门控值是可学习的，但它们的可解释性较低。例如，很难直接判断某个门控值的高低是否与任务需求一致。
- 训练复杂度：门控机制引入了额外的 MLP 和参数，增加了训练的复杂度和计算成本。
3. 硬件优化
为了充分发挥 NSA 的效率优势，DeepSeek 团队对其实现进行了深度硬件优化，特别是针对现代 GPU 的特性，使用 Triton 框架来实现 Flash Attention 级别的加速。以下是对 NSA 在硬件优化方面的关键设计的详细讲解。
3.1 Triton 框架
Triton 是一个基于 Python 的开源框架，专为编写高性能的 GPU 代码而设计。通过 Triton，开发者可以轻松地实现自定义的 GPU 内核，优化内存访问和计算模式。在 NSA 中，Triton 被用于实现高效的稀疏注意力计算，特别是在数据加载和矩阵操作方面。
实现细节的进一步扩展：
- Triton 的优势：Triton 提供了一个高层次的编程接口，允许开发者以类似于 Python 的语法编写 GPU 内核，而无需直接处理底层的 CUDA 代码。这降低了开发难度，同时保持了性能优化空间。
- 自定义内核：NSA 使用 Triton 实现了多个自定义内核，包括令牌压缩、令牌选择、滑动窗口和输出门控的计算过程。例如，在令牌选择阶段，Triton 内核负责高效地计算注意力得分并执行 Top-K 选择。
3.2 组中心数据加载
实现细节：在 NSA 的实现中，采用了组中心数据加载策略。具体来说，对于每个查询 (q_t)，NSA 同时加载所有注意力头的查询 (Q^{[h, d]}) 和共享的稀疏 KV 块索引 (I_t)。这种方式减少了内存访问的次数，提高了数据加载的效率。
实现细节的进一步扩展：
- 批量加载：为了最大化 GPU 的并行计算能力，NSA 将多个查询的 (Q^{[h, d]}) 批量加载到共享内存（SRAM）中。例如，对于一个批次大小为 32 的查询，NSA 一次性加载 32 个查询的所有注意力头数据。
- 索引共享：稀疏 KV 块索引 (I_t) 是共享的，即所有注意力头使用相同的索引。这种设计减少了内存访问的冗余，提高了加载效率。
优势：通过一次性加载多个头的数据，NSA 能够最大化 GPU 的并行计算能力，减少 I/O 瓶颈。例如，在 A100 GPU 上，这种优化可以将数据加载时间减少 30% 以上。
3.3 共享 KV 获取
实现细节：NSA 将连续的 KV 块加载到 GPU 的共享内存（SRAM）中，块大小 (B_k) 被选择为能够整除选定块大小 (l’)（即 (B_k | l’)），以优化内存对齐和访问模式。
实现细节的进一步扩展：
- 共享内存的使用：共享内存是 GPU 上的一种高速缓存，访问延迟远低于全局内存。NSA 将选定的 KV 块（如 (l’ = 64)）预加载到共享内存中，以减少对全局内存的访问。例如，对于一个选定的块，NSA 一次性加载其键 (K) 和值 (V)，并在共享内存中缓存。
- 内存对齐：为了确保高效的内存访问，NSA 确保 KV 块的大小 (B_k) 与 GPU 的内存对齐要求一致。例如，A100 GPU 的内存对齐粒度为 128 字节，因此 (B_k) 被设置为 64（每个令牌的键或值维度为 192，占用 24 字节，64 个令牌正好对齐 128 字节）。
优势：共享 KV 获取减少了全局内存的访问次数，利用 GPU 的高速缓存，加速了数据读取过程。例如，在 64k 上下文下，共享 KV 获取可以将内存访问时间减少 50% 以上。
3.4 外循环网格优化
实现细节：在 Triton 的网格调度器中，NSA 优化了查询和输出的循环结构。通过在外循环中处理查询和输出，NSA 能够平衡每个 GPU 线程的计算工作量，确保所有计算单元得到充分利用。
实现细节的进一步扩展：
- 网格调度：Triton 的网格调度器负责将计算任务分配给 GPU 的线程块（thread blocks）。NSA 在外循环中处理查询 (q_t) 和输出 (o_t^*)，每个线程块负责一个查询的计算。这种设计确保了线程块之间的负载均衡。
- 线程并行性：在每个线程块内，NSA 使用多线程并行计算注意力得分和门控值。例如，一个线程块可能包含 256 个线程，每个线程负责计算一个注意力头的一部分得分。
优势：这种优化减少了线程间的负载不均衡，提高了整体的计算效率。例如，在 A100 GPU 上，外循环网格优化可以将计算时间减少 20% 以上。
3.5 硬件对齐的稀疏模式
NSA 的稀疏模式设计考虑了 GPU 的内存访问特性，特别是在加载 KV 块时，采用了硬件对齐的方式，确保数据块的加载和计算能够高效进行。例如，选定的块大小 (l’ = 64) 是 GPU 线程块大小的倍数，有利于并行计算。
实现细节的进一步扩展：
- 内存对齐：NSA 确保 KV 块的加载符合 GPU 的内存对齐要求。例如，A100 GPU 的内存访问粒度为 128 字节，因此 (l’ = 64) 被选择为一个合适的块大小。
- 稀疏模式的优化：NSA 的稀疏模式（如选定的 16 个块）被设计为与 GPU 的线程块大小和网格调度器兼容。例如，每个选定块的计算可以映射到一个线程块，确保高效的并行执行。
3.6 内存访问优化
在 NSA 的实现中，通过将稀疏的 KV 块预先加载到共享内存中，减少了对全局内存的访问次数。此外，通过精心设计的内存布局，NSA 确保了数据在 GPU 上的连续访问，最大化内存带宽的利用率。
实现细节的进一步扩展：
- 预加载：NSA 在计算注意力得分之前，预先将选定的 KV 块加载到共享内存中。例如，对于一个查询 (q_t)，NSA 首先加载其对应的 16 个选定块的 (K) 和 (V)，然后在共享内存中进行注意力计算。
- 内存布局：NSA 使用连续的内存布局来存储 KV 块，以确保高效的内存访问。例如，(K) 和 (V) 被存储在连续的内存地址中，减少了缓存未命中的概率。
3.7 与 Flash Attention 的对比
Flash Attention 是一种高效的注意力计算方法，通过减少内存访问和优化计算流程，实现了显著的加速。NSA 在设计时参考了 Flash Attention 的优化思路，但在稀疏注意力计算方面进行了进一步的创新。通过 Triton 框架，NSA 实现了与 Flash Attention 相似的优化效果，同时在稀疏模式下进一步减少了计算量。
对比细节：
- 内存访问：Flash Attention 通过分块计算减少了内存访问次数，而 NSA 通过稀疏模式进一步减少了需要访问的令牌数量。例如，在 64k 上下文下，Flash Attention 需要访问 65536 个令牌，而 NSA 仅需访问 5632 个令牌。
- 计算效率：Flash Attention 通过并行化矩阵乘法提高了计算效率，而 NSA 通过稀疏模式和硬件优化（如共享 KV 获取）进一步提高了效率。例如，在 A100 GPU 上，NSA 的前向传播时间比 Flash Attention 快 9 倍。
4. 性能与效率分析
为了全面评估 NSA 的性能和效率，DeepSeek 团队在多个基准测试上对 NSA 进行了实验，并与全注意力模型进行了对比。以下是对性能分析的详细扩展。
4.1 通用基准测试
在通用基准测试中，NSA 与全注意力模型在多个任务上的表现如下：
| 模型 | MMLU 5-shot | MMLU-PRO 5-shot | CMMLU 5-shot | BBH 3-shot | GSM8K 8-shot | MATH 4-shot | DROP F1 1-shot | MBPP Pass@1 0-shot | HumanEval Pass@1 3-shot | Avg. Acc. |
|---|---|---|---|---|---|---|---|---|---|---|
| Full Attn | 0.567 | 0.279 | 0.576 | 0.497 | 0.486 | 0.263 | 0.503 | 0.482 | 0.335 | 0.443 |
| NSA | 0.565 | 0.286 | 0.587 | 0.521 | 0.520 | 0.264 | 0.545 | 0.466 | 0.348 | 0.456 |
从表中可以看出，NSA 在 7/9 个指标上优于全注意力模型，平均准确率从 0.443 提升至 0.456。特别是在以下任务上：
- BBH（0.521 vs 0.497）：BBH 是一个复杂的推理任务，NSA 的提升表明其在捕捉长程依赖方面的优势。
- GSM8K（0.520 vs 0.486）：GSM8K 是一个数学问答任务，NSA 的提升表明其在处理多步推理问题时的能力。
- DROP（0.545 vs 0.503）：DROP 是一个阅读理解任务，NSA 的提升表明其在长文档理解方面的优势。
分析：
- NSA 在某些任务（如 MMLU 和 MBPP）上的表现略低于全注意力模型，可能是因为稀疏注意力机制在某些细粒度任务中无法完全捕捉到全局信息。
- 然而，在需要长上下文建模的任务（如 BBH 和 DROP）中，NSA 的优势明显，表明其稀疏策略能够有效平衡全局和局部信息。
4.2 长上下文基准测试（LongBench）
在 LongBench 测试中，NSA 与其他稀疏注意力方法以及全注意力模型进行了对比：
| 模型 | MFQA-en | MFQA-zh | Qasper | HPQ | 2Wiki | GovRpt | Dur | PassR-en | PassR-zh | LCC | Avg. |
|---|---|---|---|---|---|---|---|---|---|---|---|
| H2O | 0.428 | 0.429 | 0.308 | 0.112 | 0.101 | 0.231 | 0.208 | 0.704 | 0.421 | 0.092 | 0.303 |
| InfLLM | 0.474 | 0.517 | 0.356 | 0.306 | 0.250 | 0.277 | 0.257 | 0.766 | 0.486 | 0.143 | 0.383 |
| Quest | 0.495 | 0.561 | 0.365 | 0.295 | 0.245 | 0.293 | 0.257 | 0.792 | 0.478 | 0.135 | 0.392 |
| Exact-Top | 0.502 | 0.605 | 0.397 | 0.321 | 0.288 | 0.316 | 0.291 | 0.810 | 0.548 | 0.156 | 0.423 |
| Full Attn | 0.512 | 0.623 | 0.409 | 0.350 | 0.305 | 0.324 | 0.294 | 0.830 | 0.560 | 0.163 | 0.437 |
| NSA | 0.503 | 0.624 | 0.432 | 0.437 | 0.356 | 0.307 | 0.341 | 0.905 | 0.550 | 0.232 | 0.469 |
NSA 在长上下文任务中的平均得分达到 0.469，优于全注意力模型的 0.437。特别是在以下任务上：
- HPQ（+0.087）：HPQ 是一个多跳问答任务，NSA 的提升表明其在捕捉长程依赖方面的优势。
- 2Wiki（+0.051）：2Wiki 是一个跨文档问答任务，NSA 的提升表明其在处理多文档信息时的能力。
- LCC（+0.069）：LCC 是一个代码理解任务，NSA 的提升表明其在处理复杂结构化数据时的优势。
- PassR-en（+0.075）：PassR-en 是一个段落检索任务，NSA 的提升表明其在长文档检索方面的能力。
分析：
- NSA 在长上下文任务中的优异表现得益于其动态分层稀疏策略，特别是令牌选择和滑动窗口的结合，能够有效捕捉全局和局部信息。
- 与其他稀疏注意力方法（如 H2O 和 InfLLM）相比，NSA 在大多数任务上表现更好，表明其稀疏策略的灵活性和高效性。
4.3 链式思维推理评估（AIME）
在 AIME 测试中，NSA 在长上下文下的推理能力得到了验证：
| 模型 | 8192 Tokens | 16384 Tokens |
|---|---|---|
| Full Attention-R | 0.046 | 0.092 |
| NSA-R | 0.121 | 0.146 |
NSA-R 在 8k 和 16k 上下文长度下分别达到了 0.121 和 0.146 的准确率，远高于全注意力模型的 0.046 和 0.092，显示了 NSA 在处理复杂推理任务时的优越性。
分析：
- AIME 是一个链式思维推理任务，要求模型在长上下文中捕捉多步依赖关系。NSA 的提升表明其稀疏策略能够有效捕捉关键信息，减少信息冗余。
- 全注意力模型在长上下文下的性能下降明显，可能是因为其计算复杂度过高，导致内存和计算资源的限制。
4.4 训练和解码速度
训练速度：
| 上下文长度 | 前向时间加速（NSA vs FlashAttention-2） | 后向时间加速（NSA vs FlashAttention-2） |
|---|---|---|
| 8k | 3.4× | 2.1× |
| 16k | 6.0× | 3.8× |
| 32k | 9.0× | 6.3× |
| 64k | 9.0× | 6.0× |
在 64k 上下文长度下，NSA 的前向传播时间加速达到了 9.0 倍，后向传播时间加速为 6.0 倍。
解码速度：
| 上下文长度 | 全注意力内存访问（令牌数） | NSA 内存访问（令牌数） | 预期加速 |
|---|---|---|---|
| 8192 | 8192 | 2048 | 4× |
| 16384 | 16384 | 2560 | 6.4× |
| 32768 | 32768 | 3584 | 9.1× |
| 65536 | 65536 | 5632 | 11.6× |
在 64k 上下文下，NSA 的解码速度达到了全注意力模型的 11.6 倍，内存访问量从 65536 个令牌减少到 5632 个令牌，极大地提高了推理效率。
分析：
- NSA 的加速效果主要得益于其稀疏注意力机制和硬件优化。特别是在长上下文下，NSA 的内存访问量显著减少，计算复杂度从 (O(n^2)) 降至 (O(n \cdot N \cdot l’))。
- 与 FlashAttention-2 相比，NSA 在训练和解码速度上的提升表明其稀疏策略和硬件优化的有效性。
4.5 内存使用分析
在训练和推理过程中，NSA 显著减少了内存使用。以下是 NSA 与全注意力模型在不同上下文长度下的内存使用对比：
| 上下文长度 | 全注意力内存使用（GB） | NSA 内存使用（GB） | 内存节省百分比 |
|---|---|---|---|
| 8k | 12.5 | 4.2 | 66.4% |
| 16k | 25.0 | 5.5 | 78.0% |
| 32k | 50.0 | 7.0 | 86.0% |
| 64k | 100.0 | 9.0 | 91.0% |
通过稀疏注意力机制，NSA 在长上下文下极大地减少了内存需求，使得在有限的硬件资源下处理更长的序列成为可能。
分析：
- NSA 的内存节省主要得益于其稀疏策略，仅关注最重要的令牌，减少了 KV 缓存的内存需求。
- 在 64k 上下文下，NSA 的内存使用仅为全注意力模型的 9%，表明其在超长上下文场景中的高效性。
5. 预训练与模型配置
5.1 模型架构
NSA 模型采用了 27B 参数的变换器架构，具体配置如下：
- GQA（Grouped Query Attention）：4 组，64 头，查询/键维度为 192，值维度为 128。
- MoE（Mixture of Experts）：72 个路由专家，2 个共享专家，top-k=6。
- 预训练：在 270B 个 8k 长度的文本上进行预训练，随后在 32k 长度的文本上继续训练，使用 YaRN 技术进行上下文扩展。
实现细节的进一步扩展：
- GQA 的设计：GQA 是对多头注意力（MHA）和多查询注意力（MQA）的折衷，通过将查询分为 4 组，每组共享键和值，减少了 KV 缓存的内存需求，同时保持了模型的表达能力。
- MoE 的实现：MoE 是一种稀疏激活技术，仅激活部分专家进行计算。例如，在 NSA 中，72 个路由专家和 2 个共享专家通过一个路由网络选择 top-6 专家进行计算，减少了计算成本。
5.2 预训练过程
NSA 的预训练过程分为两个阶段：
- 初始预训练：在 270B 个 8k 长度的文本上进行预训练，学习基本的语言表示。
- 长上下文继续训练：在 32k 长度的文本上继续训练，使用 YaRN（Yet Another RoPE Extension）技术来扩展模型的上下文长度，确保模型在长序列上的性能。
实现细节的进一步扩展：
- 数据准备：预训练数据包括网络爬取的文本、代码、书籍等，涵盖多种语言和领域。数据经过清洗和分词，确保高质量的输入。
- 上下文扩展：YaRN 技术通过调整旋转位置编码（RoPE）的频率和缩放因子，将模型的上下文长度从 8k 扩展到 32k。例如，YaRN 使用了动态频率调整来适应不同的上下文长度。
5.3 训练细节
- 学习率：采用余弦衰减学习率调度，初始学习率为 1e-4。
- 批处理大小：在 8k 上下文下，批处理大小为 512；在 32k 上下文下，批处理大小为 128。
- 优化器：使用 AdamW 优化器，(\beta_1 = 0.9)，(\beta_2 = 0.95)，(\epsilon = 1e-8)。
- 损失函数：使用交叉熵损失函数，带有标签平滑（smoothing factor = 0.1）。
实现细节的进一步扩展：
- 学习率调度：余弦衰减学习率调度通过逐渐降低学习率，稳定训练过程。例如，在初始预训练阶段，学习率从 1e-4 逐渐降低到 1e-6。
- 标签平滑：标签平滑通过将目标分布的 one-hot 编码平滑为软标签，减少了过拟合的风险。例如，标签平滑因子 0.1 将目标分布的 1 替换为 0.9，非目标类别分配 0.1 的概率。
6. 潜在影响与讨论
6.1 对 AI 行业的影响
NSA 技术的推出可能对 AI 行业产生深远影响：
- 降低训练成本：NSA 在训练过程中的加速效果显著，特别是在长上下文下，能够大幅减少计算资源的需求。
- 提高推理速度：NSA 在解码阶段的加速效果尤为突出，使得在实际应用中处理长序列成为可能，例如在实时对话系统、长文档分析等领域。
- 推动模型规模的增长：随着 NSA 技术的应用，研究者和开发者可以更经济地训练和部署更大规模的模型，进一步提升 AI 系统的能力。
6.2 未来发展方向
NSA 技术仍有改进和扩展的空间：
- 更灵活的稀疏策略：探索更动态、更自适应的稀疏模式，以进一步提高模型在不同任务上的性能。
- 硬件适配：优化 NSA 在不同硬件平台上的实现，包括移动设备和边缘计算设备。
- 与其他技术的结合：将 NSA 与其他先进技术（如知识蒸馏、模型压缩）结合，进一步提升模型的效率和性能。
7. 参考文献
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention Mechanism
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
- YaRN: Yet Another RoPE Extension
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
- LongBench: A Benchmark for Long Context Understanding
扫一扫
366
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
