论文笔记《MEMO: Test Time Robustness via Adaptation and Augmentation》

论文笔记《MEMO: Test Time Robustness via Adaptation and Augmentation》

Administrator 3 2025-05-14
  • 论文 - 《MEMO: Test Time Robustness via Adaptation and Augmentation》
  • 代码 - Github
  • 关键词 - 数据增强、TTA、熵、即插即用

摘要

  • 研究问题：测试时自适应的方法通常引入了一些额外的假设条件，例如需要访问多个测试样本，这限制了它们的广泛应用。
  • 本文工作
    • 目标：在训练过程中不作任何假设、并且在测试时具有广泛适用性的方法。
    • 方法：当面对一个测试样本时，对该样本进行多种数据增强，然后通过最小化模型在这些增强下的平均输出分布（或称边缘分布）的熵，来自适应地更新（全部）模型参数。
    • 原理：直观上，这个目标鼓励模型在不同的增强下做出相同的预测，从而强制模型学习这些增强所编码的不变性（invariances），同时保持预测的置信度。

1 引言

  • 目标：关注测试时鲁棒性，即利用具体的测试输入来提升模型在该样本上的预测性能。“即插即用”，能够与其它鲁棒化技术协同工作。
  • 方法 - MEMO
    • 全称：Marginal Entropy Minimization with One test point，基于单个测试点的边缘熵最小化
    • 如图1所示，当面对一个测试样本时，我们通过对该样本进行多种数据增强，并在此基础上调整模型，以鼓励模型在不同增强下做出一致的预测，从而体现出数据增强所编码的不变性。
    • 同时，我们也鼓励模型做出高置信度的预测，由此得出了我们提出的方法：最小化模型在该测试样本各种增强版本下的边缘预测分布的熵。
paper46-1.webp

2 测试时增强与自适应

数据增强：通常用于训练模型以尊重某些不变性——例如，光照或视角的变化不会改变底层的类别标签。（然而在分布偏移的情况下，模型不一定能够遵守不变性）

MEMO 基本思想：这是一种用于测试时鲁棒性的技术，它通过调整模型使其在测试输入上尊重这些不变性。

测试时鲁棒性：特指直接作用于预训练模型和单个测试输入的技术，例如“单点BN自适应”和“测试时增强”，这些都是先前的测试时鲁棒性方法的例子。

测试时鲁棒性设置问题定义：给予一个已训练好的模型 ​ f_\theta ，其参数为 ​ \theta \in \Theta 。不需要任何特殊的训练过程，也不对模型做出任何假设，除了以下两点：(1) 参数 ​ \theta 是可调整的；(2) 模型 ​ f_\theta 生成的条件输出分布 ​ p_\theta(y|\mathbf{x}) 对 ​ \theta 是可微的。一个单点 ​ \mathbf{x} \in \mathcal{X} 被提供给 ​ f_\theta ，模型必须立即预测一个标签 ​ \hat{y} \in \mathcal{Y} 。请注意，这与常规监督学习模型的标准测试推理过程完全相同，并且没有对训练过程或测试数据的可用性做出额外假设。这使得测试时鲁棒性方法成为一种简单的“即插即用”替换方案，适用于普遍存在的标准测试推理过程。假设可以访问一组增强函数 ​ \mathcal{A} = \{a_1, \dots, a_M\} ，这些函数可以应用于测试点 ​ \mathbf{x} 。作者利用这些增强函数以及下面详细说明的自监督目标，在模型对 ​ \mathbf{x} 进行预测之前对其进行调整。当给定一组测试输入时，模型会独立地对每个测试点进行调整并做出预测。


2.1 单个测试点的边缘熵最小化

给定一个测试点 ​ \mathbf{x} 和一组增强函数 ​ \mathcal{A} ，我们从 ​ \mathcal{A} 中采样 ​ B 个增强，并将其应用于 ​ \mathbf{x} ，以生成一批增强数据 ​ \tilde{\mathbf{x}}_1, \dots, \tilde{\mathbf{x}}_B 。模型关于这些增强点的平均（或边缘）输出分布由以下公式给出：

