LLM2D
面向熔蟒的随机长上下文访问通过硬件对齐分层稀疏注意机制
Random Long-Context Access for Mamba via Hardware-aligned Hierarchical Sparse Attention
作者: Xiang Hu, Jiaqi Leng, Jun Zhao, Kewei Tu, Wei Wu
发布日期: 4/24/2025
arXiv ID: oai:arXiv.org:2504.16795v1

摘要

arXiv:2504.16795v1 交叉公告类型:交叉 摘要:与变换器相比,循环神经网络(RNNs)的一个关键优势是它们的线性计算和空间复杂度使得它们能够更快地训练和推断长序列。然而,RNNs本质上无法随机访问历史上下文,简单地集成注意力机制可能会削弱它们的效率优势。为了解决这一限制,我们提出了一种新颖的注意机制——**层次稀疏注意**(HSA),它增强了RNNs的长距离随机访问灵活性,同时保持了它们在效率和长度泛化的优点。HSA 将输入分为块,选择前 $k$ 个块,并逐级聚合信息。核心创新在于基于每个块内部的细粒度词级信息来学习词块相关性。这种方法增强了在不同领域内外上下文中块选择的精确度。为了使HSA有效,我们进一步引入了一种与硬件对齐的内核设计。通过将HSA与Mamba结合,我们引入了RAMba,它在只在4K长度上下文中进行预训练的情况下,实现了在6400万个上下文中密码钥匙检索的完美准确率,并在各种下游任务上取得了显著改进,内存占用几乎保持恒定。这些结果表明,RAMba 在长上下文建模方面具有巨大的潜力。