文章目录
1 要点
题目：用于WSI分类的多模态去噪扩散预训练模型 (Multi-modal denoising diffusion pre-training for whole-slide image)
代码：https://github.com/lhaof/MDDP
研究目的：
免疫组织化学 (IHC) 染色能够提供分子标记和蛋白质表达信息，突出显示癌症区域，但其制备过程需要额外的劳动、时间和专业设备。因此，研究者们希望利用多模态训练图像来推断缺失模态的WSI标签。具体地，通过仅使用H&E染色的WSI作为输入，来捕获两种染色图像模态之间的结构相似性和染色差异。
关键技术：
- 多模态去噪扩散预训练框架：
- MDDP框架包含两个关键任务：H&E染色图像到IHC图像的翻译任务，以及IHC图像到IHC图像的重构任务；
- 翻译任务和重构任务使用基本的生成损失
- 类约束对比损失：
使用从IHC染色图像估计的先验图像级标签来对齐两个生成任务的全局语义； - 包特征增强策略：
通过所设计的预训练模型提取的特征来扩展包特征，如图3；
数据集：
- Camelyon16
- TCGA-COAD
- TCGA-NSCLC
难点：
- H&E图像和IHC属于背景知识，可以不用过多关注。可以这样简单理解：我们平时使用的WSI，例如Camelyon和TCGA等，都是这样的，而IHC图像是在H&E染色的基础上进行，但是这需要更多的成本；
- 对于去噪扩散预训练模型，其输入是H&E-IHC图像对，作者在此基础上进行图像的翻译任务和重构任务。为了对齐这两个任务的生成特征，所以设计了类约束对比损失；
- 对于我们常用的WSI分类任务，本文的方法可以看作是一种类似于Resnet的特征提取器，其所提取的特征包含了IHC信息，可以看作是已有图像的另一个模态。在本文的策略下，其这些特征可以和Resnet提取的特征拼接，并用于后续分类任务；
2 方法
2.1 多模态图像配准和标签分配
如图2(a)所示，在我们的数据获取过程中，首先对未染色的组织切片进行H&E染色，然后在去除H&E染色后进行IHC染色。这种顺序染色方法确保了不同染色的两幅图像之间在位置和结构的对应。然而，染色过程可能仍然会经历不可避免的变换，如平移或旋转，导致错位。为了解决这个问题，本文采样以下步骤：
- 使用刚性配准方法对来自同一组织切片的H&E染色和IHC染色WSI进行对齐；
- 使用Yottixel技术，根据组织分割结果和RGB直方图，从配准后的H&E染色WSI中选择有效的图像块；
- 在Yottixel中，从分割的组织区域裁剪出大小为1000×1000的正方形区域，并调整为256×256的图像块；
- 相应的IHC染色区域也以与其配对的H&E染色图像块相同的位置裁剪和调整大小；
- 在上述图像块选择程序之后，生成了大量配对的H&E-IHC染色样本 ( x h e , x i h c ) (x^{he},x^{ihc}) (xhe,xihc)；
2.1.1 伪标签分配
在IHC染色中，特定蛋白质的表达被可视化为棕色区域，而正常组织往往更白。例如，如图2(a)中的 x i h c x^{ihc} xihc行所示，左侧的补丁是阴性样本，右侧带有较大棕色区域的是阳性样本。因此，可以通过分析IHC样本中棕色区域的比例来评估特定蛋白质的表达。在本文中，棕色区域大于1%的样本定义为阳性，小于0.1%的为阴性，其余样本则被丢弃。每对H&E-IHC染色图像块被分配一个二元伪标签 ρ \rho ρ（阳性/阴性），表明蛋白质表达的水平。
2.2 多模态去噪扩散预训练架构
我们提出的多模态去噪扩散预训练框架旨在从仅含H&E染色图像的输入中，提取IHC染色引导的特征。该框架通过两个任务进行训练：
- 多模态图像到图像的翻译任务：去噪扩散模型学习将H&E染色图像映射到其IHC染色对应图像；
- 单模态图像重建任务：要求去噪扩散模型从其噪声版本中恢复IHC染色图像；
如图2(b)所示，给定一对H&E染色和IHC染色的图像块 x h e x_{he} xhe和 x i h c x_{ihc} xihc:
- 首先使用预训练编码器 τ \tau τ提取其潜在特征 z h e 0 z_{he}^0 zhe0和 z i h c 0 z_{ihc}^0 zihc0；
- 分别向潜在特征添加噪声，得到噪声特征 z h e t z_{he}^t zhet和 z i h c t z_{ihc}^t zihct；
- U-Net模型将 z h e t z_{he}^t zhet和 z i h c t z_{ihc}^t zihct作为输入，分别执行重新染色任务和重建任务。根据输入特征，去噪U-Net模型不仅能将H&E染色图像块的潜在特征转换为IHC染色图像的潜在特征，还能从噪声特征中恢复IHC染色图像的原始特征；
- 设 f t f_t ft为图像到图像翻译任务生成的特征，可以使用预训练解码器 D D D，以 f t f_t ft为输入生成合成的IHC染色图像；
2.2.1 潜在特征提取
潜在扩散模型 (LDM) 将扩散过程移入潜在空间，并在潜在特征上执行去噪扩散过程，有效地降低了时间和内存成本。具体来说，LDM采用了由编码器和解码器组成的自动编码器架构：
- 编码器 τ \tau τ 将输入图像 x x x映射为潜在特征图 z = τ ( x ) z = \tau(x) z=τ(x)，而解码器 D D D将潜在特征图重新转换为图像；
- 在训练去噪U-net之前，首先在H&E-IHC染色数据集上训练名为VQGAN的自编码器；
- 在提出框架的训练过程中，编码器 τ \tau τ和解码器 D D D的权重保持冻结状态；
- 编码器以配对的H&E和IHC染色图像块 x h e x_{he} xhe和 x i h c x_{ihc} xihc为输入，输入形状为 h × w × 3 h \times w \times 3 h×w×3，输出两个潜在特征 Z h e 0 Z_{he}^0 Zhe0和 Z i h c 0 Z_{ihc}^0 Zihc0，其形状为 h ′ × w ′ × c ′ h' \times w' \times c' h′×w′×c′。这两个潜在特征将在扩散过程中使用。
2.2.2 扩散过程
扩散过程在每个时间步
t
t
t向输入特征添加噪声，经过长时间步后得到噪声特征。我们的框架遵循一种名为BBDM的图像到图像翻译网络来执行扩散过程。在扩散过程中，噪声分别添加到两个不同任务的潜在特征中。两个任务的前向扩散过程定义如下：
q
h
e
→
i
h
c
(
z
h
e
t
∣
z
h
e
0
,
z
i
h
c
0
)
=
N
(
z
h
e
t
;
(
1
−
m
t
)
z
h
e
0
+
m
t
z
i
h
c
0
,
δ
t
I
)
,
(1)
\tag{1} q_{he \rightarrow ihc}(z_{he}^t | z_{he}^0, z_{ihc}^0) = \mathcal{N}(z_{he}^t; (1 - m_t) z_{he}^0 + m_t z_{ihc}^0, \delta_t I),
qhe→ihc(zhet∣zhe0,zihc0)=N(zhet;(1−mt)zhe0+mtzihc0,δtI),(1)
q
i
h
c
→
i
h
c
(
z
i
h
c
t
∣
z
i
h
c
0
)
=
N
(
z
i
h
c
t
;
z
i
h
c
0
,
δ
t
I
)
,
(2)
\tag{2} q_{ihc \rightarrow ihc}(z_{ihc}^t | z_{ihc}^0) = \mathcal{N}(z_{ihc}^t; z_{ihc}^0, \delta_t I),
qihc→ihc(zihct∣zihc0)=N(zihct;zihc0,δtI),(2)其中，
q
h
e
→
i
h
c
q_{he \rightarrow ihc}
qhe→ihc表示H&E特征到IHC特征翻译任务的前向转移概率，
q
i
h
c
→
i
h
c
q_{ihc \rightarrow ihc}
qihc→ihc表示IHC特征重建任务的前向转移概率。
N
\mathcal{N}
N表示高斯分布，
t
t
t是介于
0
0
0和
T
T
T之间的时间步、
m
t
=
t
/
T
m_t = t/T
mt=t/T表示添加的噪声比例、以及
δ
t
\delta_t
δt设计为KaTeX parse error: Can't use function '\(' in math mode at position 1: \̲(̲2(m_t - m_t^2)并遵循BBDM的最佳设置。在时间步
t
t
t时，两个任务的扩散特征计算如下：
z
h
e
t
=
(
1
−
m
t
)
z
h
e
0
+
m
t
z
i
h
c
0
+
δ
t
ϵ
t
,
(3)
\tag{3} z_{he}^t = (1 - m_t) z_{he}^0 + m_t z_{ihc}^0 + \sqrt{\delta_t} \epsilon_t,
zhet=(1−mt)zhe0+mtzihc0+δtϵt,(3)
z
i
h
c
t
=
z
i
h
c
0
+
δ
t
ϵ
t
,
(4)
\tag{4} z_{ihc}^t = z_{ihc}^0 + \sqrt{\delta_t} \epsilon_t,
zihct=zihc0+δtϵt,(4)其中，
z
h
e
t
z_{he}^t
zhet和
z
i
h
c
t
z_{ihc}^t
zihct分别表示在时间步
t
t
t加入噪声的H&E染色和IHC染色图像特征，以及
ϵ
t
∼
N
(
0
,
I
)
\epsilon_t \sim N(0, I)
ϵt∼N(0,I) 是来自标准正态分布的高斯噪声。
2.2.3 去噪过程
在去噪过程中，给定
z
h
e
t
z_{he}^t
zhet、
z
i
h
c
t
z_{ihc}^t
zihct，以及时间步
t
t
t，使用一个U-Net模型
μ
θ
\mu_\theta
μθ来分别预测添加到这两个输入噪声特征上的噪声。需要注意的是，
μ
θ
\mu_\theta
μθ仅依赖输入特征为不同任务 (图像到图像翻译/图像重建) 预测噪声。去噪特征可以通过以下公式计算：
f
t
=
z
h
e
t
−
μ
θ
(
z
h
e
t
,
t
)
,
f_t = z_{he}^t - \mu_\theta(z_{he}^t, t),
ft=zhet−μθ(zhet,t),
g
t
=
z
i
h
c
t
−
μ
θ
(
z
i
h
c
t
,
t
)
,
g_t = z_{ihc}^t - \mu_\theta(z_{ihc}^t, t),
gt=zihct−μθ(zihct,t),其中，
f
t
∈
R
h
′
×
w
′
×
c
′
f_t \in \mathbb{R}^{h' \times w' \times c'}
ft∈Rh′×w′×c′表示从H&E染色图像的噪声特征图生成的合成IHC虚拟染色特征图，以及
g
t
∈
R
h
′
×
w
′
×
c
′
g_t \in \mathbb{R}^{h' \times w' \times c'}
gt∈Rh′×w′×c′表示从原始IHC染色图像的噪声特征图重建的特征图。
2.3 训练策略
2.3.1 重建和重染色损失
在提出的多模态预训练网络中，有两个生成损失分别衡量重建和重染色任务的合成质量：
- 第一个损失：用于评估通过重建任务生成的去噪IHC特征是否与编码器 τ \tau τ提取的原始特征一致。这个重建损失帮助去噪U-net更好地学习IHC染色切片原始特征的分布；
- 第二个生成损失：评估通过去噪扩散模型使用H&E染色图像块进行重染色任务生成的特征是否与其对应的IHC染色图像块的特征相似。该损失使得去噪U-Net能够通过单模态的H&E染色输入预测跨模态的IHC染色引导特征。我们对生成损失使用
L
1
L1
L1损失函数：
L h e → i h c r e c = ∣ ∣ Z i h c 0 − f t ∣ ∣ , L_{he \to ihc}^{rec} = ||Z_{ihc}^0 - f_t||, Lhe→ihcrec=∣∣Zihc0−ft∣∣, L i h c → i h c r e c = ∣ ∣ Z i h c 0 − g t ∣ ∣ . L_{ihc \to ihc}^{rec} = ||Z_{ihc}^0 - g_t||. Lihc→ihcrec=∣∣Zihc0−gt∣∣.
3.3.2 类约束对比损失
每个IHC染色的图像块可以根据棕色区域的百分比分类为阳性或阴性。H&E染色的图像块与其对应的IHC染色图像块共享相同的类别。因此，提出了一种类约束对比损失，以保持重染色特征 f t f_t ft和重建特征 g t g_t gt之间的语义一致性。
对于一批重染色特征
F
=
{
f
t
0
,
f
t
1
,
.
.
.
,
f
t
b
}
F = \{f_t^0, f_t^1, ..., f_t^b\}
F={ft0,ft1,...,ftb}和对应的重建特征
G
=
{
g
t
0
,
g
t
1
,
.
.
.
,
g
t
b
}
G = \{g_t^0, g_t^1, ..., g_t^b\}
G={gt0,gt1,...,gtb}，它们共享相同的伪标签
ρ
=
{
ρ
0
,
ρ
1
,
.
.
.
,
ρ
b
}
\rho = \{\rho^0, \rho^1, ..., \rho^b\}
ρ={ρ0,ρ1,...,ρb}。同一类别的特征在潜在空间中应该是相似的。因此，定义了一个CLIP风格的对比损失来执行类约束特征对齐：
L
c
o
n
=
−
∑
i
=
0
b
[
ρ
i
⋅
log
(
e
(
ϕ
(
f
t
i
)
⋅
ϕ
(
g
t
i
)
/
σ
)
∑
j
≠
i
e
(
ϕ
(
f
t
j
)
⋅
ϕ
(
g
t
j
)
/
σ
)
)
+
(
1
−
ρ
i
)
⋅
log
(
e
(
−
ϕ
(
f
t
i
)
⋅
ϕ
(
g
t
i
)
/
σ
)
∑
j
≠
i
e
(
−
ϕ
(
f
t
j
)
⋅
ϕ
(
g
t
j
)
/
σ
)
)
]
(5)
\tag{5} L_{con} = - \sum_{i=0}^{b} \left[ \rho_i \cdot \log \left( \frac{e^{(\phi(f_t^i) \cdot \phi(g_t^i) / \sigma)}}{\sum_{j \neq i} e^{(\phi(f_t^j) \cdot \phi(g_t^j) / \sigma)}} \right) + (1 - \rho_i) \cdot \log \left( \frac{e^{(-\phi(f_t^i) \cdot \phi(g_t^i) / \sigma)}}{\sum_{j \neq i} e^{(-\phi(f_t^j) \cdot \phi(g_t^j) / \sigma)}} \right) \right]
Lcon=−i=0∑b[ρi⋅log(∑j=ie(ϕ(ftj)⋅ϕ(gtj)/σ)e(ϕ(fti)⋅ϕ(gti)/σ))+(1−ρi)⋅log(∑j=ie(−ϕ(ftj)⋅ϕ(gtj)/σ)e(−ϕ(fti)⋅ϕ(gti)/σ))](5)其中，
b
b
b是批次大小、
f
t
i
f_t^i
fti和
g
t
i
g_t^i
gti分别表示输入批次中第
i
i
i对的重染色特征和重建特征、
f
t
j
f_t^j
ftj和
g
t
j
g_t^j
gtj代表批次中的其他特征、
ϕ
\phi
ϕ是将所有特征转换为形状为
1
×
c
′
1 \times c'
1×c′的特征向量的全局平均池化操作，以及
σ
\sigma
σ是温度因子。
总损失函数由三个项组成：
L
(
X
h
e
,
X
i
h
c
,
ρ
)
=
λ
1
L
h
e
→
i
h
c
r
e
c
+
λ
2
L
i
h
c
→
i
h
c
r
e
c
+
λ
3
L
c
o
n
L(X_{he}, X_{ihc}, \rho) = \lambda_1 L_{he \to ihc}^{rec} + \lambda_2 L_{ihc \to ihc}^{rec} + \lambda_3 L_{con}
L(Xhe,Xihc,ρ)=λ1Lhe→ihcrec+λ2Lihc→ihcrec+λ3Lcon其中，
X
h
e
X_{he}
Xhe、
X
i
h
c
X_{ihc}
Xihc，以及
ρ
\rho
ρ分别表示一批H&E染色图像块、IHC染色图像块，及其对应的伪标签。
λ
1
\lambda_1
λ1、
λ
2
\lambda_2
λ2，以及
λ
3
\lambda_3
λ3分别设为10、1和0.1。
3.4 下游任务的包特征增强
对于预训练的去噪扩散网络，我们将其用作H&E染色图像的特征提取器，而不是图像生成器。考虑到去噪U-Net已经学习将H&E染色图像的特征转换为IHC染色图像的表示，假设它不仅可以提供H&E染色图像的形态和纹理特征，还可以提供表明特定蛋白质表达水平的IHC染色引导特征。
如图3所示，在分类WSI的任务中，形状为 H × W × 3 H \times W \times 3 H×W×3的H&E染色WSI被裁剪为 n n n个图像块，形成一个用于MIL的包 B B B：
- 包首先被送入一个在ImageNet上预训练的编码器 E E E，以提取通用的实例级特征 f u n i ∈ R n × C f_{uni} \in \mathbb{R}^{n \times C} funi∈Rn×C；
- 同时，包也输入到所提出的多模态框架预训练的特征提取器中，旨在提取IHC引导的表示。特征提取器由预训练的编码器
τ
\tau
τ、单步扩散过程
Ψ
\Psi
Ψ，以及去噪U-net
μ
θ
\mu_\theta
μθ组成：
- 对于包中的每个H&E染色图像块 x h e x_{he} xhe，首先使用 τ \tau τ计算特征图 z h e 0 z_{he}^0 zhe0；
- 给定一个时间步 t t t，通过单步扩散过程得到噪声特征 z h e t = z h e 0 + δ t ϵ t z_{he}^t = z_{he}^0 + \sqrt{\delta_t} \epsilon_t zhet=zhe0+δtϵt，其中 δ t \delta_t δt和 ϵ t \epsilon_t ϵt与公式(3)中的相同；
- 与公式(3)不同的是，由于仅使用单模态的H&E染色图像来解决下游任务，公式(3)中的 z i h c 0 z_{ihc}^0 zihc0被替换为 z h e 0 z_{he}^0 zhe0；
- 随后，去噪U-net模型将噪声特征和时间步
t
t
t作为输入，并输出特征表示
f
d
e
∈
R
n
×
h
′
′
×
w
′
′
×
c
′
′
f_{de} \in \mathbb{R}^{n \times h'' \times w'' \times c''}
fde∈Rn×h′′×w′′×c′′；
为了对齐两个预训练模型中不同潜在空间的特征，对预训练模型的特征应用全局平均池化操作和线性层 (如图3中GAP+Linear)。然后，我们将 f d e f_{de} fde和 f u n i f_{uni} funi连接在一起，构建增强的包级特征 f b a g ∈ R 2 n × C f_{bag} \in \mathbb{R}^{2n \times C} fbag∈R2n×C。包级特征的提取过程可以公式化为：
f b a g = F C ( ϕ ( μ θ ( Ψ ( τ ( B ) , t ) , t ) ) ) ∣ ∣ E ( B ) (7) \tag{7} f_{bag} = FC(\phi(\mu_\theta(\Psi(\tau(B), t), t))) || E(B) fbag=FC(ϕ(μθ(Ψ(τ(B),t),t)))∣∣E(B)(7)其中， ϕ \phi ϕ表示全局平均池化操作、 F C FC FC是全连接层。 τ \tau τ、 Ψ \Psi Ψ，以及 μ θ \mu_\theta μθ分别表示提出的多模态框架中的预训练编码器、单步扩散过程，以及去噪U-net。 E E E代表在ImageNet上预训练的编码器、 ∣ ∣ || ∣∣表示连接操作。然后，得到的包级特征可以用于广泛的基于MIL的WSI分类器中，以获得最终的预测结果。
去噪U-Net模型提供多尺度的特征表示。以往研究表明，利用来自不同解码器层的特征会影响下游任务的性能。因此，我们进行了实验，研究如何选择去噪U-Net的解码器层来提取特征。我们从这些层中选择最佳特征作为我们预训练模型的最终输出表示。此外，不同的时间步 t t t的选择也会影响提取特征的强度。因此，我们搜索能为下游任务带来最佳表现的 t t t。
扫一扫
2133
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
