论文
PHBA:递归前缀状态与稀疏块注意力
PHBA: Prefix-State Hybrid Block Attention
摘要
将线性序列模型与 softmax 注意力相结合的混合架构在高效的长上下文建模和精确的 token 检索之间提供了有效的平衡。诸如本机混合注意力(NHA)之类的现有设计将压缩的长期状态与滑动窗口注意力相结合,但它们的确切注意力仅限于固定的局部窗口。在这项工作中,我们引入了前缀状态混合块注意力(PHBA),它将局部滑动窗口注意力替换为 top-k 块稀疏检索,并将每个检索到的块与总结其先前上下文的紧凑前缀状态耦合。前缀状态是通过块边界处的门控线性递归构建的,并与相应的 token 块一起检索,从而允许模型将精确的远程证据与统一层内的压缩历史上下文结合起来。我们进一步开发了一种硬件感知的 Triton 实现,可以流式传输路由的 token 块和前缀状态,而无需实现大型中间张量。实验表明,与强线性和混合基线相比,PHBA 提高了长上下文和检索性能,同时保留了高效的训练和推理。
