[论文解读] Transformers from an Optimization Perspective

Yongyi Yang, Zengfeng Huang+1 more|arXiv (Cornell University)|May 27, 2022

Explainable Artificial Intelligence (XAI)被引用 6

一句话总结

本文通过推导一个能量函数，提出了一种基于优化的新型Transformer架构解释，该能量函数的最小化过程可近似表示每个Transformer层的前向传播。通过将自注意力机制和前馈网络视为在该能量函数上的交替最小化步骤，该研究建立了Transformer与展开优化之间的理论基础，从而为注意力机制提供了新见解，并为定制化模型设计提供了可能。

ABSTRACT

Deep learning models such as the Transformer are often constructed by heuristics and experience. To provide a complementary foundation, in this work we study the following problem: Is it possible to find an energy function underlying the Transformer model, such that descent steps along this energy correspond with the Transformer forward pass? By finding such a function, we can view Transformers as the unfolding of an interpretable optimization process across iterations. This unfolding perspective has been frequently adopted in the past to elucidate more straightforward deep models such as MLPs and CNNs; however, it has thus far remained elusive obtaining a similar equivalence for more complex models with self-attention mechanisms like the Transformer. To this end, we first outline several major obstacles before providing companion techniques to at least partially address them, demonstrating for the first time a close association between energy function minimization and deep layers with self-attention. This interpretation contributes to our intuition and understanding of Transformers, while potentially laying the ground-work for new model designs.

研究动机与目标

• 建立Transformer层与最小化可解释能量函数的优化过程迭代之间的一一对应关系。
• 解决将自注意力机制与前馈网络整合到统一能量基框架中的挑战。
• 在特定条件下，为Transformer前向传播作为近似能量最小化的理论提供依据。
• 通过将注意力行为与能量函数组件的选择关联，实现对模型设计的新见解。
• 通过实证验证，确认在真实世界Transformer模型中，所推导的能量函数在各层间单调递减。

提出的方法

• 推导出一个能量函数，通过使用非线性激活函数 ρ(z) = -e^(-z) 将自注意力建模为最小化问题，该函数对应于Softmax归一化。
• 提出一种新颖的交替最小化方案，将自注意力与前馈组件联合优化于单一能量函数中，以解决模块异质性问题。
• 应用邻近方法处理能量函数中的非线性激活，实现对包含非光滑项的梯度优化。
• 用数据保真项 ‖Y - X‖²_F 替代标准的Frobenius范数正则化，以避免退化表征，与实际中的残差连接保持一致。
• 在IMDB和SST2数据集上的随机初始化和训练完成的12层Transformer中，实证评估了各层间的能量曲线。
• 证明能量函数在各层间单调递减，支持理论收敛性声明。

实验结果

研究问题

• RQ1 能否构建一个单一、统一的能量函数，使得每个Transformer层对应于该函数最小化过程中的一步下降？
• RQ2 如何在保持可解释性的前提下，将自注意力与前馈网络联合建模于单一能量函数中？
• RQ3 非线性激活函数和任意权重矩阵在能量最小化框架中起什么作用，以及如何处理它们？
• RQ4 在真实世界Transformer模型中，所推导的能量函数在前向传播过程中减少的程度如何？
• RQ5 能量函数中ρ函数的其他选择是否能导致具有可解释归纳偏置的新注意力机制？

主要发现

• 所推导的能量函数在理论条件下能紧密复现Transformer中自注意力模块的行为，且能量在各层间递减。
• 实证结果表明，在IMDB和SST2数据集上的随机初始化和训练完成的12层Transformer中，能量函数在各层间单调递减。
• 能量函数的最小化框架将Softmax注意力机制解释为ρ(z) = -e^(-z)这一特定选择的结果，从而实现将注意力重新诠释为优化过程。
• ρ(z)的其他选择（如log(z+2)）可导致更强调token间差异性的注意力机制，展示了设计灵活性。
• 该框架通过数据保真项‖Y - X‖²_F支持残差连接，可防止退化表征，与标准Transformer实践保持一致。
• 理论收敛结果在温和的技术条件下成立，实证证据表明这些条件在实际设置中极有可能满足。