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%,同时在通用任务、数学任务、编码任务等标准基准上保持与现有模型相当的性能,验证了其在稀疏性、随机访问灵活性和长度泛化三大核心属性上的优势。二、创新点分层稀疏注意力(HSA)机制设计:借鉴混合专家(MoE)思想,将文本划分为固定长度块并生成地标表示,每个token通过检索分数选择top-k相关块,分别进行块内注意力计算后按检索分数加权融合,实现端到端的块选择与注意力学习,解决了现有稀疏注意力方案块选择不准确的问题。多机制融合的模型架构:结合SWA处理局部信息与HSA处理全局信息,共享中间层KV缓存压缩内存占用;采用“短序列Ro