跳转至

MedVTok: A General-Purpose Medical Visual Tokenizer

会议: NeurIPS 2026
arXiv: 2610.35842
OpenReview: https://openreview.net/forum?id=VGTDLvJS3w
代码: https://masaaki-75.github.io/medvtok
领域: 医学影像
关键词: 视觉分词器、多模态医学AI、切片差分正则化、多专家表示对齐、统一表征空间

一句话总结

MedVTok 是一个通用的多模态医学视觉分词器,通过切片差分正则化保持 3D 体积解剖连续性,结合多专家双向表示对齐注入临床语义,实现了跨 9 种成像模态与 2D/3D 维度的统一视觉 Token 空间。

研究背景与动机

多模态医学人工智能正从孤立的单一任务模型迈向通用临床系统,其核心前提是构建一个能够高效编码、迁移和复用异构临床证据的统一表征接口。在各类临床数据中,医学图像承载了极其丰富的解剖结构、病理组织形态和病灶进展信息,但成像数据在维度和物理模态上表现出极强的异构性:既包括眼底彩照(CFP)、皮肤镜、病理切片等 2D 图像,也涵盖计算机断层扫描(CT)、磁共振成像(MRI)等 3D 断层扫描体积。目前绝大多数医学视觉分词器(Visual Tokenizer)仅针对特定维度或单一模态定制,导致下游多模态模型必须自行承担繁重的模态适配与特征对齐负担。

构建跨维度与跨模态的通用医学分词器面临两大核心矛盾。一方面是跨维度的解剖结构保真度问题:传统 2D 图像仅需重建局部外观,而 3D 断层扫描具备显著的体积连续性,器官与血管在切片间平滑过渡;以往分词器采用单像素或单体素的重构损失(如 \(\ell_1\) 范数),容易产生切片间不连续或阶梯状伪影。另一方面是跨模态的临床语义捕获问题:病理图像侧重细胞级微观结构,而放射影像反映宏观器官解剖,现有分词器通常依赖单一通用医学编码器做语义对齐,容易被单一专家的特征几何所局限,抹杀特定模态独有的互补临床语义。

针对上述挑战,本文的核心设计思想是将几何维度的物理连续性约束与多模态临床语义的专家协同解耦并联合建模。核心 idea:通过无缝支持 2D/3D 任意分辨率的无网格 Patch 序列化机制,引入一阶与二阶切片差分正则化保证 3D 解剖连贯性,并设计共享空间对比与专家空间反向投影的双向多专家对齐机制,构建一个兼具像素级保真度与深层临床判别力的统一医学 Token 接口。

方法详解

整体框架

MedVTok 的整体目标是将任意维度(2D 切片或 3D 体积)和 9 种不同成像模态的医学图像映射至统一且紧凑的潜在 Token 空间。其核心处理流程包括四个阶段:维度无关的 Patch 化切片序列化、Transformer 编码器与对角块自注意力、连续潜变量重参数化与解码重构、以及多任务多层级监督。

输入张量 \(\boldsymbol{x}\in\mathbb{R}^{C\times T\times H\times W}\)(对于 2D 图像取深度 \(T=1\) 并填充至 \(P_t\),对于 3D 体积取 \(T>1\))首先被划分为不重叠的三维 Patch。所有样本的 Patch 被展平成一个无全局网格约束的长序列,并赋予记录批次与三维空间坐标的位置编码向量。由 27 层 Transformer Block 构成的编码器采用旋转位置编码(RoPE)并配合对角块注意力掩码(Block-diagonal Attention Mask),防止批次内不同样本发生信息泄漏。编码生成的 Patch Token 经由对角高斯后验重参数化为潜在编码,由对称的 Transformer 解码器还原输入数据。与此同时,Patch Token 经由均值池化头聚合为全局表征向量,与离线预计算的多个专业医学模型进行多专家双向语义对齐。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据: 2D 图像 / 3D 体积<br/>9 种医学模态"] --> B["维度无关 Patch 化<br/>坐标元组与扁平化长序列"]
    B --> C["Transformer 编码器<br/>3D RoPE 与对角块掩码自注意力"]
    C --> D["隐空间重参数化<br/>高斯分布采样与 KL 正则"]
    D --> E["切片差分正则化 SliDR<br/>一阶与二阶深度差分保真"]
    D --> F["多专家表示对齐 MERA<br/>共享空间对比 + 专家空间反投影"]
    E --> G["Transformer 解码重构"]
    F --> H["下游统一多任务 Token 接口"]

