SwAV 논문 리뷰

Unsupervised Learning of Visual Features by Contrasting Cluster Assignments

---

Summary

본 논문은 두 개의 augmented view 간 표현을 직접 비교하지 않고, prototype과의 softmax 확률 분포와 반대편 view에서 계산된 code 간의 cross-entropy loss를 통해 자기지도학습을 수행한다. 이러한 접근을 통해 큰 배치 사이즈나 별도의 momentum encoder 없이도 높은 성능을 달성하며, SimCLR의 배치 크기 제약과 MoCo의 추가 네트워크 사용 한계를 효과적으로 극복한다.

---

Method

본 논문이 제시하는 두 가지 방법: (1) Swapped Prediction과 (2) Multi-crop Augmentation에 대해서 알아보자. 또한 representation collapse를 어떻게 막는지에 관한 내용인 Equipartition constraint와 Prototypes를 통한 soft assignment 방법을 알아보도록 하자.

Swapped Prediction

여기서

• $x_1, x_2$: Augmentation $t \sim T$를 통해서 만들어진 두 이미지
• $f_{\theta}$: 이미지 처리 모델. Output은 input $x_1, x_2$에 대한 representation $z_1, z_2$ 이다.
• Prototypes $C$: Soft label을 의미한다. 학습이 가능한 matrix이며, $z_1, z_2$와 행렬곱이 적용된 후 softmax로 확률분포 $p$를 만들어준다.
• $Q_1, Q_2$: Code라고 부르며, $z_1, z_2$에서 만들어진 확률분포에 대한 실제 정답(으로 간주하는) soft assignment이다. $Q$는 Sinkhorn-Knopp 알고리즘을 통해 최적화되었으며 이때 $Q_1$은 $x_2$에 대한 타깃, $Q_2$는 $x_1$에 대한 타깃이다.

Loss의 구조를 먼저 확인해보자. Swapped prediction, 두 augment된 이미지의 표현이 서로의 code $q$와 같아지도록 설계된 것을 확인할 수 있다.

단일 표현의 예측 loss $l$은 다음과 같다:

Soft alignment 타깃인 $q$와의 cross-entropy loss를 낮추는데, 이때 $p$는 prototype인 $C$와의 행렬곱에 softmax을 적용한 확률분포를 output으로 낸다. 아직까지는 단일 데이터에 대한 loss였고, 실제로 $N$개 데이터에 대한 loss는 어떻게 다루어지는지 확인해보자. Matrix form이 어떻게 나올지 적당히 생각해보면 쉽다.

모델은 각 뷰로부터 확률분포 $p$를 출력하며, 이때 학습 대상인 타겟 $q$는 해당 뷰의 자체 라벨이 아니라 상대편 뷰에서 계산된 soft code로 설정된다. 이러한 “스왑된” 예측 구조가 SwAV의 핵심이다. 이러한 방법을 적용하게 되면 구조적으로 표현이 collapse 될 수 없는 것을 확인할 수 있다. Collapse된 표현이라고 하면, 서로 다른 augment에서 나온 코드와 일치시킬 수 없기 때문이다.

 

이번에는 Equipartition constraint과 최적화에 대해서 먼저 알아보도록 하자.

우선 간단하게 보면, $Q$의 row-wise mean과 colum-wise mean이 균일분포를 따라야 한다는 제약임을 알 수 있다. 이는 곧 $Q$를 구성하는 $q$들이 특정 프로토타입에 치우치지 않고, 집합적으로 다양성을 유지하도록 강제한다는 뜻이다.. 이것이 어떤 의미인지 당장 이해하기 어렵다면, 각 방향에서의 평균이 균일분포를 따르지 않는 경우, 즉 $q$들 간 비교를 했을 때 서로 거의 비슷한 모습을 하고 있는 경우($q$ 자체는 uniform하지 않음)을 생각해보자.

