[논문 리뷰] MoCo v3 요약, 코드, 구현

논문을 상세히 번역하고 한단어씩 해석해주는 포스팅은 많다.

나는 논문을 누구나 알아듣도록 쉽고 간결하게 전달하고자 한다.

 

MoCo v3

An Empirical Study of Training Self-Supervised Vision Transformers

CHEN, Xinlei; XIE, Saining; HE, Kaiming. An empirical study of training self-supervised vision transformers. In: Proceedings of the IEEE/CVF international conference on computer vision. 2021. p. 9640-9649.

 

논문 원문 링크

 

저자의 의도

새로운 방법을 제시하는 것은 아니다.
ViT의 self-supervised learning(SSL)에 대한 연구다.
CNN의 레시피는 상당히 성숙하고 robust한 것에 비하여 ViT는 미숙하다.
특히 SSL에서는 더욱 부족한 경향이 있다.
ViT의 부분적인 실패들을 분석해서 안정적이게 만들자.

 

기존 문제점

프리 트레이닝 부분에서 NLP와 CV가 격차가 너무 크다.

ViT의 self-supervised framework를 만들어야 한다.

 

CNN과 ViT는 완전히 다른 모델이고 기존 CNN 프레임워크는 통하지 않는다.
흥미로운 점은, 불안정한 ViT의 학습이 치명적인 실패가 아니라 약간의 퍼포먼스 저하만 생긴다는 점이다.
따라서 ViT가 불안정한지 안정한지 판단하기 매우 어렵다.
이런 문제점을 입증하기 위해 저자들은 간단한 트릭을 사용했다.
저자들은 경험적으로 patch projection layer를 고정하는 방법이 유용하다고 어필했다.

 

해결 아이디어

1. MoCo v3

MoCo v3

MoCo v1과 v2에 비해 단순성, 정확성, 확장성 간의 밸런스를 높였다.

MoCo에 대한 내용을 모르면 아래를 읽어보자.

 

MoCo v1 : https://davidlds.tistory.com/42

[논문 리뷰] MoCo v1 요약, 코드, 구현

MoCo v2: https://davidlds.tistory.com/43

[논문 리뷰] MoCo v2 요약, 코드, 구현

 

MoCo

일반적인 접근법과 동일하게 각 이미지에 랜덤 증강을 적용해 2개의 crop을 얻는다.
이 2개의 crop은 각각 쿼리 인코더, 키 인코더에 의해 인코딩 된다.

인코더의 결과물로 벡터 q와 k를 얻는다.

각각 쿼리와 키의 역할을 하고 학습의 목표는 쿼리로 키를 검색하는 것이다.
이는 contrastive loss function을 최소화하는 것으로 공식화할 수 있다.

우리는 이를 InfoNCE 라 하며 아래 Equation 1으로 나타낸다.

 

Equation 1

여기서 k+는 쿼리와 같은 이미지를 키 인코더로 인코딩한 것이다.
그리고 k-는 쿼리와 다른 이미지를 키 인코더로 인코딩한 것이다.

이번 버전부터 MoCo는 dictionary queue 방식을 포기한다.
대신 거대한 사이즈의 배치 (ex. 4096)를 채택했다.

인코더들은 ResNet 혹은 ViT를 backbone으로 사용한다.
앞선 버전의 MoCo처럼 키 인코더의 파라미터를 업데이트 한다.

(큐 인코더의 moving average를 통해서 천천히 업데이트 시킨다.)

 

2. Stability of self-supervised ViT training

프레임워크에서 간단하게 backbone을 R50에서 ViT로 바꾸고자 했다.
그런데 학습의 불안정성에 직면했는데, 이게 정확도로 판단하기 어려웠다.
왜냐하면 불안정성이 있다고 판단할 만큼 낮은 퍼포먼스가 아니기 때문이다.

 

트랜스포머를 학습해본 사람은 다 아는 문제 일 것이다.

lr이나 epoch가 심하게 틀어지지 않으면 생각보다 정확도가 높게 나온다.

그래서 진짜 최적의 lr과 epoch를 찾는게 엄청난 노가다일 수 밖에 없다.

 

이게 새로운 아키텍처를 연구하는 사람에게 심각한 문제로 다가온 것이다.

아키텍처가 잘못 됐는데 정확도가 나쁘지 않게 나오고 티가 안나는 것이다.

생각만해도 골치 아프다....... (이거 나잖아...?)

 

저자들은 kNN 커브를 그려보고 나서야 이를 증명할 수 있었다.

 

2-1. Empirical observations on basic factors

어떻게 basic factor들이 안정성에 영향을 주는지 연구했다.

 

먼저 Batch size 이다.

ViT는 기본적으로 연산량 측면에서 무겁게 설계됐다.

따라서 배치 사이즈도 더 클수록 유리하다.

과연 실제 결과도 그럴까.

 

