DLTTA: Dynamic Learning Rate for Test-time Adaptation on Cross-domain Medical Images--论文笔记

论文笔记

资料

1.代码地址

2.论文地址

3.数据集地址

OCT数据集

1论文摘要的翻译

为了有效地解决来自不同机构的医学图像在测试时间的跨域分布偏移问题，测试时间自适应(TTA)日益成为一个重要的研究课题

以前的TTA方法有一个共同的限制，即对所有测试样本使用固定的学习率。这样的实践对于TTA来说是次优的，因为测试数据可能是顺序到达的，因此分布转移的规模会经常变化。为了解决这个问题，我们提出了一种新的测试时间适应的动态学习率调整方法，称为DLTTA，该方法动态调节每个测试图像的权重更新量，以考虑其分布变化的差异。具体来说，我们的DLTTA配备了一个基于记忆库的估计方案，以有效地测量给定测试样本的差异。基于这个估计的差异，然后开发一个动态学习率调整策略，以实现每个测试样本的适当程度的适应。我们的DLTTA在视网膜光学相干断层扫描(OCT)分割、组织病理学图像分类和前列腺3D MRI分割等三个任务中得到了广泛的证明

我们的方法实现了有效和快速的测试时间自适应，并且与当前最先进的测试时间自适应方法相比具有一致的性能改进。

2论文的创新点

• 本工作提出了一种新的动态学习率调整框架，用于跨域医学图像分析的测试时间自适应。动态学习率策略，以克服测试时模型自适应推理数据的变化分布偏移。
• 本工作设计了一种基于内存库的测试样本差异估计策略来模拟适应需求，基于该策略，以在线方式为每个特定的测试样本专门调整学习率。
• 我们已经验证了我们的方法在三种不同的2D或3D模型医学成像模式上的分类和分割任务的有效性。实验结果表明，我们的DLTTA适用于不同的网络架构，并且可以始终优于当前最先进的TTA方法。

3 论文方法的概述

本节将介绍了测试时间自适应动态学习率方法，其中加入了一个memory bank 来估计推理阶段预测的差异。根据估计的差异，设计了一个动态学习率来动态调节测试时间自适应过程，以获得稳定和改进的自适应结果。图2显示了我们有效测试时间自适应框架的概述。

当将模型应用于新的测试数据时，一个问题是新的测试样本可能遵循未知的数据分布，导致模型预测的性能严重下降。以前的领域自适应方法需要组装足够数量的测试数据，而在现实世界的应用中，测试样本通常会以不同的分布偏移顺序到达。更吸引人的解决方案是测试时间自适应，旨在直接根据每个呈现的测试样本不断调整模型。具体来说，对于测试时间自适应，我们只得到一个由一组参数${\theta }^s$参数化的预训练模型$f_{{\theta }^s}$和新的测试样本$x_1，x_2，\dots ，x_{t-1}，x_t$，这些样本以不同的分布偏移顺序在线到达。
通过最小化损失函数，使用源训练数据${(x_i^s，y_i^s)}_{i=1}^{N^s}$对模型$f_{{\theta }^s}$进行优化，如下所示：$${\theta }^s=\arg \underset{\theta }{\min }\frac{1}{N^s}\sum _{i=1}^{N^s}{L}_m(f_{\theta }(x_i^s),y_i^s),(1)$$
其中$L_m$表示训练的监督损失，例如交叉熵损失。模型$f_{{\theta }^s}$在到达时间步长$t$并且在数据分布上与源训练样本不同的新测试样本$x_t$上表现不佳。然后，为了不断调整模型以实现更好的泛化性能，需要专门设计一个测试时间目标函数$L_{tt}$，以根据每个测试样本提供的分布信息更新模型参数。对于在模型更新的第t次迭代中出现的测试样本$x_t$，我们有：$${\theta }_{t+1}\leftarrow {\theta }_t-\eta (\nabla L_{tt}{f}_{{\theta }_t}(x_t)),(2)$$其中$\eta$表示测试时间自适应的学习率，${\theta }_1用$θ_s$初始化。更新后的模型$f_{θ_{t+1}}$用于获得测试样本$x_t$的预测。

3.2 Dynamic Learning Rate on Test Data

之前的测试时间自适应方法在公式（2）中采用固定的学习率η，由于自适应性能对学习率很敏感，因此需要仔细选择$\eta$。本工作认为适应步骤的静态学习率都不能准确地更新模型，以克服测试数据分布的变化。因此，我们提出了一种动态学习策略，在测试时间自适应过程中捕捉模型变化和测试数据的不同偏移程度。

