Transformers from an Optimization Perspective

前言

从优化的角度理解Transformer模型

论文百度云链接: https://pan.baidu.com/s/1fQiIXRyUw4LBhXIfqCVUKA
提取码: 3w62

加油读文章

摘要

翻译

诸如 Transformer 之类的深度学习模型通常是通过启发式方法和经验构建的。 为了补充基础研究，在这项工作中，我们关注了以下问题：是否有可能找到 Transformer 模型的能量函数，使得该能量函数的下降步骤与 Transformer 模型的前向传递相对应？ 通过找到这样的函数，我们可以将 Transformer 视为跨迭代的可解释优化过程的展开。 过去经常采用这种展开的观点来阐明更直接的深度模型，例如 MLP 和 CNN，然而，到目前为止，对于带自注意力机制（如 Transformer）的复杂模型，获得类似的等价性仍然难以捉摸。 为此，我们首先概述了几个主要障碍，然后提供了至少部分解决这些障碍的配套技术，首次证明了能量函数最小化与带自注意力机制的深层模型之间的密切联系。 这种解释有助于我们对 Transformer 的直觉和理解，同时可能为新模型设计奠定基础。

理解

建立Transformer模型的能量函数，期待能量最小化的过程对应模型的前向传播过程，已经在 CNN 和 MLP 上有了一定成果，暂时没有在Transformer 模型上的研究成果出现，这是本文的出发点。

一，介绍

尽管深度学习在实践中取得了广泛的成功 [26、27]，但模型架构通常是启发式创建，多数时候仍然难以理解，因此，最近投入了大量精力来分析和重新解释流行的深度网络结构的各种方法。 出于我们的目的，这里需要特别注意深度模型层和一些优化过程展开迭代之间的一对一关联，这些优化过程旨在最小化可解释的能量函数 [14、15、53、55、41]。 通过这种方式，前向传递至少可以近似地被视为计算，所以具有最小能量的表示，其中模型参数决定了这种能量的形式，并且可以通过反向传递来训练下游感兴趣的任务，例如分类或回归等。

站在这种所谓的展开优化视角，可以深入了解不同架构中归纳偏差的性质，同时可能作为一种指南，通过所涉及的基础能量函数的性质来指导深度网络结构的实例化。事实上，对后者的修改可以导致可预测的模型行为，这些行为在过去会被剪枝优化来保证包含有用的正则化因素 [59]，避免微妙的退化方案 [55]，或组合起来设计有效的新模型 [1、14]。

尽管取得了这些成功，但该领域的大多数前置工作要么解决了相对简单的网络结构，例如 MLP [53]，要么将问题限制在不同于大系统的单个模型组件 [41] 上。 因此，诸如广泛使用的 Transformer [49] 之类的复杂模型大多被忽略了，这在很大程度上是因为很难同时将自注意力和前馈网络 (FFN) 模块映射到一个集成的优化过程，而且可以进行最小化的单个、统一的能量函数。

我们通过以下步骤将展开的优化视角扩展到完整的 Transformer 堆栈：

• 在提供了展开优化的背景细节和现有示例之后，第 2 节形式化了将 Transformer 层映射到优化步骤的四个关键挑战，降低了一个显式表示的能量函数。 包括：(i) 处理自注意力，(ii) 集成异构 Transformer 层模块（即自注意力和 FFN），(iii) 考虑非线性激活，以及 (iv) 允许任意、非结构化权重矩阵。

• 随后，第 3 节推导了一个能量函数，该函数紧密地再现了挑战(i)中的自注意力 ，第 4 节严格推导了收敛结果，展示了如何通过一种新颖的交替最小化形式近似处理挑战 (ii)，以及第 5 节 提供通过近端方法解决挑战 (iii) 的明确细节； 对挑战 (iv) 的考虑与其他支持分析一起推迟到补充部分。

• 上述贡献最终形成了一个能量函数，当某些技术条件成立时，该能量函数被 Transformer 层最小化。 为了提供进一步的动力，第 6 节凭经验证明，即使在技术条件难以正式验证的情况下，实际数据也会将这种能量最小化。

总的来说，这些结果为理解 Transformer 模型提供了补充基础，并为未来的增强提出了合理的途径。

二，由展开的优化形成的深层架构简介

展开优化视角的基本思想是创建一对一的对应关系，对象分别是深度模型的层和模型的能量函数的最小化求解过程的迭代算法。 这样，前向传递可以被视为计算最小能量解的近似值。 同时，假设每个前向迭代都是可微的，这种情况下等能量函数参数，可以通过近似最小化步骤中的反向传递来传递一些元目标函数（例如，分类、回归等）的梯度，如此形成一个可解释的双层优化过程[11、14、24]。 总的来说，展开优化产生最小的能量表示，由此能量函数的特定结构可以在相关下游应用领域进行训练。