关键设计

1. 维度无关 Patch 化与序列拼接:消除长宽比与维度填充冗余

传统多模态分词器通常要求固定的网格尺寸,或通过将 3D 卷积权重池化到 2D、或复制 2D 切片堆叠至 3D,既割裂了不同维度的物理特性,又引入了密集的空白填充计算开销。MedVTok 采用基础 Patch 尺寸 \(P_t\times P_h\times P_w = 4\times 16\times 16\)(输入通道 \(C=3\),展平维度为 3072)。对于 2D 图像(\(T=1\)),在深度轴补零至 \(P_t=4\) 后切出 \(HW / 256\) 个 Patch;对于 3D 体积(深度 \(T\)),直接切出 \((T/P_t)\times(H/P_h)\times(W/P_w)\) 个 Patch。一个批次内所有样本生成的 Patch 被直接沿样本轴拼接为一个大扁平序列(总长度为 \(N\)),每个 Token 携带坐标元组 \((b, z, x, y)\),通过三维 RoPE 注入空间位置信息。自注意力层施加对角块掩码,完全避免不同样本间的 cross-attention。这一设计使单批次显存开销达到 \(\mathcal{O}(N)\) 严格线性,支持 2D/3D 以及不同分辨率在同批次内混合计算。

2. 切片差分正则化(SliDR):显式约束穿平面解剖形态连续性

在 CT 和 MRI 等断层扫描体积中,解剖组织、器官分界和病灶边缘在轴向切片之间具有高度的连续过渡特性。传统的单体素重构损失(\(\ell_1\) 或 \(\ell_2\) 范数)仅逐点独立考量,容易在体积重建中产生切片跳变和锯齿状假影。SliDR 不直接约束绝对强度,而是对输入真值与重构输出之间的层间过渡梯度施加一致性约束。对于沿深度轴 \(T>2\) 的三维体积,定义一阶切片差分 \(\Delta_k(\boldsymbol{x})=\boldsymbol{x}_{(k+1)}-\boldsymbol{x}_{(k)}\)(捕获相邻切片的演变步长)和二阶切片差分 \(\Delta^2_k(\boldsymbol{x})=\boldsymbol{x}_{(k+1)}-2\boldsymbol{x}_{(k)}+\boldsymbol{x}_{(k-1)}\)(捕获穿平面曲率与加速度)。SliDR 损失函数定义为:

\[\mathcal{L}_{\text{SliDR}}=\frac{1}{T-1}\sum_{k=1}^{T-1}\left\|\Delta_k(\hat{\boldsymbol{x}})-\Delta_k(\boldsymbol{x})\right\|_1+\frac{1}{T-2}\sum_{k=2}^{T-1}\left\|\Delta^2_k(\hat{\boldsymbol{x}})-\Delta^2_k(\boldsymbol{x})\right\|_1\]

相比无参考的 Total Variation(TV)或纯平滑拉普拉斯算子(容易导致真实解剖突变处过度模糊),SliDR 以原始体积的真实切片差分为目标,既严密保留了器官轮廓的平滑演变,又避免了纹理失真。

3. 多专家双向表示对齐(MERA):平衡跨模态共享空间与领域特异语义

