[논문 리뷰] MEMO: Test Time Robustness via Adaptation and Augmentation

Marvin Zhang, Sergey Levine외 1명|arXiv (Cornell University)|2021. 10. 18.

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 레이블은 필요하지 않다.
• 훈련 절차를 바꾸지 않고도 BN-통계 적응 또는 다른 강건성 방법과 MEMO를 선택적으로 결합한다.

실험 결과

연구 질문

• 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 적응과 결합하여 성능을 더욱 개선할 수 있으며, 강한 이득을 나타내는 경우가 많다.