2.1 数学表示

上面描述的双向优化过程可以更正式地表示为：

$$\begin{align*} Y^\ast(W,X)&=\mathop{argmin}_YE(Y;W,X)\tag{1}\\ W^\ast(X)&=\mathop{argmin}_W\ell(\psi[Y^\ast(W,X)])\tag{2} \end{align*}$$

理解
式子（1）是优化问题，式子（2）是深度模型解决的问题。

其中，$X\in\mathbb{R}^{n\times d}$是模型的输入，在语言模型中就是输入的语句，$n$ 是 token 长度，$d$是token维度。参考 $Y\in\mathbb{R}^{n\times d}$ 是模型的隐藏表征，我们用 $Y^\ast$ 来表示模型的输出结果（也就是Transformer 编码器的单词表证输出），并且 $Y^\ast$ 最小化最小能量函数 $E:\mathbb{R}^{n\times d}\to\mathbb{R}$，能量函数的形式由参数 $W$ 决定，对应于Transformer 中的可训练参数。$\ell$ 是元损失函数，$\psi: \mathbb{R}^{n\times d}\to\mathbb{R}^{n\times d}$ 是将模型的隐藏表达转化为模型输出的函数，当然也可能有可训练的参数。假设 $\partial Y^\ast(W,X)/\partial W$ 和 $\partial\ell/\partial Y^\ast(W,X)$ 是有意义的，当然这个形式也要包含求近似 $Y^\ast(W,X)$ 时所求的梯度，之后上文中提到的双边系统就可以用参数 $W$ 来优化。

请注意，在下文中，如果能量函数不包括令牌间的交互，我们将其视为向量函数，并且为简单起见，我们使用小写字符 $y$ 和 $x$ 来表示 $Y$ 和 $X$ 的任意行。 此外，当清楚上下文信息时，我们省略写函数参数，例如我们可以用 $E(Y)$ 来表示 $E(Y ; W, X)$。

2.2 相关工作和局限

在本节中，我们介绍了展开优化的相关工作，特别是分析了一些相关的限制，这些限制是我们努力的动力，并强调了所涉及的挑战。

优化诱导的前馈网络

各种各样的前置工作已经从展开的优化角度研究了前馈结构 [14、16、20、22、47、53]。 在这里，我们仔细研究基于近端方法的 [53]。 [53]中的基本能量函数是

$$\begin{align*} E(y)=\pmb{1}^\top\tilde{\sigma}^\ast(W^{-\top}y)-\langle f(x),W^{-\top}y\rangle-\frac{1}{2}\|y\|^2\tag{3} \end{align*}$$

其中 $\pmb{1}$ 是一个所有元素都为 1 的向量，$\tilde{\sigma}^\ast$ 是 $\tilde{\sigma}:a\to\int^a_0\sigma(r)dr$ 的凸共轭 [42],$\sigma$ 是激活函数，$f(x)$ 是输入表示 $x$ 的某种变换。 式子(3) 的近端算子 [12] 是

$$\begin{align*} y^{(t+1)}=prox_E\big(y^{(t)}\big)=W^\top\sigma(W^{(t)}+f(x))\tag{4} \end{align*}$$

这类似于 MLP 的前馈层。 请注意，只要维数适当对齐，虽然 $W$ 原则上可以是任意矩阵，但在堆叠该模型的多个层后，实际上强制约束了有效变换：

$$\begin{align*} y^{(t+2)}=W^\top\sigma[WW^\top\sigma(Wy^{(t)}+f(x))+f(x)]\tag{5} \end{align*}$$

很明显，在第一层之后，实际的前馈变换变成了 $WW^\top$，这必然是半正定的 (PSD)。 因此，在此范例中实际上层权重不可能不受约束。

优化诱导的注意力机制

也有部分有限的工作试图从展开优化的角度解释注意力机制 [15, 41]。 例如，若有如下形式的能量函数

$$\begin{align*} E(y)=-\beta^{-1}log\big(\sum^d_{i=1}e^{\beta S_iy}\big)+\frac{1}{2}\|y\|^2\tag{6} \end{align*}$$

在[41]中提出，随后使用凹凸过程[57]进行更新，有形式如下的迭代

$$\begin{align*} y^{(t+1)}=Ssoftmax(\beta SY^{(t)})\tag{7} \end{align*}$$

其中 $S\in\mathbb{R}^{n\times n}$ 对应于注意力机制中的键，而 $y$ 映射到注意力机制中的查询，$β$ 是常数标量。 但至关重要的是，这种注意力机制实际上是交叉注意力（使用 $y$ 来关注 $S$），这与 Transformers 的典型自注意力用例不一致。 此外，这项工作没有考虑聚合的 Transformer 前馈网络模块和随之而来的非线性。

