[論文レビュー] A Constrained Optimization Perspective of Unrolled Transformers

Javier Porras-Valenzuela, Samar Hadou+1 more|arXiv (Cornell University)|Jan 24, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

本論文は、 primal–dual トレーニングスキームを用いて unrolled transformers に層ごとの降下を課す制約付き学習フレームワークを導入し、ID性能を損なうことなく頑健性とOOD一般化を向上させる。

ABSTRACT

We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms. Specifically, we enforce layerwise descent constraints on the objective function and replace standard empirical risk minimization (ERM) with a primal-dual training scheme. This approach yields models whose intermediate representations decrease the loss monotonically in expectation across layers. We apply our method to both unrolled transformer architectures and conventional pretrained transformers on tasks of video denoising and text classification. Across these settings, we observe constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving in-distribution performance.

研究の動機と目的

• 層ごとの損失減少を課すことで、トランスフォーマーを降下アルゴリズムの挙動に近づけることを動機づける。
• 制約付き unrolled transformers を訓練するための primal–dual 最適化アルゴリズムを開発する。
• 収束と分布シフト下のOOD一般化に対する理論的保証を提供する。
• 摂動を受けた入力で動画ノイズ除去とテキスト分類における頑健性の向上を実証する。

提案手法

• 各トランスフォーマー層は期待損失をある因子だけ減少させることを要求する制約学習問題を定式化する（降下制約）。
• 制約問題を双対（ラグランジェ）形に変換し、 primal–dual 交互アルゴリズム（Algorithm 1）で解く。
• dual 問題は実践上ほぼ最適でほぼ実現可能な解を小さいデュアル性ギャップで得られることを証明し、層間での漸近的収束を示す。
• 必要に応じてスラック変数で実現可能性を緩和し降下制約の耐性を提供する。
• 動画ノイズ除去の unrolled transformers と、摂動埋め込みを持つテキスト分類の事前学習済み transformers の両方に本手法を適用する。
• 分布シフト下の保証と摂動に対する頑健性を議論する。

Figure 1: Layerwise descent improves OOD robustness. Left: Test loss at each layer ( $\downarrow$ lower is better). Constrained RoBERTa exhibits monotonic descent, unlike the unconstrained baseline. Right: Out-of-distribution accuracy under increasing embedding perturbation levels $\gamma$ ( $\uparr

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマー訓練で層ごとの降下制約を課し、最適化降下の挙動を模倣できるか。
• RQ2 primal–dual 制約付き訓練法は層ごとの損失減少を単調に促し、摂動に対する頑健性を改善するか。
• RQ3 制約付き unrolled transformers はID性能を維持または向上させつつ、OOD一般化を強化できるか。
• RQ4 収束および分布シフト下のOOD性能に関する理論的保証は何か。
• RQ5 摂動入力を持つ動画ノイズ除去とテキスト分類で制約付き transformers はどのように性能を発揮するか。

主な発見

• 制約付きトランスフォーマは単調な層ごとの降下を示し、制約なしのベースラインとは異なる。
• 降下制約は入力摂動に対する頑健性を高め、ほとんど損失を生じさせずにOOD一般化を強化する。
• 層が深くなると、制約学習の下でOOD頑健性が向上し、層が増えるほど制約付きモデルがしばしば制約なしモデルを上回る。
• 動画ノイズ除去（CUHK Avenue、UCSD、ShanghaiTech）とテキスト分類（IMDb、MNLI）における経験的結果は、制約付きモデルが摂動下でRMSEを低減し精度を維持することを示す。
• 本手法は収束が近似的に最適な統計的損失へ収束すること、分布シフト下で降下特性を保持すること（誤差境界で定量化可能）という理論的保証を与える。
• 大規模モデル（例：Llama 8B）へもスケール可能で、層ごとの降下を維持し頑健性を向上させつつIDタスクを損なわない。

Figure 2: Video denoising error vs. test perturbation $\gamma$ (RMSE $\downarrow$ , lower is better). Columns are datasets, rows are architectures. Solid lines are constrained models; dashed lines are unconstrained. Each plot shows RMSE over increasing test perturbation levels ( $\gamma$ ). All mode