T-Mamba: Frequency-Enhanced Gated Long-Range Dependendcy for Tooth 3D CBCT Segmentation

摘要

三维成像中的高效牙齿分割对于正畸诊断至关重要，但由于CBCT图像中的噪声、低对比度和伪影，这仍然是一个挑战。

卷积神经网络（CNNs）和变压器（transformers）都已成为图像分割中流行的架构。然而，由于内在的局部性或计算复杂性，它们在处理长距离依赖方面的有效性受到限制。

为了解决这个问题，我们提出了T-Mamba，将共享位置编码和基于频率的特征整合到视觉Mamba中，以解决空间位置保留和频率域特征增强的限制。

此外，我们还设计了一个门选择单元，以自适应地整合空间域中的两个特征和频率域中的一个特征。T-Mamba是首次将基于频率的特征引入视觉Mamba的工作。

大量实验证明，T-Mamba在公共牙齿CBCT数据集上取得了新的SOTA结果，并且相较于之前的方法在较大幅度上取得了性能提升，即 IoU + 3.63%，SO + 2.43%，DSC + 2.30%，HD - 4.39mm，ASSD - 0.37mm。

代码和模型：https://github.com/isbrycee/T-Mamba

Introduction

现代数字牙科的进化关键在于三维（3D）成像的获取与分割。这项技术在口腔及颌面学科领域有多种用途，包括正畸诊断和治疗规划。牙齿分割，即从3D数字模型中提取具有强度和密度的一组体素，是数字工作流程中的必要阶段。然而，精确的牙齿分割过程由于多种原因而具有挑战性。首先，由于3D数字模型固有的高噪声和低对比度特性，难以观察。此外，CBCT图像中金属充填物和修复体的存在引入了导致畸变的伪影，显著增加了识别牙齿的难度。最后，CBCT通常在自然咬合状态下获取，由于上下牙齿密度相似，区分它们也构成了挑战。

深度学习在医学图像分割领域得到了广泛应用，许多研究者致力于利用深度学习技术实现CBCT图像中牙齿的自动分割，如Chen等人（2023年）；Cui等人（2019年）；Cui等人（2021年）；Yin等人（2023年）。卷积神经网络（CNNs）和Transformers架构因其能够学习复杂的图像特征并提供准确的分割结果，在医学图像分割方面显示出巨大的潜力。CNNs能够捕获平移不变性并提取局部特征，而Transformers擅长捕捉全局上下文信息并提高长距离依赖的能力。由于这种互补特性，许多研究探索了通过混合网络架构将Transformers融入CNNs中。然而，Transformers的明显缺点是资源密集型，因为自注意力机制与输入大小成二次方增长，并且在处理高分辨率生物医学图像时在速度和内存使用方面构成挑战。尽管大量努力致力于降低Transformers的计算复杂性，但这往往以牺牲模型准确性为代价，如Maaz等人（2022年）；Zhai等人（2023年）。因此，在CNNs中有效增强长距离依赖仍然是一个未解决的问题。

近年来，受到Mamba在语言建模领域的成功启发，许多研究将这种成功从语言领域转移到视觉领域，旨在实现线性复杂度同时不牺牲全局感受野 Liu等人（2024）；Zhu等人（2024）。然而，我们发现在利用视觉Mamba帮助卷积神经网络（CNN）建模长距离依赖时存在两个局限性。首先，CNN可以直接处理2D或3D特征，而视觉Mamba的结构旨在处理1D特征序列。在CNN和视觉Mamba混合架构中，视觉Mamba的特征转换不可避免地导致空间位置信息的丢失。这种空间上下文在需要密集精确位置预测的任务中尤为重要。其次，由于CT和X射线等医学图像的固有成像原理，这些图像从视觉角度具有高噪声和低对比度的自然属性。对于这类图像，基于频域的特征表示更准确、独特和鲁棒。然而，CNN和视觉Mamba模型通常仅从空域提取语义特征，忽略了丰富的基于频域的信息。正如Azad等人（2021）所指出的，高频成分捕捉纹理细节，而低频成分编码形状信息。因此，将频域特征与空域表示相结合，有望增强医学图像中的图像特征提取，从而提高预测的准确性。

受到这两个局限性的启发，我们设计了一个名为T-Mamba的网络，它将我们提出的Tim（牙科视觉Mamba）块与DenseVNet Gibson等人（2018）的多尺度特征相结合，用于牙科CBCT图像分割。Tim块展现出三个关键优势：（1）它包含一个共享的双位置嵌入，用于补偿在重塑操作期间丢失的位置信息。值得注意的是，我们在每个特征尺度内使用单一位置嵌入。这不仅在相同尺度的特征图之间保证了空间位置的保留，还减少了模型参数和计算负担。（2）它提取频域中的图像特征，因此我们可以为具有高噪声和低对比度的医学图像导出更准确、独特和鲁棒的特征表示。基于不同尺度特征的性质，我们为每个尺度定制了不同的带通滤波策略。（3）它包括一个门选择单元，用于自适应地整合空域中的两个特征（正向和反向）和频域中的一个特征。

