Theseus

Pineda (Meta) 2022 · 論文

一行要約 — 微分可能な非線形最小二乗（DNLS）のための応用領域を問わないオープンソースPyTorchライブラリであり、バンドル調整のような幾何的最適化ループの内部でニューラルネットワークを学習させるための再利用可能な基盤インフラを提供する。

問題

2022年までに、いくつかの画期的なシステム（BA-Net、DROID-SLAM、gradSLAM）が、非線形最小二乗ソルバーをネットワークの内部に置き、それを通して学習することの威力を示していたが、各実装は1つのシステムのために独自に作られたものだった。論文が述べる通り、既存のDNLS実装は「応用に特化しており、効率にとって重要な多くの要素を必ずしも取り込んでいない」：疎ソルバー、バッチ処理、ベクトル化、GPUサポート、メモリ効率の良い勾配計算はプロジェクトごとに再発明されたり（あるいは省略されたり）しており、先行するDNLS研究は勾配計算のためにアンローリングしか対応していなかった。一方、成熟した古典的ソルバー（Ceres、g2o、GTSAM）は効率化の仕組みは持つが、求解を通して逆伝播する手段がなかった。Theseusはこの隙間を一度に埋め、「ロボティクスとビジョンにおけるエンドツーエンドの構造化学習」のための共有基盤となる。

手法とアーキテクチャ

双レベル最適化としてのDNLS。 内側の問題は、多様体値の変数$\theta = \{\theta_j\}$（ユークリッドベクトルまたはリー群）に関する非線形最小二乗であり、残差は重みとコストに因数分解される、$r_i(\theta^i) = w_i c_i(\theta^i)$。

$$\theta^{\star} = \operatorname*{arg\,min}_{\theta} S(\theta), \qquad S(\theta) = \frac{1}{2}\sum_i \|w_i c_i(\theta^i)\|^2 .$$

これは反復的な線形化によって解かれる：$J_i = \partial r_i / \partial\theta^i$として$\big(\textstyle\sum_i J_i^{\top}J_i\big)\,\delta\theta = \textstyle\sum_i J_i^{\top} r_i$を解き、その後リトラクション$\theta \leftarrow \theta - \delta\theta$を行う（ガウス・ニュートン法；適応的減衰を持つLMおよびDoglegも提供される）。上流のネットワークパラメータ$\phi$はコスト、重み、初期化のいずれにも入り得るため、次の双レベル設定が得られる。

$$\text{内側:}\;\; \theta^{\star}(\phi) = \operatorname*{arg\,min}_{\theta} S(\theta;\phi), \qquad \text{外側:}\;\; \phi^{\star} = \operatorname*{arg\,min}_{\phi} L(\theta^{\star}(\phi)),$$

ここで外側のループは、ソルバーを通じた$\partial\theta^{\star}/\partial\phi$を用いる通常の勾配降下である。

API（ファクターグラフ形式）。 Variable（最適化対象または補助テンソル）、CostFunction（$c_i$；解析的ヤコビアンを備えたライブラリ提供のもの — ガウス計測、再投影、相対姿勢、動作モデル、衝突など — またはインプレースのAutoDiffCostFunction）、CostWeight（$w_i$、ロバスト損失を含む）、Objective（$S$）、Optimizer、そしてTheseusLayer。そのforwardは入力テンソルを、任意のPyTorchグラフ内で最適な変数値に写す。微分可能なリー群は、指数写像・対数写像・逆元・合成を解析的接空間導関数付きの閉形式で計算し、さらに自動微分の勾配が接空間に正しく写るような射影演算子を備える（LieTorchの演算ごとのカスタムカーネルとは対照的）。微分可能な順運動学はDifferentiable Robot Modelをラップする。

効率化の仕組み。 (i) 2レベルの並列性 — DNLS問題のネイティブなバッチ処理と、同種のコスト演算の自動ベクトル化（SIMD風）。(ii) PyTorchの密なコレスキー分解を置き換える、エンドツーエンドで微分可能な疎線形ソルバー：CHOLMOD（CPU）、cudaLU（GPU上のcuSolverRFベースのバッチLU）、そしてBaSpaCho — GPU対応の新規オープンソースなバッチ超節点疎コレスキー分解であり、その疎な消去によって外部のシューア補元トリックが不要になる。線形求解$y = A^{-1}b$の逆伝播は陰関数微分を用い、$\partial f/\partial b = A^{-1}\,\partial f/\partial y$および$\partial f/\partial A = -A^{-1}(\partial f/\partial y)\,y^{\top}$であり、キャッシュされた因数分解によって逆伝播が順伝播より高速になる。

