arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2410.09397cs.LGcs.AIcs.CCcs.CL
收藏

关于注意力反向传递中细粒度I/O复杂度

On Fine-Grained I/O Complexity of Attention Backward Passes

Department of XXX, University of YYY, Location, Country(YYY大学XXX系);School of ZZZ, Institute of WWW, Location, Country(WWW研究所ZZZ学院)

Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Song Yue, Jiahao Zhang

AI总结:

本文研究了注意力机制中细粒度I/O复杂度,提出了一种在小缓存环境下优于现有方法的算法,并为高效LLM训练和推理提供了理论指导。

摘要:

大型语言模型(LLMs)在处理自然语言中的大量上下文窗口方面表现出色。然而,注意力计算的二次缩放与序列长度相关,导致了显著的效率瓶颈,需要开发I/O优化的算法。在本工作中,我们系统地检查了注意力机制中固有的I/O复杂度,特别强调在小和大缓存设置下的反向传递。通过利用红蓝 Pebble 游戏框架,我们推导出在所有缓存大小范围内I/O复杂度的紧界。我们验证了FlashAttention,当前行业标准之一,在大缓存场景中在正向和反向传递中都达到最优。相反,在小缓存环境中,我们引入了一种新的算法,其性能优于当前方法,并成功达到了理论紧界。此外,我们还扩展了研究范围,包括稀疏注意力,通过在所有缓存配置中为正向和反向传递建立细粒度的下界。最终,我们的结果巩固了关于注意力机制中I/O复杂度的理论框架,为高效LLM训练和推理系统的发展提供了关键指导。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