一篇来自 IPMI 2023 的 TTA 新思路：UPL-TTA
这次要聊的工作发表于 IPMI 2023——这个会议全称是 Information Processing in Medical Imaging，两年一届，在医学影像分析领域算是非常有特色的顶级会。这篇文章属于 Test Time Adaptation（TTA）系列，之前像 CoTTA、EcoTTA、DIGA 这些工作都是这个方向的热门。如果对 TTA 还不太熟悉，可以先翻翻那几篇。
深度学习模型在医学影像落地时，一个老生常谈的问题就是：模型在训练数据分布上表现不错，一到新的测试数据上就掉链子。为了应对这种数据偏移，一个直观的思路是把源域模型在线适应到目标域数据分布上。过去的方法大多靠熵最小化或者正则化这类手段，让模型通过无监督损失反向传播来更新参数。
但在真实临床场景里，更实际的需求是：模型能在推理时实时适应新来的单张图像，同时还得避开隐私问题，以及部署环境里算力不足的坑。TTA 就是为此而生的——遇到来自未知领域的图像时，能提升模型的鲁棒性。本文介绍的工作属于 Fully Test Time Adaptation，意味着推理时要做完整的反向传播。下表简单列了几种常见设定，其中 Fully TTA 只需要目标域数据和测试损失。
现有的 TTA 方法性能不太理想，原因主要有两个：一是未标记的目标域图像提供的监督信号太弱，二是方法本身对源域的预训练策略或网络结构有特定要求。这篇文章的目标是把源域预训练和目标域适应彻底拆开，做到高性能且通用，不依赖任何预训练策略的假设。
概述
这篇工作（论文里叫 UPL-TTA）提出了一种“不确定性感知的伪标签引导的完全测试时间自适应”方法，用于医学图像分割。它不需要预训练模型在适应目标域之前，先在源域里搞什么额外的辅助分支训练，也不绑定特定策略。
下面这张图以婴儿脑部 MRI 肿瘤分割为例，从 HASTE 序列跨域到 TrueFISP 序列，用了 UPL-TTA 之后，效果比什么都不做要好得多。
UPL-TTA 的核心思路是引入了“测试时增长”（Test-Time Growing，简称 TTG）。具体来说，在目标域里多次复制源模型的预测头部，同时给输入图像和特征图加上一系列的随机扰动（比如 Dropout、空间变换），得到多个不同的分割预测。然后把这些预测集成起来，生成目标域图像的伪标签。为了抑制那些可能不靠谱的伪标签，又引入了集成方法和 MC dropout 不确定性估计，得到一张可靠性图。可靠像素的伪标签用来监督每个预测头部的输出，而不可靠像素的预测，则通过对平均预测图做熵最小化来规范化。
具体实现
源域预训练模型
第一步还是得先把源域上的预训练模型优化好（对应图2的A部分），编码器和解码器的初始权重分别记为 θenc 和 θdec。
测试时增长（TTG）适应
TTG 的过程是这样的：对一张输入图像 x，先做一个空间变换 T，包括随机翻转、旋转 π/2、π 和 3π/2。然后经过 Dropout（特征级别的扰动）再送入解码器，最后做空间变换的逆变换得到概率图。就像前面说的，在目标域里复制源模型的预测头部，对每个头的输入图像和特征做不同的随机扰动，得到 K 个不同的分割预测。最后把 K 个头的结果集成起来。
用可靠伪标签进行监督
这一步的关键是怎么拿到可靠的伪标签。图2里那个可靠性图（Reliable map），你可以简单理解成一个掩码，用来筛选伪标签。设定一个阈值 τ，根据概率图的值来确定每个像素是否可靠——只保留伪标签里置信度高的像素。这样一来，我们就有三个目标：K 个头输出的预测图、伪标签、以及用来优化伪标签的掩码。
基于平均预测的熵最小化
熵最小化在 TTA 里很常见，但 UPL-TTA 里因为有 K 个集成，情况有点不一样。假设第 K-1 个头的预测概率是 0，第 K 个头的预测概率是 1，这两个头各自的熵都是最小的，但一平均变成 0.5，对应的熵反而大了。所以需要同时对 K 个头做熵最小化。
通过自训练进行适应
最后，优化目标由两部分损失组成：一部分来自可靠伪标签的监督损失，另一部分来自平均预测的熵最小化损失。两者加起来就是最终的训练目标。
实验
数据集信息如下：
和其他 SOTA 方法对比的结果：
可视化结果很直观：
下图展示了自训练过程中不同训练步数的伪标签变化。Epoch 0 代表“仅源域”（适应之前），n 代表在目标域验证集上表现最好的轮数。在 (c)-(g) 中，只有可靠的伪标签用颜色标了出来。
消融实验的结果如下：
总结
这篇 IPMI 2023 的工作提出了一种完全测试时间自适应的方法，能在不知道源模型训练策略的前提下，把源模型适应到未标记的目标域上。无需访问源域图像，基于不确定性的伪标签引导的 TTA 方法通过测试时增长（TTG）为目标域中的同一样本生成多个预测输出，进而产生高质量的伪标签和对应的可靠性图，为未标记目标域提供了有效的监督。那些带有不可靠伪标签的像素，则通过对复制头部的平均预测做熵最小化来进一步规范化，这其实也引入了隐式的一致规范化。在胎儿脑分割的双向跨模态 TTA 实验中，该方法优于几种最先进的 TTA 方法。未来，把该方法扩展到 3D 版本并应用到其他分割任务上，是很值得期待的方向。