4つの逆伝播モード。 アンローリング（ソルバーの反復を通した逆伝播 — 計算量・メモリが反復数に対して線形に増加し、勾配消失のリスクがある）；打ち切り微分（TBPTT、バイアスあり）；陰関数微分 — 最適性条件$g(\theta;\phi) := \nabla_{\theta} S(\theta;\phi) = 0$に陰関数定理を適用したもの：

$$\mathrm{D}_{\phi}\theta^{\star}(\bar{\phi}) = -\mathrm{D}_{\theta}^{-1} g\big(\theta^{\star}(\bar{\phi});\bar{\phi}\big)\, \mathrm{D}_{\phi} g\big(\theta^{\star}(\bar{\phi});\bar{\phi}\big),$$

実際には、解において単一のニュートンステップ$h(\theta;\phi) = \theta - [\nabla^2_{\theta}S]^{-1}_{\text{stop}}\nabla_{\theta}S$を微分することで計算される；そして直接損失最小化（DLM）は、損失を付加した内側求解を用いる有限差分ベクトル-ヤコビアン積方式である。陰関数微分とDLMのコストは反復数に依存しない。

同一のコンポーネントから構築される応用例： ポーズグラフ最適化（Welschロバストカーネルの半径を学習）、タクタイル状態推定（タクタイル画像から相対姿勢へのネットワークをエンドツーエンドで学習）、バンドル調整（外れ値ソフトカーネルの半径を学習）、動作計画（学習された初期化モデルを持つ微分可能なGPMP2）、特徴量ベースのホモグラフィ推定（ロバストなアライメント特徴のためのCNNを学習）。

実験結果

• 疎ソルバー対密ソルバー（PGO、合成Cubeデータセット、V100 32GB、内側10反復／外側20エポック、陰関数モード）: PyTorchの密ソルバーはバッチサイズ128で256姿勢を超えるとメモリ不足になり、そこでの順伝播＋逆伝播の時間は既に20.81秒であるのに対し、CHOLMODは10.96秒、cudaLUは2.86秒、BaSpaChoは2.25秒。BaSpaChoは2048姿勢まで、cudaLUは4096姿勢まで、CHOLMODはバッチ256で8192姿勢までスケールする；BaSpaChoはすべての規模で密ソルバーを上回り、最大で1桁高速である。
• Ceresとの比較（バッチPGO、10反復、256問題）: 小規模ではCeresが優位（256姿勢、バッチ16で25倍高速）だが、2048姿勢／バッチ256ではBaSpaChoがCeresより約23倍高速で、他の疎ソルバーも約4倍高速；アブストラクトが謳う「最大20倍」の順伝播高速化は、バッチ処理＋ベクトル化＋疎性の組み合わせによるものである。
• 逆伝播モード（タクタイル状態推定、100エポック）: アンローリングの逆伝播時間とメモリは内側反復数に対して線形に増加する（おおよそ34MBから262MB）が、陰関数微分／DLMは約28〜29MBで一定；陰関数微分は最良の検証損失も達成し、推奨される既定手法となっている。
• 自動ベクトル化はPGOにおいて有意な順伝播・逆伝播の高速化をもたらすが、その代償として最大で約82%（順伝播）／約55%（逆伝播）の追加メモリを要する。

SLAMにおける意義

古典的なSLAMバックエンド（g2o、Ceres、GTSAM）は高度に最適化されているが微分可能ではなく、一方で深層ネットワークは微分可能だがSLAMを扱いやすくする疎構造を捨ててしまう。Theseusはこの2つを橋渡しする：ファクターグラフ風の最適化をPyTorchに持ち込み、学習済みフロントエンドと最適化バックエンドを持つハイブリッドシステムを、真の課題損失に対してエンドツーエンドで学習できるようにする。BA-NetとDROID-SLAMが先駆けたパターンを共有基盤へと一般化しており、VIO/SLAM推定器の残差重み、ロバストカーネル、初期化ネットワークを手動で調整する代わりに学習させるための自然なツールである。

関連ノート

• BA-Net — ネットワーク層としての微分可能バンドル調整であり、直接の先行研究。
• Lietorch — PyTorchにおける微分可能なリー群演算であり、同種の問題に用いられる。
• GradSLAM — 完全に微分可能な稠密SLAMパイプライン。
• DROID-SLAM — 微分可能なBA層を中心に構築されたエンドツーエンドSLAMシステム。
• Differentiability — これらすべてのシステムの根底にある概念。
• MAP inference as sparse nonlinear least squares — 微分可能にされている古典的な問題。