MonarchAttention：零-shot 转换至快速、硬件感知结构化注意力

Transformer在 various 任务中实现了最佳性能，但由于注意力机制，其序列长度导致显著的二次复杂度。本文提出MonarchAttention——一种通过Monarch矩阵实现亚二次注意力近似的新方法。基于softmax的变分形式，我们描述了一种高效的优化算法，用于计算将softmax注意力近似投影到Monarch矩阵类上的方法，其计算复杂度为 $\Theta(N\sqrt{N} d)$，内存/IO复杂度为 $\Theta(Nd)$。与之前方法不同，MonarchAttention既是（1）可迁移的，即使在替换Transformer的每个注意力层后仍可保持最小性能损失且无需额外训练，同时（2）硬件高效，利用现代GPU上最高吞吐量的张量核心单元。经过优化后的内核，MonarchAttention在FlashAttention-2的墙时时间上实现了显著的加速：对于较短序列（$N=256$）为 $1.4\times$，中等长度序列（$N=4K$）为 $4.5\times$，较长序列（$N=16K$）为 $8.2\times$。我们在视觉和语言问题中展示了MonarchAttention的质量，表明它在各种上下文中灵活且准确地近似softmax注意力。我们的代码已公开：https://github.com/cjyaras/monarch-attention。