LUNA让注意力机制既快又准,突破了效率与精度的权衡瓶颈。
LUNA: Linear Universal Neural Attention with Generalization Guarantees
- 用可学习的核特征替代固定随机特征,实现线性计算复杂度
- 在长序列任务上达到甚至超过传统注意力模型的准确率
- 适合需要高效长序列处理的场景,如超长文本或视频建模
扩展注意力机制面临关键瓶颈:softmax注意力的$/mathcal{O}(n^2)$二次计算开销,限制其在长序列任务中的应用。尽管线性注意力将成本降至$/mathcal{O}(n)$,但通常依赖固定的随机特征映射(如随机傅里叶特征或手工设计函数)。这种对静态、数据无关核函数的依赖造成根本性权衡,迫使使用者牺牲显著模型精度以换取效率。我们提出 extsc{LUNA},一种基于核函数的线性注意力机制,消除这一权衡,在保持线性复杂度的同时匹配并超越二次注意力的精度。核心思想是让核特征映射本身可学习而非预先固定。通过参数化核函数, extsc{LUNA}学得针对特定数据和任务的特征基,克服了固定特征方法的表达局限。 extsc{Luna}采用可学习特征映射,诱导正定核并支持流式计算形式,实现序列长度上的线性时间和内存复杂度。实证评估验证该方法在多种设置下的有效性。在长序列基准测试集LRA上, extsc{Luna}在相同参数量、训练步数和近似浮点运算量条件下,成为高效Transformer中平均准确率最高的模型之一。此外, extsc{Luna}在后处理转换中表现优异:仅需替换已微调的BERT和ViT-B/16检查点中的softmax,并进行简短微调,即可恢复大部分原始性能,显著优于固定线性化方法。
原文摘要 · Abstract (English)
Scaling attention faces a critical bottleneck: the $\mathcal{O}(n^2)$ quadratic computational cost of softmax attention, which limits its application in long-sequence domains. While linear attention mechanisms reduce this cost to $\mathcal{O}(n)$, they typically rely on fixed random feature maps, such as random Fourier features or hand-crafted functions. This reliance on static, data-agnostic kernels creates a fundamental trade-off, forcing practitioners to sacrifice significant model accuracy for computational efficiency. We introduce \textsc{LUNA}, a kernelized linear attention mechanism that eliminates this trade-off, retaining linear cost while matching and surpassing the accuracy of quadratic attention. \textsc{LUNA} is built on the key insight that the kernel feature map itself should be learned rather than fixed a priori. By parameterizing the kernel, \textsc{LUNA} learns a feature basis tailored to the specific data and task, overcoming the expressive limitations of fixed-feature methods. \textsc{Luna} implements this with a learnable feature map that induces a positive-definite kernel and admits a streaming form, yielding linear time and memory scaling in the sequence length. Empirical evaluations validate our approach across diverse settings. On the Long Range Arena (LRA), \textsc{Luna} achieves state-of-the-art average accuracy among efficient Transformers under compute parity, using the same parameter count, training steps, and approximate FLOPs. \textsc{Luna} also excels at post-hoc conversion: replacing softmax in fine-tuned BERT and ViT-B/16 checkpoints and briefly fine-tuning recovers most of the original performance, substantially outperforming fixed linearizations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。