🤖 machine learning

A Constrained Optimization Perspective of Unrolled Transformers

本論文は、プライマル・デュアル学習スキームを通じて層ごとの降下制約を課す制約付き最適化フレームワークを提案しており、これによりトランスフォーマーが層を跨いで損失を単調に減少させ、インディストリビューションの性能を維持しつつ、ロバスト性と分布外汎化性能を向上させることを可能にする。

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

公開日 2026-07-08

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

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

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、生徒に複雑なパズルを解く方法を教えていると想像してください。標準的な教室（伝統的なAIトレーニング）では、あなたは生徒に最終的な答えを見せ、「この答えに近づいてください」と言います。生徒は数歩前進し、次に誤って後退し、再び前進し、激しくジグザグに動き回りながら、ようやく解決策にたどり着くかもしれません。そこに到達することはできるでしょうが、その過程は乱雑で不安定なものでした。

この論文は、異なる教え方を提案しています：「ステップ・バイ・ステップ」のルールです。

著者である Javier Porras-Valenzuela、Samar Hadou、Alejandro Ribeiro は、単に最終的な答えを見るのではなく、プロセスのあらゆる単一のステップにおいて、位置を改善することを強制すべきだと提案しています。もし生徒が、問題をより難しくしたり、正解に近づかなかったりするステップを踏んだ場合、私たちは「いいえ、やり直してください」と伝えます。

以下に、簡単な比喩を用いた彼らのアイデアの解説をまとめます。

1. 問題点：「ジグザグ」な生徒

標準的なAIモデル（Transformer）は、これらジグザグに動く生徒のようなものです。これらは多くの層（建物のフロアのようなもの）で構成されています。データが下のフロアから上のフロアへと移動するにつれ、モデルはエラーを修正しようとします。しかし、時としてモデルは中間のフロアで混乱することがあります。データを改善する前に、むしろ悪化させてしまうことがあるのです。これは、山に登ろうとしているハイカーが、頂上を見つける前に、誤って谷へと足を踏み入れてしまうようなものです。

この「ジグザグ」な動きは、モデルを脆弱にします。もしモデルに、少し異なる、あるいはノイズの多い入力（ぼやけた写真やタイポのある文章など）を与えると、道中の凹凸を処理するように訓練されていないため、完全に迷走してしまう可能性があります。

2. 解決策：「単調減少」のルール

著者らは、**制約付き最適化（Constrained Optimization）**と呼ばれる新しいトレーニング手法を導入しています。これは、建物のすべてのフロアに立つ厳格なコーチを想像すると分かりやすいでしょう。

ルール： 「あなたは、下のフロアにいた時よりも、このフロアでは少なくとも10%正解に近づいていなければならない」
比喩： ボールが丘を転がり落ちる様子を想像してください。標準的なモデルは、転がり落ち、凸凹に当たり、少し転がり上がり、再び転がり落ちるかもしれません。著者らのモデルは、完璧に滑らかで急勾配な滑り台を転がるボールのようなものです。それは常に下方向へ進まなければなりません。決して逆方向に転がり上がることはありません。

彼らはこれを、Transformerの**「アンローリング（Unrolling）」**と呼んでいます。通常、「アンローリング」とは、特定の数学的問題を解くためにニューラルネットワークを構築することを意味します。ここでは、既存の標準的なAIアーキテクチャを取り上げ、それを完璧なステップ・バイ・ステップの降下アルゴリズムのように振る舞わせるように強制しています。

3. 実装方法：「主双対（Primal-Dual）」のダンス

この厳格な「後退禁止」ルールを用いてモデルをトレーニングすることは、数学的に困難です。それは、犬に座るよう教えながら、同時に吠えないようにし、さらに特定のフェンスの中に留まるようにさせるようなものです。

著者らは、**主双対トレーニング（Primal-Dual Training）**と呼ばれる巧妙な数学的トリックを使用しています。

主（Primal - 生徒）： AIモデルは、タスク（テキストの分類やビデオのノイズ除去など）を学習しようとします。
双対（Dual - コーチ）： 別の数値セット（ラグランジュ乗数と呼ばれます）が、厳格なレフェリーとして機能します。もしAIが「後退禁止」ルールに違反するステップを踏もうとすれば、コーチはペナルティを課します。
ダンス： 彼らは共にトレーニングを行います。AIはより良くなろうとし、コーチはAIがルールに従うようにペナルティを調整します。論文によれば、この「ダンス」は非常に効率的であり、コンピュータの計算速度を大きく低下させることもありません。

4. 結果：「頑健な」モデル

論文では、主に2つのタスク、テキスト分類（言語の理解）とビデオ・デノイジング（粒子の粗いビデオのクリーニング）でこの手法をテストしました。

「ノイズ」テスト： ノイズ（テレビの砂嵐や文章内のタイポなど）を加えてテストを行いました。
- 標準的なモデル： ノイズが高くなると、標準的なモデルのパフォーマンスは崩壊しました。それは、そよ風で崩れ落ちるカードの家のようなものでした。
- 制約付きモデル： ノイズが高くなっても、このモデルは崩れませんでした。それは、風に吹かれても折れることなく、しなやかに曲がる頑丈な木のように、優雅に性能を低下させました。入力が乱れていても、モデルは機能し続けました。
「分布外（Out-of-Distribution）」テスト： モデルが一度も見たことのないデータ（例：映画レビューで学習したRoBERTAモデルを、異なる種類のテキストに対してテストする場合）を用いてテストを行いました。制約付きモデルは、標準的なモデルよりもはるかにうまく踏みとどまりました。

5. なぜこれが重要なのか（論文による主張）