门选择单元是数据依赖的，可以根据输入序列标记为三个不同的特征分配权重。

在牙齿CBCT分割任务上进行了大量实验，以验证T-Mamba的有效性。我们的T-Mamba在公共牙齿CBCT数据集上的表现大幅超越了先前最先进的结果，即IoU提高了3.63%，SO提高了2.43%，DSC提高了2.30%，HD降低了4.39毫米，ASSD降低了0.37毫米。此外，我们还进行了消融研究，以验证我们提出的Tim块中三个组件的有效性。据我们所知，T-Mamba是首次将频域特征融入到视觉Mamba框架中的开创性工作。

主要贡献可以总结如下：

1. 我们提出了T-Mamba，它将我们设计的Tim块与DenseVNet结合，用于牙齿CBCT分割的全局和局部视觉上下文建模。
2. Tim块是首次通过引入基于频率的特征，为高噪声和低对比度的医学图像提取更鲁棒和独特的表示。
3. 在没有额外修饰的情况下，T-Mamba在公共牙齿CBCT数据集上取得了新的SOTA结果，并且以较大的幅度超过了先前SOTA方法，如IoU提高了3.63%，SO提高了2.43%，DSC提高了2.30%，HD降低了4.39毫米，ASSD降低了0.37毫米。

方法

3.1 实验设计
实验设计基于随机完全区组设计（RCBD），设有三个重复。

Preliminaries for Mamba
高级状态空间模型（SSM），即结构化状态空间序列模型（S4）和Mamba，是一类系统，它们通过隐藏状态 $h(t)\in {\mathbb{R}}^N$ 将一维连续函数或序列 $x(t)\in \mathbb{R}\mapsto y(t)\in \mathbb{R}$ 进行映射。从数学上讲，这些模型通常被表述为线性常微分方程（ODEs），如公式（1）所示：

$ \begin{split} h’(t) &= \mathrm{A}h(t) + \mathrm{B}x(t), \ y(t) &= \mathrm{C}h(t), \end{split} \tag{1}$

其中参数包括 $\mathcal{A}\in {\mathbb{R}}^{N\times N}$ 作为演化参数，以及 $\mathrm{B}\in {\mathbb{R}}^{N\times 1}$、$\mathrm{C}\in {\mathbb{R}}^{1\times N}$ 作为投影参数。

基于SSM的模型作为连续时间模型，在整合到深度学习算法中时应当进行离散化。这种离散化转换对于将模型与输入数据中蕴含的底层信号的采样率对齐至关重要[1]。给定输入 $x(t)\in {\mathbb{R}}^{L\times D}$，一个在信号流程中的长度为 $L$ 的采样向量[2]，公式（1）可以使用零阶保持（ZOH）规则离散化为如下形式：

$ \begin{split} h_t &= \overline{\mathbf{A}}h_{t-1} + \overline{\mathbf{B}}x_t \ y_t &= \overline{\mathbf{C}}h_t \ \overline{\mathbf{A}} &= e^{\Delta\mathbf{A}} \ \overline{\mathbf{B}} &= \Delta\mathbf{A}^{-1}(e{\Delta\mathbf{A}} - \mathbf{I}) \cdot \Delta\mathbf{B} \ \overline{\mathbf{C}} &= \mathbf{C}, \end{split} \tag{2}$

其中 $\Delta \in {\mathbb{R}}^D$ 是时间尺度参数。

最终，模型通过在结构化卷积核 $\overline{\mathbf{K}}$ 内的全局卷积操作计算输出 $y$：

其中，$\overline{\mathbf{K}}$是一个包含$L$个元素的向量，每个元素由$\mathbf{C}$和不同次幂的$\overline{\mathbf{A}}$与$\overline{\mathbf{B}}$的乘积组成。$y$是向量$\overline{\mathbf{K}}$与向量$y_t$的Kronecker积。

T-Mamba architecture

我们通过利用Mamba的线性扩展优势，增强了卷积神经网络（CNN）的表征建模能力，并提出了一种用于图像分割的通用网络T-Mamba。T-Mamba的网络架构如图1所示，它将我们提出的Tim模块与DenseVNet在多尺度特征上进行了融合。T-Mamba包含三个特征尺度，并采用了经典的单阶段V型架构。我们在DenseVNet的每个CNN层后简单地插入Tim模块，以建模长距离依赖。T-Mamba能够在空间域和频率域捕获局部的细粒度特征和长距离依赖。我们希望它能为生物医学图像分析中高效的长距离依赖建模开辟新的途径。

