DINO v1을 리뷰하고 BYOL을 리뷰하고 다시 DINO v2를 리뷰하는 뒤죽박죽이다.
Abstract
대용량 데이터셋으로 사전훈련된 모델을 이용할 수 있음에 따라 이미지의 시각적 특징만을 이용해 원하는 작업을 위한 정보를 얻을 수 있게 되었다. 이 연구에서는 다양한 원천에서 수집된 데이터로 사전훈련된 모델이 존재한다면 충분한 정보를 제공할 수 있다는 것을 보여준다. 우리는 데이터의 측면에서, 흔히 자기주도학습에서 사용되는 정돈되지 않은 데이터 대신 dedicated, diverse, curated image dataset을 위한 파이프라인을 제시한다. 또한 모델의 측면에서는 ViT 모델을 이용해 작은 모델에 distill시킬것이다.
Introduction
우리는 이 연구에서, 만약 대용량의 데이터로 자기지도학습을 진행하면 general-purpose visual features를 획득할 수 있을지 탐구한다.(특정 task에 한정되지 않고 모델의 구조 변경 없이 모든 task에 적합한 features를 획득할 수 있는지). 우리의 기술적인 초점은 모델과 데이터의 사이즈를 키웠을 때 안정적이고 빠른 학습이 가능하도록 하는 것이다. 또한 우리는 매우 많은 uncrated images를 filter, rebalance 하기 위한 파이프라인을 제시한다.
Data Processing
LVD-142M dataset을 이용하였고 추가적인 metadata나 text등 없이 이미지만을 이용해 중복처리 등과 같은 이미지 전처리 작업을 진행했다고 한다. 코드는 깃허브에 공개되었지만 데이터셋은 아직 release되지 않은 듯 하다.
curated dataset으로는 ImageNet-22k, ImageNet-1k, Google Landmarks 등이 있고 uncurated 로는 웹 크롤링을 통해 1.2B의 데이터를 수집했다. (142M = 1억 4천만 이미지 ㄷㄷ..)
Pizzi et al. (2022)에 의한 copy detection pipeline을 통해 중복되거나 유사한 이미지를 제거하였다.
데이터 전처리를 위해 V100-32GB GPU 8개가 달린 컴퓨터 20개를 동원해 2일이 걸렸다고 한다. 클라스가 다르네
Discriminative Self-supervised Pre-training
우리는 DINO + iBOT의 손실함수를 조합하고 SwAV의 centering을 이용하여 학습을 진행하였다.
1. Image-level objective (Caron et al., 2021)
- DINO v1에서와 같이 ViT의 class token을 이용하였고 student, teacher network 두개가 존재한다. 또한 모델 붕괴를 피하기 위해 centering, sharpening을 적용하였다고 했었는데 v2에서는 softmax 이후 Sinkhorn-Knopp centering과 지수 이동 평균을 통해 \(p_t\)를 획득한다.
\(p_t\) : teacher class token으로부터 얻어진 prototype score에 softmax를 적용한 것
\(p_s\) : student class token으로부터 얻어진 prototype score에 softmax를 적용한 것
2. Patch-level objective (Zhou et al., 2022a)
- student의 입력 패치에 대해 random masking을 적용(teacher에는 적용 x). 마스킹된 student token에 student iBOT을 적용하고 teacher에도 똑같이 teacher iBOT을 적용한다(visible patch token에만).
iBOT이 뭔지 몰라서 다음 리뷰는 iBOT확정
\(i\) : 마스크 토큰에 대한 패치 인덱스
\(L_{DINO}\)와 유사하게 student parameter의 지수이동평균을 통해 teacher head를 업데이트 한다.
3. Untying head weights between both objectives
- DINO, iBOT은 학습가능한 MLP 헤드로부터 손실을 계산하는데, 가중치를 공유하는 동일 헤드를 사용했을때 성능향상이 있었다는 연구가 있었지만 우리의 경우 크기를 키웠을 때 반대의 효과가 일어났기 때문에 DINO, iBOT에 각각의 고유한 헤드를 사용하였다.
4. Sinkhorn-Knopp centering (Caron et al., 2020)
- Ruan et al.(2023)에 의하면 DINO, iBOT의 teacher softmax-centering step을 SwAV의 Sinkhorn-Knopp batch normalization으로 대체하는것을 권장했다. iBOT 다음 리뷰는 SwAV 확정ㅋㅋ
우리는 Sinkhorn-Knopp을 3번 반복한 이후, Student에만 softmax normalization을 적용하였다.
5. KoLeo regularizer (Sablayrolles et al., 2019)
\(d_{n,i}\) = \(min_{j\ne i} \parallel x_i - x_j \parallel\) = \(x_i\)와 배치 내 다른 모든 점들 간의 최소거리
모든 features들에 대해 \(l_2-norm\)을 적용한 이후 위 regularizer을 적용하였다.
6. Adapting the resolution (Touvron et al., 2019)
이미지 화소를 높이는 것은 segmentation, detection 등에서 downstream을 위한 중요한 방법이다. 화소를 높이는 것은 메모리와 시간의 소비가 매우 커지기 때문에 우리는 훈련의 마지막에서 짧게 518x518의 크기로 이미지 화소를 향상시켜보았다. 이는 UniVIT, FlexiViT에서와 비슷한 방법이다.
(놀랍게도 이건 Related work가 아닌 Pre-training을 위해 사용된 method들을 나열한 것이다. 일단 여기까지만 읽어보면 대용량의 데이터를 때려박기위해 좋다는 기술들도 다 때려 박았다. 이런 노력의 결과로 우리들은 pretrained 모델을 사용할 수 있게 되었기 때문에 무한히 감사해야함)
아래에서 부터는 대용량의 데이터와 대규모 모델을 훈련하기 위해 수정된 각종 알고리즘들이 소개된다.
Efficient implementation
1. Fast and memory-efficient attention
- 우리는 메모리 사용량과 self-attention의 속도를 향상시키기 위해 우리만의 FlashAttnetion(Dao et al.)을 구현하였다. 우리의 GPU에서는 embedding dimension per head, full embedding dimension에 대해 각 64의 배수, 256의 배수가 가장 효율적이였다. 때문에 기존 ViT-g와도 구조 상 차이가 있으며 1408 embed / 16 heads 에서 1536 embed / 24 heads로 증가시켰다. 결과적으로는 정확도 상 차이는 크지 않았으나 매우 효율적인 훈련을 진행했다고 한다.
2. Sequence Packing
- DINO 알고리즘 상 (224x224), (64x64)의 이미지가 동시에 순전파 되어야 하기 때문에 token sequence의 길이가 차이가 난다. 우리는 NLP에서 사용되던 sequence packing(Krell et al., 2022)이란걸 도입했다. 두 시퀀스를 concatenate 하여 트랜스포머 블록으로 전달하는데, self-attention에서 대각행렬에 마스크가 적용되기 때문에 서로 다른 시퀀스들 간에 attention이 되지 않는다.
3. Efficient stochastic depth
- 우리는 결과값에 마스킹하는 stochastic depth 대신 dropped residual에 대한 계산을 아예 skip 하였다. 본 연구에서는 d=40%의 drop rate을 사용하였고 batch dimension 내 \(B\) sample 을 무작위로 섞고 첫 (1-d) x \(B\) sample만큼을 잘라낸후 계산을 진행하였다.
4. Fully-Sharded Data Parallel(FSDP)
- student, teacher, optimizer first moments, optimizer second moments를 계산하기 위해 AdamW optimizer의 복제본 4개가 필요하였는데 PyTorch의 FSDP를 이용하여 메모리 효율 50%를 증가시켰다고 한다.
5. Model distillation.
- 특별한 내용은 없으니 논문 참조.
Results
실험 결과와 ablation study가 매우 많기 때문에 결과는 논문을 참조하기 바람.
