Gradient Alignment Improves Test-Time Adaptation for Medical Image Segmentation

Gradient Alignment Improves Test-Time Adaptation for Medical Image Segmentation

针对解决的问题：不同中心的医学图像之间普遍存在的领域转移的问题，阻碍了预训练模型的有效部署
原来提出的一些方法，次优优化方向和固定步长(学习率不变)，效果不是太好

[!NOTE]

次优优化方向：选择的不是最佳的方向，而是一个相对较好的方向，基于当前信息，可能并不是全局最优的方向
提出了基于梯度对齐时间自适应方法，提高优化过程中的梯度方向和学习率优化，出发点从自监督目标导出的伪梯度，我们的方法将辅助梯度与伪梯度结合在一起，以促进梯度对齐。
关键点
- 相当于可以挖掘不同梯度之间的相似性，并矫正梯度方向，利用与当前分割任务相关的经验梯度
- 根据伪梯度和辅助梯度之间的余弦相似性设计了一个动态学习率，从而能够对不同测试数据上的预训练模型进行自适应微调

主要提高泛化能力

[!NOTE]

TTA: 测试时间自适应，用在迁移学习和无监督学习中，核心思想是模型在处于推理阶段时，对模型进行自适应调整，不依赖于训练阶段的标签信息，而是根据测试阶段的输入做出动态调整
认为原来的方法的主要问题是，忽略了优化过程中的两个关键因素：步长和方向
主要方法对比，如上图：
- a图蓝色的线代表需要优化的伪梯度，而粉色的线代表经验梯度的方式，原来采取的方式是直接优化伪梯度(弱增强和强增强两种，一致性损失产生的梯度)的方式
- b图则是借助一个辅助梯度(熵损失,原有测试集中导出的梯度)，最小化辅助梯度(黄色的线)和伪优化的梯度(蓝色的线)之间的角度，而因为这两个会把粉红色的线(经验梯度)包含在其中，从而达到蓝色的线(原有方式)和粉色的线(经验梯度)两者不断靠近
还提出了一个与这两个梯度之间的角度成反比的动态学习率，以自适应地微调预训练模型。角度更大意味着冲突更大，则需要更小的学习率慢慢地来调
主要贡献;
- 提出GraTa来改进优化方向和步长
- 有不同的伪梯度和辅助梯度两种，减少任务的分歧
- 不同梯度之间角度，对两个梯度之间的余弦相似性的可变学习率，有助于动态确定自适应微调的优化步长

原有的TTA方法主要是
- 基于构建自我监督的辅助任务来微调模型
- 微调批归一化层内的仿射参数
- 这样的方法但是存在一定的问题，当监督信息可靠的时候，这些方法表现出有效适应性，不可靠的时候，性能便会下降
原有梯度对齐的方法
- 对梯度进行投影，缓解任务梯度之间的梯度冲突的问题
- 联邦学习加上隐式正则化的方式

$\text{Lent}(\theta; X_i^t) = -P_i \log P_i, \quad P_i = \text{Sigmoid}(f_\theta(X_i^t)).$
测试样本的交叉熵损失函数
$\text{Lcon}(\theta; X_i^t) = \text{lce}(\tilde{P}_i, \hat{P}_i), \quad \tilde{P}_i = \text{Sigmoid}(f_\theta(\tilde{X}_i^t)), \quad \hat{P}_i = \frac{1}{j} \sum_{j=1}^j \text{Sigmoid}(f_\theta(\hat{X}_{ij}^t)).$
弱增强和强增强两种，一致性损失产生的梯度
弱增强策略：身份映射，水平翻转，垂直翻转，旋转90度，旋转180度，旋转270度
强增强策略：亮度调整，对比度调整，伽马变化，高斯噪声和高斯模糊
$\min_{\theta} \left( \text{Lcon}(\theta; X_i^t) + \angle(\nabla \text{Lcon}(\theta; X_i^t), \nabla \text{Lent}(\theta; X_i^t)) \right)$
最后的目标参数
但是优化的时候角度不能直接微分，而且如果采用内积的方式的话，计算复杂而且不稳定，因此采取下面的这种方式

$\min_{\theta} \, \text{Lcon}(\theta'; X_i^t), \quad \text{where} \quad \theta' = \theta - \nabla \text{Lent}(\theta; X_i^t).$
但是最后在对比试验加上消融实验之后得出下面的变体的效果是最好的

微调过程
- 自己的理解是，就相当于上面的那个式子两个部分先后进行计算，然后再反向传播
- 首先计算 $\text{Lent}(\theta; X_i^t)$ ，然后反向传播，更新参数
- 然后再计算 $\text{Lcon}(\theta'; X_i^t)$ , 然后再反向传播，得到学习率，然后再更新参数 $\theta_{i+1} \leftarrow \theta_i - \eta \nabla \text{Lcon}(\theta_i'; X_i^t)$
- 然后又进行推理
参数 $\eta$
$\eta = \beta \cdot \text{Cus} \left( \frac{\nabla \text{Lcon}(\theta'; X_i^t) \cdot \nabla \text{Lent}(\theta; X_i^t)}{\|\nabla \text{Lcon}(\theta'; X_i^t)\| \, \|\nabla \text{Lent}(\theta; X_i^t)\|} \right)$
β是一个缩放因子

过程总览图