🤖 machine learning

A Constrained Optimization Perspective of Unrolled Transformers

本文提出了一种约束优化框架，通过原对偶训练方案强制执行逐层下降约束，使 Transformer 能够在保持分布内性能的同时，实现跨层损失的单调递减，并提升鲁棒性和分布外泛化能力。

原作者： Javier Porras-Valenzuela, Samar Hadou, Alejandro Ribeiro

发布于 2026-07-08

📖 1 分钟阅读☕ 轻松阅读

原作者： Javier Porras-Valenzuela, Samar Hadou, Alejandro Ribeiro

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一名学生解开一个复杂的谜题。在标准的课堂（传统的 AI 训练）中，你会向学生展示最终答案，并说：“向这个答案靠近一点。”学生可能会向前走几步，然后不小心向后退了一步，接着又向前走，如此反复地左右摇摆，直到最终偶然撞上解决方案。他们可能确实达到了目标，但他们的路径是混乱且不稳定的。

这篇论文提出了一种不同的教学方式：“步步为营”规则（The "Step-by-Step" Rule）。

作者 Javier Porras-Valenzuela、Samar Hadou 和 Alejandro Ribeiro 建议，我们不应仅仅关注最终答案，而应该强制要求学生在过程中的每一步都必须有所进步。如果学生迈出的那一步让谜题变得更难，或者没有让他们离目标更近，我们就告诉他们：“不，重试。”

以下是他们利用简单类比对这一理念进行的详细拆й解：

1. 问题所在：“之字形”学生

标准的 AI 模型（Transformer）就像那些走“之字形”的学生。它们由许多层组成（就像一栋建筑的楼层）。随着数据从底层向顶层移动，模型试图修正错误。然而，模型有时会在中间楼层感到困惑。它可能会在让数据变好之前，先让数据看起来变得更糟。这就像一名登山者，在试图爬上山顶的过程中，由于误入山谷而导致进度倒退。

这种“之字形”运动使得模型非常脆弱。如果你给模型一个略有不同或带有噪声的输入（比如一张模糊的照片或一个带有拼写错误的句子），它可能会彻底迷失方向，因为它并未被训练去应对这些路途中的颠簸。

2. 解决方案：“单调下降”规则

作者引入了一种名为**约束优化（Constrained Optimization）**的新型训练方法。你可以把它想象成一位站在每一层楼上的严厉教练。

规则： “你这一层的状态，必须比你在下一层时至少更接近解决方案 10%。”
类比： 想象一个球从山上滚下。标准模型可能会向下滚动，撞到一个凸起，向上滚一点，然后再向下滚。而作者的模型就像是在一个完美光滑且陡峭的滑梯上的球。它必须在每一刻都向下移动。它永远不会往回滚。

他们将此称为对 Transformer 进行**“展开”（Unrolling）**。通常，“展开”是指构建一个专门解决数学问题的神经网络。在这里，他们是将一个现有的、标准的 AI 架构，强制使其表现得像一个完美的、步步递进的下降算法。

3. 他们是如何实现的：“原对偶”之舞

使用这种严格的“不准后退”规则来训练模型在数学上是非常困难的。这就像是试图教一只狗坐下，同时还要确保它不吠叫，并且还要确保它待在特定的围栏内。

作者使用了一个巧妙的数学技巧，称为原对偶训练（Primal-Dual Training）：

原变量（学生）： AI 模型尝试学习任务（如文本分类或视频去噪）。
对偶变量（教练）： 一组独立的数值（称为拉格朗日乘子）充当严厉的裁判。如果 AI 尝试采取违反“不准后退”规则的步骤，教练就会施加惩罚。
舞蹈： 他们共同训练。AI 试图变得更好，而教练则调整惩罚力度，以确保 AI 遵循规则。论文声称这种“舞蹈”非常高效，且不会显著拖慢计算机的速度。

4. 结果：“稳健”的模型

论文测试了该方法在两个主要任务上的表现：文本分类（理解语言）和视频去噪（清理颗粒感视频）。

“噪声”测试： 他们通过添加“噪声”（如电视上的静电或句子中的错别字）来测试模型。
- 标准模型： 当噪声增加时，标准模型的性能大幅下降。它就像一座在微风中倒塌的纸牌屋。
- 约束模型： 当噪声增加时，这个模型并没有崩溃。它表现出优雅的降级，就像一棵在风中弯曲却不折断的坚韧树木。即使输入信息很杂乱，它依然能保持运作。
“分布外”测试： 他们在模型从未见过的数据上进行了测试（例如，一个在电影评论上训练的 RoBERTa 模型，被用于测试另一种类型的文本）。相比之下，约束模型在面对这些数据时表现得更加稳健。

5. 为什么这很重要（根据论文观点）

这篇论文并不声称这能立即治愈疾病或制造自动驾驶汽车。相反，它声称解决了一个具体的问题：鲁棒性（Robustness）。

通过强制 AI 在其思考过程的每一步都进行改进，模型变得：

更稳定： 它不会被输入的微小错误所迷惑。
更可靠： 它在处理未见过的（分布外）数据时表现更好。
更可预测： 我们确切知道模型在处理信息时的行为，因为它遵循一条严格的“向下”路径。

简而言之，作者将一个强大但有时会陷入混乱的 AI 工具，添加了一个“安全护栏”，迫使它稳定地向前推进，使其在面对现实世界的噪声和混乱时变得更加坚韧。

技术摘要：从约束优化视角看展开式 Transformer

问题陈述

