Show HN: Proxy Block-CAGE，一种新的稀疏块注意力机制

在人工智能研究领域，处理长序列的Transformer模型面临着注意力机制计算量随序列长度平方增长的问题。近日，一位名为Iqbal Addou的博士研究生在GitHub上发布了一个名为Proxy Block-CAGE的注意力原型，旨在通过稀疏化方法降低长上下文预填充阶段的计算开销。该项目不仅是一项技术探索，更是一次将ChatGPT作为研究伙伴的实践。

Proxy Block-CAGE的核心思想是将查询（Q）、键（K）和值（V）分割成固定大小的块，然后使用每个块的平均池化摘要来计算块之间的代理分数。根据这些分数，每个查询块会选择得分最高的几个键块，随后只在这些选择的块内执行标准的精确注意力计算。这种方法在保持注意力质量的同时，大幅减少了需要计算的查询-键对数量。

作者在配备NVIDIA RTX 2080 Max-Q GPU的笔记本电脑上，对一个两层字符级Transformer进行了测试。结果显示，在16,384个token的上下文长度下，Proxy Block-CAGE实现了2.61倍的加速，注意力密度仅为密集注意力的0.39%，而最终token的top-1和top-5准确率保持100%。虽然测试规模较小，但足以证明该方法的潜力。

文章详细记录了项目的演变过程。最初，作者尝试使用遗传算法（GA）为每个查询token选择稀疏的注意力掩码（称为“CAGE”），但token级别的搜索耗时巨大。随后，他将搜索单元从token升级为块，大幅减少了搜索次数，但遗传算法的计算开销仍然过高。最终，作者发现对于仅基于log-sum-exp的简单目标函数，最优解等价于简单的top-k选择，遗传算法由此被更高效的有序选择取代。这一发现促成了Proxy Block-CAGE的最终形式，即通过池化摘要的代理分数进行top-k块选择。

从数学上看，Proxy Block-CAGE的计算复杂度约为O(BHd(n²/(b_q b_k) + n m b_k))，其中b_q和b_k是块大小，m是每个查询块选择的键块数。在典型配置（b_q=b_k=32, m=2）下，16K上下文时的计算位置数仅为密集注意力的约1/100。尽管由于工程优化程度不同，实际加速比未达到理论值，但已足够表明稀疏注意力在长上下文场景中的优势。

作者强调，这只是一个可重复的研究原型，而非生产级解决方案。他希望通过公开此项目，获得来自社区的反馈，并鼓励更多研究者探索稀疏注意力中的数学与算法创新。