若仅将分词器对齐到单一医学基础模型(如 BiomedCLIP),分词器将不可避免地继承该专家的几何偏差,使病理、眼底或内窥镜等特殊模态的精细特征被泛化语义压缩掩盖。MERA 组建了一个涵盖 10 个预训练专家模型的专家库(包括 DermLIP 皮肤镜、RetCLIP 眼底、Conch 病理、PeskaVLP 内窥镜、MedImageInsight 放射等常规影像,以及 Spectre 3D-CT 和 HLIP 3D-MR)。为防止单一投影将特异信息丢弃,MERA 采用双向约束:

首先,针对模态 \(m(i)\) 的专家特征 \(\boldsymbol{e}_i = E_{m(i)}(\boldsymbol{x}_i)\),利用线性投射层 \(p_{m(i)}\) 映射至全局统一空间 \(\tilde{\boldsymbol{e}}_i = p_{m(i)}(\boldsymbol{e}_i) \in \mathbb{R}^d\),并与 MedVTok 经均值池化得到的全局表征 \(\boldsymbol{r}_i\) 计算带可学习温度系数 \(\tau\) 的对称 InfoNCE 损失:

\[\mathcal{L}_{\text{shared}}=-\frac{1}{2B}\sum_{i=1}^B\left[\log\frac{\exp(\boldsymbol{r}_i\cdot\tilde{\boldsymbol{e}}_i/\tau)}{\sum_{j=1}^B\exp(\boldsymbol{r}_i\cdot\tilde{\boldsymbol{e}}_j/\tau)}+\log\frac{\exp(\boldsymbol{r}_i\cdot\tilde{\boldsymbol{e}}_i/\tau)}{\sum_{j=1}^B\exp(\boldsymbol{r}_j\cdot\tilde{\boldsymbol{e}}_i/\tau)}\right]\]

其次,为保留各专家原生表征空间中独特的解剖细微特征,引入反向投影器 \(q_{m(i)}\) 将统一表征映射回原生专家空间 \(\tilde{\boldsymbol{r}}_i = q_{m(i)}(\boldsymbol{r}_i) \in \mathbb{R}^{d_{m(i)}}\),并施加余弦相似度损失:

\[\mathcal{L}_{\text{expert}}=\frac{1}{B}\sum_{i=1}^B\left(1-\frac{\tilde{\boldsymbol{r}}_i\cdot\boldsymbol{e}_i}{\|\tilde{\boldsymbol{r}}_i\|\cdot\|\boldsymbol{e}_i\|}\right)\]

双向对齐总损失为 \(\mathcal{L}_{\text{MERA}}=\mathcal{L}_{\text{shared}}+\mathcal{L}_{\text{expert}}\),确保了全局统一性与各临床模态细粒度特异语义的共存。

损失函数 / 训练策略

MedVTok 的总训练目标由四部分加权组成:

\[\mathcal{L}=\mathcal{L}_{\text{R}}+\lambda_{\text{KL}}\mathcal{L}_{\text{KL}}+\lambda_{\text{SliDR}}\mathcal{L}_{\text{SliDR}}+\lambda_{\text{MERA}}\mathcal{L}_{\text{MERA}}\]

其中重构损失 \(\mathcal{L}_{\text{R}} = \ell_1 + \ell_2 + \mathcal{L}_{\text{LPIPS}} + w_{\text{adv}}\mathcal{L}_{\text{adv}}\) 包含多尺度像素重构误差、感知相似度及带自适应权重的判别器对抗损失。潜变量正则化权重设为 \(\lambda_{\text{KL}}=10^{-6}\),语义对齐权重 \(\lambda_{\text{MERA}}=1\),切片差分权重 \(\lambda_{\text{SliDR}}=0.1\)。

