@VukRosic99: 长上下文Transformer面临两大瓶颈:二次注意力计算和KV缓存(在1M tokens时可达数百GB)…
摘要
MiniCPM-SALA是一款9B参数的混合注意力模型,通过在稀疏注意力和线性注意力之间交替插入(每3个线性层插入1个稀疏层)来克服长上下文Transformer的二次计算和KV缓存瓶颈。在256K tokens下,其推理速度比Qwen3-8B快3.5倍,并能在消费级GPU上支持高达1M tokens。该模型采用经济高效的持续训练方法,训练成本降低约75%。
查看缓存全文
缓存时间: 2026/07/11 15:26
长上下文Transformer面临两大瓶颈:二次注意力计算和KV缓存——在100万token时可达数百GB。稀疏注意力能降低计算开销,但存储依旧密集;线性注意力虽同时缓解两者,但会压缩上下文信息。
MiniCPM-SALA(OpenBMB,9B参数)采用混合架构:每3个线性注意力层中穿插1个稀疏注意力层,层位置自动选择,全局应用QK归一化稳定训练,线性层仅使用旋转位置编码(RoPE),每个块后添加输出门控。
经济亮点:它仅用约2万亿token,分5阶段将MiniCPM-4.0的中间检查点转化为混合模型,成本约为从头训练的25%,避开了混合模型从头训练时冷启动不稳定的问题。
效果:在标准基准测试中与Qwen3-8B持平,长上下文上显著领先(RULER-128K 89.4 vs 71.7)。仅训练至52万token,无需长度扩展技巧即可外推至200万token。推理:单卡A6000D上256K序列长度比Qwen3-8B快3.5倍;32GB显存的RTX 5090上可扩展至100万token,而Qwen3-8B在128K时就已显存溢出。
附上简要的可视化拆解——每个技巧对应一张图。左右滑动查看。
论文 - https://arxiv.org/abs/2602.11761 代码 - https://github.com/OpenBMB/MiniCPM 模型 - https://huggingface.co/openbmb/MiniCPM-SALA… 完整摘要PDF - https://gist.github.com/vukrosic/0c77bcf28e99843e4f47694b5b866b2b…
每周日我会举办AI研究实战直播(含一对一辅导):https://skool.com/become-ai-researcher-2669/about…
MiniCPM-SALA:融合稀疏与线性注意力实现高效长上下文建模
来源:https://arxiv.org/html/2602.11761
摘要
大型语言模型(LLMs)向超长上下文应用发展时,面临Transformer架构高计算与内存成本的挑战。现有的稀疏注意力和线性注意力机制虽试图缓解这些问题,但通常需要在内存效率与模型性能之间权衡。本文提出MiniCPM-SALA¹,一个9B参数的混合架构,结合了稀疏注意力(InfLLM-V2)的高保真长上下文建模能力与线性注意力(Lightning Attention)的全局效率。通过采用层选择算法以1:3的比例融合这两种机制,并利用混合位置编码(HyPE),模型在长上下文任务中保持了效率与性能。此外,我们引入了一种经济高效的持续训练框架,将预训练的Transformer模型转化为混合模型,相比从头训练可降低约75%的算力成本。大量实验表明,MiniCPM-SALA在保持与全注意力模型相当的一般能力的同时,提供了更高的效率。在单张NVIDIA A6000D GPU上,该模型在256K token序列长度下推理速度比全注意力模型快3.5倍,并支持长达100万token的上下文,而传统全注意力8B模型在此规模下因内存限制而失败。
1 引言
随着大型语言模型(LLMs)(OpenAI等,2024;Comanici等,2025;Grattafiori等,2024;Yang等,2025a;DeepSeek-AI等,2025)日益强大,其应用场景正经历深刻变革:从简单的问答(Brown等,2020)转向更高级的应用,如超长上下文的深度理解与生成(Bai等,2024, 2025;Zhou等,2025;Shao等,2024)、仓库级代码工程(Guo等,2024;Jimenez等,2024;Liu等,2024)以及面向复杂任务的长时间跨度智能体(Qian等,2024;Mialon等,2023;Li等,2026)。对于这些高级应用,模型不再局限于处理碎片化信息,而必须展现处理超长上下文的能力,例如一次性掌握完整的技术手册、分析包含数万行代码的综合项目依赖树,以及在多天的人机协作中保持连贯的任务状态和记忆。这种对整体上下文信息的追求,使得处理百万级token的能力成为先进LLMs的关键特性(Kimi Team等,2025;NVIDIA等,2025b)。
然而,作为现代LLMs基础的Transformer架构(Vaswani等,2017)在处理超长上下文时,因其核心的全注意力机制而遭遇严重的计算瓶颈。该瓶颈主要体现在两个维度:(1)计算瓶颈:标准注意力机制的计算复杂度随序列长度N呈二次增长,即复杂度为O(N²)。当上下文扩展到百万token级别时,巨大开销导致推理延迟急剧增加;(2)KV缓存内存瓶颈:在自回归生成过程中,模型必须存储所有历史上下文token的键值状态(KVs)以避免重复计算。对于典型的8B参数模型,即便使用分组查询注意力(GQA)(Ainslie等,2023),百万token所需的KV缓存也可能达到数十甚至数百GB。
为应对上述挑战,现有解决方案发展出两大范式:稀疏注意力(Yuan等,2025;DeepSeek-AI等,2025;Xiao等,2024;Zhao等,2025)和线性注意力(Yang等,2024a;Gu和Dao,2024;Peng等,2023;Yang等,2024b, 2025b)。两种范式各有其独特优势与固有局限。稀疏注意力方法通过仅计算注意力矩阵中最显著的部分(如采用滑动窗口或全局锚点)来突破计算瓶颈。然而,这些方法受限于“稀疏计算,密集存储”的困境:虽然局部计算减少了即时处理开销,但模型仍需保留完整的KV缓存以支持上下文信息检索。线性注意力利用循环公式成功将计算复杂度降低至O(N)。然而,这种极致效率是通过对上下文信息的有损压缩实现的,不可避免地会导致性能下降。
MiniCPM-SALA采用稀疏与线性注意力的混合架构(Chen等,2026),专门设计用于实现高效超长序列建模。该架构结合了InfLLM-V2(Zhao等,2025)的高保真长上下文建模能力与Lightning Attention(Qin等,2024)的全局计算效率。通过这种集成方法,模型显著降低了推理开销和内存消耗,同时解决了纯线性架构在长距离信息处理中常见的精度瓶颈。因此,MiniCPM-SALA在长上下文任务中提供了兼顾效率与高性能的平衡方案。此外,我们采用持续训练范式将预训练的Transformer模型转化为我们的混合模型。通过避免从头训练,该方法显著降低了模型开发的计算成本。虽然已有工作开始探索稀疏与线性注意力的融合(Hu等,2025;Hou等,2025;He和Garner,2025),但据我们所知,MiniCPM-SALA是首个通过大规模实验证明此类混合模型能够匹配全注意力基线性能的工作。此外,该模型在长上下文处理中展现出高效率和强性能。
综上所述,本研究的主要贡献可概括如下:
- • 我们引入了稀疏-线性混合注意力机制,整合了25%的InfLLM-V2与75%的Lightning Attention,以实现吞吐量与精度之间的平衡。通过利用稀疏注意力对局部细节的精细关注和线性注意力在广泛上下文上的O(N)效率,该架构在序列长度扩展时保持了高语义准确性。
- • 我们证明了Transformer到混合模型的范式是构建强大混合模型的一种高效策略。该方法通过对预训练权重进行架构变换,规避了冷启动训练的低效性,从而将总训练预算降低至从头训练类似模型所需的大约25%。
- • 我们采用HyPE(混合位置编码)(Chen等,2026)有效协调了短上下文与长上下文下的性能。在保持与现代全注意力模型(如Qwen3-8B)相当的一般能力(如知识、数学和编码)的同时,MiniCPM-SALA在多个长上下文基准上具有显著优势。
- • MiniCPM-SALA在长上下文场景中展现出显著的资源节省和速度优势。在NVIDIA A6000D GPU上,序列长度256K token时,MiniCPM-SALA的推理速度达到Qwen3-8B的3.5倍。此外,MiniCPM-SALA在NVIDIA A6000D和5090 GPU上均支持高达100万token的上下文长度推理,而Qwen3-8B在此长度下因内存溢出(OOM)而失败。这些结果展示了MiniCPM-SALA在边缘端信息密集型应用中的广阔前景。
2 模型开发
在本节中,我们将介绍MiniCPM-SALA的模型架构和训练策略。具体而言,我们在MiniCPM-SALA中结合了用于长上下文建模的高效稀疏注意力和用于全局效率的线性注意力。此外,我们还引入了一种高效的训练方法,该方法可以将标准Transformer模型转化为稀疏-线性混合注意力模型。
参考图注图1:MiniCPM-SALA架构。模型采用高效的混合设计,以1:3的比例结合InfLLM-V2(Zhao等,2025)和Lightning Attention(Qin等,2024)模块。基于MiniCPM-4.0(MiniCPM-Team等,2025)的中间检查点,MiniCPM-SALA经历持续训练阶段,将标准Transformer模型转化为稀疏-线性混合模型。### 2.1 模型架构
MiniCPM-SALA的整体架构如图1所示。MiniCPM-SALA采用了一种将稀疏注意力层和线性注意力层交错排列的混合架构。我们保留了Transformer架构中每个注意力块后的前馈网络(FFN)块,以确保高容量的知识表示。受近期代表性研究(如Qwen3-Next(Qwen Team, 2025)和Kimi-Linear(Kimi Team等,2025))以及我们内部小规模初步实验的启发,我们采用了1:3的混合比例:25%的层使用稀疏注意力,其余75%使用线性注意力。
这种混合配置充分利用了两种注意力机制的互补优势。线性注意力层的计算和内存复杂度相对于序列长度是恒定的,有利于高效处理长上下文。另一方面,稀疏注意力层有助于有效建模长距离依赖。我们并非简单地均匀交错这两种注意力变体,而是采用Chen等人(2026)提出的层选择机制来确定稀疏注意力模块的放置位置,这能带来更优的下游性能。
训练策略现有的混合模型训练范式通常分为两类:(1)从头训练(Zuo等,2025;Qwen Team, 2025;Kimi Team等,2025;NVIDIA等,2025b);(2)通过跨架构蒸馏将预训练的Transformer模型转化为混合模型(Wang等,2024a;Hoshino等,2025;Li等,2025;Gu等,2025)。尽管从头训练简单且提供最大的架构灵活性,但持续训练转换是一种更节省资源的替代方案,它利用了已有预训练模型中的参数继承。通过重用预训练权重和表示,持续训练方法显著降低了通常与新模型训练(de novo训练)相关的巨大计算成本,以极小的预算实现了有竞争力的性能。因此,MiniCPM-SALA采用基于转换的框架,通过持续训练将Transformer适配为高效的混合版本,同时保留其核心能力。
稀疏注意力与线性注意力对于稀疏注意力层,我们采用了InfLLM-V2(Zhao等,2025),其显著优势是不向架构引入额外参数。它固有的灵活性以及能够在密集模式和稀疏模式之间无缝切换的能力与我们的转换过程高度兼容。这种兼容性允许稀疏模块在不产生架构差异的情况下继承密集权重,从而促进稳定的训练初始化,确保向混合结构的转换顺利进行。
相似文章
@akshay_pachaar: 扩展上下文窗口不仅仅是关于更大的矩阵。在传统的Transformer中,将token数量扩大8倍会…
解释了由于注意力的二次复杂度,扩展Transformer上下文窗口所带来的内存挑战,并暗示了解决方案。
@NVIDIAAI: A long-context model's serving speed is largely decided before training starts. Attention used to be a small part of a …
NVIDIA explains how attention architecture choices (group size, head dimension, KV-cache size, parallelism) set the ceiling for long-context inference performance, with guidelines for co-designing models for faster serving.
Subquadratic AI 推出 SubQ-1.1-Small,一款采用 Smart Sparse Attention 的新模型
Subquadratic AI 推出 SubQ-1.1-Small,该模型利用 Smart Sparse Attention 在长达 1200 万 token 的上下文中实现近乎完美的长上下文检索,注意力计算量减少高达 1000 倍。它兼顾了长上下文优化与强大的通用推理能力,在 NIAH 和 RULER 等基准测试中优于基线模型。
MiniMax 稀疏注意力
MiniMax 稀疏注意力 引入了一种分块稀疏注意力机制,针对超长上下文的大语言模型实现了显著的加速。在1M上下文长度下,每个token的注意力计算减少28.4倍,在H800 GPU上预填充阶段实际速度提升14.2倍,解码阶段提升7.6倍。该方法附带了一个开源推理内核以及一个公开发布的多模态模型。
学习如何遗忘:针对长上下文稀疏注意力的微调
本文提出了一种新方法,用于对具有稀疏注意力的transformer语言模型进行微调,以实现高效的长上下文推理,通常优于使用精确注意力训练的模型,并介绍了高效的实现和一个新的开源库。