[5분 컷 리뷰] DINO v2: Learning Robust Visual Features without Supervision

 

요약

---

iBOT(masked image modeling)과 DINO(knowledge distillation with no label)방법을 혼합한 방법

 

Transactions on Machine Learning Research (01/2024)

 

Data processing: LVD-142M dataset의 생성과정

많은 양의 데이터 학습을 위해, 데이터 처리방식을 1) 라벨, 비라벨 이미지 수집, 2) 중복이미지 제거, 3) 이미지 검색

1. 데이터 수집: 데이터 수집은 라벨링된 데이터와 라벨링 안된 데이터 크게 2가지의 데이터셋을 구축했습니다. 라벨링 된 데이터는 classifcation, segmentation, depth estimation 등의 다양한 문제에 사용되는 공공데이터를 수집했습니다. 다른 한편으로는 웹크롤링을 통해서 <img>테그가 있는 경우를 모두 가져오고, PCA hash dedupliation 등을 이용해서 이미지의 중복을 제거했습니다.
2. 중복제거(de-duplication): 비슷한이미지만 모여있으면, 이미지의 다양성을 학습하기 어려우니, 유사한 이미지들을 제거하기위해서 중복제거를 image-copy-detection알고리즘을 이용해서 처리했습니다.
3. Self-supervised image retrieval: 라벨링된 데이터셋과 라벨링안된 데이터셋에서도 매우 비슷한 이미지가 있을 수 있기때문에,  이를 제거해야합니다. 첫 번째로, 이미지를 사전학습된 ViT-H/16로 "비정제 /정제 "임베딩 시키고 각 이미지를 코사인유사도로 유사성을 평가합니다. 그 후, K-means clustering을 이용해서 "비정제" 데이터를 군집화합니다. 하나의 정제 이미지의 벡터를 이용해서, 가장 가까운 N(4)의 이미지를 검색합니다. 1개의 정제데이터에 대해서, 비정제 데이터 4개의 유사한 Cosine similarity을 확인하고, 유사도가 낮으면(=이미지가 충분하지 않다면), 쿼리 이미지가 속한 클러스터에서 추가로 M개를 샘플링해서 보완합니다.

실제 구현은 vector DB로 쓰이는 Faiss 을이용해서 진행했습니다. 약 20개의 노드에서 위 과정을 분산처리하여 LVD-142M dataset을 생성합니다.

 

Discriminative Self-supervised Pre-training

이 연구에서 제안하는 방법론은 DINO와 iBOT loss을 혼합하는 방법론입니다. 이를 혼합하기 위해서, Student모델과 Teacher model은 각각 2개의 head을 가지고 있게됩니다.

1. Image-level objective: DINO v1에서 사용했던 것과 같습니다. Student, Teacher 모델이 있고, Student model이 반환하는 vector 와 teacher model이 반환하는 vector가 유사해지는 것을 목적으로합니다. (+centering + moving average에 대한 설명). 그 후, 파라미터는 EMA(Exponential moving average)로 student model에서 teacher 모델로 반영합니다.

$L_{DINO}=-\sum p_{t}log p_{s}$

2.Patch-level objective: iBOT을 이용한 자기지도 학습입니다. iBOT은 MIM(Masked image modeling0)으로 이미지를 더작은 패치로 나눈다음에 훈련하는 방법을 의미합니다. 이렇게 만든 저 작은 패치를 랜덤으로 마스킹합니다. (Fig 2. iBOT) 마스킹된 경우 마스킹 토큰만 줍니다. 

1. $\textbf{x}=\{x_{i}\}^{N}_{i=1}$: 이미지를 작은 토큰으로 생각해서, 패치단위로 N개로 잘게 쪼갭니다.
2. $\textbf {m}^{N}$: 마스킹인덱스입니다. 0 또는 1입니다. N개의 패치마다 마스킹할건지 여부를 결정합니다. $m_{i}=1$이면 마스킹된 상태입니다.
3. $\tilde{x}\triangleq \{\textbf {x}_{i} | m_{i}=1\}$: 마스크된 이미지를 의미합니다. m=1인 경우, 이미지들이 마스크 토큰으로 변경됩니다.
4. $\hat{x}\triangleq  \{\tilde{x}|(1-m_{i})x_{i} +m_{i}\textbf {e}_{[MASK]}\}$: N개의 패치가 마스크 토큰 또는 보존

학습의 목적은 마스크된 토큰($\tilde{x}$)을 복원하는 학습을 합니다. 아래의 그림 Figure 2에서, VIT(student)에는 마스크된 토큰을, teacher 모델에서는 전체의 패치를 다 본 후, masking 된 패치만 가져와 두 분포가 동일하도록 하는 것을 원합니다.

$L_{MIM} = -\sum _{i=1}^{N}m_{i}\cdot P_{\theta}^{patch}(u_{i})^{T} log P_{{\theta}}^{patch}(\hat{u_{i}})$

위 식에서 $m_{i}$가 곱해져서 반드시 마스킹된 경우만, 임베딩값을 복원하는 문제로 변경됩니다. 여기서도 DINO와 마찬가지로 softmax, centering이 쓰입니다. 논문에서는 아래와 같이 간략하게 표현했습니다. 

$L_{iBOT}=\sum_{i} p_{ti}logp_{si}$ (i:masked patch index, t: teacher model, s: student model)

Figure 2. iBOT의 MIM의 방식의 예시

 

3. 두 해드의 파라미터는 공유: Zhou et al(2022)에 논문에선서는 DINO와 iBOT의 해드 파라미터를 공유하는게 더 나은 성능을 보였으나, 본인들이 실험해보니 따로두는게 더 나았다고합니다.

4. Sinkhorn-Knopp centering: DINO에서는 teacher모델의 표현을 centering 후의 temperature scaling을 진행합니다. DINO v2에서는 SK(Sinkhorn-Knopp) batch normalization을 사용하는게 더 낫다고하여 Sinkhorn-knopp algorithm을 진행합니다.

5. KoLeo regularizer: 배치내에서 피쳐가 균등하게 만들기위해서 KoLeo regularizer을 사용합니다. KoLeo 정규화전에 각 벡터에 L2 norm을 진행합니다. 

$L_{koleo}= -\frac{1}{n}\sum_{i=1}^{n}log(d_{n,i})$,

$d_{n,i}=min_{j=i}||x_{i}-x{j}||$은 벡터i와 배치 내에 어떤 벡터든간의 최소거리를 의미합니다. 벡터를 L2 norm을 진행하면 거리가 1로 정규화되고, 각만 달라지게됩니다. 이 벡터들간의 거리가 좁을 때, 점점더 좁아지면 -log값이 매우커지니 서로 띄어놓게하기위해(uniform span) 이 정규화를 진행합니다. 

KoLeo regulariziation

6. Adapting the resolution: 작은 이미지를 패치로 나누면 더 작게되어 pixel level task(semgnetaion, detection)등에서 성능열하가 발생합니다. 그렇기에 학습후반에만 일부 큰 이미지518x518로 학습합니다.