训练采用两阶段递进策略: - Stage 1(2D 基础预训练):在 28M 张 2D 图像上训练,禁用切片差分正则化(\(\lambda_{\text{SliDR}}=0\)),仅激活 2D 专家投射器。采用多样性采样器确保每个 mini-batch 至少包含每组模态各 2 个样本,学习率由 \(10^{-4}\) 余弦衰减至 \(10^{-5}\)(100K steps)。 - Stage 2(2D/3D 联合全模态训练):从 Stage 1 检查点初始化,激活 3D 专家(Spectre 和 HLIP)以及切片差分正则化(\(\lambda_{\text{SliDR}}=0.1\))。在偶数迭代输入有效批大小为 600 的 2D 数据,奇数迭代输入有效批大小为 48 的 3D 数据(CT 与 MRI 严格均衡),学习率由 \(2\times 10^{-5}\) 衰减至 \(2\times 10^{-6}\)(100K steps)。

实验关键数据

主实验

评估涵盖 24 个下游基准。在重构任务(Token 保真度)中评估 PSNR 与 SSIM;在分类(\(k\)-NN)与检索(Recall@1)中评估判别语义;在密集预测任务中评估冻结 Token 的线性分割 Dice(%)。下表汇总了 MedVTok 与代表性基准在 9 种 2D 模态均值及 3D 断层扫描上的核心对比。

模型类型 / 方法 2D 重构 PSNR (dB) 2D 重构 SSIM (%) 3D-CT PSNR (dB) 分类准确率 (%) 检索 R@1 (%) 2D 分割 Dice (%) 3D 分割 CT Dice (%)
通用自编码: SDVAE 33.06 82.63 – 53.66 53.47 44.56 –
通用统一: AToken 38.22 91.72 35.68 67.50 69.12 73.03 36.89
医学通用: BiomedCLIP* – – – 69.45 70.52 64.63 –
医学通用: MedImageInsight* – – – 76.21 77.34 64.38 –
医学自编码: PUMIT 27.78 82.16 31.87 52.41 52.53 53.63 28.67
医学统一: MedITok 32.91 87.20 – 73.55 74.43 71.99 –
本文: MedVTok 40.18 94.16 36.50 80.23 79.18 76.07 50.13

注: 表示专用语义编码器(不可直接重构图像)。

消融实验

在 Stage 1 预训练子集(3M 图像)及 Stage 2 3D 训练中对各模块设计进行消融验证:

实验配置编号与变更 2D 图像分类准确率 (%) 2D 线性分割 Dice (%) 3D-CT 重构 PSNR (dB) 说明
E1: 无语义对齐基准(纯自编码) 52.14 43.82 – 语义严重匮乏,分类与分割性能大幅垫底
E2: 单专家共享空间对齐(仅 MedImageInsight) 68.35 65.40 – 引入单专家语义后下游性能显著跃升
E3: 多专家共享空间对齐(多专家 \(\mathcal{L}_{\text{shared}}\)) 70.18 67.82 – 专家多样性补充了不同模态的判别边界
E4: 引入专家空间反向对齐(完整 MERA) 71.71 68.95 – 保留了原生专家空间中的细粒度互补知识
E5: 聚合头将 Attention Pooling 替换为 Mean Pooling 71.60 73.21 – 均值池化更好地保留了密集预测所需的空间分布
E6: 数据规模扩展(3M 扩增至 28M) 79.52 75.84 – 大规模异构数据极大巩固了跨模态泛化能力
E7: 进入 Stage 2 联合训练(w/o SliDR) 79.80 75.90 34.25 未约束层间差分时 3D 体积重构连续性不足
E8: 完整模型(Stage 2 加入 \(\mathcal{L}_{\text{SliDR}}\)) 80.23 76.07 36.50 切片差分使 3D 重构提升 2.25 dB 且促成全局最优