この論文は、この手法がすぐに病気を治したり、自動運転車を作ったりすることを主張しているわけではありません。代わりに、特定の課題、すなわち**「ロバストネス（堅牢性）」**を解決することを目的としています。

思考プロセスにおけるあらゆるステップで改善することを強制することで、モデルは以下の特性を備えます：

より安定している： 入力の小さなエラーによって混乱することはありません。
より信頼できる： 未知のデータ（分布外データ）に対しても優れた性能を発揮します。
予測可能である： モデルが情報を処理する際、厳格な「下降」の経路を辿るため、モデルがどのように振る舞うかを正確に把握できます。

要約すると、著者らは強力ではあるものの時に混沌としたAIツールを取り上げ、そこに「安全レール」を取り付けました。これにより、モデルが着実に前進することを強制し、現実世界のノイズや混沌に対してより強靭なものにしたのです。

技術要約：展開型トランスフォーマーの制約付き最適化の観点

問題提起

近年の文献では、トランスフォーマー・アーキテクチャを、最適化問題を解く反復降下アルゴリズム（例：Yang et al., 2022; De Weerdt et al., 2023）として解釈していますが、標準的な経験リスク最小化（ERM）による学習は、このような最適化アルゴリズムに期待される単調な降下挙動を強制することには失敗することがよくあります。ERMを用いて学習した場合、展開型トランスフォーマーの中間表現は、必ずしも層を追うごとに損失関数を単調に減少させるわけではありません。この不一致は、標準的な学習が、学習されたパラメータがエネルギー最小化の条件を満たすことを保証しないこと、また、展開型アーキテクチャの学習におけるバイレベルの性質を考慮していないことに起因します。その結果、これらのモデルは非単調な損失パターンを示し、入力の摂動や分布シフトに対する堅牢性が低下する可能性があります。

手法

著者らは、トランスフォーマーの学習を制約付き学習問題として扱うフレームワークを提案しています。ネットワーク全体の出力に対する平均損失を最小化する代わりに、本手法は、各層において期待損失が単調に減少することを強制します。

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$ はステップサイズ・スケジュールです。この定式化は、トランスフォーマーを降下アルゴリズムとして振る舞わせ、中間表現が目的関数を段階的に改善することを保証します。

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）: 確率的勾配降下法を用いて、トランスフォーマーのパラメータ $T$ に関して経験的ラグランジュ関数を最小化する。
双対更新（Dual Update）: 射影勾配上昇法を用いて、ラグランジュ乗数 $\lambda$ に関して双対関数を最大化する。

著者らは、深層ニューラルネットワークが小さな双対ギャップを示すという最近の結果（Chamon et al., 2023）を活用し、双対問題が最適かつ実行可能に近い解をもたらすことを利用しています。厳密な実行可能性を確保するために、スラック変数と二次ペナルティを用いたレジリエントな緩和策が導入されています。

主な貢献

本論文は、主に4つの貢献を述べています。

制約付き学習の定式化: 著者らは、トランスフォーマーの各層が期待損失を単調に減少させるよう明示的に制約を課す学習問題を定式化しており、これは標準的なERMとは異なります。
双対学習アルゴリズムと収束保証: 著者らは原始・双対学習アルゴリズムを開発し、制約付き展開型トランスフォーマーが統計的損失の最適値に漸近的に収束するという理論的保証（定理2）を提供しています。さらに、分布シフトの下でも降下特性が維持されることを示す、分布外（OOD）汎化の保証（定理4）も確立しています。
埋め込みの摂動に対する堅牢性: テキスト分類タスク（IMDb, MNLI）において、制約付きモデルは、トークン埋め込みに対するガウス、一様、および離散的な摂動に対しても、制約なしのベースラインと比較して高い精度を維持することを実証しています。
ビデオデノイジングにおける堅牢性: ビデオデノイジング（CUHK Avenue, UCSD, ShanghaiTech）に適用した結果、本手法は、イン分布の性能を維持しつつ、様々なデータセットおよびアーキテクチャにおいて、OODノイズレベルに対してより堅牢なモデル（RMSEで測定）をもたらします。

実験結果

著者らは、従来の事前学習済みトランスフォーマー（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%）では制約なしのモデルを大幅に上回りました。
トレードオフ: 本手法は、一般に、標準的な精度を犠牲にすることなく（多くの場合0.3%未満の微小な低下）、イン分布の性能を維持しており、降下制約を課すことが堅牢性のための標準的な精度を損なわないことを示唆しています。

意義と主張

本論文は、降下制約を課すことが、トランスフォーマーの内部ダイナミクスを古典的な最適化アルゴリズムの安定特性に適合させることを主張しています。この適合は、入力の摂動や分布シフトに対する優れた堅牢性のための理論的および経験的な基礎を提供します。

著者らは、本研究を、アルゴリズムの展開（ネットワークが最適化問題を解くもの）と制約付き学習の架け橋として位置付けています。トランスフォーマーを明示的に「降下」するように訓練することで、入力の摂動に対して本質的に堅牢なモデルを実現できると論じており、この特性は、ユーザーのプライバシーのために入力の摂動が必要となるローカル差分プライバシーのアプリケーションにおいて特に価値があります。本研究は、「展開型」の視点が単なる理論的な解釈ではなく、アーキテクチャの変更を必要とせずに、汎化性能と堅牢性の具体的な向上をもたらす実用的な学習目的であることを示唆しています。

1. 問題点： 「ジグザグ」な生徒

2. 解決策： 「単調減少」のルール

3. 実装方法： 「主双対（Primal-Dual）」のダンス

4. 結果： 「頑健な」モデル