提出首个近乎线性时间的RoPE注意力反向传播算法,显著加速训练。
RoPE Attention Can Be Trained in Almost Linear Time
- 结合多项式方法与快速傅里叶变换,实现高效反向计算。
- 在有界参数条件下,反向计算复杂度达n^{1+o(1)},接近线性。
- 理论证明:无有界条件则无法实现亚二次复杂度,具必要性。
旋转位置编码(RoPE)已成为Transformer架构的重要增强,使模型能有效捕捉位置信息中的令牌关系。然而,RoPE使注意力计算更复杂,导致高效算法难以设计。此前研究仅在特定有界参数条件下实现了前向计算的近乎线性时间(复杂度为n^{1+o(1)},n为输入令牌数),但未解决反向传播问题。本文首次提出一种在有界参数下,基于RoPE注意力的反向计算近乎线性时间算法。方法建立在近期快速RoPE注意力计算基础上,创新性地融合多项式方法与快速傅里叶变换。此外,我们通过强指数时间假设(SETH)推导出下界,证明有界参数条件对实现亚二次性能是必要的。
原文摘要 · Abstract (English)
The Rotary Position Embedding (RoPE) mechanism has become a powerful enhancement to the Transformer architecture, which enables models to capture token relationships when encoding positional information. However, the RoPE mechanisms make the computations of attention mechanisms more complicated, which makes efficient algorithms challenging. Earlier research introduced almost linear time algorithms for the forward computation under specific parameter settings of bounded entries (i.e., in time $n^{1+o(1)}$ where $n$ is the number of input tokens), but has not addressed backward computation. In this work, we develop the first almost linear time algorithm for backward computations in the RoPE-based attention under bounded entries. Our approach builds on recent advancements in fast RoPE attention computations, utilizing a novel combination of the polynomial method and the Fast Fourier Transform. Furthermore, we show that with lower bounds derived from the Strong Exponential Time Hypothesis (SETH), the bounded entry condition is necessary for subquadratic performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。