SEA-RAFT

Wang 2024 · 論文

一行要約 — Simple, Efficient, Accurate RAFT: 混合ラプラス損失、初期フローの直接回帰、TartanAirによる剛体フロー事前学習、そしてアーキテクチャの単純化により、RAFTに最良の精度・効率パレートフロンティアをもたらす——Springで最先端であり、比較対象の手法より少なくとも2.3倍高速に動作する。

問題

RAFT以降のオプティカルフローの進展は主に、Transformerコストボリュームエンコーダやより大きなバックボーンといった、より重いアーキテクチャからもたらされ、それによってフローをリアルタイムシステムで使えるようにする速度が犠牲になっていた。一方、RAFT自体のレシピにも弱点がある: 標準的なL1L_1エンドポイント損失は、強いオクルージョン下の曖昧で予測不能な画素に支配される; ゼロ初期化されたフローは多くの精緻化イテレーション（学習時12回、推論時最大32回）を必要とする; そして合成データのFlyingChairs/Thingsはリアリズムと一般化を制限する。SEA-RAFTは、損失・初期化・データの変更が元のアーキテクチャをどこまで押し進められるかを問う。

手法とアーキテクチャ

SEA-RAFTはRAFTの骨格を保持する。特徴エンコーダFFとコンテキストエンコーダCC（RAFT独自のエンコーダを置き換えるImageNet事前学習済みのResNetを切り取ったもの）はI1,I2∈RH×W×3I_1, I_2 \in \mathbb{R}^{H\times W\times 3}を1/8解像度の特徴に写す; 多スケール4Dコリレーションボリュームが次のように構築される。

Vk=F(I1)∘AvgPool(F(I2),2k)⊤∈Rh×w×h2k×w2k,k=4, (h,w)=18(H,W)V_k = F(I_1) \circ \texttt{AvgPool}(F(I_2), 2^k)^{\top} \in \mathbb{R}^{h\times w\times\frac{h}{2^k}\times\frac{w}{2^k}}, \qquad k=4,\ (h,w)=\tfrac{1}{8}(H,W)

各イテレーションは現在のフローμ\muの周辺で半径r=4r{=}4の動き特徴を検索する: M=MotionEncoder(LookUp({Vk},μ,r))M = \texttt{MotionEncoder}(\texttt{LookUp}(\{V_k\},\mu,r))。そして再帰ユニット（RAFTのConvGRUを置き換える2つのConvNeXtブロック）が隠れ状態を更新し、残差を回帰する: h′=RNN(h,M,C(I1))h' = \texttt{RNN}(h, M, C(I_1))、Δμ=FlowHead(h′)\Delta\mu = \texttt{FlowHead}(h')。3つの主要な変更点は次のとおりである。

MixLap(x;α,β1,β2,μ)=α⋅e−∣x−μ∣eβ12eβ1+(1−α)⋅e−∣x−μ∣eβ22eβ2MixLap(x;\alpha,\beta_1,\beta_2,\mu) = \alpha\cdot\frac{e^{-\frac{|x-\mu|}{e^{\beta_1}}}}{2e^{\beta_1}} + (1-\alpha)\cdot\frac{e^{-\frac{|x-\mu|}{e^{\beta_2}}}}{2e^{\beta_2}}

β1=0\beta_1{=}0を固定し、第一成分がL1L_1/エンドポイント誤差指標に一致するようにし、安定性のためスケールを対数空間で回帰する（β2∈[0,10]\beta_2 \in [0,10]）。損失LMoL\mathcal{L}_{MoL}は、画素とフローの両軸で平均されたグラウンドトゥルースの負の対数尤度であり、通常の指数重み付けで各イテレーションに適用される: Lall=∑i=1NγN−iLMoLi\mathcal{L}_{all}=\sum_{i=1}^{N}\gamma^{N-i}\mathcal{L}_{MoL}^{i}。混合重みα\alphaは不確実性の出力としても機能する。

バリアント: SEA-RAFT(S)はResNet-18の最初の6層を用い、(M)はResNet-34の最初の13層を用いる; (L)は(M)を推論イテレーション12回で実行するものである。

実験結果

SLAMにおける意義

SEA-RAFTは、SLAMのフロントエンドがリアルタイムのレートで密なフローを必要とする際の実用的な選択肢である: Transformer級のレイテンシを持たないRAFT級の機構である。その要素のうち2つはSLAMの要求と直接一致する——剛体運動事前学習はSLAMが仮定するほぼ剛体な世界に一致し、ラプラス混合パラメータは確率的推定における観測共分散に自然に写る画素ごとの不確実性を与える。また、損失・初期化・データの設計がアーキテクチャの成長を上回りうることを示す事例研究でもあり、これは学習型SLAMコンポーネント全般に関連する教訓である。

関連ノート