RAFT

Teed 2020 · 論文

一行要約 — 4Dのall-pairs相関ボリュームを構築し、重み共有ConvGRUで現在の推定値周辺の相関をルックアップしながら単一の高解像度オプティカルフローフィールドを反復的に精緻化する — ECCV 2020ベストペーパーであり、現代のオプティカルフローを定義するアーキテクチャ。

問題

主流の深層フローアーキテクチャ(PWC-Netおよびその系譜)は、古典的な粗から密へのピラミッドを継承していた: 低解像度でフローを推定し、その後ワーピングして精緻化する。この設計には構造的な死角がある — 各レベルのコストボリュームは小さな探索窓しか対象にしない、小さく速く動く物体は粗い解像度で消失する、ピラミッドの早い段階での誤りは元に戻すのが難しい、そしてマルチステージのカスケードはしばしば100万回を超える学習反復を必要とする。以前の反復的精緻化スキームは反復間で重みを共有していなかった(あるいは、IRRのように大きなリカレントユニットによって制約されていた)。RAFTは次のように問う: ネットワークがすべてのピクセルペア間のマッチングコストを事前計算し、軽量な学習済みオプティマイザが必要に応じてそのボリュームをクエリして単一の高解像度フローフィールドを精緻化したらどうなるか?

手法とアーキテクチャ

3段階、すべて微分可能でエンドツーエンドに学習される:

  1. 特徴抽出: エンコーダ gθ:RH×W×3↦RH/8×W/8×Dg_\theta : \mathbb{R}^{H \times W \times 3} \mapsto \mathbb{R}^{H/8 \times W/8 \times D}(D=256D = 256、6個の残差ブロック)が両フレームを符号化する。同一アーキテクチャのコンテキストネットワーク hθh_\theta が I1I_1 のみを符号化する。両方ともペアごとに1回だけ実行される。
  2. All-pairs相関: 視覚的類似性は、すべてのピクセルペアに対して単一の行列積として事前計算される:

Cijkl=∑hgθ(I1)ijh⋅gθ(I2)klh,C∈RH×W×H×WC_{ijkl} = \sum_h g_\theta(I_1)_{ijh} \cdot g_\theta(I_2)_{klh}, \qquad \mathbf{C} \in \mathbb{R}^{H \times W \times H \times W}

その後、最後の2次元がカーネル1、2、4、8で平均プーリングされ、ピラミッド {C1,C2,C3,C4}\{\mathbf{C}^1, \mathbf{C}^2, \mathbf{C}^3, \mathbf{C}^4\} になる。I2I_2 側の次元のみをプーリングすることで、I1I_1 側の次元はフル(1/8)解像度に保たれる — 大きな変位と小さな変位の両方が捉えられ、小さく速く動く物体を失わない。ルックアップ演算子 LCL_\mathbf{C} は、現在の対応関係 x′=x+f(x)\mathbf{x}' = \mathbf{x} + \mathbf{f}(\mathbf{x}) 周辺のローカルグリッド上で各レベルをバイリニアサンプリングする:

N(x′)r={x′+dx∣dx∈Z2, ∥dx∥1≤r}\mathcal{N}(\mathbf{x}')_r = \{ \mathbf{x}' + \mathbf{dx} \mid \mathbf{dx} \in \mathbb{Z}^2,\ \lVert \mathbf{dx} \rVert_1 \le r \}

各レベル kk において N(x′/2k)r\mathcal{N}(\mathbf{x}'/2^k)_r でインデックス付けされる — 一定の半径は粗いレベルほど広い文脈をカバーする(半径4は k=4k=4 で元解像度の256ピクセルをカバーする)。 3. 反復更新: f0=0\mathbf{f}_0 = \mathbf{0} から開始し、リカレント更新演算子(わずか270万パラメータ、すべての反復で重みを共有)が相関ルックアップ、フロー特徴、コンテキスト特徴 xtx_t を受け取り、畳み込みGRUを介して残差更新 fk+1=fk+Δf\mathbf{f}_{k+1} = \mathbf{f}_k + \Delta\mathbf{f} を出力する:

zt=σ(Conv3×3([ht−1,xt],Wz)),rt=σ(Conv3×3([ht−1,xt],Wr))z_t = \sigma(\mathrm{Conv}_{3\times3}([h_{t-1}, x_t], W_z)), \qquad r_t = \sigma(\mathrm{Conv}_{3\times3}([h_{t-1}, x_t], W_r))

h~t=tanh⁡(Conv3×3([rt⊙ht−1,xt],Wh)),ht=(1−zt)⊙ht−1+zt⊙h~t\tilde{h}_t = \tanh(\mathrm{Conv}_{3\times3}([r_t \odot h_{t-1}, x_t], W_h)), \qquad h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t

この演算子は一次オプティマイザを模倣する — しかしテイラー線形化されたデータ項の代わりに、降下方向を提案することを学習する。有界な活性化関数は不動点への収束を促し、発散せずに100回以上反復を実行できる。フローは1/8解像度で予測され、各ピクセルの3x3粗い近傍上での学習された凸結合(softmaxによる重み)によってアップサンプリングされる。

教師信号は、指数的に増加する重みで推定値全系列をカバーする:

L=∑i=1NγN−i∥fgt−fi∥1,γ=0.8\mathcal{L} = \sum_{i=1}^{N} \gamma^{N-i} \lVert \mathbf{f}_{gt} - \mathbf{f}_i \rVert_1, \qquad \gamma = 0.8

学習はFlyingChairsに続いてFlyingThings、その後ベンチマークでのファインチューニングという順序で行われる。映像に対しては、ウォームスタート初期化が前フレームのフローを前方投影する。

実験結果

SLAMにおける意義

RAFTの「相関ボリューム+反復的リカレント精緻化」というレシピは、SLAMにおける学習ベースのデータ対応付けの主力となった: DROID-SLAMとDPVOは本質的に、微分可能なバンドル調整層を包み込んだRAFTスタイルの更新演算子である。その後継(シーンフロー向けのRAFT-3D、リアルタイム向けのSEA-RAFT)はフローベンチマークを席巻し、RAFTが広めたアンロールされた学習済み最適化のパターンは、現在では密な予測やSLAMシステム全般に見られる。

関連ノート