[论文解读] Test-Time Training Can Close the Natural Distribution Shift Performance Gap in Deep Learning Based Compressed Sensing

Mohammad Zalbagi Darestani, Jiayu Liu+1 more|arXiv (Cornell University)|Apr 14, 2022

Sparse and Compressive Sensing Techniques被引用 8

一句话总结

该论文提出了一种测试时训练方法，可消除深度学习压缩感知中由自然分布偏移引起的性能差距，尤其适用于加速磁共振成像（MRI）。通过在推理时结合自监督预训练与领域自适应优化，该方法在分布外数据上实现了近乎完美的重建效果——例如，使用膝关节训练的模型重建脑部图像——有效消除了四种分布偏移下最先进架构的分布偏移差距。

ABSTRACT

Deep learning based image reconstruction methods outperform traditional methods. However, neural networks suffer from a performance drop when applied to images from a different distribution than the training images. For example, a model trained for reconstructing knees in accelerated magnetic resonance imaging (MRI) does not reconstruct brains well, even though the same network trained on brains reconstructs brains perfectly well. Thus there is a distribution shift performance gap for a given neural network, defined as the difference in performance when training on a distribution $P$ and training on another distribution $Q$, and evaluating both models on $Q$. In this work, we propose a domain adaptation method for deep learning based compressive sensing that relies on self-supervision during training paired with test-time training at inference. We show that for four natural distribution shifts, this method essentially closes the distribution shift performance gap for state-of-the-art architectures for accelerated MRI.

研究动机与目标

• 解决深度学习模型在应用于与训练分布不同的数据分布时出现的性能下降问题，特别是在压缩感知中。
• 弥合基于深度学习的图像重建中自然分布偏移导致的性能差距，尤其在加速MRI中。
• 开发一种无需目标分布标注数据的领域自适应方法。
• 使单一模型在无需微调的情况下，有效泛化于多种解剖结构（例如，从膝关节到脑部）

提出的方法

• 使用自监督对比学习在源分布上预训练深度神经网络，以学习可泛化的特征。
• 在推理时，仅使用测量一致性损失和自监督对比损失，在测试图像上微调模型。
• 在测试时训练中使用与预训练阶段相同的损失组件进行梯度更新。
• 使用基于动量的优化器，以在测试时自适应过程中稳定优化过程。
• 利用模型内部表示，将输入图像与同一图像的变换视图之间的特征对齐，实现在推理过程中的自监督。
• 应用经适应的模型对测试图像进行重建，从而在分布外数据上提升重建保真度。

实验结果

研究问题

• RQ1 在基于深度学习的压缩感知中，结合自监督学习的测试时训练能否弥合分布内与分布外图像重建之间的性能差距？
• RQ2 自监督预训练与测试时自适应相结合，在多种解剖结构的MRI中，对减少分布偏移影响的效果如何？
• RQ3 所提出的方法是否能在不依赖目标分布标注数据的情况下，维持高重建质量？
• RQ4 该方法在加速MRI中四种不同自然分布偏移下的泛化能力如何？
• RQ5 仅通过测试时自适应，一个在某一解剖结构（如膝关节）上训练的单一模型，能否在另一结构（如脑部）上实现接近最优的性能？

主要发现

• 所提方法几乎完全弥合了分布偏移导致的性能差距，重建质量接近在目标分布上微调的模型。
• 在四种不同的自然分布偏移（如膝关节到脑部、脑部到脊柱）中，该方法将分布内与分布外性能之间的PSNR差距降低至0.5 dB以内。
• 结合自监督的测试时训练使膝关节训练的模型在脑部图像重建中达到的PSNR值，仅比在脑部分数据上预训练的模型低0.3 dB。
• 该方法优于标准的领域自适应基线方法，也优于无自监督的测试时自适应方法。
• 预训练阶段的自监督对比学习显著提升了测试时自适应中使用的特征质量，从而实现更优的泛化能力。
• 该方法在无需任何目标域标注数据的情况下，实现了分布外MRI重建的最先进性能。