• 差异估计的memory bank构建
  对于动态适应，了解每个更新步骤所需的适应程度至关重要，这是调整学习率的关键因素。为了衡量适应程度，我们提出了一种基于memory bank库的差异估计。构建存储体以存储由连续更新的模型提取的最新特征表示和预测掩模对。这些对反映了模型的变化，可以进一步用于计算与传入测试样本的距离，以估计分布偏移程度。通过偏移度测量，我们能够相应地调整测试时间学习率。
  具体而言，memory bank $M$包括$K$对键值对$\{(q_k，v_k){\}}_{k=1}^K$。如图2所示，键 { q k ∣ q k = h ( x t ) } \{q_{k}|q_{k}=h(x_{t})\} {qk​∣qk​=h(xt​)}是由模型的特征提取器$h$计算的特征图，值 { v k ∣ v k = g ( q k ) } \{v_{k}|v_{k}=g(q_{k})\} {vk​∣vk​=g(qk​)}对应于分类器头$g$生成的预测掩码。我们通过连续缓存新的$(q_k,v_k)$对来更新$M$。由于模型在测试时会逐步更新，因此存储体中的早期元素无法在测试数据上指示最新的模型性能。因此，我们只维护具有固定大小$K$的存储体，并在写入新对时保持先进先出（FIFO）原则。然后，为了测量新来的测试样本$x_t$的差异，从M中检索集合$\mathcal{R}=\{(q_d,v_d){\}}_{d=1}^D\subseteq \mathcal{M}$，其中包含与当前测试样本$x_t$语义相似的元素的特征和预测。这是通过基于每个关键字和查询样本$h(x_t)$的特征之间的L2距离计算D-nearest neighbors来实现的。由于键保持语义级上下文特征，支持集R保持了关于高级信息的内部一致性，例如图像中的对象类别。R中的预测集合可以用作$x_t$的参考预测，如下所示：$${\hat{p}}_t=\frac{1}{D}\sum _{d=1}^D{v}_d.\text{(3)}$$然后，图像$x_t$的预测差异由下式得出：$$\begin{array}{rlr}{L}_{\mathrm{DIV}}& =\frac{1}{2}(L_{\mathrm{KL}}({\hat{p}}_t||f_{{\theta }_t}(x_t))+L_{\mathrm{KL}}(f_{{\theta }_t}(x_t)||{\hat{p}}_t)),& \\ & & \text{(4)}\\ & \mathrm{with}{L}_{\mathrm{KL}}({\hat{p}}_t||f_{{\theta }_t}(x_t))=\sum _j{\hat{p}}_{t(j)}\log \frac{{\hat{p}}_{t(j)}}{f_{{\theta }_t}(x_{t(j)})},& \end{array}$$
  这里$L_{KL}$表示KL散度，$f_{{\theta }_t}(x_t)$表示由模型$f_{{\theta }_t}$生成的$x_t$预测。
• Adaptive Learning Rate Adjustment
  根据公式（4）测量的差异，表明大的差异表示需要很大程度上去缩小的显著差距，小差异需要较小的适应体调整。因此，我们建议根据估计的差异动态调整自监督学习率。为了更通用地描述我们的方法，我们采用分批公式来说明测试样品可能会逐个或分批到达。单个测试图像对应于批量大小1。给定一批测试数据 { x t , b } b = 1 B \{x_t,b\}^B_{b=1} {xt​,b}b=1B​在每个测试时间自适应步骤中，我们通过以下方式计算总差异：$$L_{\text{DIV}}^B=\frac{1}{B}\sum _{b=1}^B{L}_{\text{DIV},b},\text{(5)}$$其中$B$是批量大小。$L_{\mathrm{DIV}}^B$捕捉了批处理预测中分布的一般差异。此后，通过函数$H(L_{\mathrm{DIV}}^B)$获得分批动态学习率，该函数直接基于$L_{\mathrm{DIV}}^B$输出测试时间自适应任务的学习率：$$H(L_{\mathrm{DIV}}^B)=\alpha \cdot L_{\mathrm{DIV}}^B,(6)$$其中$\alpha$进一步缩放学习率，并且可以根据经验设置为用于源训练图像的模型优化的值。
  我们提出的测试时动态学习率调整可以很容易地部署到任何网络架构和自监督目标函数中，以改善测试时的适应过程。由于特征表示和预测掩码是自然计算并存储在存储体中的，因此测试时的计算成本主要来自检索用于计算参考预测${\hat{p}}_t$的支持集。我们通过保持较小的检索大小（例如8）来提高效率。在实验中分析了检索大小的影响，结果表明，相对较小的大小已经可以获得良好的结果，增加大小（例如增加到20）并没有显示出更高的性能。

