【技術追蹤】通過潛在擴散和先驗知識增強時空疾病進展模型（MICCAI-2024）

??向擴散模型中引入先驗知識，實現疾病進展預測，擴散模型開始細節作業了~

論文：Enhancing Spatiotemporal Disease Progression Models via Latent Diffusion and Prior Knowledge
代碼：https://github.com/LemuelPuglisi/BrLP

0、摘要

??本文介紹了基于潛在擴散的新型時空疾病進展模型—— Brain Latent Progression（BrLP）。BrLP 旨在預測個體在三維腦 MRI 上的疾病進展。（由當前圖像獲得疾病進展后的圖像）。

??為完成這一任務而開發的現有深度生成模型主要以數據驅動為主，面臨學習疾病進展的挑戰。BrLP 通過結合疾病模型的先驗知識來增強預測的準確性，解決了這些挑戰。

??為了實現這一點，本文建議整合一個輔助模型來推斷不同腦區的體積變化。此外，本文引入了潛在平均穩定化（Latent Average Stabilization，LAS），一種新的技術來提高預測進展的時空一致性。

??BrLP 在一個大規模數據集上進行訓練和評估，該數據集包含來自 2805 名受試者的 11730 張 T1 加權腦部 MRI 圖像，這些圖像采集自三個公開的縱向阿爾茨海默病（AD）研究。

??實驗中，在橫斷面和縱向設置中，將 BrLP 生成的 MRI 掃描與受試者實際可用的隨訪 MRI 進行比較。BrLP 在體積準確性和圖像相似性方面均顯著優于現有方法，其中在與阿爾茨海默病（AD）相關的腦區體積準確性提高了 22%，與真實掃描的圖像相似性提高了 43%。

??BrLP 能夠在受試者水平生成條件化的三維掃描，并且通過整合先驗知識來提高準確性，這在疾病進展建模方面代表了一個重大進步，并為精準醫療開辟了新的途徑。

1、引言

1.1、研究意義與當前挑戰

??（1）相比于生物標志物的表征，利用豐富的、高維的影像表示疾病進展，便于可視化和精確定位復雜結構變化的模式，從而提供對疾病動態更詳細的了解；
??（2）挑戰 1： 通過基于受試者特定的元數據進行條件化，以提高個體化程度；
??（3）挑戰 2： 在有條件的情況下使用縱向掃描；
??（4）挑戰 3： 增強時空一致性，以實現空間和時間維度上的平滑進展；
??（5）挑戰 4： 管理由于使用高分辨率三維醫學圖像而導致的高內存需求；

1.2、本文貢獻

??（1）解決挑戰 1： 提出將 LDM 和 ControlNet 結合起來，根據可用的受試者數據生成個體化的腦部 MRI；
??（2）解決挑戰 2： 提出通過使用一個輔助模型來整合疾病進展的先驗知識，該模型旨在推斷不同腦區的體積變化，從而在有條件的情況下使用縱向數據；
??（3）解決挑戰 3： 提出了 LAS，這是一種用于改善預測進展的時空一致性的技術；
??（4）解決挑戰 4： 使用腦部 MRI 的潛在表示，以減少處理三維掃描的內存需求；

??本文通過訓練 BrLP 來學習不同認知狀態個體腦部的漸進性結構變化，從而評估其性能：認知正常（Cognitively Normal, CN）、輕度認知障礙（Mild Cognitive Impairment, MCI）以及阿爾茨海默病（Alzheimer’s Disease）。

??為此，本文使用了一個大規模數據集，包含來自 2805 名受試者的 11730 張 T1 加權腦部 MRI 圖像，這些圖像來自三個公開的阿爾茨海默病縱向研究。據本文作者所知，本文是首次提出一種用于腦部 MRI 的三維條件生成模型，該模型將疾病進展的先驗知識整合到圖像生成過程中。

挑戰與貢獻對應的寫法很棒的感覺~

2、方法

2.1、Background - Diffusion Models

??原文略，可參考：【Diffusion綜述】醫學圖像分析中的擴散模型（一）中 2.2 節；

2.2、Proposed Pipeline - Brain Latent Progression (BrLP)

??BrLP 架構由四個關鍵組件構成：一個 LDM（潛在擴散模型）、一個 ControlNet（控制網絡）、一個輔助模型以及一個 LAS（時空一致性增強）模塊。這四個組件共同解決了引言部分提到的挑戰，其架構總體如圖1 所示。具體來說：