이 제약의 의미를 직관적으로 이해하기 위해, 만약 각 $q$가 서로 매우 유사한 분포(즉, 모든 샘플이 동일한 프로토타입에 높은 확률을 할당)로 수렴하는 상황을 생각해보자. 이 경우 프로토타입 행렬 $C$와의 곱 $Cq$는 항상 거의 동일한 출력을 내게 되며, 이는 결국 모든 표현 $z$가 동일한 공간으로 수축(collapse)되도록 유도한다.

반면, equipartition 제약이 존재하면 각 $q$는 서로 구별되는 방향으로 분산될 수밖에 없고, 이에 따라 $Cq$ 역시 다양한 출력을 갖는다. 결과적으로 모델은 표현 공간 전체를 활용하여 다양한 프로토타입과의 관계를 학습하게 되며, 이는 collapse를 방지하는 핵심적인 메커니즘이 된다.

위의 제약 조건을 바탕으로 다음과 같은 목적식을 최적화한다:

이 최적화 식의 해석적 해는 다음과 같다:

본 논문에서는 해석적 해를 직접 도출하지 않고 Sinkhorn–Knopp 알고리즘을 활용하여 근사한다. 이 알고리즘은 확률 행렬 $Q$를 반복적으로 정규화하여, 각 프로토타입이 배치 내에서 균등하게 선택되도록 보장하며, 이를 통해 trivial solution(모든 샘플이 하나의 프로토타입에 몰리는 현상)을 방지한다.

관련 알고리즘에 대해서 모르고 쓰는 것이기 때문에 자존심이 매우매우 상하므로 열심히 공부해서 블로그에 글을 쓰도록 하겠다.

 

Feature을 직접 비교하지 않고, prototype를 통환 soft assignment 또한 중요한 개념이다. Feature를 직접 비교하지 않고, 클러스터 중심 고정점을 따로 두어 비교하기 때문에 feature가 단일 벡터로 쏠리는 경우가 구조적으로 발생할 수 없도록 막는다. 물론 프로토타입이 여러 개라는 가정 하에서다. 단일 벡터로 feature가 쏠리는 경우, loss 또한 줄어드는 것이 아닌 상수의 lower bound를 가지게 되기 때문에, 프로토타입을 두는 것이 표현 공간 자체가 다양하게 분산되는 구조를 줌을 알 수 있다.

Multi-crop

단일 이미지가 주어졌다고 가정하자. 이 이미지에 대한 두 개의 view만 교차로 비교하는 것은 자원을 덜 쓰는 꼴이고, 그렇다고 같은 해상도를 가진 view를 제공하는 것은 비용 소모가 크다. 따라서 본 논문은 작은 해상도를 가진 $V$개의 crop를 추가적으로 두어 loss 계산에 추가한다. 수식은 다음과 같다:

비용면에서도 안정적이지만, 낮은 해상도를 가진 view로 높은 해상도를 가진 crop를 예측하는 테스크가 추가되었기 때문에, 더 표현을 깊이 학습할 수 있도록 모델을 유도할 수 있다. 이렇게 보면 참 inductive bias를 줄이기 위한(?) 각종 노력들이 정말 아름다운 것 같기도 하다. 실제로 논문에서 제시한 방법이 우수함을 확인할 수도 있다:

---

Conclusion

Supervised learning보다 SwAV 표현학습 후의 성능이 더 좋은 것을 확인해보자.

(절대 귀찮아서 experiment를 뺀 것이 아니다 ㅎㅎ)

사실 이 논문을 리뷰한 이유는, Dino에서 사용되는 각종 방법들의 시조격이 되기 때문이다. SSL과 representation learning에서 전반적으로 다루는, collapse를 막기 위한 방법들과 모델이 더 좋은 표현을 깊게 학습하도록 유도하는 방법에 집중하면서 보면 이 논문이 정말 잘 만들어졌음을 알 수 있을 것이다.