BYOL 논문 리뷰

Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning

---

Summary

Self-Supervised Learning(SSL)과 Image Representation Learning 분야에서는 오랫동안 적절한 이미지 증강(image augmentation), negative sampling, 그리고 대규모 batch size가 핵심 요소로 간주되어 왔다. 그러나 이러한 접근에는 본질적인 한계가 존재한다. 우선, 이미지 증강은 모델이 학습 손실을 지나치게 쉽게 줄일 수 있는 ‘편향된 경로(shortcut)’를 제공할 위험이 있어, 매우 정교한 설계가 요구된다. Negative sampling 역시 방대한 메모리와 계산 자원을 필요로 하며, 효과적인 negative 선택 전략을 설계하는 것이 쉽지 않다. 마지막으로, 대규모 batch size는 실험 환경의 제약을 강화하여, 연구 및 실용적 적용 모두에 큰 부담으로 작용한다.

 

이러한 한계를 극복하기 위해, 본 논문은 positive pair만으로도 collapse를 방지할 수 있는 새로운 방법론을 제안한다. 제안된 방법은 기존 접근과 달리 batch size나 이미지 증강 설정에 덜 민감하여 학습이 간편하며, 동시에 ImageNet 선형 분류 벤치마크에서 SimCLR, MoCo 등 기존 기법을 능가하는 성능을 보여준다.

---

Method

BYOL은 기본적으로 target와 online 두 개의 네트워크를 사용한다. 두 네트워크는 구조가 동일하지만, target는 $\xi$, online은 $\theta$를 weight로 사용한다는 점에서 다르며, online에서는 predictior $g_\theta$를 추가적으로 사용한다는 점에서 비대칭적인 모습을 보인다.

 

초기 input image $x$는 두 개의 서로 다른 augmentation을 사용하여 각각 $v$와 $v'$으로 변환되며, 각각의 레이어를 지나 online model의 경우 $q_\theta (z_\theta)$, target의 경우 $z_\xi '$로 변환된다. 이 두 representation의 차이를 줄이기 위해, 즉 similarity가 같아지도록 학습되는데, loss는 다음과 같다:

또한 Loss를 대칭적으로 만들어주기 위해, augment된 input $v$와 $v'$을 switch해서 각각의 모델에 넣어준다.

 

여기서 target 모델은 gradient로 직접 업데이트되지 않고, exponential moving average(EMA)를 통해 간접적으로만 갱신된다. 겉보기에 방법론은 단순하지만, 필자는 직관적으로 “이 단순한 구조가 어떻게 성능 향상으로 이어지는가?”라는 의문을 가졌다. 이에 대해 다음과 같은 두 가지 질문을 던지고, 이를 탐구하는 과정을 거치고자 한다.

1. 만약 target 모델이 이미 성능이 좋은 모델이라고 가정해 보자. Online 모델은 target 모델의 표현을 예측하도록 학습되므로, 오히려 target의 성능이 online 쪽으로 희석되어 representation 품질이 저하되고, 결과적으로 online 모델의 성능도 올라가지 않을 것처럼 보인다. 그런데 논문은 target 모델에 대해 아무런 제약도 두지 않는다. 즉, target이 “깡통 모델”이어도 상관없다는 것인데, 어떻게 그런 상황에서도 높은 성능이 가능할까?
2. 또한, 어떻게 positive pair만으로 collapse를 방지할 수 있는 구조를 설계할 수 있는가? 기존 contrastive 학습법은 negative pair가 필수적이라고 여겨졌는데, BYOL은 왜 예외적으로 동작할 수 있는가?

1. 어떻게 높은 성능이 가능한 것인가?

Augmetation이 핵심이다. 한 이미지에서 증강 $v$, $v'$을 만들어서 online/target에 따로 입력되는데, 이 때 모델이 같은 이미지에 대한 서로 다른 뷰를 align하려다 보면, 자연스럽게 view에 대해서 불변인 표현(invariant representation)을 학습하게 된다.

 

실험적으로도 증명되는데, 논문에서 target을 랜덤 고정 네트워크로 두어도, 즉 EMA를 통한 업데이트 자체를 하지 않더라도 online이 그 representation을 예측하는 과정에서 퀄리티가 높아진다. 즉 target가 좋냐 나쁘냐가 중요한 것이 아닌, 예측 과정 자체가 중요한 것이다. 마치 GAN과도 비슷하다.

나아가서 EMA로 인해서 target는 online이 학습되는 동안 부드럽게 변하는 기준점 역할이 되어, 즉 online의 과거 모습의 누적을 따라가는 안정된 거울이 되어 online이 representation을 점점 더 구조화되고 general하도록 개선하는 것이다.

 

2. Positive pair만으로 어떻게 collapse를 방지할 수 있는가?

최적의 predictor을 가정하자. 다음과 같은 식을 따를 것이다:

이때의 이상적인 손실은 다음과 같다:

즉 최적 predictor을 가정하면 BYOL loss는 target의 조건부분산의 합이 된다. Total variance 법칙에 의해

이를 $L^{*}$에 대입해보면

이 때 $\theta$관점에서 $L^{*}$을 최적화하고자 하면 아래와 같은 식이 된다.

이를 최대화한다는 것은, online 표현 $z_\theta$만으로 target의 표현을 최대한 설명할 수 있도록 하는 것이다. 만약 $z_\theta$가 collapse된 상태, 즉 상수라고 가정하면 loss가 최대화되는 지점 중 하나가 되어, 자연적으로 이러한 상황이 일어나지 못하게 된다.

 

또한 online 모델의 predictor로 인해 비대칭성이 일어나, 상수 해로 변하고자 하는 안정성을 더 깨게 하며 $\xi$는 gradient로 직접 최소화되는 것이 아닌, EMA로 느린 적합을 이루어지도록 하기 때문에 변동이 줄어든다.

---

Experiments

우선 성능부터 확인해주자.

더 중요한 부분은 Ablations 파트에 나온다.

BYOL 방법이 SimCLR보다 augmentation이나 batch size에 대해서 더욱 강건한 것을 확인할 수 있다.

아직도 신기한 것은, Constant random network, 즉 target 모델이 완전 깡통 모델인데도 성능이 증가하는 것이다. 참고로 저 target 모델은 accuracy가 1.4%라는 것이다.

---

Conclusion

Negative sample 없이도 self-supervised learning을 안정적으로, 그리고 더 우수한 성능으로 해낼 수 있다는 것을 보았다. 조건부 분산의 극대화를 통해 online 표현이 상수로 수렴되는 것을 막았으며, predictor와 EMA 업데이트가 학습의 안정성을 보장함을 확인하였다.

 

아직 나는 암것도 몰랐구나를 생각하면서 공부에 더 매진해야겠다...