优化诱导的图神经网络

从类似的优化角度 [8、31、34、37、54、55、58、60] 也开发了各种图神经网络架构。 例如，图注意力机制是使用 [55] 中的迭代重新加权最小二乘法 (IRLS) 算法推导出来的。 虽然这个结果与自注意力相关，可以将其视为全连接图上的图注意力 [5]，但它无法产生 Transformer softmax 项或组合的自注意力/前馈 Transformer 堆栈。 尽管如此，我们稍后将展示如何利用 [55] 中的想法来通过 softmax 获得自我关注。

2.3 扩展到通用Transformer的主要挑战

在本文中，我们关注比之前工作更复杂的模型，即 Transformer 编码器。 Transformer 层 [49] 通常由两个主要组件组成：自注意力层和前馈网络 (FFN) 层。 如果我们将 FFN 简化为具有非线性激活的单个线性变换，并忽略层归一化和残差连接，一个 Transformer 层将简化为基本形式：

$$\begin{align*} Y^{(t+1)}=ReLU[softmax(Y^{(t)}W_aY^{(t)\top})Y^{(t)}W_f]\tag{8} \end{align*}$$

其中 $W_a$ 和 $W_f$ 分别代表self-attention和FFN层权重矩阵。 关于 (8) 和上面讨论的先前工作的局限性，我们提出了将展开优化正式应用于 Transformer 设置的四大挑战：

1. Token-Level Interactions：如前所述，Transformer 模型不仅包括前馈过程，还包括跨 token 交互，这反映在被认为是 Transformers 的必要条件的自我注意机制上。 到目前为止，之前的工作还没有衍生出 Transformer 风格的自注意力，而是要么放松到交叉注意力 [41]，要么未能产生无处不在的 softmax 运算符 [55]。

2. 异构层类型：每个 Transformer 编码器层都由两个完全不同的组件组成：self-attention 和 FFN。 我们将这种结构称为“异构层类型”，其中每个组件都有自己的参数，可以看作是一个独特的展开优化过程。 然而，在聚合前向传递过程中是否有可能结合这两个分量的相应能量以获得具有任何类型的收敛保证（甚至近似）的统一目标仍然未知。

3. 非线性激活：至少在一些相对简单的前馈模式的孤立上下文中，可以从近端算子 [29、55] 的角度理解常见神经网络架构中使用的许多激活函数。 然而，当集成到上述异构 Transformer 层类型中时，优化过程变得相当复杂，并且不知道在包含近端步骤后是否仍然存在任何收敛特性。

4. 权重不对称：正如在 (5) 的讨论中提到的，大多数处理前馈网络的工作范围实际上仅限于具有对称（或更严格的 PSD）权重变换的模型 [1、14、53] ，这限制了由此产生的普遍性。 此外，图神经网络文献中的相关模型通常基于无向图 [54、55]，其中图传播也是对称的。 因此，尽管据我们所知之前没有讨论过，但仍然不知道如何为更一般的不对称变换构造能量函数。

在本文中，我们提出了至少部分解决挑战 1、2 和 3 的技术，而对于 4，我们将初步讨论推迟到补充，将正式调查作为未来的方向。

三，Transformer Self-Attention的一种新推导

现在，我们通过构建一个能量函数来解决挑战1，该函数的迭代优化步骤与 Transformer 风格的 softmax 自注意力完全匹配，这是惯例。

3.1 展开的优化步骤 vs 基本的 Softmax 自注意力

考虑能量函数

$$\begin{align*} E_1(Y)=\sum^n_{i=1}\sum^n_{j=1}\rho\bigg(\frac{1}{2}\|y_i-y_j\|^2\bigg)+R(Y)\tag{9} \end{align*}$$

其中 $y_i\in\mathbb{R}^{d\times1}$ 是矩阵 $Y\in\mathbb{R}^{n\times d}$ 的第 $i$ 行，$\rho: \mathbb R^+ \to \mathbb R$ 是凹非减函数，$R:\mathbb R^{n\times d} \to \mathbb R$ 是凸函数。 有趣的是，尽管由于 $\rho$ 的非凸性，$E_1$ 不一定是凸的，但在 $\rho$ 和 $R$ 的特定选择下，它可以通过类似 softmax 的结构进行优化，如下所示：

定理3.1

假设 $\rho(z)=-exp\{-z\},\quad R(Y)=\frac{1}{2}|Y|^2_{\mathcal F}$，且 $\beta_i=exp\big\{-\frac{1}{2}|y_i^{(t)}|^2\big\}$，令 $Y^{(t)}$ 表示 $Y$ 中的任意固定值，则有如下更新策略：

