正确实现分层稀疏注意力:迈向无限上下文建模

发表
Lei ZHULei ZHU 提交
作者: Xiang HuXiang Hu, Xinyu Wei, guhaoHao Gu, Minshen ZhangMinshen Zhang, Tian Liang, Huayang Li, Lei ZHULei Zhu, Yan Wang, Sirui Han, Yushi Bai, Kewei Tu, Haitao Mi, Leo Liang

摘要

AI 生成总结
分层地标稀疏注意力通过端到端学习块选择,实现高效的长上下文语言建模,在训练上下文长度之外进行外推时,其性能可与全注意力相媲美。
将现代大型语言模型(LLM)扩展到长上下文受到二次计算成本和密集注意力不良长度外推的限制。分块稀疏注意力提供了一个有前景的替代方案,但所有现有方法都因其不准确的分块选择而无法达到完全注意力。我们提出了分层地标稀疏(HiLS)注意力,这是一种分块稀疏注意力机制,它在语言建模(LM)损失下端到端地学习分块选择。HiLS分层分解注意力:每个查询独立地与每个检索到的块执行注意力以提取块特定信息,然后根据块检索分数融合结果输出。通过将检索分数纳入前向注意力计算,HiLS使用LM损失直接优化它们,从而实现端到端检索学习和原生稀疏训练。实验结果表明,HiLS-Attention在域内上下文长度上实现了与完全注意力相当甚至更好的性能。同时,HiLS-Attention将训练上下文长度外推超过64倍,检索准确率达到90%,远超完全注意力。此外,现有完全注意力模型可以通过轻量级持续预训练转换为HiLS-Attention,在保持域内性能的同时获得超长上下文外推能力。HiLS-Attention结合其稀疏的KV访问和计算,打破了通常的效率-性能权衡,使得长上下文LLM在通用长上下文任务上比其完全注意力对应物更高效、更有效。
查看 arXiv 页面查看 PDF

评论

Lei ZHULei ZHU
论文作者
论文提交者

一种具有卓越长度泛化能力的完全可微分分层稀疏注意力机制。