??（1）LDM 旨在生成符合特定協變量的隨機三維腦部 MRI 圖像；
??（2）ControlNet 則專注于將這些 MRI 圖像細化到特定受試者的解剖結構；
??（3）輔助模型利用疾病進展的先驗知識，以提高對特定腦區體積變化預測的精確性；
??（4）LAS 模塊在推理過程中被用于增強時空一致性；

??
Figure 1 | BrLP訓練和推理過程概述：

在這里插入圖片描述

2.2.1、LDM：學習腦 MRI 分布

??在文獻 [12] 的基礎上，本文訓練了一個潛在擴散模型（LDM），其目標是生成反映特定協變量 $c = ? s, v ?$ 的三維腦部 MRI 圖像。其中， $s$ 包括受試者特定的元數據（如年齡、性別和認知狀態），而 $v$ 則涵蓋了與阿爾茨海默病（AD）進展相關的指標，例如與 AD 進展相關的腦區（海馬體、大腦皮層、杏仁核、大腦白質和側腦室）的體積。

??LDM 的構建是一個兩階段的過程。首先，訓練一個自編碼器 $(\mathcal E, \mathcal D)$ （圖1中的A部分），該自編碼器旨在為數據集中的每個腦部 MRI $x$ 生成一個潛在表示 $\mathcal E(x)$ 。隨后訓練一個條件 UNet（圖1中的B部分），表示為 $_θ$ ，網絡參數為 $θ$ ，旨在估計從 $z_t$ 恢復到 $z_{t-1}$ 所需的噪聲 $_θ(z_t,t,c)$ ，通過最小化損失函數 $\mathcal L_?$ 來訓練 $_θ$ 。協變量 $c$ 作為條件通過交叉注意力機制集成到網絡中。

??生成過程首先從采樣隨機高斯噪聲 $z_T～\mathcal N(0,I)$ 開始，然后迭代地逆向執行每個擴散步驟 $z_t→z_{t?1}$ ，其中 $t = T, \dots, 1$ 。對最終步驟 $t = 1$ 的輸出 $z_0$ 進行解碼，可以得到一個符合指定協變量 $c$ 的合成腦部 MRI 圖像 $\hat x=\mathcal D(z_0)$ 。

2.2.2、ControlNet：基于受試者腦部 MRI 的條件

??LDM 僅通過協變量 $c$ 對生成的腦部 MRI 提供有限程度的控制，并且不允許根據單個解剖結構對模型進行調節。這個模塊的目的是擴展 LDM 的能力，使其能夠包含這種額外的控制功能。為了實現這一點，本文使用了 ControlNet（圖1中的C部分），這是一個設計為與 LDM 協同工作的神經網絡。

??將 ControlNet 和 LDM 概念化為一個統一的網絡 $_{θ,?}$ ，其中 $θ$ 表示 LDM 的固定網絡參數， $?$ 表示 ControlNet 的可訓練網絡參數。與 LDM 一樣， $_{θ,?}$ 仍用于預測反向擴散步驟 $z_t→z_{t?1}$ 中的噪聲 $_{θ,?}(z_t,t,c,z)$ ，現在將 $\mathcal E(x)$ 也作為條件，以涵蓋生成過程中目標腦 $x$ 的結構。

??為了訓練 ControlNet，本文使用來自同一患者在不同年齡 $A < B$ 時拍攝的腦部 MRI 對的潛在表示 $z^{(A)}$ 和 $z^{(B)}$ 。與 $z^{(B)}$ 相關的協變量 $c^{(B)}$ 已知并用作目標協變量。每次訓練迭代包括：

??（1）采樣 $t ～ U [1, T]$
??（2）執行 $t$ 步前向擴散 $z^{(B)}→z^{(B)}_t$
??（3）預測噪聲 $?_{θ,?}(z^{(B)}_t,t,c^{(B)},z^{(A)})$ 以恢復 $z^{(B)}_t→z^{(B)}_{t-1}$
??（4）最小化損失 $\mathcal L_?$

2.2.3、提出的輔助模型 - 利用疾病先驗知識

??AD 相關區域隨時間而縮小或擴大，且速率不同。基于深度學習的時空模型試圖以黑盒方式直接從腦部 MRI 中學習這些進展速率，這可能非常具有挑戰性。為了幫助這一過程，本文提出將有關體積變化的先驗知識直接整合到流程中。

??為此，本文利用輔助模型 $f_ψ$ （圖1中的D部分），能夠預測與 AD 相關的區域體積隨時間的變化，并通過進展相關的協變量 $v$ 將這些信息提供給 LDM。輔助模型的選擇針對兩種情況，使 BrLP 在橫斷面和縱向數據中都具有靈活性。

