georgehu学习生活 - 论文阅读-M3Ret: Unleashing Zero-shot Multimodal Medical Image Retrieval via Self-Supervision

原文arXiv链接M3Ret: Unleashing Zero-shot Multimodal Medical Image Retrieval via Self-Supervision

基本介绍

本文提出的框架能够横跨2D灰度图、3D灰度图与RGB图进行统一训练，并得到一个统一的表征，而无需依赖于多种对不同模态的专门设计。另外，还将训练好的模型在以图搜图（可跨模态）的任务中进行了评测。

主要贡献点：

横跨2D灰度图（X光、超声），3D灰度图（CT）、RGB（内窥镜视频），统一训练
量化研究了效果提升的原因

动机图

具体方法

MAE

Masked autoencoder，即常规的像素集masked后重建的训练方式。原输入统一为4维，即$X\in \mathbb{R}^{C\times H\times W\times S}$，(S=slice) 分为N个patch后展平并线性映射到D维空间，得到$P\in\mathbb{R}^{N\times D}$。

假设掩码比例为$\alpha$，剩下可见的patch embedding变为$P_{\text{vis}}\in\mathbb{R}^{(1-\alpha)N\times D}$，经过encoder $f_\theta$后得到$P'_{\text{vis}}\in\mathbb{R}^{(1-\alpha)N\times D'}$。再将可变mask patch embedding $P_{\text{mask}}$一起给入decoder $g_\phi$以得到重建后的掩码部分像素值$\hat{X}_{\text{mask}} = g_{\phi}\left(f_{\theta}(P_{\text{vis}}), P_{\text{mask}}\right)$

使用的loss如下，其中的$\mathcal{M}$为被掩到的patch的索引号：

$$ \mathcal{L}_{\text{MAE}} = \frac{1}{\alpha N} \sum_{i \in \mathcal{M}} \| \hat{X}_i - X_i \|_2^2 $$

在图片检索时，将经过encoder后得到的patch embedding平均作为图片的表示。

SimDINO

按照SimDINO中的方法来生成local和global视角$v_c(X), v_g(X)$。每种模态都取两个global视图，local视图在2D模态上取10个、3D和视频模态上取4个。两种视图拼接CLS后分别传入教师和学生encoder以获取CLS embedding，即$z_g = f_{\text{teacher}}^{\text{cls}}(v_g(X))$与$z_c = f_{\text{student}}^{\text{cls}}(v_c(X))$

损失中结合了coding-rate正则化项，其中的$\Gamma = \mathrm{Cov}[z_c] \in \mathbb{R}^{d \times d}$，即CLS embedding的协方差矩阵：

$$ \mathcal{L}_{\text{SimDINO}} = \frac{1}{2} \|z_c - z_g\|_2^2 - \frac{1}{2} \log \det(I + \frac{d}{\epsilon^2} \Gamma) $$

教师encoder通过EMA机制更新。

在图片检索时，将原图（非增强视图）的CLS token embedding与patch embedding的平均池化拼接，作为该图的表征。