Barlow Twins 논문 리뷰

Barlow Twins: Self-Supervised Learning via Redundancy Reduction

---

Summary

본 논문은 self-supervised learning에서 자주 발생하는 collapse 문제를 해결하기 위해, cross-correlation matrix를 활용한 redundancy reduction 접근을 제안한다. 제안된 방법은 수학적으로 representation이 무너지지 않도록 설계되었으며, 대규모 negative 샘플이나 비대칭 구조 없이도 작은 배치에서 효율적인 학습이 가능하다. 이러한 단순하고 원리적인 특성 덕분에, 본 방법은 end-to-end 학습으로 안정성과 효율성을 동시에 달성한다.

---

Method

본 논문의 핵심은 loss function이다. 간단한 구조 설명 이후 이 loss function이 어떻게 설계되어 있는지 탐구하고 이러한 접근법이 기존의 SSL 방법론들과 어떻게 다른지 살펴보도록 하자.

Input image $X$는 서로 다른 두 augmentation 방법으로 증강되며, 증강된 두 이미지 $Y^A$와 $Y^B$는 모두 같은 인코더 $f_\theta$를 통과하여 임베딩이 만들어진다. 이 때 인풋은 배치 단위로 진행되며, 같은 이미지의 상관계수는 높이고 다른 이미지의 상관계수는 낮추는 방식으로 학습된다. 즉 모델의 최종 목표는 이미지들의 representation의 correlation matrix가 항등행렬에 가까워지도록 학습하는 것이다.

---

Loss Function

Loss function은 다음과 같다.

이 때 $\mathcal{C}$는 cross-correlation matrix를 의미한다.

여기에서 $z$는 0으로 centering된 벡터이다.

 

Loss function을 분석해보면, 서로 다른 이미지에 대한 correlation이 0이 되도록 학습되기 때문에 임베딩이 서로 직교하도록 학습되는 것을 확인할 수 있다. Correlation이 0이라고 해서 직교인 것은 아니지만, 벡터는 선형이기 때문에 벡터 공간 구조에서는 직교 이외의 무상관 상태가 존재하지 않기 때문에 직교임을 확인할 수 있게 되는 것이다. (애초에 $\mathcal{C}_{ij}$가 cosine similarity를 의미하는 것을 확인하는 것으로도 충분하다)

---

Comparison with InfoNCE

InfoNCE loss는 다음과 같다. 여기서 similarity는 보통 코사인 유사도로 간주된다.

InfoNCE는 representation 벡터 간 거리를 직접 밀어내고, Barlow Twins는 벡터의 좌표축을 서로 직교하도록 배치하는 방식으로 밀어낸다. Sample(또는 배치) 수가 충분할 때 embedding 전체 분포가 퍼지도록 유도하는 것은 동일하다. 즉 샘플 분포의 엔트로피를 상승시키는 방향이라는 측면에서는 같다고 볼 수 있다.

 

그러나 배치 수가 충분하다는 것은 이론적인 상황에서 가능하며,  그렇지 않은 경우 InfoNCE에서는 문제가 발생할 수 있다. InfoNCE의 loss는 softmax 분류 문제로 해석되어, negative sample가 많아야 분포가 더 넓게 퍼질 수 있다. 그러나 배치가 작은 경우 negative가 적어서, 임베딩 공간에서 positive만 붙이고 negative는 충분히 밀어내지 못하는, trivial solution의 함정으로 빠지기 쉬워지는 것이다.

 

따라서 배치 수를 늘려야 하는데, 이는 그대로 GPU의 연산과 메모리 부담으로 이어진다. 이것이 바로 CimCLR의 단점이기도 한 것이다.

 

반면 BT loss의 경우 밀어내는 방향으로 학습하는 것이 아니라, 차원 간 직교를 하도록 처음부터 설계가 되어있어 애초부터 collapse가 될 수 없다는 것을 생각할 수 있다. Trivial solution, 즉 모든 벡터가 하나의 방향만을 가지고 있는 상황을 가정하자. 이렇게 되는 경우 diagonal term 자체는 1에 가까워져서, invariance term에 관한 loss는 압도적으로 낮아진다. 그러나 off-diagonal term인 redundancy reduction term은 크게 증가하게 된다. 즉 애초에 trivial solution으로 향하는 것은 loss 설계 자체에 모순이 되는 것이다.

 

이렇게 BT loss는 처음부터 batch size에 대해서 자유롭도록 설계되어 있어서, 논문에서 설명하듯이 batch size에 대한 부담이 적은 것이 본 논문의 장점이다.

---

논문은 추가적으로 다른 방법들, asymmetric encoder을 두는 경우와 cluster을 두는 방법과도 비교하는데, 현재 cluster 기반 방법은 주류에서 멀어져 있기 때문에 asymmetric encoder만 비교해보도록 하겠다.

 

Asymmetric encoder을 두는 대표적인 방법으로 BYOL, MoCo가 있다. 이 방법들은 BT와 마찬가지로 큰 배치를 요구하지 않는다는 측면에서는 상당한 강점을 가지고 있지만 여전히 trivial solution(사실상 불가능하지만)으로 가는 경우가 존재할 수 있음을 언급한다. 특히 문제는 target이 아닌 reference를 위한 인코더를 따로 두기 때문에 진정한 end-to-end 방법은 아니라는 지적을 한다.

 

물론 Barlow Twins도 장점만 있는 것이 아니다. 애초에 직교하도록 만드는 것에서 그 함정이 있는데, 과도한 orthogonalization의 가능성이 있다. 즉 redundancy를 지나치게 줄이면 유용한 correlation까지 사라지게 만들 수 있다는 것이다. 또 필자가 논문을 읽으면서, 그리고 experiments를 보면서 BT도 비주류같은데? 라는 생각을 버릴 수가 없었던 것이 하나 있는데, 배치가 임베딩 차원보다훨씬 크게 되는 경우에 대한 것이다.

 

예를 들어 배치가 8192이고, 임베딩이 4096 차원이라고 해보자. 이러한 경우 샘플 수가 차원보다 훨씬 많다는 뜻이 되고, 각 벡터가 선형 독립이기 위한 벡터의 최대 개수(orthogonal까지 다룰 필요 없다. Orthogonal하다는 것은 선형독립을 이미 보장하기 때문)는 벡터의 차원이기 때문에 그보다 많은 수의 벡터에 대해서는 제대로 embedding space에 위치하지 못한다는 것을 내포한다.

---

Experiments

사실 이 논문의 의의는 SOTA를 찍었다가 아니라 loss function을 이렇게 잡으면 된다 + batch size에 자유로우며 + end-to-end라는 것에만 있으니 experiments는 간단하게 보고 넘어가도록 하겠다.

앞선 추측대로 실험 결과가 나왔다. 배치 사이즈가 분명 작은 경우 효율성이라는 바운더리 내에서 더 좋은 성능을 낼 수 있는 것은 사실이지만, 배치 사이즈가 늘어나는 경우 SimCLR보다 성능이 낮게 나온다. 그러면 embedding의 dimension을 키우면 되겠네? 정확하다.

Projector Dimensionality를 늘릴수록 일관적으로 좋은 성능이 나타나는 것을 확인할 수 있다.

---

Conclusion

좋은 논문이다. 이러한 접근을 할 수 있다는 것과, batch size가 작고 GPU 자원이 한정적일 때 가장 효과를 발휘할 수 있는 방법론이라고 생각하면 될 것이다. 다만 현재 주류에서는 멀어져 있으니, 이러한 접근법이 있다는 것 정도만 확인하고 가면 될 것이다!