SiamJEPA: JEPA에서의 시암스 학생 인코더 역할 연구

자기지도 학습에서 잠재 공간 예측은 픽셀 재구성 대비 학습 효율과 의미적 초점을 제공한다. 그러나 학생 네트워크 구조가 JEPA 성능에 미치는 영향은 충분히 규명되지 않았다. 본 연구는 시암스 학생 인코더가 JEPA에 규제적 바이어스를 부여하여 학습 초기에 수렴을 가속하고 최종적인 선형 분류 성능을 개선함이 확인되었다.

전통적 joint embedding 방법은 서로 다른 시야의 표현을 동일한 잠재 공간에 정렬하여 의미적 일관성을 확보한다. JEPA는 이 방향을 잠재 예측으로 전환하여 가시적 문맥으로부터 마스크된 영역의 임베딩을 직접 예측한다. 이 접근은 픽셀 단위 재구성보다 표현의 추상화와 학습 효율을 높이는 작동 원리를 가진다. 시암스 학생 인코더는 동일한 가중치를 공유하는 두 분기를 독립 마스킹으로 운용하여 서로 다른 부분 관측에서 공통 정보를 포착하게 한다. 두 분기 사이의 KL 정합 항은 한 분기만으로도 다른 분기의 포스터리어를 근사하도록 학습 압력을 부여한다. 이로 인해 표현은 뷰 불변성을 갖추면서도 분별력을 유지하게 된다. 규제 관점에서 보면 시암스 구성은 JEPA의 주 손실인 마스크 예측 신호에 보완적 역할을 한다. 마스크 예측이 주된 학습 신호를 제공하는 동안 KL 항은 표현 붕괴를 억제하고 초기 학습 단계에서 안정적 방향으로 옵티마를 유도한다. 따라서 시암스 인코더는 최종 성능 자체뿐 아니라 학습 속도와 초기 표현 분리도에도 중요한 영향을 미친다. 실험적 관찰은 이 아이디어가 실무적 이득을 제공함을 뒷받침한다. 적절한 KL 가중치와 프리비트 설정은 초반 에폭에서 빠른 성능 향상을 유도했고 블록 마스킹과의 조합은 낮은 에폭 예산에서도 우수한 전이 성능을 산출했다. 이러한 점은 대형 자원 없이도 비교적 짧은 사전학습으로 고품질 표현을 얻을 가능성을 의미한다.

전체 접근은 두 개의 학생 인코더와 EMA 기반의 교사 인코더로 구성된 삼중 구조를 채택했다. 입력 이미지는 패치화되어 두 학생 분기에 독립적으로 마스킹된 뷰로 공급되며 교사 인코더는 마스킹 없는 전체 이미지를 처리한다. 학생 출력은 CLS 및 패치 토큰으로 분해되어 글로벌 정합과 패치 수준 예측을 위한 서로 다른 예측기에 입력된다. 핵심 메커니즘은 두 가지 손실로 구성된다. 글로벌 일치 목적은 두 학생 분기에서 추정한 잠재 분포의 포스터리어와 단일 분기 기반의 프라이어 사이의 KL 발산을 최소화하는 형태로 정의되며, 이때 프라이어 쪽에는 stop-gradient가 적용되어 안정성을 확보한다. 패치 수준 예측은 정규화된 평균제곱오차로 계산되어 EMA 교사에서 얻은 목표 임베딩을 재구성하는 것이 아니라 잠재 벡터를 직접 예측하는 데 초점을 둔다. 구현 세부에서는 ViT-Base 백본과 얕은 트랜스포머 기반 예측기를 사용했다. 마스킹은 블록 마스킹과 랜덤 마스킹을 비교했으며 프리비트와 λ_KL 하이퍼파라미터가 규제 강도를 조절했다. EMA 모멘텀은 학습 진행에 따라 0.99에서 0.9999로 단계적으로 증가시키는 스케줄이 적용되었다.

주요 벤치마크로 ImageNet 선형 프로빙을 사용했고 ViT-Base 백본에서 비교 실험을 수행했다. 동일한 400 에폭 비교에서 논문은 MAE가 Top-1 61.9%인 반면 SiamJEPA는 Top-1 70.2%를 기록했으며 에폭별 성능 추이를 통해 SiamJEPA가 적은 에폭으로 동등하거나 우수한 성능에 도달함이 확인되었다. I-JEPA는 600 에폭에서 72.9%로 더 높은 값을 보였으나 학습 설정이 다르므로 직접 비교에는 주의가 필요하다. Ablation 연구에서 λ_KL 값 증가는 선형 프로빙 성능을 일관되게 향상시켰다. 예컨대 λ_KL=0.0001 설정은 에폭 300에서 66.78%를 기록한 반면 λ_KL=0.01 설정은 같은 시점에서 69.30%를 기록했다. 또한 프리비트 설정은 0.05에서 약간의 향상이 관찰되었으나 성능 민감도는 크지 않았다. 마스킹 전략 비교에서는 블록 마스킹이 특히 적은 에폭 예산에서 유리했으며 랜덤 마스킹도 적절한 비율 조정으로 경쟁력을 보였다. 이와 함께 중간층의 평균풀링이 CLS 토큰보다 초기 에폭에서 더 높은 전이 성능을 제공했음이 보고되었다.

아키텍처는 세 가지 인코더로 구성되며 학생 인코더 두 개는 가중치를 공유하는 Siamese 형태이고 교사 인코더는 EMA로 갱신된다. 입력은 패치화되어 positional embedding이 더해진 뒤 마스킹 연산을 통해 학생 분기로 들어가며 두 학생의 마스크 집합은 서로 비교집합이 되도록 설계되었다. 학습 시 학생의 출력은 CLS와 패치 임베딩으로 분리되어 글로벌 정합을 위한 선형 예측기와 패치 예측을 위한 트랜스포머 예측기로 전달된다. 수학적 손실 구성은 두 부분으로 명확히 나뉜다. Sim-1 항은 q 및 p 분포 간의 KL 발산으로 정의되고 실제 최적화에서는 prior 쪽에 stop-gradient를 적용한 KL_sg 형태가 사용되며 free-bit 임계값을 두어 KL이 완전 수렴하는 것을 방지한다. Sim-2 항은 마스크된 패치에 대해 목표 잠재 임베딩과 예측 임베딩 간의 Frobenius 기반 평균제곱오차로 정의되어 패치 수준 재구성이 아니라 잠재 예측이 직접 학습 신호가 된다. 학습 설정은 ViT-Base, 디코더 깊이 1, 배치 유효치 8192, AdamW 옵티마이저와 기본 학습률 1.5e-4 등을 사용했다. EMA 모멘텀 스케줄은 에폭 구간에 따라 0.99에서 0.9999로 증가하도록 설계되었고 weight decay와 마스킹 비율 등 하이퍼파라미터는 광범위한 소거 연구로 선택되었다. 실험은 V100/A100/H100 GPU에서 수행되었으며 일부 결과는 300 에폭 소규모 실험으로 검증된 뒤 최종 비교는 400 에폭으로 보고되었다.

연구 자체에서 언급한 한계는 주로 상대적으로 작은 규모의 모델과 설정에서 실험이 수행되었다는 점이다. 또한 최종 성능이 구현 세부사항과 하이퍼파라미터에 민감하여 더 광범위한 스케일 업과 체계적 튜닝이 필요함이 기술되었다. 연구 범위는 이미지 기반 사전학습에 국한되어 있어 비디오나 대형 아키텍처로의 일반화 가능성은 추후 검증이 요구된다.