在大模型LLM的竞赛中“长文本能力”一直是兵家必争之地。然而现有的Transformer架构面临着一个物理定律般的诅咒随着上下文长度增加计算复杂度呈二次方爆炸KV Cache更是像无底洞一样吞噬显存。ArXiv URLhttp://arxiv.org/abs/2602.01766v1虽然市面上有各种“上下文压缩”技术但它们往往以牺牲精度为代价。有没有一种方法既能保持线性时间复杂度又能让显存占用恒定不变还能精准找回100万字之前的细节来自阿里巴巴、东北大学和清华大学的研究团队给出了答案。他们提出了CoMeTCollaborative Memory Transformer这是一种全新的架构只需在32k长度的文本上进行微调就能完美泛化到1M100万Token的超长上下文且显存占用几乎不随长度增加核心痛点Transformer的“记忆危机”标准的Transformer依赖于KV Cache来存储历史信息。这导致了两个致命问题显存爆炸处理1M长度的文本KV Cache的大小会变得天文数字单张GPU根本装不下。计算龟速注意力机制的计算量随着长度呈二次方增长O(N2)O(N^2)O(N2)推理速度极慢。现有的解决方案要么是“有损压缩”丢失细节要么是基于RNN的线性注意力容易遗忘关键信息。CoMeT则走出了一条新路它引入了一套协同记忆机制像人类大脑一样既有“短期记忆”负责即时细节又有“长期记忆”负责关键线索。CoMeT的秘密武器双重协同记忆CoMeT的设计理念非常精妙它作为一个轻量级的“插件”模块可以无缝集成到预训练模型中。其核心在于将输入文本切分成一个个块Chunk并利用两套记忆系统来管理上下文1. 临时记忆Temporary MemoryFIFO队列这相当于我们的“工作记忆”。CoMeT使用一个先进先出FIFO的队列来存储最近几个块的精细信息。作用确保模型对最近发生的事件有高保真的感知不会因为压缩而丢失细节。机制随着新信息的进入最旧的信息被移出但这并不意味着遗忘因为它们会被提炼进入全局记忆。2. 全局记忆Global Memory门控更新这是CoMeT的“长期存储”。它不像传统RNN那样简单地覆盖旧状态而是引入了类似LSTM的门控更新机制Gated Update Rule。作用专门捕捉长距离依赖。门控机制会智能地判断哪些新信息需要写入哪些历史核心信息必须保留公式全局状态的更新公式为Sτ1ig⊙Sτi(1−g)⊙S~τ1i\mathbf{S}^{i}_{\tau1}\mathbf{g}\odot\mathbf{S}^{i}_{\tau}(\mathbf{1}-\mathbf{g})\odot\tilde{\mathbf{S}}^{i}_{\tau1}Sτ1ig⊙Sτi(1−g)⊙S~τ1i其中g\mathbf{g}g是门控权重。这有效防止了重要历史信息被后续无关信息“冲刷”掉。如上图所示在处理每一个文本块时模型会同时参考全局记忆Gτi\mathbf{G}^{i}_{\tau}Gτi和临时记忆Tτi\mathbf{T}^{i}_{\tau}Tτi从而兼顾了全局大局观和局部细节。32k训练1M推理惊人的泛化能力CoMeT最令人印象深刻的特性是其强大的外推能力。研究人员仅在32k长度的上下文上对模型进行了微调然后直接在1M100万Token的长度上进行测试。结果如何请看下图的“大海捞针”Passkey Retrieval测试全绿的准确率图1(a)显示无论密钥Passkey隐藏在1M上下文的哪个位置CoMeT都能准确找回准确率接近100%。线性的推理时间图1(b)显示随着长度增加CoMeT的推理时间仅呈线性增长而全注意力机制Full Attention则是指数级暴涨。恒定的显存占用图1©最为震撼无论上下文是10万还是100万CoMeT的显存占用始终保持在一条水平线上约10GB而Full Attention在128k时就已经OOM显存溢出了。训练加速层级流水线并行为了让CoMeT能在超长上下文上高效训练团队还提出了一种层级流水线并行策略Layer-level Pipeline Parallelism。传统的上下文并行Context Parallelism在处理序列块时Worker之间存在严重的等待时间气泡。而CoMeT的新策略将不同层的计算流水线化使得训练速度相比朴素方法提升了2.7倍。这意味着仅用16张80GB的GPU就能高效微调128k长度的模型。实验表现不仅快而且准在SCROLLS长文本基准测试中CoMeT的表现超越了其他高效方法如LongLLMLingua、Transformer-XL等并在摘要任务上达到了与全注意力基线相当的水平。特别是在真实世界的应用场景——用户行为序列QA中CoMeT展现了巨大的实用价值。面对电商场景下成千上万条用户交互记录CoMeT能精准捕捉用户的兴趣点性能显著优于工业界常用的xRAG和简单的截断策略。总结CoMeT的出现打破了长文本处理中“效率”与“性能”不可兼得的魔咒。对于开发者它是一个即插即用的模块改造成本极低。对于资源受限的场景它让在单卡上处理百万级上下文成为可能。随着AI Agent和长文档分析需求的爆发CoMeT这种“恒定显存、无限长度”的架构或许正是通往下一代高效大模型的关键钥匙。
阿里CoMeT：32k训练解锁100万长文，显存恒定，告别KV Cache爆炸
发布时间：2026/6/10 12:45:58
在大模型LLM的竞赛中“长文本能力”一直是兵家必争之地。然而现有的Transformer架构面临着一个物理定律般的诅咒随着上下文长度增加计算复杂度呈二次方爆炸KV Cache更是像无底洞一样吞噬显存。ArXiv URLhttp://arxiv.org/abs/2602.01766v1虽然市面上有各种“上下文压缩”技术但它们往往以牺牲精度为代价。有没有一种方法既能保持线性时间复杂度又能让显存占用恒定不变还能精准找回100万字之前的细节来自阿里巴巴、东北大学和清华大学的研究团队给出了答案。他们提出了CoMeTCollaborative Memory Transformer这是一种全新的架构只需在32k长度的文本上进行微调就能完美泛化到1M100万Token的超长上下文且显存占用几乎不随长度增加核心痛点Transformer的“记忆危机”标准的Transformer依赖于KV Cache来存储历史信息。这导致了两个致命问题显存爆炸处理1M长度的文本KV Cache的大小会变得天文数字单张GPU根本装不下。计算龟速注意力机制的计算量随着长度呈二次方增长O(N2)O(N^2)O(N2)推理速度极慢。现有的解决方案要么是“有损压缩”丢失细节要么是基于RNN的线性注意力容易遗忘关键信息。CoMeT则走出了一条新路它引入了一套协同记忆机制像人类大脑一样既有“短期记忆”负责即时细节又有“长期记忆”负责关键线索。CoMeT的秘密武器双重协同记忆CoMeT的设计理念非常精妙它作为一个轻量级的“插件”模块可以无缝集成到预训练模型中。其核心在于将输入文本切分成一个个块Chunk并利用两套记忆系统来管理上下文1. 临时记忆Temporary MemoryFIFO队列这相当于我们的“工作记忆”。CoMeT使用一个先进先出FIFO的队列来存储最近几个块的精细信息。作用确保模型对最近发生的事件有高保真的感知不会因为压缩而丢失细节。机制随着新信息的进入最旧的信息被移出但这并不意味着遗忘因为它们会被提炼进入全局记忆。2. 全局记忆Global Memory门控更新这是CoMeT的“长期存储”。它不像传统RNN那样简单地覆盖旧状态而是引入了类似LSTM的门控更新机制Gated Update Rule。作用专门捕捉长距离依赖。门控机制会智能地判断哪些新信息需要写入哪些历史核心信息必须保留公式全局状态的更新公式为Sτ1ig⊙Sτi(1−g)⊙S~τ1i\mathbf{S}^{i}_{\tau1}\mathbf{g}\odot\mathbf{S}^{i}_{\tau}(\mathbf{1}-\mathbf{g})\odot\tilde{\mathbf{S}}^{i}_{\tau1}Sτ1ig⊙Sτi(1−g)⊙S~τ1i其中g\mathbf{g}g是门控权重。这有效防止了重要历史信息被后续无关信息“冲刷”掉。如上图所示在处理每一个文本块时模型会同时参考全局记忆Gτi\mathbf{G}^{i}_{\tau}Gτi和临时记忆Tτi\mathbf{T}^{i}_{\tau}Tτi从而兼顾了全局大局观和局部细节。32k训练1M推理惊人的泛化能力CoMeT最令人印象深刻的特性是其强大的外推能力。研究人员仅在32k长度的上下文上对模型进行了微调然后直接在1M100万Token的长度上进行测试。结果如何请看下图的“大海捞针”Passkey Retrieval测试全绿的准确率图1(a)显示无论密钥Passkey隐藏在1M上下文的哪个位置CoMeT都能准确找回准确率接近100%。线性的推理时间图1(b)显示随着长度增加CoMeT的推理时间仅呈线性增长而全注意力机制Full Attention则是指数级暴涨。恒定的显存占用图1©最为震撼无论上下文是10万还是100万CoMeT的显存占用始终保持在一条水平线上约10GB而Full Attention在128k时就已经OOM显存溢出了。训练加速层级流水线并行为了让CoMeT能在超长上下文上高效训练团队还提出了一种层级流水线并行策略Layer-level Pipeline Parallelism。传统的上下文并行Context Parallelism在处理序列块时Worker之间存在严重的等待时间气泡。而CoMeT的新策略将不同层的计算流水线化使得训练速度相比朴素方法提升了2.7倍。这意味着仅用16张80GB的GPU就能高效微调128k长度的模型。实验表现不仅快而且准在SCROLLS长文本基准测试中CoMeT的表现超越了其他高效方法如LongLLMLingua、Transformer-XL等并在摘要任务上达到了与全注意力基线相当的水平。特别是在真实世界的应用场景——用户行为序列QA中CoMeT展现了巨大的实用价值。面对电商场景下成千上万条用户交互记录CoMeT能精准捕捉用户的兴趣点性能显著优于工业界常用的xRAG和简单的截断策略。总结CoMeT的出现打破了长文本处理中“效率”与“性能”不可兼得的魔咒。对于开发者它是一个即插即用的模块改造成本极低。对于资源受限的场景它让在单卡上处理百万级上下文成为可能。随着AI Agent和长文档分析需求的爆发CoMeT这种“恒定显存、无限长度”的架构或许正是通往下一代高效大模型的关键钥匙。
