SiamJEPA: On the Role of Siamese Student Encoders in JEPA

기술 요약: SiamJEPA: JEPA에서 시암 학생 인코더(Siamese Student Encoders)의 역할

문제 정의
I-JEPA 및 V-JEPA와 같은 합동 임베딩 예측 아키텍처(Joint Embedding Predictive Architectures, JEPAs)는 픽셀을 재구성하는 대신 마스킹된 영역의 잠재 임베딩을 예측함으로써 자기 지도 학습(self-supervised representation learning)을 위한 유망한 프레임워크로 부상했습니다. 그러나 기존 JEPA 모델들은 일반적으로 학생 네트워크에 단일 인코더를 사용하는 경례가 있습니다. 반면, 시암 인코더(이중 학생 네트워크)는 다른 자기 지도 학습(SSL) 패러다임(예: BYOL, SimSiam)이나 뇌 모사 모델인 PhiNets의 핵심 요소입니다. 하지만 JEPA의 잠재 공간 예측 프레임워크 내에서 이들의 구체적인 역할과 효용성은 아직 충분히 탐구되지 않았습니다. 본 논문은 다음과 같은 질문을 다룹니다: 시암 학생 인코더는 합동 임베딩 예측 아키텍처(JEPA) 내에서 어떤 역할을 수행하는가?

방법론: SiamJEPA
저자들은 JEPA 프레임워크에 시암 학생 인코더를 통합한 마스크 잠재 예측 아키텍처인 SiamJEPA를 제안합니다. 이 아키텍처는 다음으로 구성됩니다:

1. 시암 학생 인코더 (Siamese Student Encoders): 입력 이미지의 두 가지 독립적인 마스킹 뷰(view)가 가중치를 공유하는 두 개의 학생 인코더($f$)에 의해 처리됩니다. 지름길 학습(shortcut learning)을 방지하기 위해 마스킹 세트($M_1$과 $M_2$)는 서로 겹치지 않습니다.
2. EMA 교사 네트워크 (EMA Teacher Network): 학생 가중치($f$)의 지수 이동 평균(EMA)인 $f_{ema}$가 표준 JEPA 및 BYOL 패러다임과 유사하게 타겟 네트워크 역할을 합니다.
3. 예측 네트워크 (Predictor Networks):
   • 전역 정렬 (Global Alignment, $h$): 선형 예측기가 두 시암 브랜치의 전역(CLS) 표현을 정렬합니다.
   • 잠재 예측 (Latent Prediction, $g$): 확률적 예측기(트랜스포머 기반)가 마스킹된 토큰의 잠재 표현을 예측합니다. 이는 두 뷰 모두에 조건화된 사후 분포(posterior distribution)와 단일 뷰에 조건화된 사전 분포(prior distribution)를 활용합니다.
4. 손실 함수 (Loss Functions): 모델은 두 가지 상호 보완적인 목적 함수를 최적화합니다:
   • Sim-1 (KL Divergence): 사후 분포(두 뷰로부터 추론됨)와 사전 분포(단일 뷰로부터 추론됨) 사이의 쿨백-라이블러(Kullback–Leibler) 발산을 최소화합니다. 이는 인코더가 뷰 변화에 강건한 표현을 학습하도록 장려합니다. 사전 브랜치에는 스톱 그래디언트(stop-gradient) 연산자가 적용됩니다.
   • Sim-2 (NMSE): 예측된 마스킹된 토큰과 교사 네트워크로부터의 실제 마스킹된 토큰 사이의 정규화된 평균 제곱 오차(normalized mean squared error)를 최소화합니다.
5. 통합 공식 (Unified Formulation): 이 프레임워크는 시암 일관성 손실의 영향력을 조절하는 정규화 파라미터($\lambda_{KL}$)를 포함합니다. $\lambda_{KL} = 0$으로 설정하면 단일 인코더 JEPA가 되어 시암 구성 요소에 대한 체계적인 절제 실험(ablation)이 가능합니다.

주요 기여

