[論文レビュー] Explicitly disentangling image content from translation and rotation with spatial-VAE

Tristan Bepler, Ellen D. Zhong+3 more|arXiv (Cornell University)|Sep 25, 2019

Advanced Electron Microscopy Techniques and Applications被引用数 34

ひとこと要約

Spatial-VAEは、生成器をピクセル座標に依存させることにより、VAE内で回転/平行移動を内容から分離し、姿勢推定を明示的に行い、変換された画像、EMデータ、銀河データにおける再構成を改善する。

ABSTRACT

Given an image dataset, we are often interested in finding data generative factors that encode semantic content independently from pose variables such as rotation and translation. However, current disentanglement approaches do not impose any specific structure on the learned latent representations. We propose a method for explicitly disentangling image rotation and translation from other unstructured latent factors in a variational autoencoder (VAE) framework. By formulating the generative model as a function of the spatial coordinate, we make the reconstruction error differentiable with respect to latent translation and rotation parameters. This formulation allows us to train a neural network to perform approximate inference on these latent variables while explicitly constraining them to only represent rotation and translation. We demonstrate that this framework, termed spatial-VAE, effectively learns latent representations that disentangle image rotation and translation from content and improves reconstruction over standard VAEs on several benchmark datasets, including applications to modeling continuous 2-D views of proteins from single particle electron microscopy and galaxies in astronomical images.

研究の動機と目的

• 意味的内容が回転や平行移動といった姿勢変換と分離された画像の潜在表現を学習する動機付け。
• 空間的に条件付けられたVAEを提案し、回転と平行移動を内容と共に教師なしで推定する。
• ベンチマークおよびドメインデータセット（MNIST の派生、cryo-EM シミュレーション、Galaxy Zoo）で、再構成の改善と真の内容の分離を示す。

提案手法

• 座標 x^i（回転/変換版を含む）と潜在 z を入力とするMLPによって画像尤度 p_g(y^i|x^i,z) を計算する空間生成ネットワークを定義する。
• 座標空間で直接回転と平行移動をモデル化するよう座標を変換し、p_g(y|z,θ,Δx) が x^i R(θ) + Δx に依存するようにする。
• VAE目的関数の下で、近似后验 q(z,θ,Δx|y) = N(μ(y),σ(y)^2 I) を介して回転 θ、平行移動 Δx、そして非構造的潜在 z を推定する。
• 期待対数尤度と z、θ、Δx のKL項を組み合わせた変分下界を用い、θ に一様事前分布を希望する場合には θ のKL の代替を提供する（式3）。
• 推論時に θ と Δx が変動できる学習スキームを提供し、変換された要因と内容因子のエンドツーエンド学習を可能にする。
• 空間VAEを通常のVAEおよび姿勢固定のバリアントと実験的に比較し、分離性と再構成の改善を示す。

実験結果

リサーチクエスチョン

• RQ1 教師なしで連続的な姿勢要因（回転と平行移動）を意味的内容から分離するようにVAEを訓練できるか？
• RQ2 画像がランダムに回転・平行移動される場合に、座標ベースの生成と姿勢推定を明示的にモデル化することは再構成を改善するか？
• RQ3 このフレームワークは、姿勢に不変でありつつ、シミュレートされた cryo-EM 投影や銀河画像などのドメインデータセットで真の内容の変動を回復できるか？
• RQ4 学習された潜在空間は内容と回転/平行移動を分離し、ドメインを跨いだ意味のある補間を反映できるか？

主な発見

• 空間-VAEは、変換された MNIST 派生における ELBO で通常の VAE を上回り、特に非構造的潜在次元が小さい場合に優れる。
• モデルは content (z) を回転 θ と平行移動 Δx から分離することを学習し、ポーズが content 表現から取り除かれた学習データマニフォールドによって示される。
• シミュレートされた cryo-EM データでは、自由度のある潜在変数1つと姿勢推定を組み合わせた空間VAEが、向きが要因化されたコンフォメーション変動を捉え、通常のVAEを上回る。
• Galaxy Zoo および EM データセットで、空間-VAE は回転/平行移動に依存しない変換不変表現を学習し、銀河の形態とタンパク質のコンフォメーションを独立して反映する。
• 潜在空間の補間は、姿勢が分離された状態で内容の滑らかな遷移を示し、因子の頑健な分離を示唆している。