[논문 리뷰] SeLa

개요

제목: SELF-LABELLING VIA SIMULTANEOUS CLUSTERING AND REPRESENTATION LEARNING

논문 링크: https://arxiv.org/abs/1911.05371

관련 글:

• 2023.02.26 - [paper review] - [논문 리뷰] DeepCluster (Deep Clustering for Unsupervised Learning of Visual Features)
•  2023.02.18 - [paper review] - [논문 리뷰] Sinkhorn Distances: Lightspeed Computation of Optimal Transport

ImageNet 같은 dataset에 classification task로 pre-trained network가 다른 task에도 좋은 성능을 보여준다. 따라서 자동으로 label을 얻고 이 label로 classification을 하면 좋은 self-supervised learning기법을 얻을 수 있을 것이다. 이것은 DeepCluster 방법과 비슷하다. label을 얻고 그 label을 예측하는 방식들은 collapse(모든 data에 대해서 같은 label을 예측하는 현상)를 피하기 위해서 트릭들을 사용해야 하는데, DeepCluster같은 논문들은 별로 최적화되어있지 않은 방법을 사용하였고, 이 논문에서 이런 부분들을 개선하였다고 생각해도 좋다.

방법

Classification 간단 리뷰

• Classification 관련 notation:
  \( \Phi \): neural net
  \(I\): image
  \( \Phi(I) = x \in \mathbb{R}^{D}\): feature
  N: 데이터 수, K개의 label
  \( h: \mathbb{R}^D \rightarrow \mathbb{R}^K\): classification head
  \( y_1, ..., y_N \in \{ 1,...,K\}\): \( I_1,...,I_N\)에 대응되는 labels
• 따라서 class score은 \( p(y = \cdot | I_i) = \textrm{softmax}(h \circ \Phi(I_i))\)\)  (논문에서 notation을 조금 혼란스럽게 사용한다: x와 I를 혼용해서 사용하는데 I로 해석해야 맞는 곳들이 있다. 이 포스트에서는 I로 바꿔서 사용했는데 논문을 읽을 때 참고.)
• Cross entropy loss를 사용한다:
  \[ E(p|y_1, ..., y_N) = -\frac{1}{N} \sum_1^N \log p(y_i | I_i) \quad (1) \]
• 보통 label이 존제하는 dataset에서 이 방식을 사용하지만 label이 없는 경우 사용하는 방법을 이 논문에서 제시한다

Self-label 얻는 방법

• semi-supervised learning같이 일부 label이 이미 존재하면 위에 식을 최소화하면 collapse하지 않을 수도 있지만 완전히 unsupervised setting에서는 모든 이미지에게 동일한 label을 줘서 최소화할 수 있다.
• 이것을 방지하기 위해서 다음가같이 loss를 적어보자:
  \[ E(p,q) = -\frac{1}{N} \sum_{i=1}^{N}\sum_{y=1}^{K} q(y|I_i) \log p(y|I_i) \quad (2)\]
  
  • \( q(y|I_i) = \delta(y-y_i)\)이면 식 (1)과 (2)는 동일하다
• 이제 collapse를 예방하기 위해서 K개의 class로 분류할 때 각 class에 동일하게 많은 데이터가 할당되어야 하는 조건을 추가하자:
  \[ \min_{p,q} E(p,q), \forall y: q(y|I_i) \in \{ 0,1\} \textrm{and} \sum_{i=1}^{N}q(y|I_i) = N/K \quad (3) \]
  
  • 다시 말해서 각 class마다 N/K개의 data가 할당된다
• 위 문제는 언뜻 보면 아주 어려워보이지만 optimal transport 문제다:
  \( P_{yi} = p(y|I_i)/N\)는 neural net이 예측한 y,I에 대한 joint probability matrix로 해석할 수 있고 (shape: K x N),
  \(Q_{y,i} = q(y|x_i)/N\)는 예측해야 하는 joint probability matrix라고 해석할 수 있다. 