??對于在年齡 $A$ 時有單次掃描的受試者，本文采用回歸模型來估計年齡 $B$ 時的體積變化 $\hat v^{(B)} = f_ψ(c^{(A)})$ 。對于可以訪問在年齡 $A_1, ..., A_n$ 時的 $n$ 次過去訪問的受試者，使用疾病進程映射 (DCM) 來預測 $\hat v^{(B)} = f_ψ(c^{(A_1)}, ... ,c^{(A_n)})$ ，該模型專門用于疾病進展。DCM 旨在根據受試者可用的體積變化歷史提供更準確的軌跡。雖然使用 DCM 作為潛在的解決方案，但任何合適的疾病進展模型都可用于 BrLP。

2.2.4、推理過程

??將 $x^{(A)}$ 定義為年齡為 $A$ 的受試者的輸入腦部 MRI，已知受試者特定元數據 $s^{(A)}$ 和從 $x^{(A)}$ 測量的進展相關體積 $v^{(A)}$ 。如圖1中的E部分總結，為了推斷年齡 $B > A$ 時的腦部 MRI $x^{(B)}$ ，執行 6 個步驟：

??（1）使用輔助模型預測進展相關體積 $\hat v^{(B)} = f_ψ(c^{(A)})$ ；
??（2）將此信息與受試者特定的元數據 $s^{(B)}$ 連接，形成目標協變量 $c^{(B)} =?s^{(B)},\hat v^{(B)}?$ ；
??（3）計算潛在表示 $z^{(A)} = \mathcal E(x^{(A)})$ ；
??（4）對隨機高斯噪聲 $z_T～\mathcal N(0,I)$ 進行采樣；
??（5）通過預測噪聲 $_{θ,?}(z_t,t,c^{(B)},z^{(A)})$ 來運行反向擴散過程，以反轉每個擴散步驟 $t = T, \dots, 1$ ；
??（6）使用解碼器 $\mathcal D$ 在圖像域中重建預測的腦 MRI $x^{(B)} = \mathcal D(z_0)$ ；

??這個推理過程被總結為一個緊湊的符號 $\hat z^{(B)} = \mathcal I(z_T, x^{(A)},c^{(A)})$ 和 $\hat x^{(B)} = \mathcal D(\hat z^{(B)})$ 。

2.2.5、通過提出的潛在平均穩定化（LAS）增強推理

??初始值 $x_T～\mathcal N(0,I)$ 的變化可能導致推理過程產生的結果出現輕微差異。當對連續時間步進行預測時，這些差異尤為明顯，表現為不規則的模式或進程的非平滑過渡。因此，本文引入了 LAS（圖1中的F部分），這是一種通過平均推理過程的不同結果來提高時空一致性的技術。

??特別地，LAS 基于這樣的假設：預測值 $\hat z^{(B)} = \mathcal I(z_T, x^{(A)},c^{(A)})$ 偏離理論平均值 $\mu ^{(B)} = \mathbb E[\hat z^{(B)}]$ 。為了估計期望值 $\mu ^{(B)}$ ，本文建議重復推理過程 $m$ 次，并對結果進行平均：

在這里插入圖片描述

??與之前類似，本文對預測的掃描進行解碼 $\hat x^{(B)} = \mathcal D(\mu ^{(B)})$ 。整個推理過程（其中 $m = 4$ ）在消費級 GPU 上每 MRI 需要 4.8 秒。

3、實驗與結果

3.1、數據

??收集了一個大型數據集，包含來自 2,805 名受試者的 11,730 張 T1 加權腦部 MR I掃描圖像，這些數據來自多個公開的縱向研究：ADNI 1/2/3/GO (1,990名受試者）、OASIS-3 (573名受試者）和 AIBL (242名受試者）。每位受試者至少有兩張 MRI 圖像，且每次掃描均在不同的訪問中獲取。

??所有數據集中均可以獲得年齡、性別和認知狀態。平均年齡為 74±7 歲，53% 的受試者為男性。根據最終訪視結果，43.8% 的受試者被歸類為正常對照組，25.7% 表現出或發展為輕度認知障礙，30.5% 表現出或發展為阿爾茨海默病。

??將數據隨機分為訓練集（80%）、驗證集（5%）和測試集（15%），且各組之間沒有重疊的受試者。驗證集用于訓練過程中的提前停止。每張腦部 MRI 圖像均經過預處理，包括：N4 偏場校正、顱骨剝離、到 MNI 空間的仿射配準、強度歸一化以及重采樣至 1.5mm3。用于評估進展的相關協變量體積是使用 SynthSeg 2.0 計算得出，并以總腦體積的百分比表示，以考慮個體差異。