3.3 学习过程和实施细节

3.3.1 学习过程

执行测试时间自适应的总体程序总结如下。用从源训练数据中学习的参数初始化要调整的模型。对于即将到来的测试样本，该模型首先执行前向传递以获得一对特征表示和预测掩码，从而从存储体中检索语义相似的元素，并计算特定测试样本的差异和动态学习率。然后，通过使用具有测试时间目标函数的导出学习率来更新模型，以实现所需的自适应。在一步自适应之后，模型执行另一个前向传递，以获得当前测试样本的预测，并更新存储体的新元素对。该模型随着每个顺序到来的测试样本的前一个过程迭代更新。在构建存储体之前，学习率保持为初始值。DLTTA的伪码如算法1所示

3.3.2实现细节

我们的方法中网络架构的选择是灵活的。在不失泛化性的情况下，我们采用U-Net进行视网膜OCT图像分割，使用在ImageNet上预训练的DenseNet-121进行组织病理学图像分类，并采用3D U-Net进行前列腺MRI图像分割。我们的2D U-Net和3D UNet的编码器包含4个卷积BatchNorm ReLU块，这些块不断对图像分辨率进行下采样，并将特征通道的尺寸加倍。然后，特征由瓶颈层处理并馈入解码器，解码器有4个转置卷积块，用于对中间特征图进行上采样。所有2D卷积使用核大小3×3，3D卷积使用核尺寸3×3×3。编码器特征在每个阶段跳过连接到解码器，每个块的特征信道尺寸为[16,32,64128]。对于分割网络，我们将编码器和瓶层视为特征提取器，将解码器视为预测头。DenseNet-121的架构遵循Torchvision库中的实现，该库由4个密集连接的块组成，然后是全局池和完全连接的分类层。对于每项任务，我们的方法和其他比较方法的网络骨干是相同的，以确保公平的比较。
对于使用源图像的模型训练，使用Adam优化器从头开始训练模型100个迭代周期，并将学习率初始化为1e-3和3e-4，用于分割和分类任务。对于我们的方法和所有比较基线的测试时间自适应，我们对每批测试数据进行了一步自适应，批大小为1用于分割，200用于分类任务。我们为最新的KB自适应步骤存储了键和值对。我们根据经验将$\frac{K}{B}$设置为20用于分割，设置为4用于分类。检索大小D对于分割任务设置为8，对于分类任务设置为12。从所有方法的测试数据中收集BN统计数据。该框架使用Pytorch 1.7.0实现，并在一个NVIDIA TitanXp GPU上进行训练。

4 论文实验

4.4.1 数据集和评价指标

我们首先使用光学相干断层扫描（OCT）数据集验证了我们的方法在视网膜层分割方面的有效性，这些数据集显示了由不同的海德堡Spectralis扫描仪和Cirrus扫描仪引起的分布偏移。然后，我们在更大的肿瘤和正常组织分类组织病理学图像数据集上评估我们的方法，这些数据集来自不同医院，具有不同的协议和患者群体。我们还使用从不同医疗中心收集的3D MRI图像验证了我们的前列腺分割方法。

• 数据集
  Retinal OCT Image Segmentation
  我们遵循文献，使用从海德堡Spectralis扫描仪获取的公共OCT数据集作为源训练数据，并使用从Cirrus扫描仪获得的OCT图像公共数据集作为偏移测试数据集。源训练数据集由35个3D体积组成，但由于切片之间的物理距离较大，根据，3D体积被分为588个2D切片用于训练，147个切片用于验证，980个切片用于测试。任务是分割八个不同的视网膜层，并为模型训练提供基础真实数据。测试数据集包含来自6名受试者的48个切片Cirrus 2D切片被重新采样，以具有与Spectralis扫描相同的物理切片内分辨率。两个数据集中的所有切片都经过了视网膜展平预处理，并被裁剪和调整为128×1024。为了进行评估，我们遵循之前的工作，使用相同的两个OCT数据集采用流行的Dice评分来评估分割性能。
  Histopathological Image Classification.
  我们使用大型Camelon17数据集来评估我们的方法，该数据集包含来自五家不同医院的400000多个96×96大小的组织病理学图像块。在我们的实验中，从三家医院收集的302436个组织补丁被用作源训练数据，从另外两家医院收集的34904个补丁和85054个补丁分别被视为两个看不见的测试数据集，即看不见数据集A和看不见数据集B。任务是预测给定的组织区域是否包含任何肿瘤组织。我们使用ImageNet的数据统计将每个图像块归一化为零均值和单位方差。采用五种最常用的分类任务评估指标，包括准确性、敏感性、特异性、AUC和F1评分，来评估不同的方法。
  前列腺MRI图像分割
  我们通过多中心3D MRI数据进一步验证了我们在前列腺分割任务中的方法。我们使用从六个临床中心收集的三个公共数据集NCI-ISBI13、I2CVB和PROMISE12的前列腺T2加权MRI体积。根据收集数据集的临床中心，将数据划分为六个数据集A至F。数据集A用作训练数据，其他数据集用作不同的看不见的测试数据集。对于数据预处理，每个MRI体积的强度值归一化为零均值和单位方差，并在轴向平面上调整为384384。裁剪大小为8080*80的补丁作为网络输入。采用通用的评价指标Dice评分来比较不同方法的性能。

