Large Language Models for Multimodal Deformable Image Registration
https://github.com/ninjannn/LLM-Morph
The challenge of Multimodal Deformable Image Registration (MDIR) lies in the conversion and alignment of features between images of different modalities. Generative models (GMs) cannot retain the necessary information enough from the source modality to the target one, while non-GMs struggle to align features across these two modalities.
In this paper, we propose a novel coarse-to-fine MDIR framework, LLM-Morph, which is applicable to various pre-trained Large Language Models (LLMs) to solve these concerns by aligning the deep features from different modal medical images. Specifically, we first utilize a CNN encoder to extract deep visual features from cross-modal image pairs, then use the first adapter to adjust these tokens, and use LoRA in pre-trained LLMs to fine-tune their weights, both aimed at eliminating the domain gap between the pre-trained LLMs and the MDIR task. Third, for the alignment of tokens, we utilize other four adapters to transform the LLM-encoded tokens into multi-scale visual features, generating multi-scale deformation fields and facilitating the coarse-to-fine MDIR task.
Extensive experiments in MR-CT Abdomen and SR-Reg Brain datasets demonstrate the effectiveness of our framework and the potential of pre-trained LLMs for MDIR task.
多模态可变形图像配准(MDIR)的挑战在于不同模态图像之间特征的转换和对齐。生成模型(GMs)无法从源模态保留足够的信息到目标模态,而非生成模型(non-GMs)则难以跨这两种模态对齐特征。
在本文中,本文提出了一种新颖的粗到细MDIR框架LLM-Morph,该框架适用于各种预训练的大型语言模型(LLMs),通过对齐来自不同模态医学图像的深度特征来解决这些问题。具体而言,本文首先利用卷积神经网络(CNN)编码器从跨模态图像对中提取深度视觉特征,然后使用第一个适配器调整这些标记,并在预训练的LLMs中使用LoRA微调其权重,两者都旨在消除预训练LLMs与MDIR任务之间的领域差距。第三,为了对齐标记,本文利用其他四个适配器将LLM编码的标记转换为多尺度视觉特征,生成多尺度变形场,并促进粗到细的MDIR任务。
在MR-CT腹部和SR-Reg脑部数据集上进行的大量实验证明了本文的框架的有效性以及预训练LLMs在MDIR任务中的潜力。
Introduction
多模态可变形图像配准(Multimodal Deformable Image Registration,MDIR)将来自不同成像模式的两幅图像中的类似解剖结构进行对齐,从而提供互补的解剖和功能见解。这一能力对于各种应用至关重要,包括精确的肿瘤定位、器官移植评估以及解剖结构的分割和功能分析。
传统的MDIR方法使用迭代优化算法进行图像对齐,但往往面临耗时问题和局部最优解。最近,基于深度学习的方法,包括卷积神经网络(CNNs)和Transformer,显著提高了计算效率并减少了计算时间,从而促进了MDIR领域的发展。一些MDIR工作使用生成模型(GMs),如SymReg-GAN、DualStream-GAN和TarGAN,通过单模态相似性度量在不同模态之间转换图像,从而避免了多模态相似性度量的挑战。然而,当使用GMs将源信息转换为目标时,关键的源特定信息可能在转换过程中丢失,这可能会损害扭曲图像的质量。
相比之下,非生成模型(non-GMs)用于预测明确的变形场而不进行模态转换。然而,大多数non-GMs预测的是单尺度变形场,无法处理大变形预测。此外,non-GMs采用基于信息的相似性度量,如互信息和MIND,来学习不同模态之间体素的对应关系。尽管这些相似性度量可以独立于模态特定信息进行测量,但基于深度学习的模型在跨不同模态统一特征方面仍面临挑战。因此,它们有限的表示能力更适合于单模态而非多模态数据,这限制了MDIR任务的性能。
最近,大型语言模型(LLMs)因其丰富的语料库知识和多任务能力而广受关注。在语言-视觉任务中,LLMs能够理解文本信息并将其与视觉数据相关联,展示了其在语言和视觉模态之间理解和对齐特征的能力。最近的研究表明,采用轻量级微调策略,如Adapter和低秩适应(LoRA),可以提高预训练LLMs在不同模态任务中的通用性。这些策略显著提高了LLMs的适应性。LLMs出色的跨模态能力为解决GMs和non-GMs在MDIR任务中面临的挑战提供了新的视角。
为了解决生成模型(GMs)在采样过程中导致的体素丢失问题,以及非生成模型(non-GMs)在统一不同模态特征时面临的困难,本文提出了一种基于预训练大型语言模型(LLMs)的由粗到细的多模态可变形图像配准(MDIR)方法,称为LLM-Morph。
通过利用预训练的大型语言模型作为中间模态,本文旨在通过预训练的大型语言模型实现两个模态图像在同一语义层面的深度特征对齐。然而,目前大多数大型语言模型主要集中在自然语言处理或视觉-语言任务上,与MDIR任务存在领域差异,阻碍了大型语言模型直接执行MDIR的能力。
为了解决这个问题,本文构建了LLM-Morph,它包含一个可训练的基于卷积神经网络(CNN)的编码器、五个适配器以及两个带有预训练大型语言模型层的模块。首先,CNN编码器从一对不同模态的图像中提取深层特征。接下来,在大型语言模型外部,这些特征通过第一个适配器被标记化并调整到大型语言模型所需的维度,通过投影这些标记来解决领域差异。为了进一步缩小这一领域差异并增强特征对齐,本文在大型语言模型内部使用LoRA来微调预训练权重。在解码阶段,本文使用四个适配器将大型语言模型对齐的多模态标记同时映射到每个阶段所需的维度,然后将这些标记恢复到每个分辨率级别的视觉特征。这些视觉特征直接用于生成变形场,从而实现从粗到细的多尺度MDIR。
此外,本文还研究了不同层和各种大型语言模型对该任务特异性和性能的影响。据本文所知,这是首次将大型语言模型应用于MDIR任务的工作。
本文的主要贡献如下:
• 本文开发了一种新颖的多尺度配准框架LLM-Morph,利用大型语言模型来促进多模态图像特征的对齐。
• 本文分别使用适配器和LoRA来调整大型语言模型外部和内部的特征和预训练权重,以消除预训练大型语言模型与MDIR之间的领域差异。
• 本文进行了广泛的实验,包括测试LLM-Morph中不同预训练层的大型语言模型的性能。此外,本文还测试了一些预训练的大型语言模型对MDIR性能的影响。
Method
MDIR的优化问题表述
MDIR(多模态可变形图像配准)可以基于一对不同模态图像之间的相似性度量来构造成一个优化问题。这两张图像都在三维空间域内定义。主要的优化目标是确定一个最优的变形场,该变形场能够最大限度地使两张图像对齐。这一优化问题可以表述为:
其中,I_m 和 I_f 分别代表输入的运动图像和固定图像,ϕ 表示变形场,它描述了从 I_m 到 I_f 的空间像素变换的幅度和方向。∘ 代表插值操作,I_m∘ϕ 表示通过变形场 ϕ 对运动图像进行变换后的图像。本文根据等式(1)以半监督的方式训练本文的模型。具体来说,本文采用Dice损失作为 Esim(⋅,⋅),用于计算经过变换的分割图 S_m∘ϕ 和固定图像的分割图 S_f 之间的相似性。R(⋅) 是L2正则化项,用于增强变形场的平滑性;而 λ 是一个超参数,用于平衡相似性损失和正则化项的作用。因此,本工作的半监督训练损失函数为 L=E_sim(S_m∘ϕ,S_f)+λR(ϕ)。
LLM-Morph的架构
本文提出的LLM-Morph的总体架构如图1所示。LLM-Morph由基于CNN的特征提取器、两个LLM编码块(LEBs)和四个解码分支组成。与传统非生成模型架构不同,本文在编码器底部引入了两个LEB作为中间模态,以对齐两种不同视觉模态下的深层特征。每个LEB包含一个适配器Adapter0,它由线性投影组成,用于将这些特征调整到与预训练LLM输入维度相匹配的维度,从而使预训练的LLM能够更好地理解和发现这些特征之间的潜在对应关系。在解码阶段,本文使用四个适配器模块以分层方式逐步处理和转换特征,并将LLM对齐的深层特征调整到每个解码阶段所需的特征维度。
在编码阶段,本文采用DualPyramid方法(Kang等人,2022)对运动图像 Im 和固定图像 If 进行编码,并生成五对不同尺度的特征对。这些特征分别表示为 和 Fif
,其中 F_im 和 F_if 是提取的特征,并且按照因子
进行下采样,i 的取值范围从1到5,代表从全分辨率阶段到1/16分辨率阶段的特征。为了使视觉特征满足预训练LLM的输入要求,本文将上述特征进行标记化(tokenization),即将视觉特征展平为固定长度 L=d×h×w 的标记,其中 d、h、w 分别是每个图块的深度、高度和宽度。
解码阶段
解码阶段由四个子阶段组成:S4、S3、S2和S1,分别对应从1/8分辨率到全分辨率阶段的配准过程。在每个阶段,都使用一个适配器来将已对齐的深层特征转换为能够产生多尺度变形的特征。具体来说,每个阶段的适配器负责将上一个LEB输出的特征映射到该阶段所需的特征维度,从而逐步从粗到细地恢复特征,允许对S3、S2和S1阶段生成的变形场进行精确调整。
Adapter_i 的输出形状为(B, L, C_j),其中B是批量大小,L是标记的长度。在本工作中,B和L是固定的。C_j代表Adapter_i的输出维度,其中j ∈ {4, 3, 2, 1}。从S4到S1,每个阶段Adapter_j的输出通道数C_j是前一阶段的4倍(j: 4 → 1)。
这种设置确保了在每个解码阶段,适配器都能有效地将已对齐的标记恢复为3D视觉特征,使它们的大小与 F_im 和 F_if 相匹配。
最后,通过四个阶段的适配器处理,得到了相应的1/8分辨率F^4、1/4分辨率F_3、1/2分辨率F^2和全分辨率F^1,以供后续处理。
在S4阶段,将F_4、F^4_m和F^4_f沿通道维度进行拼接,并通过一个卷积(卷积核大小为3,步长为1)生成融合特征F'_4。该融合特征通过上采样提高分辨率,并传递到S3阶段,其中通过融合特征F'_3生成初始变形场ϕ_3,并将其分辨率加倍以调整半分辨率下的运动图像I^2_m。变形后的图像I^2'_m,以及ϕ^_3和F'_3,被用于提取下一个变形场。
在S2阶段,由F'_2生成的变形场ϕ_2与ϕc_3相结合,重复S3阶段的扭曲过程。S1阶段遵循与S2阶段相同的步骤。最后,生成全分辨率的ϕ用于空间扭曲。随后,使用空间变换网络(STN)通过变形场ϕ对运动图像的分割图进行扭曲。这使得能够计算损失函数L,从而指导LLM-Morph学习其权重。
LLM编码块
LLM-Morph的关键组件是提出的LEB(LLM编码块),如图1所示,每个LEB都嵌入了一个预训练的LLM(大型语言模型)。提出的LEB通过构建两个连续LEB之间的相关性来编码标记化特征并对不同模态进行对齐。
在被预训练的LLM编码之前,视觉特征会被调整以满足LLM的输入要求。首先,图像特征被转换为称为“标记化深层特征”的标记格式,然后通过一个专门设计的适配器进行处理。
LEB中的这些适配器被统一定义为Adapter_0。每个Adapter0由两个配备有实例归一化和LeakyReLU的线性投影组成。
Adapter_0的目的是将标记化特征映射到LLM所需的通道数,并将特征与所使用LLM的编码空间对齐。之后,每个LLM都会进行微调,以最大化预训练权重对MDIR(多模态图像配准)任务的适应性。由此,构成了编码标记化特征的LEB。
LoRA微调
低秩适应(LoRA)在自注意力机制的查询(Q)、键(K)和值(V)组件中引入了小而低秩的矩阵,从而能够在不显著增加参数的情况下对模型进行调整。它可以表示为
其中,W ∈ R^d×k 是从预训练模型中获得的Q、K、V权重矩阵,d是输出特征的维度,k是输入特征的维度,x是输入标记,Wr = BA是权重矩阵的低秩更新,其中B ∈ R^d×r,A ∈ R^r×k,且秩r远小于d和k的最小值。
一般来说,W继承自预训练模型,这些权重不会接收梯度更新,以保持模型从大规模数据中学习的知识,而W_r = BA,即矩阵B和A在微调过程中会接收梯度更新,从而使模型能够适应新任务的需求,同时不会显著增加参数数量。
本文使用LoRA微调方法来更新和微调每个LEB块内LLM预训练层中的Q、K和V矩阵。具体来说,本文添加一个低秩权重Wr来实现LLM对于跨模态配准任务的微调。这种方法在解析和对齐来自各种医学成像模态的图像时特别有效,从而提高了MDIR的性能。