[CVPR 2021]Exploring Simple Siamese Representation Learning

안녕하세요 정의철 연구원입니다. 이번에 제가 리뷰할 논문은 2021 CVPR에 기재된 ‘Exploring Simple Siamese Representation Learning’이란 논문이고 흔히 Simsiam이라고 불리는 방법론입니다. 그럼 바로 리뷰 시작하겠습니다.

Abstract

Siamese 네트워크는 여러 모델에서 비지도 학습을 위한 시각적 표현 학습에 일반적인 구조가 되었습니다. 이러한 모델은 하나의 이미지에 두 가지 augmentation간의 유사성을 최대화하며, collapsing solution을 피하기 위한 특정 조건에 따라 작동합니다. 본 논문에서는 단순한 Siamese 네트워크조차도 다음과 같은 것들을 사용하지 않고도 의미 있는 표현을 학습할 수 있다는 결과를 보고합니다: (i) negative sample pairs, (ii) large batches, (iii) momentum encoders. 실험 결과, loss와 structure에는 collapsing solutions이 존재하지만, stop-gradient 작업이 collapsing를 방지하는 데 중요한 역할을 합니다. 본 논문에서는 stop-gradient에 내제된 가설을 제시하고, 이를 확인하는 개념 증명 실험을 추가로 보여줍니다.

Introduction

그동안 un-/self- supervised representation learning에서의 대부분의 모델들은 일반적으로 Siamese 네트워크의 특정 형태를 포함합니다. Siamese 네트워크란 두 개의 입력에 대해 가중치를 공유하는 동일한 신경망 아키텍처를 말합니다. 이 네트워크는 두 입력 데이터에 대해 비슷한 항목을 식별하거나 비교하는 데 사용되는 구조입니다. 최근 방법들은 입력을 동일 이미지의 두 가지 증강으로 정의하고, 유사성을 최대화하며 다양한 조건에 따라 이를 수행합니다.

하지만 Siamese 네트워크의 주요 문제로는 모든 출력이 상수로 “collapsing”하는 것이 있습니다. Siamese 네트워크가 collapsing되지 않도록 방지하는 여러 일반적인 전략들이 있는데 대표적으로 SimCLR, SwAV, BYOL 등이 있습니다. SimCLR에서 사용되는 Contrastive learning은 Negative pairs간에는 거리를 떨어뜨리고, 동일 이미지에서 온 positive pairs는 거리를 가깝게 끌어당기도록 학습됩니다. SwAV 방법론은 Clustering을 이용해서 collapsing을 방지하는데 모델은 서로 다른 증강된 이미지를 유사성에 따라 클러스터로 그룹화하고, 각 클러스터에 대한 특징을 추출하여 학습을 한다고 합니다.BYOL의 경우 Positive Pair만 사용하지만 momentum encoder도 사용하기에 collapsing을 방지할 수 있다고 합니다.본 논문에서는 위에 언급된 방법들을 사용하지 않고도 단순한 Siamese 네트워크를 구성하여 잘 작동할 수 있다고 저자는 말하고 있습니다.

저자의 모델은 Negative pairs나 모멘텀 인코더를 사용하지 않고도 한 이미지의 두 뷰의 유사성을 직접 최대화합니다. 또한 큰 배치 사이즈에 의존하지 않고 일반적인 배치 크기로 task를 진행합니다. SimSiam은 개념적으로 간단함으로 인해 기존의 여러 방법과 관련시킬 수 있는 중심 역할을 할 수 있습니다. 예를 들어, SimSiam의 방법은 “모멘텀 인코더 없는 BYOL”로 생각할 수 있습니다. BYOL과는 달리 SimCLR 및 SwAV처럼 SimSiam은 두 branch 사이에서 가중치를 직접 공유하므로 “Negative pairs가 없는 SimCLR” 및 “온라인 클러스터링 없는 SwAV”로도 볼 수 있습니다. 이처럼 SimSiam은 각 방법과 관련이 있으며, 핵심 구성 요소 중 하나를 제거함으로써 그들 사이의 관계를 보여줍니다. 그럼에도 불구하고 SimSiam은 collapsing을 유발하지 않고 좋은 성능을 보여줍니다. 저자는 collapsing solution이 존재하지만, stop-gradient 가 이러한 해결책을 방지하는 데 중요하다는 것을 실험적으로 보여줍니다

Method

