技術要約：主・双対推論を用いた制約付き拡散モデル

問題提起

標準的な拡散モデルは、非制約の設定（画像や音声の合成など）において、複雑で高次元な分布からのサンプリングに長けている。しかし、サンプリング過程において統計的（平均的）制約を強制するための堅牢なメカニズムを欠いている。無線リソース割り当てやポートフォリオ管理といった多くの実世界の最適化問題では、決定論的な解ではなく、期待値において制約を満たす（例：ユーザーごとのレート保証やリスク予算）確率的なポリシーが最適となる。

制約付きサンプリングに対する既存のアプローチは、以下のいずれかに依存している：

1. 点別制約（Pointwise constraints）： 個々のサンプルに対して個別に制約を強制する（例：投影や境界反射による）。これは過度に保守的であり、最適性を損なう可能性がある。
2. 双対学習（Dual Training: DT）： サンプリング前に最適なラグランジュ乗数（双対変数）を推定し、それを固定する。この手法は脆弱である。なぜなら、最適な乗数は学習中の分布下での計算困難な期待値によって定義されるため、新しい問題インスタンスや制約の変化に対して、毎回高価な再学習を必要とするからである。さらに、固定されたサンプラーは、生成中に観察された制約違反に対して適応することができない。

手法：主・双対推論（Primal-Dual Inference: PDI）

本論文では、双対変数を事前計算されたパラメータとしてではなく、逆拡散プロセスと共に進化するサンプラーの動的な状態として扱うフレームワークである**主・双対推論（PDI）**を提案する。

1. 理論的定式化

問題は、ラグランジュ双対領域におけるエントロピー正則化最適化として記述される。最適分布 $\mu^*$ は、最適な双対変数 $\lambda^*$ によってインデックス付けされたギブス分布であることが示される：
$$\mu^*(x) \propto \exp\left(-\frac{1}{\beta}(f_0(x) + \lambda^{*\top}f(x))\right)$$
ここで、$f_0$ は目的関数、$f$ は制約を表し、$\beta$ は温度である。

2. PDIアルゴリズム

PDIは、すべての推論ステップ $t$ において、主方向のデノイジング（逆拡散ステップ）と双対上昇（乗数の更新）を交互に行う。

• 主ステップ (PDI-P): サンプラーは、現在の 双対変数 $\lambda_t$ に条件付けられたスコア場を用いて、ノイズを含むサンプル $x_t$ を更新する。
• 双ステップ (PDI-D): 双対変数は、デノイズされたサンプルの制約違反の推定値に基づく勾配上昇によって更新される。極めて重要な点は、制約違反が真のサンプルではなく、ツイディの公式（Tweedie's formula）（事後平均推定値）を用いて推定されることであり、これにより真のデータにアクセスすることなく更新が可能になる。
  $$\lambda_{t+1} = \left[ \lambda_t + \eta_t \mathbb{E}_{x_{t+1}} \left[ f(\hat{y}_0) \right] \right]_+$$
  ここで、$\hat{y}_0$ はクリーンなサンプルのツイディ推定値である。

3. スコアネットワークの学習

この動的な結合を可能にするために、著者らは単一の双対条件付きスコアネットワーク $s_\theta(x_t, t, \lambda_t)$ を学習させる。標準的なモデルは単一のターゲット分布に対して学習されるが、このネットワークは、推論軌跡中に遭遇する双対変数の範囲によってインデックス付けされた、ギブス・スコア場のファミリーを学習する。学習手順は以下を組み合わせる：

• 活用（Exploitation）: 現在のモデルを展開して、$(x_t, \lambda_t)$ のペアの軌跡を収集する。
• 探索（Exploration）: 潜在的な乗数の空間をカバーするために、事前分布から双対変数をサンプリングする。
• モンテカルロ推定: 回帰ターゲットに対して、ギブス分布の真のスコアを近似するためにMC推定器を使用する。

主な貢献

本論文は、5つの主要な貢献を述べている：

