[論文レビュー] A Diffusion Model Predicts 3D Shapes from 2D Microscopy Images

Dominik Waibel, Ernst Röell+3 more|arXiv (Cornell University)|Aug 30, 2022

Cell Image Analysis Techniques被引用数 4

ひとこと要約

本論文では、2次元蛍光顕微鏡画像を入力として、3次元細胞形状を現実的に予測する拡散モデルDISPRを紹介する。この手法は、1枚の2次元蛍光画像を条件として用いることで、赤血球の形態特徴分類タスクにおける少数クラス分類性能を著しく向上させ、データ拡張として合成された3次元再構成を用いることで、マクロF1スコアを55.2%から72.2%に向上させた。

ABSTRACT

Diffusion models are a special type of generative model, capable of synthesising new data from a learnt distribution. We introduce DISPR, a diffusion-based model for solving the inverse problem of three-dimensional (3D) cell shape prediction from two-dimensional (2D) single cell microscopy images. Using the 2D microscopy image as a prior, DISPR is conditioned to predict realistic 3D shape reconstructions. To showcase the applicability of DISPR as a data augmentation tool in a feature-based single cell classification task, we extract morphological features from the red blood cells grouped into six highly imbalanced classes. Adding features from the DISPR predictions to the three minority classes improved the macro F1 score from $F1_ ext{macro} = 55.2 \pm 4.6\%$ to $F1_ ext{macro} = 72.2 \pm 4.9\%$. We thus demonstrate that diffusion models can be successfully applied to inverse biomedical problems, and that they learn to reconstruct 3D shapes with realistic morphological features from 2D microscopy images.

研究の動機と目的

• 3次元共焦点顕微鏡の遅さと高コストな問題に対処し、2次元画像から高速かつ現実的な3次元形状予測を可能にすること。
• 単一細胞形態分類におけるデータ不足とクラス不均衡を克服し、データ拡張のための合成3次元形状を生成すること。
• 拡散モデルの分布的出力特性を活用し、1枚の2次元入力から多様で現実的な3次元再構成を生成すること。
• 不均衡な生物学的データセットにおける下流分類性能向上の観点から、拡散ベースの3次元形状予測の有効性を評価すること。

提案手法

• ノイズ除去のための拡散確率的モデルを訓練し、Tステップにわたり3次元細胞ボリュームに徐々にガウスノイズを加える前向きのノイズ付与プロセスを逆方向に実行する。
• モデルは2次元蛍光顕微鏡画像とセグメンテーションマスクを条件として用い、ノイズ除去プロセス中に事前知識として機能する。
• 推論時、純粋なノイズから出発し、Tステップのノイズ除去を実行し、t=0ステップで3次元形状予測を生成する。
• ノイズスケジュールを最適化することでサンプリング効率を向上させる変分下界損失を用いてモデルを訓練する。
• 3次元ボリューム内の長距離依存関係をモデル化するため、スキップ接続とアテンション機構を備えたU-Netベースのアーキテクチャを採用する。
• 少数クラスのデータ拡張のために、1枚の2次元画像あたり5つの3次元再構成を予測することで、合成3次元形状を生成する。

Fig. 1 : Diffusion steps of one forward pass through DISPR visualized during inference between step $t=999$ with the largest amount of Gaussian noise and the DISPR prediction of $t=0$ , for a stromatocyte cell (see section 2 ). Note that the zoom level in the visualizations is not uniform.

実験結果

リサーチクエスチョン

• RQ1 拡散モデルは1枚の2次元顕微鏡画像から現実的な3次元細胞形状を生成できるか？
• RQ2 拡散モデルの分布的出力は、不均衡な単一細胞分類タスクにおけるデータ拡張を改善するか？
• RQ3 DISPRが生成した3次元形状の形態的特徴は、真値と比較してどれほど現実的で正確か？
• RQ4 DISPRで生成されたデータを用いることで、少数クラス分類におけるマクロF1スコアはどの程度向上するか？

主な発見

• DISPRで生成された3次元形状の特徴を用いた訓練データのデータ拡張により、赤血球形態分類のマクロF1スコアが55.2% ± 4.6%から72.2% ± 4.9%に向上した。
• 重み付きF1スコアは85.7% ± 2.1%から88.6% ± 1.1%に上昇し、全体の分類性能が向上したことが示された。
• データ拡張後、キラトサイトの正しく分類された数は14から24に、マルチロベート細胞は1から8に増加した。
• 最先端の決定的モデルと比較して、DISPRが生成した形状は相対的表面粗さと曲率誤差が低く、より高い形態的忠実度を示した。
• 生成された形状の分布は、真値データセットのクラス固有の形態的特徴分布を適切に捉えており、その妥当性が裏付けられた。
• 推論時間（1予測あたり4分45秒）がやや長くても、SHAPRのような決定的オートエンコーダーよりも多様で現実的な3次元再構成を生成でき、性能を上回った。

Fig. 2 : DISPR is trained to denoise the 3D volume $\mathbf{x}_{b,t}$ containing added stochastic Gaussian noise, to obtain $\mathbf{x}_{b,t-1}$ , thus reversing the noising step $q$ . In each forward pass $p_{\theta}$ , we constrain our 3D model with one 2D image $b$ containing a fluorecent image a