Tim block

原始的Mamba块是为1-D序列设计的，这不适用于需要空间感知理解的任务。Vision Mamba Zhu等人，2024年提出的Vim块在Mamba块上集成了双向序列建模。基于这个强大的设计，我们通过在Vim块中引入三个组件来进一步增强视觉特征建模，这对于2D和3D牙齿分割任务至关重要：1) 共享双向位置编码补偿；2) 基于频率的带通滤波；3) 门控选择单元。我们的牙齿视觉Mamba（Tim）块如图1所示。

共享双向位置编码补偿。 我们的T-Mamba网络将卷积层的局部特征提取能力与SSMs捕捉长距离依赖的能力相结合。卷积层通常处理2-D或3-D特征图，而我们设计的Tim块专注于处理1-D序列。因此，将高维特征重塑为1-D特征标记是必要的。然而，这一过程不可避免地会导致关键位置信息的丢失，这对于密集预测任务是至关重要的。为了减轻这一点，我们采用共享位置嵌入来补偿在重塑操作中丢失的位置信息。

具体来说，对于一个形状为（B, C, H, W, D）的输入特征，我们首先将其展平为1-D特征标记，其形状为（B, L, C），其中L = H × W × D，然后向特征标记中添加一个可学习的位置嵌入，其形状为（C, L），以保留位置信息。遵循[transformer is all u need]，我们通过正弦函数初始化位置嵌入：

KaTeX parse error: {split} can be used only in display mode.

其中pos是沿L的位置，i是沿C的索引。关于Tim块的输出，我们还需要将1-D特征标记重塑为高维特征，以供下一次卷积操作使用。为了进一步强化1-D特征标记中的空间信息，在将其重塑回高维特征之前，之前使用过的位置嵌入被再次添加到1-D特征标记中。请注意，在每个特征尺度上我们只使用单一的位置嵌入。这种做法确保了在同一尺度特征图中空间位置保持不变，同时减少了模型参数和计算负担。我们认为，在Tim块的输入和输出中添加共享双向位置嵌入，可以显著保留高维特征的位置信息。通过消融实验，我们验证了使用共享双向位置编码相比于使用单一位置嵌入能带来更高的性能。

基于频率的带通滤波。傅里叶域在提取基于频率的图像信息分析中起着主要作用，很明显，即使是在质量较差的X射线和CT图像中，也可以在频率域中捕捉到细微的细节和轮廓。Said等人，2004；余毅等人，2021年和Azad等人，2021；李等人，2024年。卷积层具有强烈的纹理诱导偏差，它们倾向于学习基于纹理的特征。以频率的形式表示对象可以减少纹理偏差的影响，因为只有高频部分负责纹理信息（如边界），而低频部分可能与形状有关。受到这一点的启发，我们通过增强频率域中的特征提取来演进Vim的结构。具体来说，我们首先将序列标记X转换到傅里叶域，然后使用可学习的权重参数提取频率特征，并实施带通滤波，最后进行逆傅里叶变换以恢复信号。之后，通过带有最大池化操作的激活$Z^{\prime }$来聚合频率特征。整个过程可以表述为：

$\begin{split}\mathcal{F}{freq}=IFFT(Bandpass(W{f}(FFT(X))))* Maxpool(Z^{\prime}),\ Bandpass=\begin{cases}&X*(|X|<s_{low}),x\in low-level=“” features,\=“” &x*(s_{low}<|x|

其中$S_{low}$、$S_{high}$是带通滤波的阈值。在实验中，我们设置$S_{low}$=0.1和$S_{high}$=0.9。低级别、中级别、高级别分别表示我们网络中的三种不同特征尺度。

门选择单元。演进的Vim块包括空间域中的两个特征（前后方向）和频率域中的一个特征。我们设计了一个门选择单元，旨在自适应地融合这些独特的特征。输入嵌入序列首先下采样到一个固定维度，如2048，并通过MLP（多线性投影）进行投影，并通过一个全连接层预测与三个特征相对应的三个比例。之后，通过三个特征的加权和得到$f_{fuse}$，然后通过一个线性层进行投影，门选择单元输出$f_{fuse}$的总和、残差信息以及共享位置嵌入。

KaTeX parse error: {split} can be used only in display mode.

请注意，内嵌公式和单独成行的公式未翻译，并且内嵌公式与文本之间保留了空格。

其中，shared pos 是前面提到的共享位置嵌入。门选择单元（Gate Selection Unit）是数据依赖的，因为这三个权重系数是从源数据 X 计算得出的，然后这些权重系数被用来更新 X 的三种不同形式的特征。因此，门选择单元可以自适应地根据不同的输入调整这三种形式特征的组合，从而获得更好的特征表示。