揭示注意力反向传播的输入输出复杂度，为高效大模型训练提供理论依据。
On Fine-Grained I/O Complexity of Attention Backward Passes
- 用红蓝石子游戏框架推导注意力反向传播的输入输出复杂度界。
- 发现FlashAttention在大缓存下已最优，小缓存下新算法逼近理论极限。
- 首次建立稀疏注意力的细粒度复杂度下界，适用于各类缓存场景。
大规模语言模型在处理长文本上下文方面表现出色，但注意力计算随序列长度呈二次增长，带来显著效率瓶颈，亟需输入输出（I/O）优化算法。本文系统研究了注意力机制中反向传播的I/O复杂度，覆盖小缓存与大缓存两种情形。基于红蓝石子游戏框架，我们推导出全缓存尺度下的紧致下界。验证表明，当前工业标准FlashAttention在大缓存下对前向与反向传播均达到最优。而在小缓存环境下，我们提出一种新算法，优于现有方法并逼近理论最优。此外，我们将研究扩展至稀疏注意力，建立了所有缓存配置下前向与反向传播的细粒度下界。最终结果构建了注意力I/O复杂度的完整理论体系，为高效大模型训练与推理系统开发提供关键指导。
原文摘要 · Abstract (English)
Large Language Models (LLMs) exhibit exceptional proficiency in handling extensive context windows in natural language. Nevertheless, the quadratic scaling of attention computation relative to sequence length creates substantial efficiency bottlenecks, necessitating the development of I/O-optimized algorithms. In this work, we conduct a systematic examination of the I/O complexity inherent in attention mechanisms, with a specific emphasis on the backward pass under both small and large cache settings. By leveraging the red-blue pebble game framework, we derive tight bounds for I/O complexity across the full spectrum of cache sizes. We validate that FlashAttention, one of the current industry standards, achieves optimality in the large-cache scenario for both forward and backward passes. Conversely, for small-cache environments, we introduce a novel algorithm that outperforms contemporary methods and successfully attains theoretical tight bounds. Furthermore, we expand our investigation to include sparse attention by establishing granular lower bounds for both forward and backward passes across all cache configurations. Ultimately, our results solidify the theoretical framework regarding I/O complexity in attention mechanisms, providing critical guidance for the development of efficient LLM training and inference systems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。