[论文解读] 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 将旋转/平移从内容中解耦，通过使生成器依赖像素坐标，实现对姿态的显式推断，并在变换图像、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）上展示改进的重建和真实内容的解耦。

提出的方法

• 定义一个空间生成器网络，其中图像似然 p_g(y^i|x^i,z) 由一个多层感知机（MLP）计算，该 MLP 接受坐标 x^i（及其旋转/变换版本）和潜在变量 z。
• 在坐标空间直接对坐标进行变换，以建模旋转和平移，因此 p_g(y|z,θ,Δx) 取决于 x^i 经过 R(θ) 的旋转再加上 Δx。
• 通过近似后验 q(z,θ,Δx|y) = N(μ(y),σ(y)^2 I) 在 VAE 目标下推断旋转 θ、平移 Δx 和无结构潜在 z。
• 使用一个变分下界，将期望对数似然与 z、θ、Δx 的 KL 项相结合；当需要均匀先验时，为 θ 提供 KL 的替代项（式3）。
• 提供一个训练方案，在推断阶段允许 θ 和 Δx 变化，从而实现对变换与内容因子的端到端学习。
• 在实验中将 spatial-VAE 与普通 VAE 以及具有固定姿态的变体进行比较，以证明解耦与重建改进。

实验结果

研究问题

• RQ1 是否可以在无监督的情况下训练 VAE 将连续的姿态因子（旋转和平移）与语义内容解耦？
• RQ2 在图像被随机旋转和平移时，显式建模基于坐标的生成和姿态推断是否能改善重建？
• RQ3 在诸如模拟 cryo-EM 投影和银河图像等领域数据集中，框架是否能在保持对姿态不变的同时恢复真实的内容变异？
• RQ4 所学习的潜在空间是否能够将内容与旋转/平移分离，并在跨域上体现有意义的插值？

主要发现

• Spatial-VAE 在经过变换的 MNIST 变体的 ELBO 上优于普通 VAE，尤其是在无结构潜在维度较小时。
• 模型学会将内容（z）与旋转 θ 和平移 Δx 分离，通过学习的数据流形来证明姿态从内容表征中移除。
• 在模拟的 cryo-EM 数据中，带有一个无约束潜在变量加上姿态推断的 spatial-VAE 捕捉到构象变异且方向被因子化，优于普通 VAE。
• 在 Galaxy Zoo 和 EM 数据集上，Spatial-VAE 学会了对变换不变的表征，能够独立于旋转/平移反映星系形态和蛋白质构象。
• 潜在空间中的插值显示内容的平滑过渡，同时姿态保持解耦，表明因素分离稳健。