[论文解读] Using VAEs to Learn Latent Variables: Observations on Applications in cryo-EM

Daniel G. Edelberg, Roy R. Lederman|arXiv (Cornell University)|Mar 13, 2023

Computational Physics and Python Applications被引用 5

一句话总结

本文研究了变分自编码器（VAEs）在冷冻电镜（cryo-EM）应用中摊销推理的性质，特别关注连续构象异质性的建模。研究发现，CryoDRGN 中的 VAE 编码器在旋转或平移后的粒子图像上泛化能力较差——表明其对噪声特征产生了过拟合，而显式潜在变量估计的性能与之相当或更优，这挑战了关于科学成像中 VAE 泛化能力的既有假设。

ABSTRACT

Variational autoencoders (VAEs) are a popular generative model used to approximate distributions. The encoder part of the VAE is used in amortized learning of latent variables, producing a latent representation for data samples. Recently, VAEs have been used to characterize physical and biological systems. In this case study, we qualitatively examine the amortization properties of a VAE used in biological applications. We find that in this application the encoder bears a qualitative resemblance to more traditional explicit representation of latent variables.

研究动机与目标

• 评估 VAE 中摊销编码器在未见或经过增强的冷冻电镜粒子图像上的泛化能力。
• 比较基于 VAE 的摊销推理与显式潜在变量估计在冷冻电镜数据中的性能表现。
• 探究 CryoDRGN 中的编码器是否更倾向于对噪声或结构伪影过拟合，而非捕捉真实的构象变异性。
• 评估数据不变性（如旋转、平移）对高噪声生物成像中 VAE 泛化能力的影响。
• 探讨在冷冻电镜等科学应用中，VAE 的行为是否更接近显式潜在变量模型，而非通常所假设的那样。

提出的方法

• 使用 CryoDRGN 框架在冷冻电镜粒子图像上训练 VAE，采用共享编码器网络为每个输入生成高斯潜在分布。
• 通过平面内旋转和平移进行数据增强，利用冷冻电镜数据中已知的不变性作为未见数据的代理，以测试泛化能力。
• 比较增强图像与原始图像的 VAE 编码器潜在表示，以评估其一致性和泛化性能。
• 通过为每张粒子图像单独拟合潜在分布，实现显式变分推理，避免共享编码器参数。
• 使用潜在空间结构的定性比较以及在不同增强视图下的表现一致性来评估性能。
• 在附录 B 中通过 MNIST VAE 进行消融实验，为标准深度学习设置提供基线比较。

Figure 1: Experimental setup of a standard VAE. $x_{i}$ is a datapoint used as input, $z_{i}$ is the associated latent space point, $\hat{x}_{i}$ is the reconstructed output of $x_{i}$ . $\mathcal{L}$ represents the calculation of the ELBO between $x$ and $\hat{x}_{i}$

实验结果

研究问题

• RQ1 VAE 编码器在冷冻电镜中对训练期间未见的旋转或平移粒子图像是否具有良好的泛化能力？
• RQ2 通过共享编码器实现的摊销推理性能与显式潜在变量估计相比如何？
• RQ3 VAE 编码器在冷冻电镜数据中对噪声或结构伪影的过拟合程度有多大，而非捕捉真实的构象状态？
• RQ4 冷冻电镜 VAE 的泛化失败是由于高噪声水平和图像可变性，还是架构限制所致？
• RQ5 在该场景下，显式潜在变量估计是否能产生定性上更优或更鲁棒的表示，优于摊销 VAE？

主要发现

• CryoDRGN 中的 VAE 编码器在旋转或平移后的粒子图像上泛化能力较差，表明其对自然数据增强的泛化性能不佳。
• 增强图像的编码器潜在表示与原始图像相比存在显著差异，提示其对噪声或图像特异性特征产生过拟合。
• 不使用共享编码器参数的显式潜在变量估计，在相同任务中表现出定性上相似或更优的性能。
• CryoDRGN 缺乏架构上的不变性（如无旋转不变性）可能加剧了泛化问题，尽管此类不变性并不能保证完全解决该问题。
• VAE 编码器在冷冻电镜中的行为比预期更接近显式潜在变量模型，挑战了摊销推理优势的既有假设。
• 与 MNIST VAE 不同，冷冻电镜场景下即使模型训练成功，高噪声和图像可变性仍可能破坏 VAE 的泛化能力。

Figure 2: The CryoDRGN Pipeline for heterogeneous reconstruction. $x_{i}$ is the cryo-EM image datapoint as input. $\omega_{i}$ , $t_{i}$ , $h_{i}$ are the given rotation, translation, and contrast transfer function of the input image. These variables with latent variable $z_{i}$ are used as input t