Every Token Counts: Generalizing 16M Ultra-Long Context in Large Language Models

文章主要内容与创新点总结

一、主要内容

本文围绕大语言模型的超长上下文建模问题展开研究，核心目标是构建具备高效长期记忆能力的“能记忆的机器”。现有Transformer架构因全注意力机制的二次计算复杂度和长度泛化能力不足，难以处理超长序列，而现有稀疏注意力、循环架构等方案存在信息瓶颈、随机访问缺失或块选择不准确等缺陷。

研究团队提出融合分层稀疏注意力（HSA）与滑动窗口注意力（SWA）的HSA-UltraLong模型，包括0.5B稠密模型和8B混合专家（MoE）模型（激活参数1B），在8万亿 tokens 数据上完成训练，并通过预热、预训练、长上下文中期训练、退火和有监督微调的五阶段训练流程优化模型性能。

实验表明，该模型在8K预训练窗口、32K中期训练窗口的基础上，能成功外推至1600万token的超长上下文，在多数上下文检索任务中准确率超90%，同时在通用任务、数学任务、编码任务等标准基准上保持与现有模型相当的性能，验证了其在稀疏性、随机访问灵活性和长度泛化三大核心属性上的优势。

二、创新点

1. 分层稀疏注意力（HSA）机制设计：借鉴混合专家（MoE）思想，将文本划分为固定长度块并生成地标表示，每个token通过检索分数选择top-k相关块，分别进行块内注意力计算后按检索分数加权融合，实现端到端的块选择与注意力学习，解决了现有稀疏注意力方案块选择不准确的问题。
2. 多机制融合的模型架构：结合SWA处理局部信息与HSA处理全局信息，共享中间层KV缓存压缩内存占用；采用“短序列Ro