AI 自动翻译
arXiv：2610.31911v1 发布类型：新
摘要：随着大型语言模型（LLMs）处理的序列长度不断增长，注意力机制已成为主要的计算瓶颈。稀疏注意力方法大多面向推理阶段，但许多现有方法仍采用密集训练，或无法支持稀疏训练；加之其额外的复杂性，使得稀疏性难以转化为实际的运行时间加速，因此在实践中密集注意力仍被广泛使用。本文提出了 LazyFlash，一种前向密集、后向稀疏的注意力机制，其核心基于这样一个观察：注意力机制的训练过程主要由反向传播主导。LazyFlash 保持前向传播完全密集，因此训练得到的模型仍然是标准的密集模型，现有的 vLLM 等系统无需任何修改即可直接服务。在反向传播阶段，LazyFlash 提出了一种梯度感知的分块稀疏化设计：根据查询块的注意力输出梯度范数进行排序，并跳过那些对参数更新贡献可忽略不计的块，从而将密集的反向传播转化为分块稀疏的形式。通过引入梯度总量预算，可根据每个注意力头和每一步训练动态调整被跳过的计算量；同时，重要性采样机制确保了期望意义上的密集梯度得以保留。这些设计均集成于分块化的 FlashAttention 反向传播内核中，其中被跳过的块既不会被加载也不会被计算，从而将梯度稀疏性切实转化为实际的运行时间加速。针对长上下文文本与语音类大型语言模型的实验表明，LazyFlash 在保持下游任务性能的同时，相较于 FlashAttention 实现了最高 5.7 倍的注意力反向传播加速，并带来了 1.7 倍的端到端训练加速。