$$\begin{align*} y^{(t+1)}_i=\frac{\sum^n_{j=1}\beta_j exp\{y_i^{(t)\top}y_j^{(t)}\}y^{(t)}_j}{\sum^n_{j=1}\beta_j exp\{y_i^{(t)\top}y_j^{(j)}\}},\;\forall i\tag{10} \end{align*}$$

满足 $E_1(Y^{(t+1)})\leq E_1(Y^{(t)})$，等式成立的条件是当且仅当 $Y^{(t)}$ 是 $E_1$的固定点。

值得一提的是，虽然不是很明显，但更新步骤 (10) 可以通过最大-最小化 (MM) 算法生成 [48]，其中最大化步骤产生凸上界，最小化步骤沿上界的梯度下降。 因此，本次更新的核心本质上是凸函数上的梯度步，这将与后续章节的讨论相关； 请参阅补充文件中的证明和更多详细信息。

remark 3.2

定理 3.1 有几个值得注意的地方。 首先，虽然在定理 3.1 中我们采用特定形式的 $\rho$ 和 $R$ 来恢复 softmax 算子，但是通过这些函数的其他选择，相应的展开优化算法可以生成不同类型的注意机制。 其次，获得（10）依赖于梯度步长的特定选择； 然而，对于更广泛的选择，由此产生的收敛更新会导致残余连接作为优化轨迹的自然副产品（见补充）。 最后，定理 3.1 可以很容易地修改，以适应完全 Transformer 连接受某些图结构约束的情况，如 [10, 18] 中所示（同样，详见补充）。

3.2 拓展到包含可训练参数

在聚合成矩阵形式后，我们到目前为止已经证明了迭代：

$$\begin{align*} Y^{(t+1)}=softmax_{\beta}\big(Y^{(t)}Y^{(t)\top}\big)Y^{(t)}\tag{11} \end{align*}$$

将减少（或保持不变）来自 (9) 的能量，其中 $softmax_{\beta}(y)_i = \frac{\beta_i exp\{y_i\}}{\sum_j \beta_j exp\{y_j\}}$ 表示具有重新加权系数向量 $\beta$ 的 softmax 运算符。 如果 $\beta_i$ 独立于 $i$，即$|y_i^{(t)}|$ 是常数（它可以通过层归一化强制执行，如前所述，层归一化也可以包含在我们的框架中），那么这个重新加权的 softmax 相当于 Transformers 中使用的规范 softmax。

现在考虑重新参数化 $Y = ZW_a$，其中 $W_a \in \mathbb R^{d\times d}$ 是可逆矩阵。 由此得出 $Z^{(t+1)}W_a = softmax_{\beta} (Z^{(t)}W_aW_a^\top Z^{(t)\top}) Z^{(t)}W_a$，导致修改后的更新规则：

$$\begin{align*} Z^{(t+1)} = softmax_\beta (Z^{(t)}W_a^sZ^{(t)\top})Z^{(t)} \tag{12} \end{align*}$$

其中 $W_a^s = W_aW_a^\top$，我们采用上标“$s$”表示该矩阵是对称的。 总的来说，本节的结果直接解决了挑战 1，在有和没有可训练参数的情况下（对称权重的主要挥之不去的限制被降级为挑战 4）密切再现 softmax 风格的自注意力。

四，通过交替不精确最小化组合Transformer组件

接下来我们将解决挑战2 和异构 Transformer 层类型。 带着这个目标，我们首先介绍一个通用的优化场景。 具体来说，我们提出以下问题：
给定两个（凸）目标 $f(y)$ 和 $g(y)$，在什么条件下或者在什么程度上，将交替使用单独的梯度步骤，也就是说 $f$ 和 $g$ 优化聚合函数 $f + g$ ？

我们将这种优化策略称为交替不精确最小化 (alternating inexact minimization, AIM)，我们将很快展示它收敛到包含 $f + g$ 最佳点的有限半径球的条件。 稍后我们将讨论这些结果如何有助于解决挑战2。

我们预先强调，如此定义的 AIM 与文献中通常所说的交替最小化完全不同 [19, 33]。 后者指的是在固定其他变量的情况下，以交替方式一次最小化具有多个变量的统一目标函数的场景，这样可以轻松实现下降。 相比之下，我们的 AIM 场景涉及多个带有共享变量的客观项，我们使用相同的变量交替最小化每个项，这是一个更具挑战性的分析过程。

4.1 一般交替不精确最小化

给定两个目标 $f, g:\mathbb R^d \to \mathbb R$，AIM 形式化为算法 1，其中 $\alpha_1$ 和 $\alpha_2$ 构成步长。 我们现在将研究算法 1 如何与由 $h(y) = f(y) + g(y)$ 定义的组合目标的最小化相关。

