Sparse Chunky Attention: 一种基于分块思想且跳略非面相邻块的线性稀疏注意力机制

我们提出了Sparse Chunky Attention(稀疏分块注意力)来加速注意力计算。我们应用分块思想，将高维的嵌入空间分割成多个边长为chunk_size的超立方体块。仅计算Q所在块及其面相邻的块内的所有K-V对的注意力。尽管舍去了其他非面相邻块内的K-V对，但在1024维且chunk_size=512时随机KVQ实验中得到了82%的成功率。

摘要

我们提出了Sparse Chunky Attention(稀疏分块注意力)来加速注意力计算。我们应用分块思想，将高维的嵌入空间分割成多个边长为chunk_size的超立方体块。仅计算Q所在块及其面相邻的块内的所有K-V对的注意力。尽管舍去了其他非面相邻块内的K-V对，但在1024维且chunk_size=512时随机KVQ实验中得到了82%的成功率。

背景

目前LLM全部在使用的Transformer架构，由于传统注意力的缺陷，导致时间复杂度是平方级别，各家厂商也推出了自己的稀疏注意力来加快推理/训练。当然，他们多是通过稀疏掩码等方式选择性跳略。

介绍

把高维空间切分为多个边长为chunk_size的超立方体块。当Q查询时把Q所在的块及其面相邻的块内所有KV对取出并进行注意力运算。

缺点

显然，当Q在块的角或边缘上会忽略非面相邻的相邻块内的KV对。但端到端模型的梯度下降可能可以将Q的锁在块中心。

初步实验

在1024维chunk_size=512时得到了82%的成功率（Q是从已有K中取出并加上随机小扰动生成）。

未来规划

极度感谢@ElandaDRM朋友的5070显卡支持。我将在中秋节左右检验我在第二段的推论并训练一个小规模模型。

结尾

附随机Q查询实验的代码，仅依赖numpy库。
致谢@ElendaDRM的支持。
感谢@hotpad100c的初步评审。