• 식 (3)에서 \( q \in \{ 0, 1\} \) 조건은 까다롭기 때문에 transportation polytope로 조건을 완화한다 (이전 포스트 Sinkhorn distance 참고). Transportation polytope란:
  \[ U(r,c) = \{ Q\in \mathbb{R}^{K\times N}_+ | Q \mathbb{1} = r, Q^T \mathbb{1} = c \} \quad (4)\]
  • 여기서 r, c는 모든 component를 더하면 1이 되는 확률 벡터로
  • 따라서 row의 합이 marginal probability r이 되고, column의 합이 marginal probability c가 되는 행렬이라고 생각 가능
  • SeLa에서는 r, c를 다음과 같이 잡는다: \( r = \frac{1}{K} \mathbb{1}\), \(c = \frac{1}{N} \mathbb{1}\)
  • Q가 transportation polytope라는 조건은 식 (3)에서 q가 discrete한 0, 1값만 갖게 했던 조건이 \( \sum_{y} q(y|I_i) = 1\)으로 완화하는 효과를 갖는다
• 식 (3)을 다음과 같이 적을 수 있다:
  \[ E(p,q) + \log N = \langle Q,-\log P \rangle \quad (5) \]
  여기서 \( \langle \cdot \rangle\) 은 Frobenius dot product (그냥 elementwise product하고 더한 것)을 의미한다. 따라서 Q를 transportation polytope로 relax했을 때 eq (3)의 task는 다음과 같다:
  \[ \min_{Q \in U(r,c)} \langle Q, -\log P \rangle \quad (6)\]
• 위 문제를 푸는 방식은 알려져 있으나 행렬 사이즈가 커지면 푸는데 오래 걸린다. 그리고 항상 discrete solution을 준다 (Q를 continuous한 transportation polytope로 조건을 완화했지만 solution은 polytope의 vertex에 나온다, i.e. discrete solution)
• 다음과 같이 KL divergence로 regularization term을 주면 solution을 구하는 것도 편해지면서 discrete solution을 방지할 수 있다 (Sinkhorn distance 포스트 참고):
  \[ \min_{Q \in U(r,c)} \langle Q, -\log P \rangle + \frac{1}{\lambda} KL(Q||rc^T)\quad (7)\]
  
  • 다음 꼴로 Q를 구할 수 있다는 것이 알려져 있다:
    \[ Q=\textrm{diag}(\alpha) P^\lambda \textrm{diag}(\beta) \quad (8)\]
  • 여기서 \( P^\lambda \)는 행렬 P의 원소들을 각각 \( \lambda\)로 exponentiation한 것을 뜻함
  • 위에 꼴로 문제가 풀린다는 것을 염두에 두면 Sinkhorn-Knopp 알고리즘을 이해하는 것은 어렵지 않다: 행렬 \(P^\lambda \)의 row sum이 r이 되고 column sum이 c가 되게 \(\textrm{diag}(\alpha), \textrm{diag}(\beta) \)를 구하면 된다. 이것을 만족하기 위해서 시도해 볼 수 있는 것은 iterative하게 row sum이 r이 되도록 \( \alpha\)를 설정하고 column sum이 c가 되도록 \( \beta\)를 설정하는 것을 반복하는 것이다. 재미있게도 이런 간단한 방법이 통한다.

따라서 다음과 같은 알고리즘이 완성된다:

참고로 step 2에서 Sinkhorn-Knopp 알고리즘을 돌려서 Q를 찾고 이것을 step 1에서 사용한다고 적혀있는데 실제로 코드를 보면 step 2에서 Q를 찾고, 이 Q에서 각 data가 가장 높은 확률로 갖는 label을 찾는 방식으로 discrete label을 얻는다 (해당 깃헙 코드 링크: https://github.com/yukimasano/self-label/blob/609bd5842633ec9933b936a3a37323f40d4528fa/sinkhornknopp.py#L136)

(Notation관한 언급: 논문에서 계속 확률들을 feature에 conditional하게 적는데, e.g. \( p(y|x_i)\) 이렇게 하면 말이 안된다. step 1같은 경우 모델을 업데이트 해야하는데 확률들을 feature x에 conditional하게 해석하면, feature이 주어진 상태에서 모델을 업데이트 한다는 이상한 상황이 된다. 이미지가 주어졌을 때 확률, i.e. \( p(y|I_i)\) 로 해석해야 맞고 이번 포스트에서 모두 이미지에 conditional하게 바꿨다.)

Mutual Information과 연관성

• Data index를 uniform distribution을 가지는 random variable로 해석해 보자 (uniform distribution은 일종의 가정): \( p(i) = 1/N\)
• posterior distribution을 이미지가 아니라 data index에 conditional하도록 생각해 보자 (그냥 관점의 차이): \( p(y|x_i) = p(y|i), q(y|x_i) = q(y|i)\)
• 그럼 식 (5)를 다음과 같이 적을 수 있다:
  \[ E(p,q) + \log N = -\sum_{i=1}^{N} \sum_{j=1}^{K} q(y,i) \log p(y,i) = H(p,q) \quad (9)\]
• p=q일 때 (9)의 minimum이 되고, 이 경우 (9)는 그냥 entropy \( H_q (y,i)\)가 된다. 
• Marginal distribution 이 q(i) = 1/N, q(y) = 1/K 이라는 가정을 했기 때문에 marginal entropy 도 \(H_q(i) = \log N,  H_q(y) \log K\)로 고정된다. 따라서 다음과 같은 식을 얻을 수 있다

• 여기서 \( I_q(y,i)\) 는 label y와 data index i 사이의 mutual information.
• 위에서 언급한 알고리즘에서는 equipartitioning이 되어있지만 이 조건을 완화하고 mutual information을 최대화하는 방식으로 문제를 풀 수도 있다. \(I_q(y,i) = H(y)-H(y|i)\)이기 때문에 \( H(y|i)=0\)이 solution. 
  • H(y|i)=0은 각 데이터 i에 한 개의 label이 고정된다는 뜻
  • H(y)=ln K이기 때문에 equipartion이 된다는 뜻
• 결론적으로 SeLa 알고리즘은 equipartition을 하면서 data index와 label 사이의 mutual information을 최대화하는 것으로 해석 가능.
• Collapsed solution은 y와 i 사이에 mutual information이 없기 때문에 maximum mutual information이라는 task는 collapsed solution으로 수렴하지 않음

Multiple heads:

• 여러 개의 classification head를 사용하면 더 좋은 결과를 얻을 수 있다고 한다.
• 구체적으로 neural net backbone에 \( h_1, ..., h_T\)개의 classification head를 달고 각 head의 output을 가지고 Eq (6)을 optimize하는 방식으로 적용한다.

추가 노트

• 참고로 SeLa 알고리즘 (step 1& step 2)는 동일한 objective를 optimize한다
• DeepCluster같은 경우 1. feature space에서 K-means 2. classifier을 학습하는 스텝을 반복하는데 두 스텝들이 동일한 objective를 optimize하지 않기 때문에 convergence의 문제가 있다 (i.e. converge 하는가?)

실험

optimal 세팅 찾는 실험

세팅:

• task: 학습된 feature에 linear classifier 학습
• data: ImageNet LSVRC-12
• architecture: AlexNet

결과:

• Self labeling steps: step 1과 step 2의 비율은 1:2정도가 적당했다고 한다 (step 1을 160번 반복했을 때) (table 1; #opts가 학습할 때 step 2를 반복한 수)
• Number of clusters K: 1k~3k 정도가 적당하다
• Number of heads: 10개 정도가 적당하다

SeLa [K x T]의미: K는 cluster의 수, T는 head의 수; c3, c4, c5는 alex net의 몇번 째 convolutional layer에서 feature을 뽑아왔는지 의미.

다른 방법과 비교:

Linear classification task & weighted KNN. 당시 sota 찍음: