LASER: Attention with Exponential Transformation

📄 arXiv: 2411.03493v2 📥 PDF

作者: Sai Surya Duvvuri, Inderjit S. Dhillon

分类: cs.LG, cs.CL

发布日期: 2024-11-05 (更新: 2025-07-13)

备注: ICML 2025

💡 一句话要点

提出LASER注意力机制，通过指数变换提升梯度信号，改善Transformer学习效率。

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

关键词: 注意力机制 Transformer 梯度消失 指数变换 大型语言模型 深度学习 序列建模

📋 核心要点

  1. Transformer的softmax注意力机制存在梯度消失问题，影响了模型训练效率。
  2. LASER注意力机制通过指数变换，增强了梯度信号，从而改善了学习过程。
  3. 实验表明，LASER在多种任务和模型上均有提升，包括LLM、ViT、Conformer和BERT。

📝 摘要（中文）

Transformer在序列相关任务中取得了显著成果，这主要归功于其基于softmax的点积注意力机制，能够从序列的任何部分检索信息。本文分析了注意力机制中通过softmax操作反向传播的梯度，发现这些梯度通常很小。这种不良的梯度信号反向传播会导致注意力操作之前的参数学习效率低下。为此，我们提出了一种新的注意力机制LASER，通过分析表明其可以产生更大的梯度信号。我们展示了LASER注意力可以通过对现有注意力实现的少量修改来实现。我们在参数高达77亿的自回归大型语言模型（LLM）上进行了实验，在下游评估中，LASER相比标准注意力平均提高了高达1.44%，微调改进了1.65%。此外，LASER还在各种任务（视觉、文本和语音）中展示了泛化性能的提升：在Imagenet上的Vision Transformer（ViT），在Librispeech语音转文本上的Conformer以及具有22亿参数的BERT。

🔬 方法详解

问题定义：Transformer模型中的softmax注意力机制在反向传播过程中存在梯度消失的问题。具体来说，softmax函数容易导致梯度变得非常小，尤其是在输入值差异较大时，这会阻碍模型有效学习注意力机制之前的参数，降低了训练效率和最终性能。

核心思路：LASER的核心思路是通过引入指数变换来增强梯度信号。具体来说，LASER使用一种新的注意力计算方式，该方式在计算注意力权重时，对输入进行指数变换，从而使得梯度在反向传播时能够保持较大的值，避免梯度消失，从而提升学习效率。

技术框架：LASER注意力机制可以很容易地集成到现有的Transformer架构中。它替换了标准的softmax注意力模块，而不需要对整个模型架构进行大的改动。整体流程与标准注意力机制类似：首先，输入Query、Key和Value；然后，使用LASER计算注意力权重；最后，将注意力权重应用于Value，得到最终的注意力输出。

关键创新：LASER最重要的创新在于其注意力权重的计算方式，它使用指数变换来放大梯度信号。与传统的softmax注意力相比，LASER在理论上和实验上都证明了其具有更大的梯度，从而能够更有效地训练模型。

关键设计：LASER的关键设计在于指数变换的具体形式。论文中给出了具体的数学公式，通过调整指数函数的参数，可以控制梯度放大的程度。此外，为了保证数值稳定性，LASER还采用了一些技巧，例如对输入进行归一化处理，避免指数爆炸。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LASER注意力机制在多个任务上均取得了显著的性能提升。在大型语言模型上，LASER相比标准注意力平均提高了高达1.44%，微调改进了1.65%。在Imagenet上的Vision Transformer（ViT），在Librispeech语音转文本上的Conformer以及具有22亿参数的BERT模型上，LASER也展示了泛化性能的提升。

🎯 应用场景

LASER注意力机制具有广泛的应用前景，可以应用于各种基于Transformer的模型中，包括自然语言处理、计算机视觉和语音识别等领域。它可以提升大型语言模型的训练效率和性能，改善图像分类、目标检测等视觉任务的效果，并提高语音识别的准确率。此外，LASER还可以应用于其他序列建模任务，例如时间序列预测和生物信息学等。

📄 摘要（原文）

Transformers have had tremendous impact for several sequence related tasks, largely due to their ability to retrieve from any part of the sequence via softmax based dot-product attention. This mechanism plays a crucial role in Transformer's performance. We analyze the gradients backpropagated through the softmax operation in the attention mechanism and observe that these gradients can often be small. This poor gradient signal backpropagation can lead to inefficient learning of parameters preceeding the attention operations. To this end, we introduce a new attention mechanism called LASER, which we analytically show to admit a larger gradient signal. We show that LASER attention can be implemented by making small modifications to existing attention implementations. We conduct experiments on autoregressive large language models (LLMs) with upto 7.7 billion parameters with an average improvement of upto 1.44% over standard attention on downstream evaluations and 1.65% finetuning improvements. Additionally, LASER demonstrates generalization performance improvement across a variety of tasks (vision, text and speech):Vision Transformer (ViT) on Imagenet, Conformer on the Librispeech speech-to-text and BERT with 2.2 billion parameters.