近日,DeepSeek公司又推出了名为NSA的先进技术,该框架通过创新的稀疏注意力机制,显著提高了大型语言模型(LLM)在处理长上下文时的效率和性能。紧随其后,Kimi公司也不甘落后,推出了MoBA框架,该框架能够将上下文长度扩展到10M,从而支持更复杂的任务和长文本分析。本文将深入探讨这两个框架的关键差异和各自的优势,以及它们对未来人工智能研究和应用可能产生的影响。
技术报告地址:kimi moba: https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf
deepseek nsa: https://arxiv.org/pdf/2502.11089
其中kimi开源了代码,项目地址:https://github.com/MoonshotAI/MoBA
一、背景
在大型语言模型的军备竞赛中,“上下文窗口长度"已成为衡量模型能力的黄金标准。从GPT-4的128k到Claude 3的1M,再到Kimi近期宣称的10M级窗口,这场竞赛背后隐藏着严峻的技术挑战:传统的全注意力机制(Full Attention)在计算复杂度上呈现O(N²)的爆炸式增长。当处理百万级token时,单次推理的显存占用将超过当前最强GPU(如H100的80GB显存)的物理极限,这使得扩展上下文窗口成为"不可能三角”——既要保持模型性能,又要控制计算资源消耗,还要实现商业可行性。
这解释了为何DeepSeek、Kimi等头部玩家纷纷聚焦稀疏注意力算法(Sparse Attention)的研究。这场技术博弈的本质,是探索如何在人类认知效率与机器计算效率之间找到最优解。
两种方法均提出稀疏注意力算法以降低计算复杂度并扩展上下文处理能力,但在实现路径上存在显著差异。
NSA(Nested Sparse Attention)通过动态分层稀疏策略,采用"粗粒度区域筛选-细粒度特征关联"的双阶段机制,首先对输入特征进行空间维度的区域级压缩,进而在筛选出的关键区域内执行细粒度令牌级注意力计算,这种层级化稀疏架构有效平衡了计算效率与特征捕获能力。
MOBA(Mixture-of-Blocks Attention)则受混合专家(MoE)范式启发,将块级稀疏注意力与动态路由机制相结合,通过可学习的门控网络对输入序列进行分块,并基于内容相似性动态选择最相关的键-值块进行注意力聚合,这种模块化设计显著提升了长程依赖建模的灵活性。下文将详细介绍这两种方法的内容。
二、NSA(Nested Sparse Attention)
NSA采用动态分层稀疏策略,通过结合粗粒度的token压缩与细粒度的token选择,既保留了全局上下文感知能力,也确保了局部精度。我们的方法通过两项关键创新推进了稀疏注意力设计:
(1) 通过算术强度平衡的算法设计实现显著加速,并针对现代硬件进行了实现优化;
(2) 支持端到端训练,在保持模型性能的同时减少预训练计算量。
▲图1|完全注意力机制和NSA的效果与性能对比图©️【深蓝AI】编译
如图1所示,实验表明采用NSA预训练的模型在通用基准测试、长上下文任务和基于指令的推理中均保持或超越全注意力模型。同时,NSA在处理64k长度序列时,在解码、前向传播和反向传播阶段均实现了对全注意力模型的显著加速,验证了其在模型全生命周期中的高效性。
▲图2|NSA架构概览图©️【深蓝AI】编译
如图2所示,NSA通过将键和值组织成时间块,并经由三个注意力路径进行处理,有效降低单次查询的计算量:压缩的粗粒度标记、选择性保留的细粒度标记,以及用于局部上下文信息的滑动窗口。随后通过专用内核实现最大化实际效率。NSA针对前述关键需求引入两大核心创新:(1) 硬件对齐系统:优化块状稀疏注意力以实现张量核心的高效利用和内存访问优化,确保算术强度的平衡性;(2) 训练感知设计:通过高效算法和反向算子实现稳定的端到端训练。这种优化使NSA既能支持高效部署,又能实现端到端训练。
现代稀疏注意力方法在降低Transformer模型的理论计算复杂度方面取得了显著进展。然而,大多数方法主要在推理阶段应用稀疏性,同时保留预训练的全注意力主干网络,这种做法可能引入架构偏差,从而限制了其充分挖掘稀疏注意力优势的能力。为了充分发挥自然稀疏模式下注意力机制的潜力,研究人员提出将完全注意力机制中的原始键值对替换为针对每个查询 q𝑡 生成的更为紧凑且信息密度更高的键值对集合 𝐾𝑡, 𝑉𝑡。具体而言,将优化后的注意力输出形式化定义如下:
然后可以因此设计不同的映射方案以获得不同类别的