• SiamJEPA 제안: 시암 학생 인코더를 JEPA로 확장한 통합 잠재 예측 프레임워크를 제안하며, 이를 시암 기여도가 정규화 파라미터를 통해 제어되는 연속적 확장형으로 공식화했습니다.
• 체계적 조사: 시암 학생 인코더가 표현 공간을 제한하는 효과적인 정규화 도구로서 작용함을 입증함으로써, JEPA 내에서의 시암 학생 인코더에 대한 최초의 체계적인 연구를 수행했습니다.
• 학습 효율성: 시암 인코더를 사용한 잠재 예측이 재구성 기반 방식(예: MAE)보다 훨씬 더 높은 학습 효율성을 보이며, 훨씬 적은 에포크(epoch)만으로도 경쟁력 있는 성능을 달착할 수 있음을 입증했습니다.

실험 결과
실험은 ViT-Base 백본을 사용하여 ImageNet-1K 데이터셋에서 선형 프로빙(linear probing) 평가를 통해 진행되었습니다:

• 성능 대비 에포크 (Performance vs. Epochs): SiamJEPA는 400 에포크 후 **70.7%**의 Top-1 선형 프로빙 정확도를 달성했습니다. 이는 400 에포크에서 61.9%를 기록한 MAE를 크게 상회하며, 1600 에포크 동안 학습된 MAE(68.0%)와도 견줄 만한 수준입니다.
• 시암 인코더의 역할: KL 정규화 가중치($\lambda_{KL}$)에 대한 절제 실험 결과, $\lambda_{KL}$을 높이면(예: 0.01 또는 0.03) 거의 0에 가까운 설정($\lambda_{KL} = 0.00001$, 사실상 단일 인코더 JEPA)에 비해 성능이 일관되게 향상되었습니다.
• 수렴 속도: 더 높은 $\lambda_{KL}$을 가진 모델은 더 빠르게 수렴하여, 약한 정규화 설정이 400 에포크에서 달성한 성능 수준에 200 에포크 만에 도달했습니다.
• 마스킹 전략: 블록 마스킹(Block masking)이 특히 학습 초기 단계에서 무작위 마스킹(random masking)보다 더 효과적인 것으로 나타났으나, 세심한 튜닝을 거친 무작위 마스킹 또한 경쟁력 있는 결과를 보였습니다.
• 표현 추출 (Representation Extraction): 10번째 레이어(중간층)의 평균 풀링(mean pooling)이 초기 학습 단계에서는 약간 더 나은 결과를 보였으나, 12번째 레이어(최종층)가 결국 이를 따라잡거나 능가했습니다. 이는 서로 다른 레이어가 학습 단계에 따라 서로 다른 전이 가능한 의미론(semantics)을 포착함을 시사합니다.

의의 및 주장
본 논문은 시암 학생 인코더가 단순한 구조적 선택이 아니라, 예측적 표현 학습을 위한 중요한 **귀납적 편향(inductive bias)**이라고 주장합니다. 연구 결과는 다음과 같은 통찰을 제공합니다:

1. 시암 아키텍처는 표현의 분리 가능성(separability)을 개선하고 학습을 가속화하며, 특히 학습 초기 단계에서 효과적인 정규화제로 작용합니다.
2. 시암 학생 아키텍처를 도입하는 것은 자기 지도 학습을 개선하는 단순하면서도 효과적인 방법이며, 재구성 기반 방식보다 더 효율적인 대안을 제공합니다.
3. 결과는 JEPA 기반 모델 설계에 대한 새로운 통찰을 제공하며, 시암 인코더를 사용하는 뇌 모사 PhiNet 아키텍처가 잠재 예측을 통해 일반적인 이미지 표현 학습에 효과적으로 적응될 수 있음을 시사합니다.

저자들은 절대적인 최첨단(SOTA) 성능에 대해서는 겸허한 태도를 유지하며, 주요 목표가 벤치마크 점수를 극대화하는 것이 아니라 JEPA 내에서 시암 인코더의 역할을 검증하는 것임을 명시했습니다. 또한, 향후 연구를 위해 하이퍼파라미터 최적화 및 더 큰 모델(ViT-Large/Huge)로의 확장이 필요함을 인정했습니다.