在下文中，我们假设 $f、g$ 都是 Lipschitz 连续且强凸的 [42]，分别具有 Lipschitz 常数 $L_f$ 和 $L_g$，以及凸性参数 $c_f$ 和 $c_g$。 因此，$h$ 也将是 Lipschitz 光滑且强凸的，具有 Lipschitz 常数 $L_h$ 和凸性参数 $c_h$。 我们将 $f、g$ 和 $h$ 的最优点分别表示为 $y_f^\top、y_g^\top$ 和 $y_h\top$。

整合算法1中的两个步骤，也就是：

$$\begin{align*} y^{(t+1)} = y^{(t)}-\alpha_1 \nabla f(y^{(t)})-\alpha_2 \nabla g[y^{(t)}-\alpha_1 \nabla f(y^{(t)})] \tag{13} \end{align*}$$

而目标 $h$ 上的规范梯度下降步骤是 $y^{(t+1)} = y^{(t)}−\alpha_2[\nabla f(y^{(t)}) + \nabla g( y^{(t)})]$ 。 比较两个更新规则很明显，算法 1 可以被视为具有步长 $\alpha_2$ 和噪声因子的噪声梯度下降步骤，

$$\begin{align*} \Delta_t = \nabla h(y^{(t)})-\frac{\alpha_1}{\alpha_2}\nabla f(y^{(t)})-\nabla g[y^{(t)}-\alpha_1\nabla f(y^{(t)})] \tag{14} \end{align*}$$

在某些情况下，噪声梯度下降是一个经过充分研究的问题 [4, 36]。 然而，在我们的特定场景中，我们需要一种新颖的、不同的（实际上更严格的）界限，而不是简单地应用现有结果所能得到的界限。

具体来说，我们证明当 $\delta(y^{(t)}) = |\Delta_t| / |\nabla h(y^{(t)})|$ 有界时，(13) 保证下降目标 $h$ 如下：

定理4.1

当 $\alpha_1 \leq \alpha_2 \leq L_h^{-1}$ ，假设 $y^{(t)}$ 和 $y^{(t+1)}$ 的关系由(13)给出，而且有 $\delta(y^{(t)})\leq \mathcal L$ 和 $\mathcal L = \frac{\alpha_2}{\alpha_2-\alpha_1+\alpha_2\alpha_1L_g}$，那么有 $h(y^{(t+1)})\leq h(y^{(t)}) $。

Remark 4.2

尽管 $\mathcal L$ 的定义可能看起来相当复杂，但当 $\alpha_1、\alpha_2、L_f$ 或 $L_g$ 之一足够小时，$\mathcal L$ 表现为 $\Omega([\alpha_1\alpha_2L_fL_g]^{−1})$ 。

我们进一步考虑如何解释约束 $\delta (y) ≤ \mathcal L$ 以及 (13) 可以优化 $h$ 的区域，如下所示：

引理4.3

可令 $S(\mathcal L) = \{y|\frac{|y-y_f^\ast|}{|y-y_h^\ast|}\leq\frac{c_h\mathcal L}{L_f} \}$，若有 $y^{(t)}\in S(\mathcal L),\;\delta(y)\leq\mathcal L $。

结合定理 4.1 和引理 4.3，我们因此可以得出结论，当 $y^{(t)} \in S(\mathcal L)$ 时， $h(y^{(t+1)}) ≤ h(y^{(t)})$ 。 请注意，引理 4.3 中给出的 $S(\mathcal L )$ 的边界称为阿波罗圆[43]。 当 $\mathcal L≤\frac{L_f}{c_h}$ 时，$S(\mathcal L)$ 是一个以 $y_f^\ast$ 为中心的球，当 $\mathcal L\geq \frac{L_f}{c_h}$ 时，$S(\mathcal L)$ 是整个不包括以 $y_h^\ast$ 为中心的球的空间。 图 1 分别提供了每个案例的二维可视化。 此外，请注意，可以通过切换 $f$ 和 $g$ 的角色（因为这两个过程是交替的）并进一步限制例外区域来完成相同的分析，尽管这里我们为简单起见省略了这一点。

Remark 4.4

我们上面的发现可以总结如下：对于足够小的 $\alpha_1$ 和 $\alpha_2$ 值，算法 1 减少了组合目标 $h$，至少假设 $y$ 距离最优点 $y_h^\ast$ 有一定距离。