batch size

1k, 2k 배치는 kNN 커브가 완만하게 올라간다.
4k 배치에 도달하면 이제 학습이 불안정해진다.
중간중간에 구덩이인 'dips'가 발생한다.

4k 배치는 2k 배치와 최종 acc 측면에서 거의 차이가 없어 보인다.
하지만 다음 섹션에서 나오는 테스트를 해보면 문제가 발견됐다.

결론은 차이가 거의 없어 보이지만 학습이 잘못됐다.

(이는 Appendix에서 설명하겠다고 했는데 모든 논문의 버전에 appendix가 사라져있다.)

(아마도 robustness 문제를 설명하지 않았을까 싶지만... 다시 실험해보니 아니였는지... 아무튼 내용을 지워버려 알 길이 없다.)

6k 배치를 보면 왜 ViT가 골치아픈 모델인지 볼 수 있다.
구덩이가 매우 깊고 많은 것을 볼 수 있다.

저자들의 가정은 학습을 하다가 이번 로컬 옵티멈을 빠져나오기 위해

부분적으로 재시작 해버리고 'jump out'하여 새로운 전략을 모색한다는 것이다.

보통 이런 경우 발산하여 아예 망해버리고 acc가 낮게 나와야 한다.
그런데 결과를보면 acc 측면에서 3%로 작은 차이가 난다.

그러니까 acc만 보면 되는게 아니라 dips가 생기는지 다 그려봐야 한다.........ㅇㅁㅇ

 

 

다음으로 learning rate 이다.

learning rate

여기서 best lr은 1.0e-4 이다.
lr이 작은 경우는 학습이 더 안정적이지만 언더피팅 된다.

lr이 커지면 불안정성을 보인다.
lr이 큰 경우는 학습이 진행되며 dips가 더 자주 깊게 발생한다.

 

3. A trick for improving stability

간단한 트릭을 이용해 안정성을 향상시킨다.

 

A trick for improving stability

앞서 관찰한 dips가 발생하는 이유는 갑작스러운 gradient 업데이트인 'spike' 때문이다.

위 그림은 iteration이 진행되면서 레이어들의 gradient 변화에 대한 그래프 이다.
전체 레이어를 모두 관찰해본 결과 가장 첫 레이어에서 많이 발생했다.

 

첫 레이어는 patch projection layer 이다.
저자들은 이런 현상을 방지하기 위해 patch projection layer를 고정했다.
구현은 random patch projection을 만들어 값을 고정하였다.

이제 이 해결책의 실험 결과를 보자.

 

Comparisions

(patch projection을 학습하는 경우) vs (고정하는 경우) 두가지를 비교한다.

고정한 경우 매우 안정적이고 스무스하게 학습된다.
모든 lr 조건에서 더 좋은 최종 acc가 나왔다.

이걸 다른 프레임워크인 SimCLR, BYOL에 적용해봤는데 거기서도 효과가 있었다.

다시말해 ViT의 SSL 상황에서 보편적인 문제점을 발견한 것이다.

 

결과 분석

[Implementation details]

AdamW, step weight decay, warm up, contrastive loss(τ=0.2), Linear probing은 SGD

 

Table 2

Table 3

 

사용한 ViT들의 세부 속성과 특징들이다.

 

1. Self-supervised learning frameworks

Table 4

Table 4를 보자.

모든 framework에 random projection trick 사용해서 평가했다.

MoCo v3가 다른 framework 보다 우수하다.
R50 에서는 BYOL이 더 높게 나왔다.

 

2. Comparisons with prior art

Table 1

Table 1을 보자.

Self-supervised transformers 사이의 비교다.

ViT의 모든 종류를 MoCo v3로 학습했다.

결과는 iGPT 보다 매우 가벼우며 퍼포먼스도 더 좋다.

아래 'mask patch prediction'은 그냥 일반 ViT를 JFT-300M에서 프리트레이닝한 모델이다.
용어에서 혼동을 야기할 수 있으나 auto encoder가 아니다.

프리트레이닝 후 파인튜닝한 경우도 MoCo v3가 더 우수하다.

 

Figure 8

SimCLR이나 BYOL으로 학습한 거대 ResNet 보다 MoCo v3 ViT가 더 퍼포먼스가 좋다.

 

Table 5

이는 프리트레이닝 후 웨이트 frozen 없이 전체를 학습한 것이다.
일반 ViT와 DeiT 보다 MoCo v3가 퍼포먼스가 더 좋다.

 

Table 6

4가지 down stream task을 평가했다.
ImageNet classification supervised 보다도 높게 나왔다.

 

코드 및 구현

기존 MoCo와 다른 점은 backbone을 인스턴스해서 사용한 것 뿐이므로 생략함.

 

오피셜 깃허브 코드

 

 

관련 논문 리스트 (스크롤 내려서 Paper List 참고)