SimSiam 아키텍처(Figure 1)는 이미지 x로부터 두 가지 무작위로 증강된 이미지 x1과 x2를 입력으로 받습니다. 이 두 이미지는 백본(ResNet)과 projection MLP 헤드로 구성된 인코더 네트워크 f의 입력으로 들어갑니다. 인코더 f는 두 이미지 간에 가중치를 공유합니다. 하나의 이미지의 출력을 변환하고 다른 이미지와 일치시키는 predictor MLP 헤드인 h가 이 출력을 입력으로 받습니다. 두 출력 벡터를 p1 = h(f(x1)) 및 z2 = f(x2)로 표기하고, 두 벡터의 negative cosine similarity를 최소화합니다. negative cosine similarity의 대칭적인 학습을 위해 p2, z1의 loss도 구해주게 되며 최종적인 loss는 다음과 같이 정의됩니다.

그리고 저자의 방법론이 작동하는 데 중요한 구성 요소 중 하나는 stop-gradient 연산입니다. 이를 위의 수식에 다시 적용시키면 각 항은 다음과 같이 표현됩니다.

이는 이 항에서 z2를 상수로 취급한다는 것을 의미합니다. 마찬가지로, (2)의 형태는 다음과 같이 표현됩니다.

여기서 x2의 인코더는 첫 번째 항에서 z2로부터 gradient를받지 않지만 두 번째 항에서 p2로부터 gradient를 받습니다.

위의 수식은 매우 간단하기에 SimSiam의 의사 코드도 다음과 같이 간단하게 표현됩니다.

논문에서 설정된 Baseline settings은 다음과 같습니다.

• Optimizer : 저자는 pre-training에 SGD를 사용하고 lr = 0.05로 설정하여 lr×BatchSize/256으로 learning rate을 사용합니다(배치 크기는 기본적으로 512입니다). learning rate은 cosine decay schedule를 사용하고 weight decay는 0.0001이고, SGD 모멘텀은 0.9입니다.

• Projection MLP : Projection MLP는 각 fully-connected (fc) 레이어에 BN이 적용되며, output fc에는 ReLU가 없습니다. hidden fc는 2048차원이고 MLP는 총 3개의 레이어를 가지고 있습니다.

• Prediction MLP : Prediction MLP(h)는 hidden fc 레이어에 BN이 적용되고 output fc에는 BN, ReLU가 없습니다. 이 MLP는 총 2개의 레이어를 가지고 있습니다. h의 입력 및 출력(z와 p)의 차원은 d = 2048이며, h의 hidden 레이어의 차원은 512입니다.

Experimental setup

저자는 1000-class ImageNet training set 에 대해 unsupervised pre-training을 수행합니다. 이후 pre-trained 된 representations을 사용하여 supervised linear classifier를 훈련시키고 validation set를 이용하여 평가합니다.

Empirical Study

이 섹션이 SimSiam의 실증 연구라고 할 수 있고 모델이 non-collapsing 솔루션에 기여할 수 있는 요소에 집중합니다.

Stop-gradient

Figure 2는 “Stop-gradient” 사용 여부에 따른 비교를 제시합니다. Figure 2의 가장 왼쪽 첫 번째 결과는 training loss를 보여줍니다.

Stop-gradient를 사용하지 않으면 옵티마이저가 빠르게 degenerated solution을 찾아내고 최소 손실 -1에 도달합니다. degeneration이 collapsing으로 인한 것인지를 보여주기 위해 저자는 ‘ℓ2-normalized output z/‖z‖₂‘의 표준 편차(std)를 연구합니다. 출력이 상수 벡터로 붕괴되면 모든 샘플에 대해 각 채널의 표준 편차는 제로여야 합니다. 이는 Figure 2 (가운데)의 빨간 곡선에서 관찰할 수 있습니다. 반면, 출력 z가 평균이 제로인 isotropic Gaussian distribution를 가진다면, z/‖z‖₂의 표준 편차는 1/√d임을 보여줄 수 있습니다. Figure 2 (가운데)의 파란 곡선은 Stop-gradient를 사용할 때 std 값이 거의 1/√d에 가까움을 보여줍니다. 이는 출력이 붕괴되지 않고있음을 보여줍니다. Figure 2 (오른쪽)은 k- nearest-neighbor(kNN) 분류기의 검증 정확도를 그립니다. 이 kNN 분류기는 진행 상황을 모니터할 수 있습니다. Stop-gradient를 사용하면 kNN 모니터가 지속적으로 향상되는 정확도를 보여줍니다. Figure 2의 테이블에서는 선형 평가 결과가 나타납니다. SimSiam은 67.7%의 상당한 정확도를 달성합니다. 이 결과는 5회 시도의 std에 의해 안정적으로 보여집니다. Stop-gradient만 제거하면 정확도가 0.1%가 되며, 이는 ImageNet에서의 랜덤 수준 추측입니다.

Discussion

