学习跳跃块：自我发现的超度量路由用于硬件加速稀疏注意力

Reddit r/artificial 论文
sparse-attention ultrametric block-sparse triton-kernel transformer long-context hardware-acceleration

摘要

本文介绍了动态超度量注意力（Dynamic Ultrametric Attention），这是一个框架，其中Transformer在训练期间学习每头块稀疏路由拓扑，然后在推理时将这些拓扑卸载到自定义的Triton块稀疏内核上，与密集注意力相比，实现了高达28倍的加速和98.4%的内存减少。

摘要。标准的密集自注意力在序列长度上呈二次方扩展，为长上下文Transformer造成了难以处理的内存和计算瓶颈。我们引入了动态超度量注意力（Dynamic Ultrametric Attention），这是一个框架，其中Transformer通过Gumbel-Sigmoid深度门在训练期间自主地学习每头块稀疏路由拓扑，然后在推理时将这些学习到的稀疏模式直接卸载到自定义的Triton块稀疏内核上。路由拓扑来源于一个超度量（树结构）距离矩阵，该矩阵编码了标记位置之间的层次关系。在涵盖Dyck-k括号语言、Long Range Arena ListOps基准、自回归服务以及自然语言建模的九项实验中，我们证明了：(1) 动态门有机地发现了层级专业化——早期层专门进行层次解析，后期层专注于密集聚合——无需任何架构约束；(2) 学习到的稀疏图无损地转移到块稀疏Triton内核上，该内核会跳过非参与块的整个SRAM加载；(3) 由此产生的系统在2048个标记时实现了比PyTorch密集注意力快11.59倍的实际推理加速，在8192个标记时扩展到28倍，内存减少98.4%；(4) 稀疏PagedAttention解码内核通过条件性地跳过KV缓存块加载，实现了比密集解码高8倍的有效内存带宽；(5) 当加入局部滑动窗口时，该架构在真实自然语言（莎士比亚）的所有层上保持>88%的稀疏性，同时将交叉熵损失从10.9降低到1.55。据我们所知，这是首次展示LLM自主学习其硬件最优稀疏模式，并将其桥接到物理加速内核上，无需事后剪枝或蒸馏。https://github.com/sneed-and-feed/adelic-spectral-zeta/blob/main/papers/learning_to_skip_blocks.md
查看原文

相似文章

用于可扩展视觉 Transformer 的弹性注意力核 [R]

Reddit r/MachineLearning

本文介绍了一篇关于视觉 Transformer 弹性注意力核的新论文，提出了一种核心-外围块稀疏注意力结构，与 DINOv3 等密集自注意力方法相比，该结构提高了可扩展性和准确率。

Block-Based Double Decoders

arXiv cs.LG

提出了一种基于块的雙解碼器（block-based double decoders），这是一种使用双重因果块注意力掩码的新型Transformer架构，结合了解码器仅训练效率与编码器-解码器推理效率，实现了强大的扩展性能并减少了KV缓存内存。