1. 鞍点定式化: 制約付きサンプリングを、平均制約を持つサドルポイント問題としてラグランジュ双対領域にキャストし、最適分布が双対変数によってインデックス付けされたギブス分布であることを特定した。
2. PDIアルゴリズム: 主方向のデノイジングと双対上昇が共同で進化する逆拡散サンプラーを導入した。乗数は、デノイズされたサンプルから推定された制約違反に基づいて、毎ステップ更新される。
3. 一般化スコアネットワーク: 双対変数に条件付けられた単一のネットワークを学習させることで、固定された単一のターゲットではなく、ギブス分布のファミリー全体に汎化することを可能にした。
4. 収束解析: 時間平均された双対イテレートが、最適な乗数の近傍に $O(1/\sqrt{T})$ の速度で収束することを証明した。また、残留する双対ミスマッチが終端分布に与える影響を抑え、初期の高ノイズステップが誤差を減衰させ、後のステップが解を洗練させることを示した。
5. 実証的検証: ガウス混合サンプリング、無線電力制御、およびポートフォリオ管理という3つの異なる確率的最適化問題において、本手法の有効性を実証した。

実験結果

著者らは、PDIを、投影拡散モデル（PDM）、主・双対ランジュバン（PDL）、拡散事後サンプリング（DPS）、および双対学習（DT）を含むベースラインと比較評価した。

• ガウス混合モデル: PDIは、制約の完全な実現可能性を維持しながら、より良い目的関数値（無制約の最適値に近い値）を達成する。PDMやPDLは、個々のサンプルに対して制約を強制するため（これにより保守的な解となる）、PDIは期待値において制約を満たすことで、より多様で最適なサンプルを可能にする。
• 無線リソース割り当て: PDIは、平均総レート（最適性）と実現可能性（ユーザーの最小レート制約の充足）の間の最良のバランスを実現する。PDIは、実現可能性の指標（例：5パーセンタイルレート）においてPDMやPDLを大幅に上回り、高い平均レートも維持している。また、DTと比較して、分布外（OOD）の制約シフト（例：異なる $r_{min}$ 値）に対する堅牢性も示している。
• ポートフォリオ管理: PDIは、期待リターン、リスクの実現可能性、およびポートフォリオの多様化（エントロピー）の間の最良のトレードオフを提供する。PDIは、無制約モデル（リスク予算を違反する）や、過度に保守的なベースライン（リターンを犠牲にする）よりも優れた性能を示す。

アブレーション研究:

• 固定された乗数: 推論中に双対変数を固定する（最終イテレートまたは時間平均のいずれか）と、フルPDIダイナミクスと比較して劣った性能となる。これは、乗数の経路依存的な進化が極めて重要であることを裏付けている。
• ノイズスケジュール: ノイズスケジュールの選択は極めて重要である。低SNR（信号対雑音比）を長く維持するスケジュール（例：標準的なコサイン型）は、逆プロセスが誤差に敏感になる前に双対変数が安定することを可能にする。クリーンなサンプルを再構成しすぎるスケジュールは、双対ミスマッチを増幅させ、制約違反を招く。
• 双対学習 vs PDI: DTは有益な双対変数に到達できるが、そのスコアモデルは学習軌跡上の「平均化された」場であり、推論中の特定の制約残差に適応できないため、PDIの性能には及ばない。

意義と主張

PDIは、拡散ベースの生成モデリングと制約付き確率的最適化の間の溝を埋めるものである。その意義は以下の点にある：

• 動的な適応: 双対乗数を事前計算して固定する必要がある従来の手法とは異なり、PDIは乗則を推論時の状態として扱うため、サンプラーが制約違反に対して即座に反応することを可能にする。
• 効率性と堅牢性: 単一の分布ファミリーに対して単一のネットワークを学習させることで、PDIは新しい問題インスタンスや制約の変化に対して、双対学習（DT）で必要とされる高価な再学習を回避する。
• 理論的保証: 本研究は、時間平均された双対変数に対する収束保証と、逆プロセスの安定性に関する境界を提供し、制約付き拡散サンプリングの理論的基礎を与える。

著者らは、収束保証が最終的な主方向の分布に直接ではなく、時間平均された双対イテレートに適用されること、および本手法が問題固有のハイパーパラメータ（例：エントロピー正則化強度 $\beta$ および双対事前分布）のチューニングを必要とすることを、限界事項として述べている。それにもかかわらず、PDIは、解が確率的である必要がある制約付き最適化問題において、最適な分布からサンプリングするための堅牢なフレームワークとして提示されている。