[논문 리뷰] Barlow Twins

개요

논문 링크: https://arxiv.org/abs/2103.03230 (Barlow Twins: Self-Supervised Learning via Redundancy Reduction)

Self-supervised learning 기법 중 하나. 이미 여러 가지 방법들이 있는데, contrastive 방법들 (예: [논문 리뷰] SimCLR ), clustering 방법들(예: [논문 리뷰] SeLa,[논문 리뷰] SwAV), prediction network를 사용하는 방법들 ([논문 리뷰] BYOL,[논문 리뷰] SimSiam)같은 방법들과는 다른 방법을 제시한다. 밑에 그림과 같이 이미지를 두 가지 방법으로 augmentation을 해줘서 embedding z를 얻고, cross-correlation matrix를 계산한다. 이 cross-correlation matrix가 identity matrix가 되도록 학습하는 것이 Barlow Twins에서 제시하는 방법이다.

Barlow Twins 방법의 motivation은 다음과 같다: augmentation에 대해서는 feature들이 invariant 해야 하기 때문에 cross correlation matrix의 diagonal component가 1이 되는 것이 좋다. 또한 z의 다른 component 사이에는 중복되는 정보가 없어야 효율적으로 embedding 되기 때문에 diagonal component가 0이 되도록 학습한다. 이러한 redundancy reduction 하는 아이디어는 Barlow의 1961년 논문 (Possible Principles Underlying the Transformations of Sensory Messages)에서 출발하였다고 한다:

In his influential article Possible Principles Underlying the Transformation of Sensory Messages (Barlow, 1961), neu- roscientist H. Barlow hypothesized that the goal of sensory processing is to recode highly redundant sensory inputs into a factorial code (a code with statistically independent components)

학습할 때 large batch size, prediction network, momentum encoder, Sinkhorn-Knopp algorithm, stop-gradient같은 방법들이 필요 없는 간단한 방법이지만 좋은 성능을 보인다.

---

방법

알고리즘

• \( X \): batch of images
• \( f_\theta\): neural network
• \( Y^A, Y^B\): \( \mathcal{T}\)에서 augmentation transformation을 샘플링해서 X를 변환시켜 준 batch of images.
• \( Z^A, Z^B\): \( Y^A, Y^B\)를 neural network에 통과시켜서 얻은 batch of features. 참고로 batch dimension으로 mean=0이 되도록 normalize 한다
• Loss:
  \[ \mathcal{L}_{BT} = \sum_{i} (1-\mathcal{C}_{ij})^2 + \lambda \sum_{i,j\neq i} \mathcal{C}_{ij}^2 \quad (1)\]
  \[ \mathcal{C}_{ij} = \frac{\sum_b z^A_{b,i} z^B_{b,j}}{\sqrt{\sum_b (z^A_{b,i})^2} \sqrt{\sum_b (z^B_{b,i})^2}} \quad (2) \]
  • \( \mathcal{C}\): cross-correlation matrix (사람마다 약간씩 다르게 사용하는데 여기서는 batch에 대해서 normalized 된 covariance를 cross-correlation으로 사용했다)
  • Eq (1) 에서 첫 번째 항은 augmentation에 대해서 invariant 하게 enforce 한다
  • Eq (1)에서 두 번째 항은 embedding 되는 vector의 component끼리 independent 하도록 enforce 한다 (redundancy reduction)
  • b: batch 안에 있는 sample index
  •  i,j: feature vector dimension index
• 알고리즘:

Information Bottleneck principle과 연관

• 핵심 내용은 아니라서 스킵해도 괜찮다. 개인적으로 재미있어서 자세하게 설명함.
• Information Bottleneck principle에 의하면 좋은 representation은 데이터에 대해서 정보를 최대한 많이 담고 있어야 하지만 관심 있는 정보와 무관한 정보(다시 말해서 데이터의 distortion; augmentation이라고 보면 된다)는 담고 있지 않아야 한다 (augmentation에 대해서 invariant 해야 한다는 뜻)
• 밑에 그림이 참고가 될 것이다

논문 Fig. 6

• Fig. 6에서 말하는 information bottleneck principle을 다음과 같은 식으로 적을 수 있다:
  \[ IB_{\theta} = I(Z_\theta,Y) - \beta I(Z_\theta,X) \quad (5)\]
  • I: mutual inforamtion
  • \( \beta > 0\): tradeoff coefficient
  • 설명: 개인적으로 \( \beta = 1 + \lambda\)로 두면 더 이해하기 쉬운 것 같다: \( IB_{\theta} = (I(Z_\theta,Y) - I(Z_\theta,X)) -  \lambda I(Z_\theta,X) \). 괄호 안에 있는 항은 distorsion이 있던 없던 representation Z가 갖는 information이 같다고 해석이 되고, 두 번째 항은 representation이 data에 있는 information을 최대한 보존하는 역할을 한다.  (이렇게 decompose 한다고 식의 의미는 바뀌지 않지만 이해하기 편한 것 같다)
• Mutual information의 성질 \( I(X,Y)=H(X)-H(X|Y)\) 의하여 식 (5)는 다음과 같다 (함수 f는 deterministic 하기 때문에 \( H(Z_\theta |Y)=0\)인 사실을 사용했다):
  \[ IB_\theta = [H(Z_\theta)] - \beta [H(Z_\theta) - H(Z_\theta |X)] \quad (6)\]