\bar{p}_\theta(y|\mathbf{x}) \triangleq \mathbb{E}_{a \sim \mathcal{U}(\mathcal{A})}[p_\theta(y|a(\mathbf{x}))] \approx \frac{1}{B} \sum_{i=1}^B p_\theta(y|\tilde{\mathbf{x}}_i), \tag{1}

其中期望是对均匀采样的增强 ​ a \sim \mathcal{U}(\mathcal{A}) 计算的。

我们希望从这个边缘分布中得到什么性质？为了回答这个问题，考虑数据增强在训练期间通常扮演的角色。对于每个训练样本，模型 ​ f_\theta 使用输入的多种增强形式进行训练，使得模型 ​ f 遵守增强与标签之间的不变性。因此，作者希望在测试时实现类似的自监督信号，而无需任何真实标签。也就是说，在调整后：

  1. 模型 ​ f_\theta 的预测应该对测试点的增强版本保持不变；
  2. 模型 ​ f_\theta 应该对其预测充满信心，即使对于高度增强的测试点版本也是如此。

为了使模型更自信地进行预测，可以从假设“类之间的真正决策边界位于数据空间的低密度区域”来证明其合理性。带着这两个目标，我们提出通过增强的边缘输出分布的熵来调整模型（见公式 (1)），即：

\ell(\theta; \mathbf{x}) \triangleq H(\bar{p}_\theta(\cdot|\mathbf{x})) = -\sum_{y \in \mathcal{Y}} \bar{p}_\theta(y|\mathbf{x}) \log \bar{p}_\theta(y|\mathbf{x}). \tag{2}

请注意，这一目标不同于优化模型预测分布的平均条件熵，即：

\ell_\text{CE}(\theta; \mathbf{x}) \triangleq \frac{1}{B} \sum_{i=1}^B H(p_\theta(\cdot|\tilde{\mathbf{x}}_i)). \tag{3}

一个在增强之间自信但不同地进行预测的模型可以最小化公式 (3)，但无法最小化公式 (2)。优化公式 (2) 同时鼓励了置信度和不变性，因为当模型无论增强如何都输出相同的（自信的）预测时，​ \bar{p}_\theta(\cdot|\mathbf{x}) 的熵被最小化。

算法 1 展示了整体方法 MEMO 在测试时的自适应过程

paper46-2.webp

3.2 将 MEMO 与先前方法结合使用

MEMO 的另一个优势是它可以与其它应对分布偏移（distribution shift）的方法协同工作。具体来说，MEMO 可以与先前用于训练鲁棒模型和自适应模型统计信息的方法结合使用，从而综合各项技术带来的性能提升。

1. 使用预训练的鲁棒模型

  • 由于 MEMO 不对模型的训练过程做出任何假设，也不对其进行修改，因此在已经预训练好的鲁棒模型上进行自适应是非常直接的。
  • 测试时所使用的增强集合 ​ \mathcal{A} 并不一定要与训练过程中使用的增强策略一致。为了简单和高效，作者采用的是一些易于采样、并能直接作用于输入 ​ \mathbf{x} 的增强方式。
  • 在我们的实验中，我们选择使用 AugMix [13] 作为我们的增强集合，因为它满足上述所有性质，并且在实际应用中能够产生显著的多样性。

  1. 调整 BN 统计量

Schneider 等人表明，即使仅使用单个测试样本，部分调整模型中每个批归一化（BN）层激活的估计均值和方差，在某些情况下仍然可以有效处理分布偏移。在这种情况下，为了防止对测试点过拟合，从该点估计得到的通道级均值和方差 ​ [\mu_\text{test}, \sigma^2_\text{test}] 会根据一个先验强度 ​ N 与训练期间计算的均值和方差 ​ [\mu_\text{train}, \sigma^2_\text{train}] 进行混合。也就是说，对于 ​ \nu \in \{\mu, \sigma^2\} ，

\nu \triangleq \frac{N}{N+1} \nu_\text{train} + \frac{1}{N+1} \nu_\text{test}.

这种方法也很容易与 MEMO 结合使用：在计算模型输出分布时直接使用调整后的 BN 统计量。根据 Schneider 等人的建议，作者在下一节的所有实验中设置 ​ N = 16 。


TTA 论文笔记