技術要約：円錐最適化器のスケーラブルなディープ・アンフォールディング（Deep Unfolding）

問題提起

ディープ・アンフォールディング（DU）は、圧縮センシングやロボティクスを含む様々な領域において、反復最適化アルゴリズムを加速させる手法として有効であることが証明されている。しかし、ロボティクスの応用分野（共分散ステアリングなど）で一般的な大規模半定値計画問題（SDP）へのDUの拡張は、フルアップデート型のADMMベースの円錐ソルバー（例：COSMO）をアンローリングする際に直面する、以下の2つの具体的な障害によって阻まれてきた。

1. 線形システム解法におけるメモリのボトルネック： 各反復における線形システムの解法をバックプロパゲーション（誤差逆伝播）する場合、通常、係数行列を明示的に形成する必要がある。ADMMの場合、これは $A^\top A$ という項を伴う。制約行列 $A$ が疎（スパース）であっても、$A^\top A$ はしばしば密（デンス）になり、$O(n^2)$ のメモリ複雑度を招く。これにより、直接分解を行うとメモリ不足に陥る大規模な問題サイズへのスケールが妨げられる。
2. PSD射影における数値的不安定性： 正定値（PSD）錐への射影には固有値分解が必要である。この操作に対する標準的なバックプロパゲーションでは、$\lambda_i = \lambda_j$ の場合に $1/(\lambda_i - \lambda_j)$ という項が含まれる。固有値が一致する場合（等方的な共分散ターゲットを用いる際など、学習中に頻繁に発生する）、これらの項は数値的不安定性や未定義の勾配（NaN）を引き起こす。

手法

著者らは、これらの2つの障害に対処するカスタム・バックワード・ルールを用いた、COSMOソルバーのためのディープ・アンフォールディング・フレームワークを提案している。これにより、大規模な円錐問題に対するハイパーパラメータ・ポリシーおよびウォームスタートの学習が可能となる。

1. 線形システム解法の行列フリー（Matrix-Free）な暗黙的微分

$O(n^2)$ の密な $A^\top A$ 行列を形成することによるメモリコストを回避するため、著者らは共役勾配法（CG）に基づく線形システム解法に対して、行列フリーの暗黙的微分ルールを採用している。

• フォワードパス： ソルバーは間接モードを使用し、行列ベクトル積（matvecs）のみに依存するCGを用いて、縮小された正規方程式 $(P + \sigma I + A^\top D(\rho)A)\tilde{x} = \xi$ を解く。
• バックワードパス： CGの反復をアンローリングしたり、密なヤコビアンを形成したりする代わりに、著者らは随伴システムに基づくルールを導出する。随伴システム $M\tilde{\lambda} = -\nabla_{\tilde{x}}L$ を同じCG法を用いて解くことで、ハイパーパラメータ $\rho$ に関する勾配 $\nabla_\rho L = (A\tilde{\lambda}) \odot (A\tilde{x})$ が計算される。
• 結果： このアプローチでは、密な係数行列を明示的に形成することが決してないため、メモリ複雑度が $O(n^2)$ から $O(n + \text{nnz}(A))$ に削減され、従来の手法が失敗する規模でのバックプロパゲーションが可能になる。

2. PSD射影のための安定したバックワード・ルール

重複する固有値の下での固有値分解の勾配の不安定性を解決するために、著者らは標準的なオートグラッド（自動微分）を、Dalečii–Krein (DK) 表象に基づくバックワード・ルールに置き換えている。

• メカニズム： PSD射影は、主要な行列関数 $T = U f(\Lambda) U^\top$ （ここで $f(x) = \max(x, 0)$）として扱われる。DK定理は、特異項 $1/(\lambda_j - \lambda_i)$ が分子 $f(\lambda_j) - f(\lambda_i)$ によって乗じられる形式の微分公式を提供する。
• 安定性： 固有値が一致する場合（$\lambda_i = \lambda_j$）、分子は分母と同じ速度で消失するため、除去可能な特異点が $f'(\lambda_i)$ という極限として定義される。これにより、重複する固有値があっても勾配は有限かつ定義された状態に保たれる。
• 効率性： このルールは、フォワードパスで計算された固有値分解を再利用するため、Newton-Schulz反復などの代替案とは異なり、漸近的な計算コストが増加することはない。

主な貢献

1. スケーラブルな線形システム微分： メモリ使用量を $O(n^2)$ から $O(n)$ に削減する行列フリーの暗黙的微分ルールにより、大規模なSDPに対するフルアップデート型円錐ソルバーのディープ・アンフォールディングを可能にした。
2. 数値的に安定したPSD射影： 重複する固有値の下でも定義が維持される、Dalečii–Krein定理に基づくバックワード・ルールを提供した。これは、共分散ステアリングやその他のSDP問題の学習において極めて重要である。
3. 実証的検証： これらのルールにより、軽量なハイパーパラメータ・ポリシーやウォームスタートを学習するためにCOSMOソルバーをアンローリングできることを示した。学習されたアプローチは、スタンドアロンのSDP/SOCP、および非線形共分散ステアリングのための逐次凸計画法（SCP）のサブルーチンとして評価されている。

結果

著者らは、学習されたCOSMOソルバーを、適応型ヒューリスティックを備えた標準的なCOSMOソルバー、および最先端の内部点法ソルバーであるClarabelと比較し、複数の問題クラスにおいて評価を行った。

• スタンドアロンのSDPおよびSOCP： 学習されたポリシーは、問題クラスに応じて大幅な高速化を実現した：
  • Lovász $\vartheta$ SDP： 実時間（wall-clock time）で最大 50.4倍 の高速化。
  • Max-cut SDP： 2.62倍 の高速化。
  • ロバスト・カルマンフィルタ： 4.91倍 の高速化。
  • リアプノフLMI： 1.16倍 の高速化。
• 非線形共分散ステアリング (SCP)： 非線形システム（二重積分器およびユニサイクル・ダイナミクス）のSCPにおけるサブルーチンとして使用した場合、学習されたアプローチは、バニラなCOSMOと比較して 30倍 の高速化を実現した。
• メモリのスケーリング： 実験により、ベースラインの手法が $n \approx 20\text{--}50k$ 変数でGPUメモリ（24 GB）を使い果たす一方で、提案された行列フリー手法は線形にスケールし、$n=50k$ でも300 MB未満に抑えられることが示された。

意義と限界

本論文は、これらの貢献により、これまでメモリと数値的安定性の制約によって制限されていた大規模な円錐最適化へのディープ・アンフォールディングの適用が可能になったと主張している。学習されたポリシーは、固定された計算予算が要求されるリアルタイム・ロボティクスへの応用において、大幅な加速を提供する。

著者らは以下の特定の限界についても言及している：

• スケーラビリティ： 線形システム解法は行列フリーであるが、PSD射影は依然として各半定値ブロックのサイズ $p$ に対して $O(p^3)$ の密な固有値分解を必要とする。スケーラビリティの向上は、変数や制約の数に対して実現されるものであり、ブロックサイズ自体に対してではない。
• 精度領域： 本手法は、リアルタイム制御において典型的な、低〜中程度の精度（$\epsilon \approx 10^{-3}$）を対象としている。一次近似的手法が収束しにくい非常にタイトな許容誤差において、内部点法（Clarabelなど）を凌駕することを主張するものではない。
• 学習のオーバーヘッド： 問題ファミリーごとに個別のポリシーを学習する必要があり、学習時には（約100個の問題クラスあたり）事前のコストが発生する。現在の研究では、ファミリー間の転移や、学習時よりも大幅に大きなインスタンスへのスケーリングについては実証されていない。