Multi-Granularity Cross-modal Alignment for Generalized Medical Visual Representation Learning

Authors:
Fuying Wang,Yuyin Zhou,Shujun Wang,Varut Vardhanabhuti,Lequan Yu

📝Introduction

随着大规模标记数据的增长，深度学习技术促进了医学图像理解的进步，然而这类标记数据是十分昂贵的，并且很耗费时间。因此，直接从医学图像中直接学习逐渐变成了主流，即致力于从医生详细的医学报告中学习通用的诗句表征，并将其应用于下游的计算机视觉任务。
先前的工作关注于病理学往往只发生在整个图像的局部很小的位置，因此基于 attention-based 的对比学习策略被提出来学习局部表征。
医学图像和医学报告展示出了多粒度的语义关系，如 disease-level, instance-level, and pathological regionlevel。现有的方法并没有将这些维度全部考虑在内，没有充分利用报告中的信息，进而导致了不充分的模型表征能力。

本文提出了 MGCA(Multi-Granularity Cross-modal Alignment)来充分利用这种报告和图像的多粒度关系，通过三种对齐：instance-wise alignment，token-wise alignment 和 disease-level alignment 来提升视觉表征的泛化性。

📝Method

报告和图像首先通过 Text Encoder(BERT)和 Image Encoder(ViT)编码为 Image Representation(绿色) 和 Report Representation(紫色).

Instance-wise Image-Text Alignment(ITA): 整个框架的核心主力，用于报告-图像对映射到相近的空间，而不成对的映射到较远的距离。首先使用非线性映射层将文本和图像embeding 映射到标准化的低维空间，然后计算每对的余弦相似度。然后使用对称的标准化 InfoNCE losses 来最大限度地保留潜在空间中真实对之间的互信息。这部分可以表示为：

sim (x_{v, i}, x_{t, i}) = {\tilde{v}}_{i}^{⊤} {\tilde{t}}_{i}, where {\tilde{v}}_{i} = g_{v} (v_{i}), {\tilde{t}}_{i} = g_{t} (t_{i}) .

ℓ_{i}^{v 2 t} = - \log \frac{\exp (sim (x_{v, i}, x_{t, i}) / τ_{1})}{\sum_{k = 1}^{B} \exp (sim (x_{v, i}, x_{t, k}) / τ_{1})}, ℓ_{i}^{t 2 v} = - \log \frac{\exp (sim (x_{t, i}, x_{v, i}) / τ_{1})}{\sum_{k = 1}^{B} \exp (sim (x_{t, i}, x_{v, k}) / τ_{1})}

这部分的 loss 为：

L_{ITA} = \frac{1}{2 N} \sum_{i = 1}^{N} (ℓ_{i}^{v 2 t} + ℓ_{i}^{t 2 v}),

Cross-attention-based Token-wise Alignment(CTA): 病理只占整个图像的一小部分，报告中只有少数疾病标签描述了关键的医疗状况，基于这个情况，本文提出使用基于双向交叉注意力的 Token-wise 对齐。首先将将文本和图像 embeding 映射到标准化的低维空间，然后通过 cross attention 来计算视觉和文本 tokens 的软匹配，即让视觉 tokens 去关注所有的文本 tokens，然后计算它对应跨模态文本 embedding

o_{i}^{j} = \sum_{k = 1}^{N} O (α_{i}^{j 2 k} (V {\tilde{z}}_{i}^{k})), α_{i}^{j 2 k} = softmax (\frac{{(Q {\tilde{r}}_{i}^{j})}^{⊤} (K {\tilde{z}}_{i}^{k})}{\sqrt{d}})

然后让文本 tokens 尽可能接近他对应的 cross-modal 文本 embedding，然后远离其他的。考虑到不同的视觉 tokens 有不同的作用，因此为每个视觉 tokens 都设置了一个权重，文本部分相应的 loss 可以写为：

L_{LIA} = - \frac{1}{2 N S} \sum_{i = 1}^{N} \sum_{j = 1}^{S} w_{i}^{j} (\log \frac{\exp (sim ({\tilde{r}}_{i}^{j}, o_{i}^{j}) / τ_{2})}{\sum_{k = 1}^{S} \exp (sim ({\tilde{r}}_{i}^{j}, o_{i}^{k}) / τ_{2})} + \log \frac{\exp (sim (o_{i}^{j}, {\tilde{r}}_{i}^{j}) / τ_{2})}{\sum_{k = 1}^{S} \exp (sim (o_{i}^{j}, {\tilde{r}}_{i}^{k}) / τ_{2})}),

视觉部分也进行类似的计算，这部分总体损失为文本和视觉部分的 loss 的均值。

Crossmodal Prototype Alignment(CPA): ITA and CTA 将不同的实例都作为了负样本对，忽视了不同病理之间存在着相似的高级语义信息。对于疾病级别的对齐，本文通过强制执行跨模式聚类分配一致性来利用医学图像和放射学报告之间的受试者间关系对应关系。具体来说，Sinkhorn-Knopp聚类算法通过将文本和视觉分为 K 类，得到两个软簇分配代码，然后分别计算文本和诗句部分的 CE Loss

p_{v, i}^{(k)} = \frac{\exp ({\tilde{v}}_{i}^{⊤} c_{k} / τ_{3})}{\sum_{k^{'}} \exp ({\tilde{v}}_{i}^{⊤} c_{k^{'}} / τ_{3})}, p_{t, i}^{(k)} = \frac{\exp ({\tilde{t}}_{i}^{⊤} c_{k} / τ_{3})}{\sum_{k^{'}} \exp ({\tilde{t}}_{i}^{⊤} c_{k^{'}} / τ_{3})}

ℓ ({\tilde{v}}_{i}, q_{t, i}) = \sum_{k = 1}^{K} q_{t, i}^{(k)} \log p_{v, i}^{(k)}, ℓ ({\tilde{t}}_{i}, q_{v, i}) = \sum_{k = 1}^{K} q_{v, i}^{(k)} \log p_{t, i}^{(k)} .

L_{CPA} = \frac{1}{2 N} \sum_{i = 1}^{N} (ℓ ({\tilde{v}}_{i}, q_{t, i}) + ℓ ({\tilde{t}}_{i}, q_{v, i})) .

最后，总的 loss 为三部分 loss 和：

L = λ_{1} * L_{ITA} + λ_{2} * L_{CTA} + λ_{3} * L_{CPA}