4.2 与之前先进的方法比较

我们的动态学习策略可以部署到不同的测试时间自适应方法中，以提高模型更新的有效性。对于每项任务，我们将动态学习率策略部署到三种具有代表性的测试时间自适应方法中，包括TTT，Tent，以及ATTA。我们还将我们的方法与PTBN进行了比较，PTBN从测试数据中重新估计了批归一化统计数据。为了提供强有力的基线，我们采用了UDA方法进行比较，这些方法同时利用源训练数据和整个测试数据集进行自适应。对于视网膜OCT图像分割，我们包括中用于调整OCT图像的三种UDA方法，即图像转换方法UDA-ST和CycleGAN，以及广泛使用的输出空间调整方法UDAS。对于组织病理学图像分类，我们与UDA SwAV进行了比较，UDA SwAV[41]是使用我们的数据集进行组织病理学影像分类的最佳UDA模型。对于前列腺分割，我们还使用UDAS进行比较。OCT图像上的ATTA、UDA-ST结果和看不见的组织病理学数据集A上的UDA-SwAV结果分别直接参考[21]和[41]，因为他们的方法和我们的方法使用了相同的数据集、网络骨干和数据分割。其他结果是通过基于发布的代码重新实现而获得的，其中网络骨干网对所有比较方法都是一致的。此外，W/o自适应模型表示直接应用用训练数据集优化的模型来获得测试样本的预测。

• Results for Retinal Layer Segmentation
  表1显示了Cirrus图像OCT测试数据集视网膜层分割的比较结果。我们可以看到，所有测试时自适应和UDA方法都比“W/o自适应”基线有所改善，表明了模型自适应对测试样本数据分布的好处。我们的DLTTA配备了动态学习率策略，持续改进了不同的测试时间自适应方法TTT、Tent和ATTA，用于分割所有视网膜层。这些结果证明了我们的动态学习率在改进测试时间模型更新以克服测试数据分布变化方面的有效性，并验证了我们的静态学习率策略可以应用于不同的测试时间自适应方法，以提高自适应性能。通过部署到ATTA，我们的方法实现了最高的平均Dice值83.41%，在适应过程中优于所有使用源训练数据的UDA方法。Plain Tent的性能与UDA-ST相当，但进一步添加我们的动态学习率策略将平均Dice从81.50%提高到82.63%，优于UDA-ST。UDAS的性能优于Plain Tent，但略低于ATTA。当ATTA配备动态学习率策略时，它的性能进一步优于UDAS。这表明了动态测试时间自适应的巨大潜力，并表明不一定需要源训练数据来实现良好的自适应性能。可能的原因是，动态测试时间自适应直接基于每个特定测试样本的分布信息更新模型，而不是试图为整个训练和测试数据找到域不变空间，当训练和测试图像在数据分布上存在显著差异时，这可能很困难。
  
• 组织病理学图像分类结果
  表II显示了肿瘤和正常组织分类的大型组织病理学影像数据集的比较结果。结果报告在两个不同的测试数据集上。我们可以看到，即使该模型是在三家医院的图像上训练的，覆盖了多个数据分布，但该模型在两个具有分布偏移的看不见的数据集上的表现仍然很差，在看不见数据集A上的准确率仅为63.15%。这表明了模型适应测试数据分布的必要性。与视网膜层分割类似，我们的动态学习率策略也在组织病理学图像分类任务中持续改进TTT、Tent和ATTA，表明我们的方法在测试时间适应目标、分析任务和网络骨干方面具有普遍适用性。通过部署到Tent，我们的方法对两个测试数据集的准确率分别为91.49%和88.93%，与利用训练和测试数据的UDA-SwAV方法的强基线相当。这些改进得益于我们的自适应学习率调整，根据每个测试样本的估计差异明确调整在线适应过程。

5论文内一些重要的图

• 可以和一些实验、方法进行呼应，帮助你以后快速回忆起这篇论文内容核心。
• 你的论文方法概述也可以围绕这个图讲