背景:
皮质表面重构(cortical surface reconstruction)是指从脑磁共振成像(MRI)中提取大脑皮层内部(白质)和外部(颅内膜)表面的3D网格。这些表面不仅提供了3D可视化,还提供了皮层的形态学特征的估计。传统的神经像素分析管道(如FreeSurfer)通过一系列的处理步骤从MRI中提取皮质表面,但是它们的准确性有限且需要几个小时来处理一个单独的MRI扫描。随着几何深度学习的最新进展,许多快速的基于深度学习的方法被提出用于学习隐式表面表示或显式网格变形以实现皮质表面重构。这些方法提高了准确性并将处理时间减少到几秒钟。然而,这些方法大多依赖成人MRI作为训练数据,而新生儿脑部MRI则具有更低的分辨率和对比度,且新生儿脑部表面在不同的孕周年龄(PMA)之间具有显著的形状和尺度变化。此外,新生儿头部较小,因此新生儿的皮质皱襞比成人更窄,更容易受到部分体积效应的影响并产生表面自相交。
图1。不同经后年龄的新生儿皮质表面。
图2。左图:来自ADNI数据集[20]的成人大脑MRI。右图:来自dHCP数据集[9]的新生儿脑MRI。
解决方案:
本工作提出了条件时间注意网络(Conditional Temporal Attention Network,CoTAN)。CoTAN采用注意力机制学习条件时变速度场(CTVF)以实现新生儿皮质表面重构。给定一个输入的脑部MRI,CoTAN首先预测不同分辨率的多个稳定速度场(SVFs)。与将所有SVFs集成在一起的方法不同,CoTAN学习条件时间注意力映射以专注于不同时间步长和新生儿孕周年龄的特定SVFs。CTVF由学习到的SVFs的加权和表示,因此单个CoTAN模型足以模拟新生儿皮质表面的大型变形和变化。在dHCP新生儿数据集上的评估结果显示,与最新方法相比,CoTAN在几何准确性、网格质量和计算效率方面表现更好。注意力映射的可视化显示,CoTAN可以在没有中间约束的情况下自动学习粗到细的变形。
Methodology
**微分同胚表面变形**
定义微分同胚表面变形 \(\phi_t : \mathbb{R}^3 \times \mathbb{R} \rightarrow \mathbb{R}^3\) 作为一个流常微分方程(ODE):
\[
\frac{\partial \phi_t}{\partial t} = v_t(\phi_t), \quad \phi_0 = \text{Id}, \quad t \in [0, T],
\]
其中 \(v_t\) 是一个随时间变化的速度场(TVF),Id 是恒等映射。给定一个初始表面 \(S_0 \subset \mathbb{R}^3\),定义 \(x_t := \phi_t(x_0)\) 作为在可变形表面 \(S_t = \phi_t(S_0)\) 上点的轨迹,随着时间 \(t \in [0, T]\) 的变化。然后,流方程可以重写为 \(\frac{d}{dt}x_t = v_t(x_t)\),具有初始值 \(x_0\)。根据ODE解的存在性和唯一性定理,如果 \(v_t(x)\) 关于 \(x\) 是Lipschitz连续的,那么轨迹 \(x_t\) 将不会相交,从而有效地防止表面自交叉。通过积分ODE,得到一个微分同胚 \(\phi_T\),将 \(S_0\) 变形为一个流形表面 \(S_T\),上面的点为 \(x_T = \phi_T(x_0) = x_0 + \int_0^T v_t(x_t) \, dt\)。
图3。所提出的CoTAN框架的体系结构。给定一个输入的3D大脑MRI,CoTAN使用U-Net来预测每个分辨率水平r的M SVFs。学习注意力图来关注随输入整合时间t变化的特定SVFs,并以新生儿受试者的年龄a为条件。进一步的条件反射可以用最小的努力来实现,例如,生理性别、诊断等。对于每个时间步长t,CTVF vt用所有R×Msvf的加权和表示。通过整合CTVF,CoTAN将输入的初始表面变形为预测的皮层表面。
图4。从初始模板网格到皮质表面的微分变形。
**条件时间注意力网络(CoTAN)**
CoTAN的架构概览如图3所示。CoTAN首先从一个3D脑部MRI体积预测多个SVF。使用一个3D U-Net来提取具有 R 分辨率级别的特征图,每个级别将输入大小缩小 \(2^{r-R}\) 倍,其中 \(r = 1, ..., R\)。然后对多尺度特征图进行上采样,并学习每个分辨率的 M 个体积SVF。设 V 表示所有 R × M 离散SVF。连续的多分辨率SVF \(v : \mathbb{R}^3 \rightarrow \mathbb{R}^{R \times M \times 3}\) 可以通过 \(v(x) = \text{Lerp}(x, V)\) 获得,其中 \(\text{Lerp}(\cdot)\) 是三线性插值函数。每个元素 \(v_{r,m} : \mathbb{R}^3 \rightarrow \mathbb{R}^3\) 是一个分辨率 r 和 m 的SVF。注意 \(v(x)\) 是Lipschitz连续的,因为 \(\text{Lerp}(\cdot)\) 是连续的和分段线性的。
CoTAN采用通道式注意力机制集中特定SVF,因为整合所有 R × M SVF耗时较多。在输入时间 \(t\) 和新生儿年龄 \(a\) 的条件下,学习一个注意力图来聚焦特定SVF。进一步的条件设定可以轻松实现,例如生物性别、诊断等。对于每个时间步长 \(t\),CTVF \(v_t\) 由所有 R × M SVF 的加权和表示。通过整合CTVF,CoTAN将输入的初始表面变形为预测的皮层表面。注意力不仅取决于时间 \(t \in [0, T]\),还取决于受试者的信息。为了模拟婴儿大脑之间的高变异性,将扫描时新生儿的产后月龄 \(a \in \mathbb{R}\) 作为这项工作的条件变量。注意,不使用自注意力模块
来学习关键和查询对。相反,学习一个概率注意力图来衡量每个SVF的重要性。更准确地说,如图3所示,使用一个全连接网络(FCN)来编码输入时间 \(t\) 和 PMA \(a\) 成一个 \(R \cdot M\) 维的特征向量。经过重塑和softmax激活后,FCN 学习条件时间注意力图 \(p(t, a) \in \mathbb{R}^{R \times M}\),满足 \(\sum_{r=1}^R \sum_{m=1}^M p_{r,m}(t, a) = 1\) 对任何 \(t\) 和 \(a\)。然后,通过计算所有SVF的加权和来预测条件时变速度场(CTVF):
\[
v_t(x; a) = \sum_{r=1}^R \sum_{m=1}^M p_{r,m}(t, a) \cdot v_{r,m}(x).
\]
CTVF 对整合时间和受试者年龄具有适应性,能够处理新生儿皮层表面的大变形和变异。这种注意力机制鼓励 CoTAN 学习从粗到细的表面变形,通过关注不同分辨率的 SVF。
为了将初始表面 \(S_0\) 变形为目标表面,通过前向欧拉方法整合流 ODE(1)与CTVF。对于 \(k = 0, ..., K-1\),表面点通过 \(x_{k+1} = x_k + h v_k(x_k; a)\) 更新,其中 \(K\) 是总的整合步骤数,\(h = T/K\) 是步长,\(T = 1\)。对于每一步 \(k\),只需要重新计算注意力图 \(p(hk, a)\) 并相应地通过方程(2)更新 CTVF \(v_k(x_k; a)\)。CoTAN 只整合单个CTVF,与直接整合多个SVF相比,节省了大量运行时间。
**新生儿皮层表面重建**
在dHCP新生儿数据集上训练了两个CoTAN模型:1) 将初始表面变形为白质表面;2) 将白质表面扩展为软脑膜表面,如图4所示。对所有受试者使用相同的初始表面(图4中最左边的),该表面是通过对Conte-69表面图集进行迭代拉普拉斯平滑创建的。通过dHCP结构新生儿管道生成伪地面真实(GT)表面,该管道已通过临床专家的质量控制进行了完全验证。
对于白质表面重建,考虑了在以前的工作中广泛使用的损失函数:Chamfer 距离损失 \(L_{cd}\) 计算两个点云之间的距离,网格拉普拉斯损失 \(L_{lap}\) 规范化网格的平滑度,法线一致性损失 \(L_{nc}\) 限制两个相邻面之间法线的余弦相似性。最终损失由 \(L = L_{cd} + \lambda_{lap}L_{lap} + \lambda_{nc}L_{nc}\) 加权。首先使用相对较大的权重 \(\lambda_{lap}\) 和 \(\lambda_{nc}\) 对模型进行预训练,以进行正则化。Chamfer距离是在预测和伪GT表面的顶点之间计算的。这些确保了在训练期间初始表面可以稳健地变形。然后,使用较小的权重对CoTAN进行微调,以
来学习关键和查询对。相反,学习一个概率注意力图来衡量每个SVF的重要性。更准确地说,如图3所示,使用一个全连接网络(FCN)来编码输入时间 \(t\) 和 PMA \(a\) 成一个 \(R \cdot M\) 维的特征向量。经过重塑和softmax激活后,FCN 学习条件时间注意力图 \(p(t, a) \in \mathbb{R}^{R \times M}\),满足 \(\sum_{r=1}^R \sum_{m=1}^M p_{r,m}(t, a) = 1\) 对任何 \(t\) 和 \(a\)。然后,通过计算所有SVF的加权和来预测条件时变速度场(CTVF):
\[
v_t(x; a) = \sum_{r=1}^R \sum_{m=1}^M p_{r,m}(t, a) \cdot v_{r,m}(x).
\]
CTVF 对整合时间和受试者年龄具有适应性,能够处理新生儿皮层表面的大变形和变异。这种注意力机制鼓励 CoTAN 学习从粗到细的表面变形,通过关注不同分辨率的 SVF。
为了将初始表面 \(S_0\) 变形为目标表面,通过前向欧拉方法整合流 ODE(1)与CTVF。对于 \(k = 0, ..., K-1\),表面点通过 \(x_{k+1} = x_k + h v_k(x_k; a)\) 更新,其中 \(K\) 是总的整合步骤数,\(h = T/K\) 是步长,\(T = 1\)。对于每一步 \(k\),只需要重新计算注意力图 \(p(hk, a)\) 并相应地通过方程(2)更新 CTVF \(v_k(x_k; a)\)。CoTAN 只整合单个CTVF,与直接整合多个SVF相比,节省了大量运行时间。
图5。由不同年龄的dHCP结构管道[24]和CoTAN生成的新生儿皮质表面的可视化。CoTAN在表面网格和相应的大脑MRI图像中都显示出更好的解剖准确性。
实验部分主要介绍了CoTAN在dHCP新生儿数据集上的实验结果。首先介绍了实验细节,包括数据集的分割、模型的参数设置和评估指标。然后给出了与其他方法比较的实验结果,包括几何准确性和网格质量的比较。最后,给出了模型的运行时间、GPU内存消耗和参数数量等信息。实验结果显示,CoTAN比其他方法在几何准确性和网格质量方面有显著提高,同时运行时间更快、内存消耗更低。在几何准确性方面,CoTAN的ASSD和HD90指标都比其他方法有显著改善;在网格质量方面,CoTAN产生的自交叉面比其他方法少得多,且产生的自交叉面主要是由三角网格表示和ODE积分引起的。CoTAN还在运行时间和内存消耗方面有所提升,只需要0.21秒就可以提取每个半球的皮质表面,比最佳基线方法快了2倍。
图6。由CoTAN预测的白质表面重建的注意图的可视化。左:不同时间和年龄的条件时间注意图pr(t、a)。右图:由CoTAN学习到的从粗到细的白质表面变形。
总之,实验结果表明,CoTAN是一种高效准确的新生儿皮质表面重建方法。
Source
[1] Ma, Q., Li, L., Kyriakopoulou, V., Hajnal, J. V., Robinson, E. C., Kainz, B., & Rueckert, D. (2023, October). Conditional Temporal Attention Networks for Neonatal Cortical Surface Reconstruction. In International Conference on Medical Image Computing and Computer-Assisted Intervention (pp. 312-322). Cham: Springer Nature Switzerland.