이 실험은 collapsing 솔루션이 존재한다는 것을 보여줍니다. collapsing은 최소 가능 손실과 상수 출력으로 확인할 수 있습니다. collapsing 솔루션이 존재한다는 것은 저자의 방법이 아키텍처 설계만으로는 (ex: preductor, BN, ℓ2-normalized 등) 붕괴를 방지하는 데 충분하지 않다는 것을 의미합니다. 비교에서 모든 이러한 아키텍처 설계가 변경되지 않았지만, Stop-gradient를 제거하면 collapsing이 방지되지 않습니다. Stop-gradient의 도입은 해결되고 있는 또 다른 최적화 문제가 있음을 의미합니다. 이 부분은 뒷 부분의 Hypothesis에서 더 자세히 알아보겠습니다.

Predictor

Table 1에서 예측 MLP h의 효과를 연구합니다. Table 1a에서는 모델에서 h를 제거하면 모델이 작동하지 않음을 확인할 수 있습니다. 또한 Table 1b에서는 h를 무작위로 초기화시키고 고정할때에도 모델이 작동하지 않음을 확인할 수 있는데 이는 collapsing과는 관련이 없습니다. 이는 모델이 훈련이 수렴되지 않고 손실이 높은 상태로 유지된다고 저자는 말하고 있고 Predictor MLP h는 representation에 적응하도록 훈련되어야 함을 확인할 수 있습니다. Table 1c에서는 h에 constant learning rate를 decay없이 사용하면 잘 작동하고 baseline 보다 더 좋은 결과를 냄을 확인할 수 있는데 이는 h가 latest representation에 적응하도록 훈련되어야 하므로 representation이 충분히 훈련되기 전에 수렴을 강제할 필요가 없기 때문일 수 있다고 말하고 있습니다.

Batch Size

Table 2는 배치 크기가 64에서 4096까지 변화할 때의 결과를 보여준 것입니다. 저자의 모델은 이와 같은 다양한 배치 크기 범위에서 잘 작동하는 것을 확인할 수 있습니다. 앞의 서론에서 SimCLR과 SwAV 같은 모델은 모두 잘 작동하려면 큰 배치 (예: 4096)가 필요하다고 말씀드렸습니다. 하지만 그에 반해 Simsiam은 배치 크기가 128 또는 64인 경우에도 성능이 괜찮은 모습을 보여주고 있습니다.

Batch Normalization

able 3은 MLP 헤드에서 BN의 구성을 비교합니다. Table 3a에서는 MLP 헤드의 모든 BN 레이어를 제거합니다. 이 변형은 collapsing를 유발하지 않지만 정확도가 낮습니다 (34.6%). 낮은 정확도는 최적화의 어려움 때문일 가능성이 높습니다. hidden 레이어에 BN을 추가하면 (Table 3b), 정확도가 67.4%로 증가합니다. proj MLP의 출력에 BN을 추가하면 정확도가 68.1%로 향상됩니다 (Table 3c). 반면 pred MLP h의 출력에 BN을 추가하면 (Table 3d) 잘 작동하지 않습니다. 요약하면 BN이 적절하게 사용될 때 최적화에 도움이 된다는 것을 보여줍니다. 그러나 BN이 collapsing를 방지하는 데 도움이 되는지에 대한 증거는 없다고 저자는 말하고 있습니다.

Similarity Function

Simsiam은 코사인 유사도 함수 외에도 크로스 엔트로피 유사도를 사용해도 작동할 수 있음을 보여줍니다. 이경우 식 D는 다음과 같이 수정 됩니다.

여기서 softmax 함수는 채널 차원을 따라 적용됩니다. softmax의 출력은 각 d pseudo-categories에 속할 확률로 생각할 수 있습니다. 이때 비교 결과는 다음과 같습니다.

• 코사인 유사도: 68.1%
• 크로스 엔트로피: 63.2%

크로스 엔트로피 변형은 Collapsing 없이 수렴할 수 있다는 것을 보여주고 단순히 코사인 유사도에 의해 Collapsing을 예방할 수 있다는 것은 않음을 보여줍니다.

Symmetrization

지금까지 실험은 대칭화된 손실을 기반으로 진행되었고 SimSiam의 Collapsing 방지는 대칭화에 의존하지 않음을 관찰했습니다.

대칭화는 정확도 향상에 도움이 되지만 Collapsing 방지와는 관련이 없습니다.

Summary

다양한 설정에서 저자는 Empirical하게 SimSiam이 Collapsing 없이 의미 있는 결과를 생성할 수 있다는 것을 보여주었습니다. 배치 크기, 배치 정규화, 유사도 함수 및 대칭화는 정확도에 영향을 미칠 수 있지만, 이들이 붕괴 방지와 관련이 있다는 증거는 없습니다. 이로써 주요한 역할을 하는 것은 주로 Stop-gradient라고 저자는 말하고 있습니다.

Hypothesis

이 섹션에서는 SimSiam이 Expectation-Maximization (EM)과 유사한 알고리즘이라고 가정합니다. 이 알고리즘은 두 개의 변수 집합을 내재적으로 포함하며 두 가지 sub-problem를 해결합니다. 이를 위해, SimSiam의 objective를 수식으로 재정의하면 다음과 같습니다.

