LazyFlash: Forward-Dense, Backward-Sparse Attention for Accelerated Long-Context Fine-Tuning

📄 arXiv: 2610.31911v1 📥 PDF

作者: Keqi Deng, Ruchao Fan, Jinyu Li

分类: cs.CL

发布日期: 2026-10-22

💡 一句话要点

提出LazyFlash以加速长序列上下文微调

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长序列处理 稀疏注意力 模型训练加速 梯度感知 自然语言处理 语音识别

📋 核心要点

  1. 现有的注意力机制在处理长序列时面临计算瓶颈，稀疏注意力方法在训练时的应用仍然有限。
  2. LazyFlash通过前向密集和后向稀疏的设计，优化了训练过程，特别是通过跳过对参数更新贡献小的块来加速后向传播。
  3. 实验结果显示，LazyFlash在保持性能的同时，后向传播速度提升最高可达5.7倍，端到端训练速度提升1.7倍。

📝 摘要（中文）

随着大型语言模型（LLMs）处理越来越长的序列，注意力机制成为主要的计算瓶颈。现有的稀疏注意力方法主要针对推理，而许多方法在训练时仍然采用密集计算，导致稀疏训练难以实现实际的时间加速。本文提出LazyFlash，一种前向密集、后向稀疏的注意力方法，基于注意力训练主要受后向传播支配的观察。LazyFlash在前向传播中保持完全密集，使得训练后的模型仍然可以被现有系统如vLLM无缝使用。对于后向传播，LazyFlash提出了一种梯度感知的块稀疏化设计，通过注意力输出梯度范数对查询块进行排序，跳过那些对参数更新贡献微小的块，从而将密集的后向传播转变为块稀疏。实验表明，LazyFlash在长上下文文本和语音LLMs上保持下游性能，同时在FlashAttention上实现了最高5.7倍的后向速度提升和1.7倍的端到端训练速度提升。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在长序列处理中的计算瓶颈，尤其是在训练阶段，现有方法多采用密集计算，难以实现稀疏训练的速度优势。

核心思路：LazyFlash的核心思想是保持前向传播的密集计算，以确保模型兼容性，同时通过后向传播的块稀疏化设计来减少计算量，从而加速训练过程。

技术框架：LazyFlash的整体架构包括前向传播和后向传播两个阶段。在前向传播中，使用标准的密集注意力机制；在后向传播中，采用梯度感知的块稀疏化策略，动态调整计算量。

关键创新：LazyFlash的主要创新在于其梯度感知的块稀疏化设计，通过对查询块的注意力输出梯度进行排序，跳过对参数更新贡献小的块，从而实现了真正的计算加速。

关键设计：在后向传播中，LazyFlash引入了梯度质量预算，根据每个头和训练步骤动态调整跳过的计算量，同时使用重要性采样来保持密集梯度的期望值。

📊 实验亮点

LazyFlash在长上下文文本和语音LLMs上的实验结果显示，后向传播速度提升最高可达5.7倍，相较于FlashAttention，端到端训练速度提升1.7倍，且保持了下游任务的性能。

🎯 应用场景

LazyFlash的研究成果在自然语言处理、语音识别等领域具有广泛的应用潜力，尤其是在需要处理长上下文的任务中。其加速效果能够显著提升模型训练效率，降低计算资源消耗，推动更大规模模型的实际应用。

📄 摘要（原文）

As large language models (LLMs) process increasingly long sequences, attention becomes a major computational bottleneck. Sparse attention mostly targets inference, yet many methods still train densely or cannot support sparse training, and their added complexity makes sparsity hard to turn into wall-clock speedups, leaving dense attention widely used in practice. This paper introduces LazyFlash, a forward-dense, backward-sparse attention method built on the observation that attention training is dominated by the backward pass. LazyFlash keeps the forward pass exactly dense, so the trained model remains a standard dense model that existing systems such as vLLM can serve without any change. For the backward pass, LazyFlash proposes a gradient-aware block sparsification design that ranks query blocks by their attention-output gradient norms and skips those contributing negligibly to parameter updates, turning the dense backward into a block-sparse one. A gradient-mass budget adapts the amount of skipped computation to each head and training step, while importance sampling preserves the dense gradient in expectation. These designs are implemented inside the tiled FlashAttention backward kernel, where skipped blocks are never loaded or computed, translating gradient sparsity into real wall-clock acceleration. Experiments on long-context text and speech LLMs show that LazyFlash maintains downstream performance while achieving up to 5.7x attention-backward speedup over FlashAttention and 1.7x end-to-end training speedup.