AdaSplash-2：更快速的可微稀疏注意力

稀疏注意力被用于缓解Transformer的二次计算成本，这是长上下文训练的主要瓶颈。α-entmax注意力是softmax的一种可微稀疏替代方式，能够实现依赖输入的稀疏性，但计算归一化阈值τ所需的开销使其效率落后于softmax。本文提出AdaSplash-2，通过一种基于直方图的初始化方法，将计算τ所需的迭代次数通常减少到1–2次。关键做法是即时计算注意力分数的粗粒度直方图，并将其存储在片上SRAM中，以获得更准确的初始化，加速前向与反向计算。结合能够低开销跳过全零块的稀疏感知GPU实现，当块稀疏度为中高水平、例如超过60%时，AdaSplash-2的单步训练时间达到或优于FlashAttention-2；这类稀疏度在长上下文中较常出现。下游任务中，使用该高效α-entmax注意力训练的模型在短上下文下达到softmax基线的表现，在长上下文下取得明显改善。