为了说明这个结论，我们提出了一个合成示例 $f(Y ) = |SY |^2_{\mathcal F} + |Y − B_1|^2_{\mathcal F}$ 和 $g(Y)=|YW|^2_{\mathcal F} +|Y −B_2|^2_{\mathcal F}$。注意这里我们将变量展开到矩阵 $Y$ 并让 $f$ 和 $g$ 分别由 $Y$ 的左变换和右变换组成，带有一个额外的偏置项以防止退化解 ($Y_\ast =0$)。我们随机设置 $\{S,W,B_1,B_2\}$ 的每个条目并以固定步长执行算法 1 并投影 $Y^{(t)}$ 的轨迹 到带有 PCA 的二维空间进行可视化。 $Y^{(t)}$ 的轨迹显示在图 3 中，组合目标 $h$ 跨迭代显示在图 2 中。从这些图中，我们的理论预测的行为可以得到验证： 当 $Y^{(t)}$ 与 $Y_h^\ast$ 有足够的距离时 （能量比较高），优化轨迹向 $Y_h^\ast$ 靠拢（能量下降），当 $Y^{(t)}$ 足够接近 $Y_h^\ast$ 后（且 $h$ 处于比较低的水平），能量震荡在一定范围内的最优解。

4.2 异构层类型模型的展开优化

我们现在回到特定于 Transformer 的模型。 令 $E_2(Y ) = \frac{1}{2} Tr(Y W_f Y^\top) + \frac{1}{2} |Y |^2_{\mathcal F}$ 并考虑组合能量 $E(Y ) = E_1(Y W_a) + E_2(Y )$（类似于上一节中的 $h$ ），或者 执行以下 AIM 更新：

$$\begin{align*} U^{(t)}=\;& softmax_{\beta} (Y^{(t)}W_a^sY)Y^{(t)} \tag{15}\\ Y^{(t+1)}=\;& U^{(t)}-\alpha_2\frac{\partial E_2}{\partial Y}|_{Y=U^{(t)}}=U^{(t)}W_f^s \tag{16} \end{align*}$$

其中 $W_f^s =(1−\alpha_2)I−\alpha_2\frac{W_f+W_f^\top}{2}$ 在 softmax 之后提供了额外的线性变换。
从第 3 节中，我们知道 (15) 本质上是在主要化之后使用逐行梯度下降步骤来优化 $E_1$。 同样，(16) 也是步长为 $\alpha_2$ 的梯度步。 因此，组合规则属于算法1的范畴，从4.1节的分析，我们可以得出如下结论：

命题4.5

如果 $Y^{(t+1)}$ 通过 (15) 和 (16) 以输入 $Y^{(t)}$ 计算，则 $E(Y^{(t+1)}) \leq E(Y^{(t)})$，当 $Y^{(t)} \notin \mathcal S$，其中 $\mathcal S$ 是一个有限半径的球，包含 $Y_{\hat E}^\ast$，是 $\hat E$ 的最优点，$E$ 的凸上界。

请注意，在组合 (15) 和 (16) 之后，更新聚合更新规则已经与我们来自 (8) 的目标非常相似，唯一的区别是缺少非线性和对称权重（注意我们使用对称权重矩阵 $W_a^s$ 和 $W_f^s$ 而不是 $W_a$ 和 $W_f$），分别对应 2.3 节中的挑战 3 和 4； 接下来我们解决前者。

五，展开范式中集成的非线性激活

之前已经考虑过在展开的优化设置中处理非线性激活 [16、47、51、55]。 然而，之前的工作在很大程度上依赖于近端算子来创建与简单线性滤波器匹配的非线性激活，并且分析无法转移到本文中的 Transformer 案例。 虽然我们将在本节中将注意力集中在添加 ReLU 激活上，但为了简化说明以及它们在 Transformer 模型中的普遍性，我们的结果可以推广到更广泛的选择。【意思是也可以考虑其他的激活函数？】

与第 4 节类似，我们首先通过将 ReLU 激活引入算法 1 来研究优化问题的一般形式，将此修改形式化为具有近端步骤的 AIM，然后分析适当约束下的收敛性。 稍后我们将这些结果应用到 Transformer。

5.1 近端交替不精确最小化

已经确定 [29, 47] 可以将 ReLU 激活建模为近端算子，如下所示：

$$\begin{align*} ReLU(y)=argmin_z\frac{1}{2\lambda}\|z-y\|^2+\phi(z) \tag{17} \end{align*}$$

其中 $\phi$ 是第一象限的指示函数，且 $\phi(z)=\begin{cases}&+\inf \quad &if\;z<0\\&0 \quad &if\; z\geq0\end{cases}$ 。此外，
在第 4 节中，我们演示了算法 1 的步骤如何一起形成损失 $h(y) = f(y) + g(y)$ 的不精确梯度下降迭代，其中噪声项 $\Delta_t$ 在（14）中定义。 在这里，通过在算法 2 中添加近端步骤，我们获得了近端梯度下降的不精确版本。 事实上，算法 2 的一轮相当于

$$\begin{align*} y^{(t+1)}=argmin_z\frac{1}{2\lambda}\|z-y^{(t)}+\alpha_2\nabla h(y^{(t)})-\alpha_2\Delta_t\|^2+\phi(z) \tag{18} \end{align*}$$

