基于幂律引导的动态分层以提高注意力效率
机器学习
2025-06-06 v1
摘要
在GPU上使用大型语言模型进行推理仍然具有挑战性，主要由于注意力计算中High带宽存储器（HBM）与SRAM之间的数据传输存在内存带宽限制。近似注意力方法通过减少计算和内存开销来解决此问题，但常依赖昂贵的top-k操作，在GPU上表现不佳。我们提出了SiftAttention，一种新的近似注意力方法，用基于阈值值的元素级过滤操作取代top-k步骤。我们的直觉基于这样一个经验观察：注意力得分的τ分位数随序列生成步骤呈可预测的幂律分布。利用这一洞察，我们的方法在每个生成步骤动态估计每个提示的阈值。只有超过此阈值的注意力得分及其对应的值向量才会被加载/用于计算注意力输出，从而减少HBM和SRAM之间的数据传输。我们的评估表明，SiftAttention在保持模型质量方面优于现有的近似注意力方法，同时在加载值向量时降低了内存带宽使用。
引用
@article{arxiv.2506.05300,
title = {Power Law Guided Dynamic Sifting for Efficient Attention},
author = {Nirav Koley and Prajwal Singhania and Abhinav Bhatele},
journal= {arXiv preprint arXiv:2506.05300},
year = {2025}
}