이전 게시글 중 DINO에서는 Self-supervised learning은 NLP 뿐만 아니라 CV에서도 적절한 전략을 잘 사용한다면 기존 ViT/CNN 구조에서 발견하지 못한 유의미한 visual feature를 획득할 수 있음이 증명되었다는 점을 소개했었다. Supervised learning은 label이 존재하는 형태의 학습에서 손실 함수로 적용되는 objective value가 명확하다는 점, 그렇기 때문에 학습되는 encoder 및 decoder의 hypothesis를 정확하게 align할 수 있다는 장점이 있었으나 특정 task에 최적화된 파라미터는 label space가 조금만 달라지거나 input image의 domain이 조금만 달라지더라도 optimal point에서 크게 벗어날 수 있기 때문에 일반화된 성능을 보여주기 힘들다는 명확한 한계점이 존재했다. 특히 classification과 같이 이미지를 전반적으로 해석하는 task에서는 큰 문제가 없지만 segmentation과 같이 입력 이미지와 동일한 resolution에서 픽셀별 prediction이 진행되는 high-level task에서는 더 큰 문제를 불러오게 된다. NLP가 상대적으로 성능 수렴을 빠르게 달성하고 거대 모델의 property를 찾거나, tuning 방법과 관련된 연구가 진행된 바탕에는 바로 supervised learning에서 벗어났다는 사실이 존재한다. 즉 task 마다 직접 생성해주는 ground truth는 high level로 올라갈수록 cost가 높아진다는 superficial한 단점 말고도 궁극적으로 얻고자 하는 robust한 visual feature를 얻는 과정에 악영향을 준다는 문제가 있다.
DINO 첫번째 논문을 간단하게 요약하면 EMA 구조를 가져가면서 batch size에 따라 overfitting/underfitting되지 않도록 무관한 안정적인 학습을 위해 model output의 entropy를 조절하는 centering/sharpening 작업을 도입했었다. DINO에서는 ‘ ViT를 위한 SSL, 그리고 이를 통해 획득할 수 있었던 visual feature의 특징에 대해 서술했다.’라고 한다면 DINOv2는 ‘이러한 SSL 구조를 어떻게 확장시킬지 데이터/모델링 관점에서 서술했다.’고 요약할 수 있다.
SSL이 가지는 장점을 보여주는 하나의 qualitative 예시는 위의 그림에서 확인할 수 있다. 첫번째 칼럼에서 ‘새’와 ‘비행기’는 low level에서 비슷한 semantic 특징을 가지지만 엄연히 다른 도메인에 속한다. 그럼에도 불구하고 각 이미지의 pixel 사이의 PCA를 수행하고 가장 큰 값을 가지는 3개의 component를 visualize하면 비행기의 날개 부분이 새의 날개 부분에 매칭되거나 몸통 부분은 몸통 부분에, 꼬리는 꼬리 부분에 매칭되는 것을 확인할 수 있다(색을 보면). 마찬가지로 세번째 칼럼에서 하나의 말이 존재하는 이미지, 여러 말이 존재하는 이미지와 같이 이미지 인스턴스 내에 특정 물체를 나타내는 feature가 많이 존재하는 상황에서도 이러한 correspondence가 잘 유지된다던지, sketch(drawing)과 같이 input에 대한 natural shift가 발생한 상황에서도 PCA 결과가 합리적인 것을 볼 수 있다. 학습된 모델 스스로가 같은 종류의 instance를 포함하는 여러 도메인 사진에 대해서도 correspondence를 잘 인지할 수 있고, 이는 곧 추가적인 fine-tuning 작업 없이도 학습된 visual feature를 다양한 downstream task/dataset에 활용할 수 있음을 보여준다. 이를 흔히 표현하는 방식으로는 “Out of Box model”라고 부른다. 한국말로 “우물 안 개구리 형태의 모델에서 벗어남”이라고 하고 싶다.
대용량의 데이터에 대해 SSL을 수행하기에 앞서 paper에서는 data curation(데이터 정제) 작업의 중요성을 언급한다. 특정 dataset을 semantic하게 분석했을 때 bias가 존재한다면 다양한 이미지에 대한 visual feature를 뽑을 때 bias가 포함될 수 있기 때문이다. 이러한 ‘문맥상’의 정규화에 대한 중요성은 이미 SSL이 만연하게 적용된 NLP에서는 어느 정도 당연하게 인식하고 있는 사실이며, computer vision에서도 비슷한 맥락의 효과를 얻고자 한다면 data curation이 필수적이라는 것이다. 사실 생각해보면 supervised learning에서도 이와 같은 data curation이 필요하기는 하지만 label space 상에서 카테고리 간의 비율만 얼추 맞으면 학습 수렴에 큰 문제가 없었기 때문에 치명적으로 작용하는 문제가 아니었으나, self-supervised/unsupervised learning에서는 모델은 오로지 ‘이미지’에만 의존한 학습을 하기 때문이라고 볼 수 있다. 데이터를 처리한 과정은 다음과 같다. 위의 framework figure를 참고하면서 보면 이해하기가 편하다.
데이터셋 정제와 더불어 학습법도 DINO에 비해 일부 추가된 점이 있는데, 각 요소들에 대해 간단히 요약하면 다음과 같다.
DINO 원래 논문에서 사용하기로는 student/teacher network에 각각 local/global feature를 구분해서 넣고, 나오는 output에 대해 consistency loss를 cross entropy term으로 적용했었다. Loss에 대한 최적화는 student에만 적용하고 teacher는 EMA로 파라미터 업데이트하는 것까지 동일하게 사용하였다.
Student model로 들어가는 일부 patch를 랜덤하게 마스킹하고, 각 mask patch 위치의 feature 간의 cross-entropy loss를 추가하였다. Mask에 의한 augmentation 효과가 더해졌다고 보면 된다.
image/patch loss를 적용할 때 같은 head(classifier)를 사용하면 image-level loss는 overfitting되고 patch-level loss는 underfitting되는 문제가 발생하였고, 각 loss가 적용되는 헤드의 분리를 통해 이 문제를 해결할 수 있었다고 한다.
DINO에서 teacher softmax-centering하는 방식을 SWaV로 바꾸게 된다.
SWaV 논문에서는 prototype $C$가 일종의 code book 역할이 되는 $Q$로 향하는 경로로서 학습이 진행된다. 서로 다르게 augmentation된 이미지는 각자 prototype에 의해 코드북으로 매핑이 진행되고, 각자 본인의 코드북을 예측하는 것이 아니라 다르게 augmentation된 이미지에 의해 매핑된 코드북을 예측한다. 이로써 일반적인 contrative learning을 대체할 수 있다는 논리 전개가 가능하다. Collapse (모든 latent인 $z$가 동일한 code $q$로 매핑되는 케이스)를 막기 위해 code book $Q$는 각 배치 단위에서 구성되는 모든 샘플들을 각각의 코드북에 균등하게 배분하는 과정을 거친다. 즉 $B$만큼의 배치 사이즈로 $K$개의 코드북에 매핑될 때, 각 iteration 마다 코드북 하나는 최소한 $B/K$ 만큼 선택될 수 있어야한다는 조건이 필요하다. 해당 조건 내에서 최적화 문제를 풀어내는 과정을 쭉 요약하는 것이 SWaV 에 대한 내용이다. 이 논문에서 제안된 방법은 centering을 위한 빌드업이었다.
Centering은 단일 헤드의 학습 불안정성을 해소하기 위한 일종의 앙상블 장치였다. 과거의 prediction 정보를 accumulation 함에 따라 이전 input들의 정보가 이후 input의 prediction을 보다 bias되지 않게 해줄 수 있었다. 이러한 방법 대신 representation에 head를 $m$개씩 달아두고, 해당 예측 정보들에 대한 SWaV에 weight를 주어 앙상블하는 방법을 제안한 것이 바로 Weighted Ensemble Self-supervised learning이다. 엔트로피(예측의 확실성 지표)에 따른 weight가 가장 좋은 성능을 보였고, centering을 SWaV 방식에 여러 head로부터의 앙상블로 대체한 것이 DINO의 representation 성능을 증가시키는 것을 확인할 수 있었다.
배치 내에서 각 샘플들이 embedding되는 point 간의 간격을 동일하게 유지하고자 하는 정규화 term이다. 예컨데 embedding point $n$개 모두에 대해 가장 가까운 다른 point와의 거리를 log 분포로 나타내면, 이 분포가 균등하면 균등할수록 엔트로피는 커질 것이고 regularizer는 reduce 가 되게끔 역수를 취해 사용한다. 물론 embedding point 간의 거리 자체가 확률 분포를 표방하는 값으로 간주되므로 정규화 작업 전 normalization을 통해 scale을 조정해준다. 즉 거리를 균등하게 함으로써 필요한 정보량을 최소로 하고자 하는 것이다.
[ \mathcal{L}_{\text{koleo}}= -\frac{1}{n}\sum_{i=1}^n \log (d_{n, i}),~~d_{n, i} = \min_{j \neq i} \parallel x_i -x_j \parallel ]
보다 높은 해상도의 이미지를 사용했을 때 segmentation/detection과 같은 pixel level task에서의 성능이 올라간다. 이는 small object의 경우 low resolution에서 백본에 연산을 돌리면 작은 물체의 feature가 일종의 노이즈처럼 사라지는 현상이 발생되기 때문이다. 그렇다고 해서 무작정 고차원의 이미지를 가지고 모델을 학습시키는 건 메모리나 시간이 투머치로 소모적이기 때문에 권장되지는 않는다.
따라서 DINOv2에서는 pretraining 마지막 일부만 이미지의 해상도를 $518 \times 518$로 증가시켜서 학습시키는 전략을 사용했다.
DINO-v2는 일종의 테크니컬 리포트다. 사실 앞부분만 보아도 데이터 정제 과정이나 학습에 사용한 프레임워크를 방법론으로 제시했다기보단 기존의 연구로부터 이어지는 여러 insight 및 approach를 사용해서 좋은 모델을 만들어보겠다는 노력이 보이기 때문이다.
Flash attention에서 제안한 방법을 사용하면 하드웨어에 최적화해서 적용되기 때문에 더 빠르고 효율적인 연산이 가능하다. 이 논문 저자의 경우 논문에 나온 FlashAttention을 직접 구현하여 사용하였다. 사실 아직 본인은 FlashAttention 논문 자체는 이해하지 못하고 있다. 하드웨어 요소에 대한 이해가 부족한데, 어떻게 하면 이쪽으로 지식 및 기술 스택을 쌓을 수 있으려나…??
이전의 implementation에서는 서로 다른 patch token 수를 가지게 되는 global crop/local crop이 서로 따로 forward passing 및 backward passing 과정을 거쳤었다. 그러나 새롭게 구현된 버전에서는 이를 동시에 수행함으로써 그만큼의 연산량을 줄일 수 있었다.
Stochastic depth는 네크워크 깊이에 따른 layer dropout을 수행하여 레이어 간의 의존성 문제를 해결하고 학습 시 각 layer의 feature map을 빠르게 최적화하는 것이 주된 contribution이었다. DINO-v2에서는 이를 다르게 수행하는데, 레이어를 skip한다는 개념이 아니라 batch 단위로 랜덤하게 들어오는 샘플들을 각 레이어에서 drop out rate $d$에 따라 $(1-d)$만큼의 batch만 block에 통과시키는 전략을 취한다. 어차피 unsupervised setting이기 때문에 batch 순서에 따른 label space의 영향을 무시할 수 있으며, 학습마다 네트워크의 모든 block에 대해 stochastic하게 drop-out을 해줄 경우 오히려 학습 시간이 늘어날 수 있기 때문에 이를 최소화한 전략으로 보인다.
AdamW로 EMA 구조를 가지는 모델을 최적화할 때, model로 하여금 4개의 replica가 필요하다. Student, teacher와 이에 추가로 Adam/AdamW 최적화에 사용되는 first momentum 그리고 second momentum이 필요하기 때문이다. 보다 큰 모델을 사용할 때 memory footprint가 급격히 증가할 수 밖에 없는데, 이를 해소하기 위해 data parallel을 사용하였다. 그리고 이렇게 replica를 gpu에 분리를 할 때 GPU memory가 분리된다는 점에서 또다른 장점이 생기는데, 이는 weight 저장 자체는 float32로 하고 최적화 시 gpu 간 통신에서는 float16으로 부동 소숫점 절반을 날려버려도 학습 성능의 큰 저하 없이 메모리를 줄일 수 있다는 것이다. 원래대로라면 GPU 갯수가 증가하더라도 메모리 총합은 같은게 DDP의 특징이었는데, FSDP를 사용하면 communication 단의 메모리를 절반으로 줄여버리니까, 결론적으로는 보다 많은 GPU를 분리해서 사용할수록 학습 전체 메모리는 줄어드는 장점이 생긴다.
큰 모델이 가지고 있는representation을 효과적으로 작은 모델에 넘겨줄 때, 작은 모델을 scratch부터 학습시키는 것보다는 큰 모델의 prediction에 align하는 distillation 학습법이 효과적인 것은 어느 정도 알려진 사실이다. 따라서 DINO-v2에서도 결론적으로 작은 모델에 넘겨주는 방식을 distillation으로 했는데, 이때 기존 학습 framework인 EMA는 그대로 가되, 약간의 차이가 발생한다. 우선 self-distillation 구조가 아니므로 teacher는 학습된 large model을 frozen한 채로, 작은 모델을 student로 잡아 spare EMA를 수행한다. 또한 앞서 설명했던 stochastic depth, masking 같은 방법론은 제외한다.
DINO-v2는 베이스라인 논문인 iBOT(image BERT 논문)을 기준으로 짜잘한 방법들이 추가로 사용되었는데, 그래서 제안한 방법들이 얼마나 효과적인지를 task를 고정한 채로 ablation을 진행한 표이다. 이외에 여러 실험 결과들에 대한 내용이 페이퍼에 있는데, 대부분 성능이 올라갔다는 점을 드러내고 있다.
DINO-v2 논문을 보면 주인공이 모든 무술을 연마하여 마스터해버리는 무협 영화(?)장르가 생각난다. 메타의 이 연구는 단순히 scaling-up하는데 목적을 두지 않고 보다 효율적인 학습법과 그러면서도 어떻게 좋은 representation을 얻을 수 있는지 다양한 방법들을 적용해보고 실험해본 결과물로 보인다. 어느새 학계에서 성능 좋은 베이스라인 모델을 만드는 것은 불가능에 가까운 게 아니라 불가능이 되어버렸다. 이제는 논문을 쓰는 과정에서 집중해야 할 곳들은 이런 SSL 자체보다는 학습된 representation을 어떻게 활용하고, mapping하고 혹은 tuning할 것인가에 있어 보인다.