与精确版本 $y^{(t+1)} = argmin_z P(z;y^{(t)})$ 相比，这是噪声为 $\Delta_t$ 的典型近端梯度步长的不精确版本，其中 $P(z;y)$ 是近端问题

$$\begin{align*} P(z;y^{(t)})=\frac{1}{2\lambda}\|z-y^{(t)}+\alpha_2\nabla h(y^{(t)})\|^2+\phi(z) \tag{19} \end{align*}$$

虽然过去研究了各种形式的不精确近端梯度下降 [13、17、44、52]，但现有工作仍然假设恒定噪声 [13]、随机噪声 [52] 或递减(收敛)噪声 [17、44] . 至关重要的是，我们所知道的任何先前工作都不适用于我们的案例，其中噪声可能会随着迭代而增加。 此外，现有文献中的分析主要关注收敛到一个固定点，而在我们的场景中，我们反而考虑进入围绕特定点形成的特定区域。

除了像第 4 节中那样限制 $\delta(y^{(t)})$ 之外，我们还需要类似限制当前位置 $y(t)$ 和梯度 $\alpha_2\nabla h(y^{(t)})$，定义为 $\mathfrak D(\xi_1,\xi_2) =\frac{1}{|\xi_1|^2}\sum^d_{i=1}min(\xi^2_{2,i}-\xi^2_{1,i},0)$。直观地，$\mathfrak D(\xi_1;\xi_2)$ 被定义为每一项为负，但只有当 $\xi_1$ 和 $\xi_2$ 都很大时才接近于 $0$。 然后我们有以下内容：

定理 5.1

如果 $\alpha_1 \leq \alpha_2 \leq L^{−1}_h,\;\mathfrak D(\alpha_2\nabla h(y^{(t)});y^{(t)})\geq -\kappa$，对于任何 $\kappa \in (0,1)$，并且$\delta(y^{(t)})\leq\mathscr{L}’$，其中 $\mathscr{L}’=\frac{\alpha_2c_P\lambda\sqrt{1-\kappa}}{\sqrt{2}(\alpha_2-\alpha_1+\alpha_1\alpha_2L_g)}$，我们有 $h(y^{(t+1)}) +\phi(y^{(t+1)})\leq h(y^{(t)})+\phi(y^{(t)}) $。

直观地，定理 5.1 表明保证 $h(y)+\phi(y)$ 下降的区域是 $S(\mathscr{L}’)$ 与引理 4.3 中定义的 $S$ 的交集，并且区域 $\mathcal{T}(\kappa) = \{y|\mathscr{D}(\alpha_2\nabla h(y);y) \geq −\kappa\}$。 虽然 $\mathcal{T}(\kappa)$ 的形状通常仍然难以指定，但我们注意到当 $\alpha_2 \to 0$ 或 $\kappa\to 1$ 时，$\mathcal{T} (\kappa)$ 趋向于整个空间。我们用二维合成示例，其中能量函数为 $h(y) = |W_y|^2_{\mathcal F} + |y − b|^2$，$W$ 和 $b$ 的条目是随机生成的。 参见图 4 使用不同值 $\kappa$ 的可视化，它表明当 $\kappa$ 足够小时，$\mathcal{T}(\kappa)$ 几乎是整个空间（原点周围的一小块区域除外），这确保了 $h(y) + \phi(y)$ 在大多数情况下的下降。

5.2 在Transformer模型中嵌入非线性

通过在前面的能量上加上惩罚项，并分别从 3.1 和 4.2 节中得到 $E_1$ 和 $E_2$，我们最终得到 Transformer 总能量 $E(Y) = E_1(Y) + E_2(Y) + \phi(Y)$。 并且 $E(Y)$ 的展开优化落入算法 2 的范围，使得前面的分析和定理 5.1 适用，我们可以得出结论，聚合更新规则

$$\begin{align*} Y^{(t+1)} = ReLU[softmax_{\beta}(Y^{(t)}W^s_{\alpha}Y^{(t)\top})] \tag{20} \end{align*}$$

是 $E$ 的下降算法，除了一个有限测度集区域。 此外，如果我们修改底层梯度下降算法使步长 $\alpha_1$ 和 $\alpha_2$ 足够小，异常区域的大小将趋于零，并且更新规则将配备残差项（见补充；类似地 关于 $\beta$ 的讨论）。 除了这些问题之外，展开的更新 (20) 与 (8) 的唯一区别在于它对对称权重的依赖。 这对应于我们尚未完全解决的挑战 4，尽管已经证明如果我们扩大表示维度 [23、56]，对称权重可以模拟不对称权重。

六，实践验证