尽管近期的文献将 Transformer 架构解释为求解优化问题的迭代下降算法（例如 Yang et al., 2022; De Weerdt et al., 2023），但通过经验风险最小化（ERM）进行的标准训练往往无法强制执行此类优化器所预期的单调下降行为。在使用 ERM 进行训练时，展开式 Transformer 的中间表示并不一定会随着层数的增加而单调降低损失函数。这种不一致性源于标准训练并不能保证学习到的参数满足能量最小化所需的条件，也未考虑到展开式架构训练的双层性质。因此，这些模型可能会表现出非单调的损失模式，并在面对输入扰动和分布偏移时表现出鲁棒性下降。

方法论

作者提出了一个将 Transformer 训练视为约束学习问题的框架。该方法并非仅仅最小化整个网络输出的平均损失，而是强制要求每一层的预期损失相对于前一层 $l-1$ 呈单调递减。

1. 约束公式化

其核心贡献是将学习问题表述为一个约束优化问题，其中要求第 $l$ 层的输出相对于前一层 $l-1$ 将预期损失 $f$ 降低一个因子 $(1 - \alpha_l)$ ：
$\begin{aligned} T^* = \arg\min_{T} \quad & \mathbb{E}[f(X, \Phi(X; T))] \\ \text{subject to} \quad & \mathbb{E}[f(X, \Phi_l(X; T))] \leq (1 - \alpha_l) \mathbb{E}[f(X, \Phi_{l-1}(X; T))], \quad \forall l \end{aligned}$
这里 $\Phi_l$ 表示第 $l$ 层的输出， $\alpha_l$ 是步长调度。该公式强制 Transformer 表现得像一个下降算法，确保中间表示能够逐步改善目标函数。

2. Primal-Dual（原对偶）训练算法

为了解决这个非凸约束问题，作者在对偶域采用了原对偶训练方案。他们构建了一个拉格朗日函数：
$\mathcal{L}(T, \lambda) = \mathbb{E}[f(X, \Phi(X; T))] + \sum_{l=1}^L \lambda_l \mathbb{E}[f(X, \Phi_l(X; T)) - (1 - \alpha_l)f(X, \Phi_{l-1}(X; T))]$
训练过程（算法 1）在以下两者之间交替进行：

原问题更新（Primal Update）： 使用随机梯度下降法，针对 Transformer 参数 $T$ 最小化经验拉格朗日函数。
对偶问题更新（Dual Update）： 使用投影梯度上升法，针对拉格朗日乘子 $\lambda$ 最大化对偶函数。

作者利用了近期研究表明深度神经网络具有较小对偶间隙（duality gaps）的结果（Chamon et al., 2023），使得对偶问题能够产生接近最优且接近可行解的结果。为了确保严格的可行性，文中引入了一个带有松弛变量和二次惩罚项的弹性松弛机制。

核心贡献

本文概述了四个主要贡献：

约束学习公式化： 作者制定了一个学习问题，其中 Transformer 层被显式约束为单调降低预期损失，这与标准的 ERM 不同。
对偶训练算法与收敛保证： 他们开发了一种原对偶训练算法，并提供了理论保证（定理 2），证明约束展开式 Transformer 在渐进意义上收敛至统计损失的最优值。他们进一步建立了分布外（OOD）泛化保证（定理 4），证明了下降特性在分布偏移下依然得以维持。
对嵌入扰动的鲁棒性： 在文本分类任务（IMDb, MNLI）中，作者证明了与无约束基准相比，约束模型在受到 Token 嵌入的高斯、均匀及离散扰动时，能保持更高的准确率。
视频去噪中的鲁棒性： 在视频去噪任务（CUHK Avenue, UCSD, ShanghaiTech）中，该方法产生的模型对 OOD 噪声水平表现出更强的鲁棒性（以均方根误差 RMSE 衡量），同时保留了分布内性能。

实验结果

作者在常规预训练 Transformer（DistilBERT, RoBERTa, Llama 3.1）以及展开式架构（Deep Unfolded Sequential Transformer - DUST, Unrolled Transformer - UT）上评估了其方法。

单调下降： 约束模型表现出平滑、单调的逐层损失下降，而无约束基准则表现出非单调模式并伴有峰值（例如 RoBERTa 在第 19 层出现的峰值）。
OOD 鲁棒性：
- 文本分类： 随着扰动水平（ $\gamma$ ）的增加，约束模型的性能下降更加平缓。例如，在针对 IMDb 的 RoBERTa 模型中，约束模型在高噪声水平下的准确率明显高于无约束版本。
- 视频去噪： 在各种数据集和架构下，约束模型在测试扰动水平较高（ $\gamma \geq 1.0$ ）时，始终实现了更低的 RMSE。
- LLM 微调： 在 Llama 3.1 8B 上，约束模型在保持分布内性能（在无噪声时为 AlpacaEval 50% 胜率）的同时，在噪声环境下（ $\gamma=1.5$ ）显著优于无约束模型（69.93% 胜率）。
权衡（Trade-offs）： 该方法通常在几乎不产生性能损失（通常 < 0.3%）的情况下保持分布内性能，这表明强制执行下降约束并不会以牺牲标准准确率为代价来换取鲁棒性。

重要性与主张

本文声称，强制执行下降约束使 Transformer 的内部动力学与经典优化算法的稳定性属性相一致。这种一致性为提高对抗输入扰动和分布偏移的鲁棒性提供了理论和实证基础。

作者将这项工作定位为算法展开（即网络求解优化问题）与约束学习之间的桥梁。他们认为，通过显式地训练 Transformer 进行下降，可以获得本质上对噪声更具鲁棒性的模型，这一特性对于涉及局部差分隐私（其中输入扰动对于用户隐私是必要的）的应用场景尤为重要。这项工作表明，“展开”视角不仅是一种理论解释，更是一种实用的训练目标，能够在无需改变架构的情况下，带来泛化能力和鲁棒性的实质性提升。