关键发现

  • 重构与语义的权衡边界:随着 \(\lambda_{\text{MERA}}\) 从 0 增至 1,下游分类准确率和分割 Dice 获得质的飞跃(分类提升近 28%);但当权重过度提高至 2 时,PSNR 发生接近 3 dB 的剧烈衰减,表明超额语义约束会压制底层的空间纹理重建,\(\lambda_{\text{MERA}}=1\) 为最优平衡点。
  • 均值池化优于注意力池化:针对全局表征聚合头,注意力池化(Attention Pooling)倾向于将整图压缩为少数显著语义点,破坏了密集 Patch 的空间连续性,导致分割 Dice 暴跌近 4.3%;改用简单的均值池化加线性投影,显著提升了密集预测能力。
  • SliDR 的解剖结构特异性:在深度正则化对比中,SliDR 明显优于全变差(TV)和拉普拉斯平滑。TV 与拉普拉斯直接对预测体做无监督平滑,抹杀了真实的断层解剖突变;而 SliDR 与真值切片差分强对齐,在 AMOS-MM 3D 分割中助力整体 Dice 达到 50.13%(大幅超越 AToken 的 36.89% 与 PUMIT 的 28.67%)。

亮点与洞察

  • 真正实现 2D/3D 原生混合批处理:基于三维坐标 RoPE 和对角块掩码自注意力,将异构医学图像全部转换为一维 Patch 序列,单批内可自由混排不同分辨率与不同深度的 2D/3D 数据,无须任何填充补齐操作。
  • MERA 的双向互补设计思路:传统特征蒸馏只做正向投影,容易使分词器沦为特征的降维容器;MERA 通过“正向共享空间拉近 + 反向专家空间还原”形成闭环,不仅把多专家的知识融于一炉,更守住了各专科模态的特异细粒度结构。
  • 零文本依赖却击败图文多模态模型:在 LLaVA-Med 视觉问答和 LaMed 放射学报告生成中,MedVTok 在预训练阶段完全不依赖图文配对语料,仅凭自监督重构与多视觉专家对齐,下游下游生成的 GREEN 分数(0.413 vs. 0.386)和 RaTEScore 均全面优于依赖数百万报告文本训练的 MedITok。

局限与展望

  • 模态广度覆盖仍有空白:当前训练集中在 9 种主流临床影像(CT、MRI、X-ray、病理、眼底等),对于脑电图(EEG)等生理信号、常规临床肉眼照片或医学教科书复合插图尚未涵盖。
  • 专家模型处于冻结状态:目前 MERA 所采用的专家编码器特征均为离线预计算,未能实现分词器与多专家的端到端端联合反向微调,计算预算充足时联合调优或能激发出更深层的跨域协同。
  • 3D 训练规模仍受限于公开体积数据:相较于 28M 级别的 2D 图像,3D 体积数据仅收集了 70K 例,长尾解剖部位与罕见病灶的三维覆盖度仍有待进一步扩充。

相关工作与启发

  • vs MedITok: MedITok 仅支持 2D 图像,且依赖单一 BiomedCLIP 进行图文对齐;MedVTok 统一了 2D/3D 原生输入,采用 10 位专科视觉专家的双向 MERA 对齐,无需文本标注即可产生更高精度的临床语义。
  • vs PUMIT / MedVAE: PUMIT 与 MedVAE 属于纯自编码分词器,缺乏深层语义引导,在分类与密集分割任务中表现匮乏;MedVTok 证明了高保真重构与判别性语义表征并非互斥,二者可在一个统一的 Transformer 架构中兼得。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出针对医学 3D 体积的切片差分正则化(SliDR)与跨模态双向多专家表征对齐(MERA)。
  • 实验充分度: ⭐⭐⭐⭐⭐ 跨 9 种模态、覆盖 24 个下游基准(重构、分类、检索、分割、生成、VQA、报告生成),包含严谨消融与跨数据集评测。
  • 写作质量: ⭐⭐⭐⭐⭐ 结构清晰,方法逻辑严密,图表对比翔实,数学表达与实验协议规范。
  • 价值: ⭐⭐⭐⭐⭐ 为通用多模态医学大模型提供了一个标准、通用且高效的视觉 Token 前端基础设施。