尽管我们严格推导了收敛标准，即 Transformer 层将明确指定的能量函数下降到最优解周围的区域，但不可否认，该分析依赖于难以在现实世界数据集上正式验证的条件。 然而，我们的结果仍然适用于有针对性的经验证实，由此我们可以检查所提出的能量是否确实在典型基准的 Transformer 正向传递期间确实下降了。

为此，我们实现了一个 Transformer 模型，满足已知的限制，如对称权重。 我们将该模型应用于两个基准测试，IMDB [35] 和 SST2 [46]，它们都是依赖 Glove-840b-300d [39] 作为词嵌入的常用情感分类数据集。 图 5 和图 6 显示了 Transformer 每一层的输出能量（如 (8) 中所定义）平均超过测试集中 200 个随机选择的样本。 图 5 使用随机初始化的权重，而图 6 涉及使用 SGD 和学习率 0.01 训练 2000 步的权重。 此外，对于经过训练的模型，我们将 $E_1$ 中的项 $|Y|^2_{\mathcal F}$ 更改为 $|Y − X|^2_{\mathcal F}$ 以避免退化表示（有时会发生在经过训练的变形金刚 [45] 中），并指出这种修改同样有效 我们的理论涵盖了这一点，并在最终的 Transformer 架构中产生了一种常用的残差连接形式。 详情见补充。

从这些数字可以清楚地看出，即使使用真实世界的数据，我们得出的 Transformer 能量（平均）跨层单调递减，与我们的分析预测相符。 此外，对于 12 层，代表在实践中并不少见的深度范围（例如 BERT），模型没有进入波动区域。 此外，这种观察甚至适用于由实践中使用的许多/大多数典型组件组成的训练有素的 Transformer 模型，即具有自注意力的层、非线性变换后的线性变换和残差连接等。因此，尽管看似复杂， 引理 4.3 和定理 5.1 采用的条件在许多实际环境中仍然可能成立。

七，结论

虽然我们在这里的贡献主要是理论性质的，但仍然有一些可能具有实际意义的重要信息。 首先，由于我们的首要目标是使用展开的优化视角尽可能接近地再现 Transformer 层，因此在构建核心底层能量函数时做出了非常具体的设计选择。 然而，在实践中，我们可以自由选择替代能源，这些替代能源可以导致不同形式的定制自注意力，这在特定应用程序的基础上可能是有利的。
作为后者的一个简短的代表性示例，请考虑以下内容：规范的 Transformer 使用 softmax 对注意力系数进行归一化。 然而，多项工作质疑 softmax 归一化的适当性，例如 [32、40]。 在我们的优化展开视角下，我们可以看到 Transformer 中的 softmax 归一化源自特定选择的 $\rho$ 函数，即 $\rho(z) = −e^{−z}$（第 3.1 节和定理 3.1）。 但是，如果我们选择不同的 $\rho$，就会产生一种新的标准化方法，而不是具有可解释属性的 softmax。 例如，如果 $\rho(z) = log(z + 2)$，则注意力系数将表现为（为简单起见，我们假设所有 $y_i$ 都具有以下单位范数）：

$$\begin{align*} a_{i,j} = \frac{1}{2-y_i^\top y_j}(\sum^n_{k=1}\frac{1}{2-y_i^\top y_k})^{-1} \end{align*}$$

其中 $a_{i,j}$ 是第 $i$ 个和第 $j$ 个 token 之间的注意力系数。 由于 $log(z + 2)$ 对于 $z \in [0, 1]$ 的增长速度比 $−e^{−z}$ 慢，因此相关的注意力公式往往会鼓励标记之间出现更多不同的表示。 这是因为在新能源中，较大的 $|y_i − y_j |$（意味着不同的 $y_i$ 和 $y_j$ ）对能量的贡献比以前少，因此优化过程减少它的动力较小。 相反，如果我们选择在规定范围内增长快于 $−e^{−z}$ 的 $\rho$（例如 $\rho(z) = log(z + 1)）$，那么派生模型可能会鼓励令牌之间有更多相似的表示。

这个总体框架还有其他具有实际意义的分支。 例如，注意力权重的实际分布可能会被产生它们的能量函数的属性更好地理解或影响。 此外，特别是对于数据有限且因此自由模型参数较少的机制，展开视角可用于设计具有与下游任务一致的归纳偏差的架构，以帮助弥补模型灵活性较低的问题。

最后，我们注意到本文中介绍的技术在某种意义上可以被视为构建更广泛的展开式体系结构系列的通用工具。 事实上，许多/大多数深度学习模型，包括具有前馈结构 [30]、令牌级交互 [50]、残差连接 [21] 和异构层类型的模型，都可以通过使用展开优化的镜头重新解释 我们引入的框架，至少要达到一些潜在的约束，比如对称权重（可以用其他方式处理）。