• θ : encoder network F𝐹의 parameter
• η : stop-gradient를 설명하기 위한 변수
  • θ처럼 network의 parameter 형태로 표현할 필요가 없기 때문에 η로 간소화하여 표현
  • 𝜂𝑥 : representation of image x
• T : augmentation

이 형식을 사용하여 다음을 해결하는 것을 고려합니다

이 형식을 k-means clustering과 동일하다고 생각할 수 있습니다. θ는 clustering center이 되고, 𝜂𝑥는 sample x에 할당된 vector (x의 representation)이라고 볼 수 있는 것입니다.

위의 optimization proble를 k-means clustering과 동일하게 변수 하나씩을 고정해 가며 alternating algorithm으로 해결할 수 있습니다.

Solving for θ

θ에 대한 optimization (7)은 SGD를 이용해 해결할 수 있고 이 때 η의 경우 constant로 취급되므로 η에 대한 gradient는 stop됩니다.

Solving for η

η에 대한 optimization (8)은 MSE를 최소화하는 문제이므로 아래의 식과 같이 𝜂𝑥는 x의 모든 augmentation에 대한 기댓값으로 나타내어질 수 있습니다.

One-step alternation

SimSiam은 (7)과 (8) 사이의 one-step alternation으로 근사될 수 있습니다. 그리고 Augmentation을 한번만 사용하면, (9)는 다음과 같이 근사시킬 수 있습니다.

이를 (7)에 대입하면 다음과 같이 식이 변형됩니다.

결과적으로 η는 식에서 사라졌고, θt+1를 얻기 위해서는 θt가 필요하다는 사실을 수식으로 확인할 수 있습니다. 그리고 여기서 θt와 관련 있는 부분이 stop-gradient operation을 의미하게 됩니다.

Predictor

위 가정은 predictor h가 포함되어 있지 않기 때문에 predictor h를 고려하여 수식을 다시 정리해 보겠습니다.

SimSiam에서는 h가 다음을 minimize하기 위해 사용되고

h의 optimal solution은 다음을 만족해야합니다.

(10)에서 expectation term을 무시했지만 h를 사용하는 것이 이 gap을 채워줄 수 있고 실제로 Expectation을 계산하는 데에는 computation이 많이 들기 때문에, h가 expectation을 approximate할 수 있도록 대신 사용하는 것입니다.

Proof of concept

Multi-step alternation

저자들은 k SGD step마다 θ를 update하는 실험을 진행합니다.

여기서 “1-step”은 SimSiam을 의미하며 “1-epoch”은 한 번의 에포크에 필요한 k steps를 나타냅니다. 위 표를 통해 step이 많아지더라도 collapsing이 발생하지 않는 것을 확인할 수 있습니다

Discussion

지금까지의 가설과 실험은 SimSiam의 optimization이 무엇인지에 관한 것입니다. 하지만 어떻게 collapsing을 방지할 수 있는지에 대해서는 다루지 못했습니다.

이에 대한 저자들의 생각은 다음과 같습니다. 랜덤으로 초기화된 네트워크의 출력인 초기화된 η가 상수일 가능성은 적고, alternating optimizer가 모든 x에 대한 gradient를 계산하는 것이 아니라, 파편화된 ηx에 대한 gradient를 계산하기 때문에, collapsing 현상이 발생하기 어려웠을 것이라고 추측합니다.

Comparisons

Result Comparisons

ImageNet

ImageNet linear evaluation 성능입니다. 100 epoch에서는 좋은 성능을 보이지만, 200, 400, 800 epoch에서는 BYOL에 비해 꽤 낮은 성능을 보이고 있습니다. 이에 대해서는 momentum encoder와 large batch size가 필요하지 않다는 점에서 간단함과 memory-efficient한 학습 방식이라고 저자는 말하고 있습니다.

Transfer Learning

ImageNet 성능에 비해 SimSiam의 transfer learning 성능은 좋은 모습을 보입니다. 추가로 lr = 0.5 및 wd = 1e-5와 같이 parameter를 수정하였을때 ImageNet classification 성능에는 별 차이가 없었지만 transfer learning 성능은 크게 상승하였다고 합니다.

Conclusion

저자는 간단한 디자인의 Siamese 네트워크를 만들었습니다. 이러한 minimalist한 method에도 경쟁력을 가지는 이유는 기존 방법들의 Siamese 구조가 그들의 효과성에 주요 이유일 수 있다는 것을 시사하고 있습니다. Siamese 네트워크는 invariance을 모델링하는 효과적인 도구이며, 이는 representation learning을 하는데 효과적이라고 저자는 말하고 있습니다.