[论文解读] On-the-Fly Test-time Adaptation for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Guo, Pengfei+2 more|arXiv (Cornell University)|Mar 10, 2022

Domain Adaptation and Few-Shot Learning被引用 4

一句话总结

该论文提出了一种即用型测试时自适应（On-the-Fly Test-Time Adaptation），一种零样本、基于事件的医学图像分割框架，可在无需反向传播的情况下，将预训练模型适配到单张测试图像。通过采用基于领域码驱动的自适应批归一化（adaptive batch normalization）的自适应UNet，该方法在2D和3D领域偏移场景下均实现了最先进性能，尽管临床约束更严格，仍优于现有测试时自适应方法。

ABSTRACT

One major problem in deep learning-based solutions for medical imaging is the drop in performance when a model is tested on a data distribution different from the one that it is trained on. Adapting the source model to target data distribution at test-time is an efficient solution for the data-shift problem. Previous methods solve this by adapting the model to target distribution by using techniques like entropy minimization or regularization. In these methods, the models are still updated by back-propagation using an unsupervised loss on complete test data distribution. In real-world clinical settings, it makes more sense to adapt a model to a new test image on-the-fly and avoid model update during inference due to privacy concerns and lack of computing resource at deployment. To this end, we propose a new setting - On-the-Fly Adaptation which is zero-shot and episodic (i.e., the model is adapted to a single image at a time and also does not perform any back-propagation during test-time). To achieve this, we propose a new framework called Adaptive UNet where each convolutional block is equipped with an adaptive batch normalization layer to adapt the features with respect to a domain code. The domain code is generated using a pre-trained encoder trained on a large corpus of medical images. During test-time, the model takes in just the new test image and generates a domain code to adapt the features of source model according to the test data. We validate the performance on both 2D and 3D data distribution shifts where we get a better performance compared to previous test-time adaptation methods. Code is available at https://github.com/jeya-maria-jose/On-The-Fly-Adaptation

研究动机与目标

• 解决模型在训练与测试数据之间存在领域偏移时医学图像分割性能下降的问题。
• 开发一种可临床部署的自适应方法，避免反向传播，并且在推理时无需访问完整的目标分布。
• 实现每张图像的零样本、即用型自适应，符合隐私保护和计算资源有限等真实临床约束。
• 设计一种框架，利用从预训练编码器提取的领域码动态调整模型特征。
• 在显著领域偏移条件下，于多样化的2D和3D医学图像分割基准上验证该方法。

提出的方法

• 提出一种新的自适应设定：即用型测试时自适应，其中自适应为基于事件的零样本设定，推理过程中不进行模型更新。
• 引入自适应UNet，通过在标准UNet中加入基于领域码的自适应批归一化层，实现特征的自适应调整。
• 利用大规模医学图像语料库上预训练的编码器，为每张测试图像生成领域码，以捕捉成像模态和扫描仪特异性特征。
• 利用领域码调节编码器和解码器中批归一化层的运行统计量，实现基于实例的特征自适应。
• 采用来自预训练模型的领域先验，指导自适应过程，而无需源数据或测试时训练。
• 对每张图像单独应用自适应模型，并在每次推理后重置自适应状态，确保隐私保护和低资源消耗。

实验结果

研究问题

• RQ1 是否能在推理过程中完全避免反向传播的前提下，实现高性能的测试时自适应？
• RQ2 当无法获取完整目标分布时，每张图像的零样本、即用型自适应方法的有效性如何？
• RQ3 基于学习到的领域码的自适应批归一化是否能有效应对2D和3D医学影像中的领域偏移？
• RQ4 在更严格的临床部署约束下，该方法是否优于现有测试时自适应基线方法？
• RQ5 该方法在不同模态偏移（如不同MRI序列或眼底成像设备之间）下的泛化能力如何？

主要发现

• 在2D眼底图像领域偏移上，即用型自适应结合自适应UNet的Dice分数达到63.14，优于TENT（57.10）和RN+CR（59.15）。
• 在3D MRI分割任务中，T1→T1ce偏移下的Dice分数分别为60.66/58.73/39.30，优于TENT（55.47/52.89/39.17）和RN+CR（57.30/54.68/40.01）。
• 在2D和3D设置下的9种领域偏移场景中，该方法在直接测试和所有基线方法上均表现出一致的性能提升。
• 定性结果表明，自适应UNet生成的分割掩码更接近真实标签和理想预测结果，且过分割和误分类现象更少。
• 尽管在2D偏移中表现优于3D偏移，可能由于跨模态MRI序列偏移的复杂性更高，但该方法仍表现出强劲性能。
• 消融实验确认，基于领域码的自适应至关重要，若移除该机制，性能退化程度与非自适应模型相当。