[論文レビュー] 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

ひとこと要約

本稿は、低温電子線画像法（cryo-EM）応用における変分オートエンコーダー（VAEs）のアンモライズド推論特性を調査する。特に、連続的なコンformationalな不均一性をモデル化する場合に焦点を当てる。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のアンモライズドエンコーダーが、トレーニング時に見られなかった、または拡張されたcryo-EM粒子画像に対して一般化できるかどうかを評価すること。
• cryo-EMデータにおいて、VAEベースのアンモライズド推論と明示的変分推定による潜在変数の性能を比較すること。
• CryoDRGNのエンコーダーが、真のコンformational変動を捉えるのではなく、ノイズや構造的アーチファクトに過剰適合していないかを調査すること。
• データ不変性（例：回転、平行移動）が、高ノイズの生物学的画像におけるVAEの一般化に与える影響を評価すること。
• cryo-EMのような科学的応用におけるVAEが、一般的に想定されているよりも明示的潜在変数モデルに近い挙動を示すかどうかを検討すること。

提案手法

• CryoDRGNフレームワークを用いて、各入力に対してガウス分布の潜在変数を出力する共有エンコーダーネットワークを用いて、cryo-EM粒子画像にVAEを訓練した。
• cryo-EMデータに内在する既知の不変性を活用し、平面内回転および平行移動によるデータ拡張を実施し、未学習データの代理として一般化をテストした。
• 拡張された画像のVAEエンコーダーの潜在表現を、元の画像のそれと比較して一貫性および一般化能力を評価した。
• 共有エンコーダーパrametersを避けることで、個々の粒子画像ごとに別々の潜在分布をフィッティングする明示的変分推論を用いた。
• 拡張されたビュー間での潜在空間構造と一貫性の定性的比較を用いて性能を評価した。
• 付録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 トレーニング時に見られなかった回転または平行移動された粒子画像に対して、cryo-EMにおけるVAEエンコーダーは一般化できるか？
• RQ2 共有エンコーダーによるアンモライズド推論と、cryo-EMにおける明示的潜在変数推定の性能は、どのように比較できるか？
• RQ3 VAEエンコーダーが、真のコンformational状態を捉えるのではなく、ノイズや画像固有の特徴に過剰適合する程度はどの程度か？
• RQ4 cryo-EM VAEにおける一般化の失敗は、高ノイズレベルや画像変動性に起因するのか、あるいはアーキテクチャ的制限に起因するのか？
• RQ5 明示的潜在変数推定は、この文脈において、アンモライズドVAEよりも定性的に優れた、もしくはより頑健な表現を生み出せるか？

主な発見

• CryoDRGNにおけるVAEエンコーダーは、回転または平行移動された粒子画像に対して一般化がうまくいかない。これは、自然なデータ拡張に対して一般化が悪いことを示唆する。
• 拡張された画像のエンコーダー潜在表現は、元の画像のそれと顕著に異なる。これは、ノイズや画像固有の特徴に過剰適合している可能性を示唆する。
• 共有エンコーダーパrametersを用いない明示的潜在変数推定は、同じタスクにおいてアンモライズドVAEと同等またはそれ以上の性能を示した。
• CryoDRGNにアーキテクチャ的不変性（例：回転不変性）が欠如していることが、一般化の悪化に寄与している可能性があるが、このような不変性が問題を完全に解消するとは限らない。
• VAEエンコーダーのcryo-EMにおける挙動は、予想よりも明示的潜在変数モデルに近い。これは、アンモライズド推論の利点に関する仮定に疑問を呈する。
• 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