[논문 리뷰] MeanShift

개요

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

관련 글: [논문 리뷰] BYOL, [논문 리뷰] SimSiam

Self-distillation 기반의 self-supervised learning 기법이다. BYOL을 고도화한 방법이라고 봐도 될 것 같다. BYOL에서 한 이미지가 있으면 두 가지 방법으로 augmentation을 해서 각각 online encoder, target encoder을 통과한다. 그리고 online encoder을 통과한 embedding은 predictor을 통과해서 target encoder을 통과한 embedding을 예측하게 된다. MeanShift에서는 target encoder의 embedding을 그대로 예측하지 않고, embedding 에 가까운 다른 데이터의 (예전에 계산했던) embedding들을 평균내서 사용한다. 밑에 그림이 잘 설명해 준다.

방법

• \( T_1, T_2, f, g, h\)는 위에 그림 참고. 기본적으로 BYOL과 동일하기 때문에 자세한 설명은 생략. Target encoder f는 BYOL처럼 exponential moving average로 구한다. Target network를 통과한 embedding의 nearest neighbor들을 평균 내서 online network & predictor의 예측 target으로 사용하는데, nearest neighbor들을 현실적인 계산량으로 구하기 위해서 memory bank를 사용한다.
• 우선 \( u = \frac{f(T_1(x))}{||f(T_1(x))||}\), \( v =  \frac{h(g(T_2(x)))}{||h(g(T_2(x)))||} \)를 계산한다.
• u는 memory bank에 추가하고 (first-in-first-out queue), k 개의 nearest neighbors들을 구해서 embedding \( \{ z_j\}^k_{j=1}\)를 구한다. 참고로 z에 u도 포함된다.
• Loss:
  \[ L = \frac{1}{k}\sum_{j=1}^k dist(v,z_j)\]
  • dist = mean square error
  • v,z는 normalize되어있기 때문에 L을 최소화하는 것은 cosine similarity를 최대화하는 것과 같은 결과를 갖는다
  • 최종 loss function은 모든 input image에 대해서 L을 더한 값
  • k=1 -> BYOL과 동일
• MoCov2, SimCLR, BYOL같은 방법들은 강한 augmentation을 필요로 하는데 MeanShift는 조금 다르다
  • 너무 강한 augmentation을 사용하면 target network를 통과한 embedding u의 nearest neighbor들은 u와 지나치게 달라서 학습에 도움이 안 될 수도 있다.
  • 따라서 target network를 통과하게 될 이미지들에게는 weak augmentation을 사용하고, online network를 통과할 이미지들에게는 strong augmentation을 사용하는 (semi-supervised learning에서 하는 것과 비슷한) 전략을 사용할 수 있는데, 이게 성능을 높이는데 좀 도움이 된다고 한다. 이 전략을 w/s version이라고 부른다 (weak/strong).
  • BYOL도 w/s 전략을 사용하면 성능이 조금 좋아졌다고 한다
  • 비슷하게 w/w, s/s 버전들도 생각해 볼 수 있다.

추가 구현 디테일

• MeanShift: momentum = 0.99, k=5, memory bank size = 1.024M. Memory bank size=128K를 사용해도 성능 저하가 없었다고 한다.
• BYOL-asym: self-supervised learning들은 계산량이 많이 요구된다. 많은 실험을 하기 위해서 계산량을 줄인 버전들로 실험을 했고, fair comparison을 위해서 BYOL도 다음과 같이 간단화했다고 한다. BYOL처럼 loss를 symmetrize하지 않았다. Batch size는 4096 -> 256. LARS 대신 SGD 사용. 200 epoch 학습
• Strong augmentation: MoCo-v2를 따라 했다 (grayscale, color jitter, horizontal flip, and Gaussian blur)
• Weak augmentation: random crop (224x224, area ratio: 0.2~1.0 사이), random horizontal flip (0.5 확률)
• Network: BYOL과 거의 동일. Encoder: ResNet50. Projection layer: 2 layer MLP (2048 -> 4096 -> BatchNorm & ReLU -> 512). Prediction layer: 2 layer MLP (512 -> 4096 -> BatchNorm & ReLU -> 512). Self-supervised training 후 projection, prediction layer을 제거한 online network의 backbone 사용.
• SGD: lr = 0.05, momentum = 0.9, weight decay = 1e-4, 200 epoch 학습

결과

ImageNet에서 linear evaluation. 참고: MSF = MeanShift

오른쪽 그림은 t-SNE 결과. 학습이 되면 cluster이 잘 형성 되는 것을 볼 수 있다. medical dataset같이 강한 augmentation을 사용할 수 없는 경우 MeanShift가 사용하기 좋다고 한다 (SimCLR, BYOL같은 경우 w/w의 performance가 안 좋다; MeanShift같은 경우 성능 저하가 작다)

Pascal VOC object detection:

추가적인 실험들은 논문 참고.