[论文解读] MEMO: Test Time Robustness via Adaptation and Augmentation

Marvin Zhang, Sergey Levine+1 more|arXiv (Cornell University)|Oct 18, 2021

Advanced Neural Network Applications参考文献 38被引用 85

一句话总结

MEMO 在测试时通过对单个测试输入进行增广来调整一个预训练的概率模型，最小化跨增广的边缘熵以强化不变量和信心，在 ImageNet-C/R/A 及 CIFAR-10 变体上实现更好的鲁棒性。

ABSTRACT

While deep neural networks can attain good accuracy on in-distribution test points, many applications require robustness even in the face of unexpected perturbations in the input, changes in the domain, or other sources of distribution shift. We study the problem of test time robustification, i.e., using the test input to improve model robustness. Recent prior works have proposed methods for test time adaptation, however, they each introduce additional assumptions, such as access to multiple test points, that prevent widespread adoption. In this work, we aim to study and devise methods that make no assumptions about the model training process and are broadly applicable at test time. We propose a simple approach that can be used in any test setting where the model is probabilistic and adaptable: when presented with a test example, perform different data augmentations on the data point, and then adapt (all of) the model parameters by minimizing the entropy of the model's average, or marginal, output distribution across the augmentations. Intuitively, this objective encourages the model to make the same prediction across different augmentations, thus enforcing the invariances encoded in these augmentations, while also maintaining confidence in its predictions. In our experiments, we evaluate two baseline ResNet models, two robust ResNet-50 models, and a robust vision transformer model, and we demonstrate that this approach achieves accuracy gains of 1-8\% over standard model evaluation and also generally outperforms prior augmentation and adaptation strategies. For the setting in which only one test point is available, we achieve state-of-the-art results on the ImageNet-C, ImageNet-R, and, among ResNet-50 models, ImageNet-A distribution shift benchmarks.

研究动机与目标

• 激发并研究不依赖训练时改变或需要大量测试批次的测试时鲁棒化方法。
• 提出一种即插即用的 MEMO 方法，在测试时使用单个测试样本对所有模型参数进行自适应。
• 通过边缘熵最小化，鼓励预测在不同增广之间保持不变，同时保持自信。
• 展示与现有鲁棒性技术和 BN 自适应结合时的兼容性及可叠加的增益。

提出的方法

• 给定一个训练好的概率模型 f_theta，呈现一个单独的测试输入 x 以及一组增广 A。
• 对 x 进行 B 次增广采样，形成一组增广输入；通过对增广结果 p_theta(y|a(x)) 取平均来估计边缘输出分布。
• 将 MEMO 损失定义为边缘分布的熵 H(bar{p}_theta(.|x))；通过梯度下降更新 theta 以最小化该边缘熵（每个测试点一个梯度步）。
• 自适应后，使用更新后的参数对原始输入 x 进行预测；在测试时不需要 ground-truth 标签。
• 可选地将 MEMO 与 BN-统计自适应或其他鲁棒性方法结合，而不改变训练过程。

实验结果

研究问题

• RQ1 在不对训练过程作出假设、也不需要访问测试批次的情况下，测试时自适应是否有效？
• RQ2 对单个测试点的增广复制品最小化边缘熵是否能提升对分布迁移的鲁棒性？
• RQ3 MEMO 如何与现有鲁棒性技术（例如 BN 自适应、AugMix、MoEx）在不同模型架构和数据集上相互作用？
• RQ4 是自适应还是增广是观测收益的主导因素，增广选择和样本数量的作用是什么？

主要发现

• MEMO 在具有挑战性的分布转移基准测试中，相较于标准评估实现了 1–8% 的准确率提升。
• 在 ImageNet 单点测试中，MEMO 在 ResNet-50 上达到现状最佳结果，并增强对 ImageNet-C、ImageNet-R、ImageNet-A 的鲁棒性。
• MEMO 提升了 ResNet 和视觉变换器模型的鲁棒性，在若干基准上优于以往的增广/自适应策略。
• 消融研究表明增广的一致性保持以及维持预测置信度都是 MEMO 的重要组成部分。
• MEMO 可以与预训练的鲁棒模型和 BN 自适应结合以进一步提升性能，即使对使用大量数据增广训练的模型也显示出显著收益。