Test-Time Training with Masked Autoencoders(2022) 논문리뷰

해당 논문은 지난 번에 리뷰했던 Test Time Training with self-supervision의 후속논문으로 Rotate방식에서 MAE방식으로 변경한 것이 차이점이며 Test Time Training이 왜 작동하는지 조금 더 자세하게 설명했다.

https://arxiv.org/pdf/2209.07522.pdf

 

Introduction

Generalization은 지도학습에서 중요한 문제이다. 기본적으로 Test할 때 train set 과 test set이 same distribution 에서 나온 것으로 가정되어서 성능을 평가하지만 실제로는 그렇지 않은 경우가 많이 때문이다. 따라서 훈련되는 모델들은 possible distribution에 대하여 robust할 필요가 있지만 possible future은 굉장히 넓기 때문에 수행하기 힘들다.

 

Test Time training(TTT)은 test input으로부터 test distribution에 대한 hint를 받을 수 있다는 점에서 시작하여, one-sample ㅣlearning problem이라고 한다. 그러나 해당 test input에 대한 label은 없다. 그렇기 때문에 이미지 자체 내에서 label을 만드는 self supervision방법을 사용하며 이는 Training과정에서도 main task와 같이 훈련된다. 그리고 사용하는 self supervision방법이 MAE이다.

 

Related work

Unsupervised domain adaption -- 가장 쉬울 수 있는 방법은 train 과정에서 test distribution에 대한 unlabeled data를 노출시키는 것이다. 그러나 이러한 방법은 train 과정에서 훈련한 데이터들에 대해서만 효과적이며 상당히 제약적(restrictive)이라고 할 수 있다.

 

이에 대해 TTT에 대해서 생각해보면 UDA에서의 one-sample unsupervised domain adaption이라고 볼 수 있다고 한다. 왜내하면 one target sample만 사용하기 때문이다(물론 test과정에서만 사용하고, 업데이트한 gradient는 데이터마다 초기화된다)

 

MAE는 encoder, decoder를 훈련하게 되는데 훈련이 끝난 후(pre-train)에는 encoder만 downstream task에 사용하게 된다. downstream task과정에서 훈련하는 방법은 2가지가 있는데, 1) Fine tuning: encoder, head 둘 다 훈련. 2) Linear probing: encoder is frozen, head만 훈련

 

Method

구조는 Y-shape을 가진다고 한다.

figure는 이전 논문의 구조를 가져온 것이며 여기서 Rotate prediction을 MAE로 바꾸었다고 생각하면 된다.

 

위 figure는 참조 영상을 참고하여 만든 것이며 Encoder, Decoder는 ViT-Large이다.

 

그리고 훈련 방식에 있어서 3가지 방식이 존재하는데 일단 f, g는 ImageNet 1K로 pre-trained 된 상태이다. 

1) Fine-tuning: f, h 모두 훈련하는 방식이다. 그러나 TTT방식에서는 잘 작동하지 않는다고 한다.

2) ViT probing: f를 frozen하고 h만 훈련하는 방식이다. 이러한 방법은 fine-tuning, joint training방식보다 fewer parameter를 가진다고 한다.

3) Joint training: h, j, g모두 훈련하는 방식이다. 그러나 이것도 ViT probing방식보다 좋지 않다고 한다.

 

따라서 ViT probing방식으로 MAE를 적용한다면 다음과 같다.

여기서 $ f_{0}$ 는 pre-trained 된 것이다.

 

이렇게 훈련된 model을 가지고 TTT를 수행할 경우 업데이트 되는 모델은 f, g 이며 다음과 같이 나타낼 수 있다.

$ f_{x}, g_{x}$로 나타낸 이유는 데이터 한개마다 업데이트 되며 다른 데이터일 경우 다시 $ f_{0}, g_{0}$으로 돌아가기 때문이다. 기존의 supervised 혹은 self-supervised방식과는 다르게 Training, Test 모두 모델을 업데이트 하기 때문에 약간 헷갈릴 수 있긴하다. 

 

최적화 방법으로는 기존 MAE에서는 AdamW를 사용하지만 여기서는 SGD를 사용해 최적화를 수행한다. 그 이유에 대해서는 epoch마다 편차가 심했다고 하며 그에 반해 SGD는 안정된 모습을 보인다.

 

Result

여러 corruption에 대해서 가장 좋은 모습을 보여준다. 그러나 한계점으로는 한 개의 데이터에 대해서 20번에 대한 훈련을 수행하면서 classification을 했기 때문에 속도가 굉장히 느리다는 단점이 있으며 이를 해결하는 것이 큰 문제라고 할 수 있으며 해당 MAE를 사용한 방식이 얼마나 일반화를 가질 수 있는지는 더 연구해야한다는 단점을 언급한다.

 

추가로 왜 TTT가 좋은지는 bias-variance tradeoff를 언급하며 설명하는데 이를 eigenvalue, PCA를 통해 설명하는데 조금 더 공부가 필요할 거 같다.

 

Reference

https://www.youtube.com/watch?v=zOJ68s0F6JY