• positive constant factor은 중요하지 않기 때문에 \( IB_\theta / \beta \rightarrow  IB_\theta\)로 다시 정의하면 
  \[ IB_\theta = H(Z_\theta |X) + \frac{1-\beta}{\beta}H(Z_\theta)  \quad (7)\]
• Z가 Gaussian distribution을 갖는다고 가정하자 (사실 Z의 conditional probability도 Gaussian이라고 가정한다). Gaussian distribution의 entropy는 상수를 좀 무시하면 log det (C), C=covariance matrix가 된다 (유도과정이 어렵지는 않다; 구글링해도 쉽게 찾을 수 있다). 따라서 다음과 같이 적을 수 있다:
  \[ IB_\theta = E_X \log |C_{Z_\theta |X}| + \frac{1-\beta}{\beta} \log |C_{Z_\theta}| \quad (8) \]
  • \( \beta \le 1\)인 경우 Z를 X와 무관하게 상수가 되도록 잡으면 variance -> 0이 되면서 information bottleneck loss (IB)가 최소화된다. 별로 의미 없는 상황.
  • \( \beta > 1\)인 경우 두 번째 항 앞에 붙는 개수(\( \frac{1-\beta}{\beta}\))는 음수라서 상수 \( -\lambda\)로 교체
  • 이렇게 해서 학습하면 SoTA까지는 안된다고 한다
  • 우선 Information bottleneck loss 에서 \( -\lambda \log |C_{Z_\theta}| \)만 고려해 보자. 이 항을 최소화하는 것은 entropy (i.e. information)을 최대화하는 것과 같다. 이 항을 cross-correlation matrix의 Frobenius norm을 최소화하는 objective로 바꿔보자. 그러면 loss function은 바뀌지만 하려는 것은 어떻게 보면 똑같다: Z가 batch dimension으로 normalize 되어있다고 가정하면 (cross-correlation matrix는 데이터를 noramlize 하던 안 하던 같기 때문에 상관없다) 새로운 objective는 cross-correlation matrix의 off-diagonal component를 최소화하는 것과 같다 (diagonal component는 normalization 때문에 1이니까). 하지만 correlation을 없애는 것은 entropy을 최대화하는 것과 같다 (entropy는 확률분포가 얼마나 random 한 지 나타내는 지표; correlation이 없으면 더 random). 따라서 cross-correlation matrix의 Frobenius norm을 최소화하는 loss로 교체해도 하려는 목표는 비슷하다.
  • 새로운 Eq (8) (두 번째 항을 Frobenius norm으로 교체한 loss)과 Eq (1)의 두번째 항이 같으려면 Eq (1)에서 auto-correlation을 사용해야 한다 (Eq 1은 서로 다른 augmentation을 사용한 A, B 사이의 correlation; Eq 8은 A, A 사이의 correlation). 이 두 가지 방법은 실험적으로 차이가 별로 없다고 한다.
  • Eq 8의 첫 번째 항은 X가 주어졌을 때 Z가 최대한 random하지 않는 loss이다. neural net 은 deterministic하기 때문에 distorsion 에 대해서 Z가 바뀌지 않는 것을 추구하는 loss라고 해석 가능하고, 따라서 Eq 1의 첫번째 항과 목표가 동일하다는 것을 알 수 있다.

Implementation Details

BYOL이랑 비슷.

• Augmentation: random cropping, resizing to 224 × 224, horizontal flipping, color jittering, converting to grayscale, Gaussian blurring, and solarization
  • 처음 두 개는 항상 사용하고 마지막 5개는 random 하게 적용
• Network: encoder (ResNet 50; feature dim 2048) + projector network (MLP: 2048->8192->8192->8192, 중간에 batch norm & ReLU 사용)
• optimizer: LARS, 1000 epochs, batch size 2048 (256으로도 잘 학습됨)
• learning rate: neural net의 weight는 0.2 * (batch size / 256); bias & batch norm parameter은 0.0048 * (batch size / 256)
• 10 epoch동안 warmup, 그 이후 cosine decay scheduler을 써서 1000배 낮춰준다. 
• \( \lambda = 5 \times 10^{-3}\) 
• weight decay = \( 1.5 \times 10^{-6} \)
• bias, batch normalization parameter에 대해서는 LARS & weight decay를 사용하지 않는다

---

실험 결과

전반적으로 다른 SoTA들과 비교할만한 결과들을 보였다 (SwAV, MoCo-v2, BYOL, SimSiam)

Barlow Twins에서 사용한 ResNet encoder의 output에 linear classifier을 붙인 결과 (ImageNet):

결과가 좋은 편이다

ResNet encoder + classifier을 사용해서 fine-tuning한 결과 (ImageNet)

다른 dataset에 대해서 classification 결과 (ImageNet 에 대해서 pre-training; encoder은 학습하지 않고 linear classifier만 학습)

Detection에 대한 transfer learning (ImageNet pre-trained, detection에 fine tuning)

스킵한 내용

많은 실험들에 대해서 궁금하면 논문의 setcion 4 Ablations 참고

InfoNCE와 연관성, IMAX 기법에 대한 코멘트, 등등은 section 5 Discussion 참고.