[논문 리뷰] Test-Time Training with Masked Autoencoders

Yossi Gandelsman, Yu Sun외 2명|arXiv (Cornell University)|2022. 09. 15.

Domain Adaptation and Few-Shot Learning인용 수 36

한 줄 요약

논문은 Masked Autoencoders (MAE)를 테스트 시 학습(TTT)으로 활용하여 self-supervised 재구성을 통해 각 테스트 입력에 모델을 적응시키고 시각 벤치마크에서 분포 이동에 대한 로버스트성을 향상시키며, ImageNet-C 및 다른 데이터셋에서 실증적 이득을 제공하고 이 접근법의 선형 모델 바이어스-분산 분석을 제시합니다.

ABSTRACT

Test-time training adapts to a new test distribution on the fly by optimizing a model for each test input using self-supervision. In this paper, we use masked autoencoders for this one-sample learning problem. Empirically, our simple method improves generalization on many visual benchmarks for distribution shifts. Theoretically, we characterize this improvement in terms of the bias-variance trade-off.

연구 동기 및 목표

• 보이지 않는 분포 이동에서도 강건한 일반화를 유도하고 테스트 시 self-supervision으로 모델을 적응시키려는 동기 부여.
• per-sample 적응을 위한 자기지도 신호를 생성하기 위해 마스킹된 자동인코딩을 자기지도 학습 과제로 활용합니다.
• 다양한 분포-이동 벤치마크(ImageNet-C, ImageNet-A, ImageNet-R, Portraits)에서 TTT-MAE를 평가하고 이론적 특성을 분석합니다.
• 훈련 시 설계 선택(미세 조정, 프로빙, 공동 학습)을 비교하고 테스트 시 적응을 위한 실용적이고 효과적인 설정을 식별합니다.

제안 방법

• TTT 프레임워크에서 자기지도 구성요소로 MAE를 채택하고 Y자형 아키텍처(인코더 f, 자기지도 헤드 g, 주요 작업 헤드 h)를 사용합니다.
• MAE로 사전 학습된 인코더 f0와 디코더 g0를 사용하고, 각 테스트 입력에 대해 마스킹된 패치의 재구성 손실을 최소화하도록 테스트 시 최적화를 수행하여 f_x와 g_x를 얻고, h∘f_x를 예측합니다.
• 강력한 베이스라인으로 ViT 프로빙(고정된 f, 학습 가능한 헤드 h)을 기본 훈련 시나리오로 채택하고, 미세 조정 및 공동 학습과 비교합니다.
• 테스트 입력당 20회의 TTT 단계에 대해 MAS(마스크 75%)와 비망각 증강을 사용하여 시작점 f0, g0에서 SGD로 학습합니다.
• 이미지넷-C 레벨-5(및 부록의 다른 레벨)에서 평가하고, 손상 특이적 증강을 사용하지 않고도 베이스라인 대비 개선을 보고합니다.
• 선형 설정에서 PCA 유사 자동인코딩과 같은 TTT가 고정된 모델보다 바이어스-분산 트레이드오프를 더 잘 설명한다는 이론적 통찰을 제공합니다.

실험 결과

연구 질문

• RQ1 다양한 분포 이동 하에서 왜곡 특이 신호에 의존하지 않고도 마스크된 자동인코더를 이용한 테스트 시 학습이 비전 모델의 강건성을 향상시킬 수 있는가?
• RQ2 MAE 기반 TTT는 회전 예측 기반의 TTT 및 다른 훈련 시 설계(미세 조정, 프로빙, 공동 학습)와 비교하여 벤치마크에서 어떤 차이를 보이는가?
• RQ3 선형 모델 설정에서 TTT-MAE의 효과를 바이어스-분산 트레이드오프 측면에서 이론적으로 어떻게 설명할 수 있는가?

주요 결과

• TTT-MAE는 베이스라인 ViT 프로빙 구성에 비해 ImageNet-C 레벨-5에서 정확도를 크게 향상시킵니다.
• TTT-MAE는 회전 예측 기반 TTT 및 베이스라인 모델들보다 대부분의 손상 유형에서 더 나은 성능을 보입니다.
• 훈련 시 설계 선택이 중요합니다: MAE 사전학습을 이용한 ViT 프로빙이 미세 조정이나 공동 학습에 비해 가장 강한 성능을 냅니다.
• 테스트 시 고정된 20스텝 예산으로 SGD를 사용한 최적화에서 검증 기반 조기 중단 없이 지속적인 개선이 나타납니다.
• TTT-MAE는 분포 이동 하에서 ImageNet-A, ImageNet-R, Portraits 데이터셋에서도 이득을 제공합니다.