Why Prototypes Collapse: Diagnosing and Preventing Partial Collapse in Prototypical Self-Supervised Learning

1. 문제 정의 (Problem Definition)

이 논문은 프로토타입 기반 자기지도학습 (Prototypical Self-Supervised Learning, SSL) 방법론에서 발생하는 부분 프로토타입 붕괴 (Partial Prototype Collapse) 현상을 규명하고 해결하는 것을 목표로 합니다.

• 현상: 프로토타입 기반 SSL (예: DINO, DINOv2 등) 은 학습able 한 벡터 (프로토타입) 를 사용하여 데이터 포인트를 클러스터링합니다. 그러나 많은 방법론에서 학습 초기에 여러 개의 프로토타입이 거의 동일한 표현 (representation) 으로 수렴하는 '부분 붕괴'가 발생합니다.
• 원인: 기존 연구들은 이를 해결하기 위해 프로토타입 수를 과다하게 늘리거나 (over-parameterization), KoLeo 와 같은 ad-hoc 정규화 항을 추가하는 방식을 사용했습니다. 하지만 이는 근본 원인을 해결하지 못하고 증상만 완화하는 접근이었습니다.
• 핵심 가설: 저자들은 인코더 (Encoder) 와 프로토타입을 동일한 손실 함수 (loss) 하에서 공동 최적화 (joint optimization) 하는 것이 붕괴의 주원인이라고 주장합니다. 공동 최적화는 초기 학습 단계에서 손실을 빠르게 줄이기 위해 프로토타입이 중복된 표현으로 이동하는 '숏컷 학습 (shortcut learning)'을 유도합니다.

2. 제안 방법론 (Methodology)

저자는 공동 최적화의 문제를 해결하기 위해 인코더 학습과 프로토타입 추정을 완전히 분리 (Fully Decoupled) 하는 새로운 훈련 전략을 제안합니다.

• 완전 분리 훈련 (Full Decoupling Strategy):

  • 기존 방식: 인코더와 프로토타입을 하나의 손실 함수로 동시에 업데이트.
  • 제안 방식: 두 객체를 별도의 목적 함수 하에서 번갈아 가며 업데이트합니다.
    1. 프로토타입 업데이트: 인코더의 손실과 무관하게, 현재 인코더가 생성한 잠재 특징 (latent features) 을 기반으로 프로토타입을 추정합니다.
    2. 인코더 업데이트: 고정된 프로토타입을 사용하여 인코더의 일관성 손실 (consistency loss) 을 최소화합니다.

• 온라인 가우시안 혼합 모델 (Online Gaussian Mixture Model, GMM):

  • 분리된 프로토타입 추정을 위해, 프로토타입을 가우시안 혼합 모델의 성분 (means) 으로 모델링합니다.
  • EM (Expectation-Maximization) 스타일의 온라인 업데이트: 배치 단위로 들어오는 데이터에 대해 GMM 파라미터를 점진적으로 업데이트합니다.
  • 핵심 기법:
    • Responsibility-weighted Forgetting: 대규모 프로토타입 세트에서 일부 성분이 업데이트를 받지 않을 때 발생하는 불안정성을 방지하기 위해, 할당 책임 (responsibility) 에 기반한 망각 (forgetting) 메커니즘을 적용합니다.
    • Deterministic Annealing: 할당의 날카로움을 조절하여 학습 안정성을 높입니다.
    • Split-Resurrect: 과도하게 큰 가중치를 가진 클러스터를 분할하고, 소멸된 클러스터를 재생성하여 프로토타입의 다양성을 유지합니다.

3. 주요 기여 (Key Contributions)

1. 체계적인 분석: DINO 계열뿐만 아니라 다양한 프로토타입 기반 SSL 프레임워크 (CARP, CAPI, iBOT 등) 에 대해 부분 붕괴 현상이 광범위하게 존재함을 정량적으로 증명했습니다.
2. 근본 원인 규명: 인코더와 프로토타입의 공동 최적화가 붕괴를 유발하는 '숏컷 학습' 메커니즘임을 이론적 및 실험적으로 규명했습니다.
3. 새로운 프레임워크 제안: 프로토타입 추정을 인코더 학습에서 완전히 분리하는 Decoupled Training을 도입했습니다. 이는 명시적인 정규화 없이도 붕괴를 방지하고 일관되게 다양한 프로토타입을 생성합니다.

4. 실험 결과 (Results)

• 프로토타입 다양성: 제안된 분리 훈련 방식 (CARP + Decoupling) 은 학습 내내 프로토타입 붕괴가 전혀 발생하지 않았습니다. 반면, 기존 DINO 는 초기 프로토타입의 1.5% 만이 고유하게 남는 등 심각한 붕괴를 보였습니다.
• 다운스트림 성능:
  • ImageNet-1k: k-NN 평가에서 CARP 대비 1.8%p 향상 (67.7% → 69.1%) 을 기록했습니다. 선형 평가 (Linear Evaluation) 에서도 유사하거나 더 나은 성능을 보였습니다.
  • 불균형 데이터 (Long-tailed Distribution): iNaturalist 2018 데이터셋에서 'Head', 'Medium', 'Tail' 클래스 전반에 걸쳐 성능이 향상되었습니다. 특히 기존 방법들이 Tail 클래스에서 어려움을 겪는 반면, 분리 훈련은 모든 클래스 빈도에서 강건한 성능을 보여주었습니다.
  • 전송 학습 (Transfer Learning): 다양한 벤치마크 (Aircraft, Cars, Flowers 등) 에서 우수한 전송 학습 성능을 입증했습니다.
• 효율성: 메모리 사용량이 감소했습니다 (대규모 배치 크기에서 더 두드러짐). 이는 프로토타입 최적화를 위한 계산 그래프를 저장할 필요가 없기 때문입니다. 학습 시간 증가량은 미미했습니다 (0.2 시간 증가).

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 프로토타입 기반 자기지도학습의 핵심 병목 현상인 '부분 붕괴'를 공동 최적화 (Joint Optimization) 의 관점에서 재해석하고, 이를 분리된 최적화 (Decoupled Optimization) 로 해결함으로써 새로운 패러다임을 제시했습니다.

• 이론적 통찰: 프로토타입이 인코더의 손실 최소화에 종속될 때 발생하는 중복 표현 문제를 해결함으로써, 더 풍부하고 다양한 표현 학습이 가능함을 보였습니다.
• 실용적 가치: 복잡한 정규화 항이나 과도한 하이퍼파라미터 튜닝 없이, 단순하고 원리 기반 (principled) 인 GMM 기반 업데이트만으로 붕괴를 방지하고 모델의 강건성 (Robustness) 을 높일 수 있음을 입증했습니다.
• 미래 방향: 특히 불균형한 데이터 분포를 가진 현실 세계의 데이터셋에서 프로토타입 기반 SSL 의 성능을 극대화할 수 있는 기초를 마련했습니다.

요약하자면, 이 연구는 "프로토타입과 인코더를 함께 학습시키는 것이 오히려 학습을 방해한다" 는 역설적인 통찰을 바탕으로, 분리된 학습 전략을 통해 더 나은 표현 학습을 가능하게 하는 획기적인 방법을 제시한 논문입니다.