[2026-1] 김지원 Self-Supervised Learning from Images with aJoint-Embedding Predictive Architecture
논문 제목 : Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
논문 링크: https://arxiv.org/pdf/2301.08243
배경
I-JEPA(Image-based Joint-Embedding Predictive Architecture)는 '라벨 없이 이미지의 semantic representation을 어떻게 학습할 것인가'라는 자기 지도 학습 문제를 다룸
학계에서는 두 가지 방법으로 이 문제를 다뤄왔음:
1. View-invariance(SimCLR, BYOL, DINO)
같은 이미지를 증강(crop, color jitter 등)으로 여러 view로 변형한 뒤 변형된 view의 임베딩이 서로 비슷해지도록 학습했음
높은 퀄리티의 semantic representation을 얻지만, 어떤 증강 방법을 사용하는가가 곧 '무엇을 같다고 볼 것인가'에 대한 강한 편향이 됨
예를 들어 색을 무작위로 바꾸면 '색은 중요하지 않다'는 편향이 생기고, 이는 분류에는 도움이 되지만 색이 실제로 중요한 다른 과제에는 오히려 해가 됨
또한, 이렇게 학습된 모델은 다른 태스크에 일반화되기 어려움
2. Generative(BERT 기반의 BEiT 등)
이미지의 일부 patch를 masking하고 가려진 부분을 픽셀 또는 토큰 단위로 복원하도록 학습함
그 과정에서 'masking된 부분의 정확한 픽셀 값'까지 맞히려다 보니 배경의 노이즈 같은 저수준 디테일에 모델 용량을 많이 쓰게됨
그 결과 linear probing 같은 평가에서 view-invariance 방법보다 의미적 수준이 떨어지는 경향이 있음
I-JEPA는 이 두 방법의 장점만 가져오고자 함
Augmentation 없이, Masking만으로 학습하되 픽셀이 아닌 표현공간에서 예측하는 것
(즉, 가려진 공간의 정보를 예측하는 것)
아키텍쳐
모델의 context-encoder, target-encoder, predictor는 모두 Vision Transformer(ViT) 구조를 사용
모델의 작동 방식은 아래와 같음:
(target)
원본 이미지를 N개의 겹치지 않는 patch로 나누고 target encoder에 통과시켜 patch 별 표현을 얻음
이후 패치 임베딩으로부터 M개 샘플(overlap이 있을 수 있음)을 샘플링
이때 샘플은 종횡비 (0.75,1.5), 크기 스케일 (0.15,0.2) 범위 내에서 무작위로 샘플링함
이 샘플들에 해당하는 임베딩 벡터들이 target이 됨
(context)
context의 경우 처음부터 이미지의 일부만 보고 인코딩 함
일반적으로 이미지는 scale(0.85,1.0), aspect ratio는 정사각형에 가깝게 고정된 블록 하나를 샘플링
context-encoder를 통해 이미지의 일부를 인코딩
(prediction)
context-encoder의 임베딩 결과와 맞춰야 하는 target block의 마스크 토큰(target block의 위치 정보)을 함께 입력받음
(Loss)
$$ \frac 1 M \Sigma_{i=1}^M\Sigma_{j\in B_i}||\hat s_{y_j}-s_{y_j}||_2^2 $$
주의할 점은 Predictor와 Context Encoder는 역전파를 통해 직접적으로 훈련이 되지만
target encoder는 그렇지 않다는 점임
target encoder는 매 스텝마다 아래 수식을 통해 context encoder를 따라감
$$ m\cdot\bar\theta + (1-m)\cdot\theta \to \bar \theta $$
이를 통해 collapse(두 인코더가 입력과 무관하게 항상 같은 상수 벡터를 뱉는 것)을 방지함
실험 정리
1) ImageNet-1K Linear Evaluation
Self-supervised 사전학습 후 인코더를 고정하고 linear classifier만 학습시킴
I-JEPA는 augmentation 없이 학습한 방법 중 최고 성능인 79.3%를 기록하며
같은 조건인 MAE(77.2%), data2vec(77.3%)를 넘어섬
해상도 448에서 학습한 버전(81.1%)은 augmentation을 쓰는 DINO, iBOT과도 맞먹음
2) ImageNet-1% Low-shot 평가
라벨의 1%(클래스당 약 12~13장)만 사용해 fine-tuning 또는 linear probing으로 적응시킨 결과
I-JEPA는 MAE보다 훨씬 적은 epoch로도 앞서고, data2vec과 동급 성능을 더 적은 계산량으로 달성
해상도를 키운 버전은 augmentation을 쓰는 MSN까지 넘어섬
3) 다운스트림 분류 전이학습 - Linear Probe
사전학습된 인코더를 고정한 채 CIFAR100, Places205(장면 분류), iNaturalist18(세밀한 종 분류)로 전이한 결과
I-JEPA는 CIFAR100, Places205에서 augmentation을 쓰는 DINO보다 높은 점수를 냄
하지만 iNat18에서는 DINO, iBOT에 뒤쳐지는데 이는 augmentation 기반 방법이 미세한 지역적 차이를 구분하는 데는 여전히 강점이 있다는 뜻으로 해석됨
4) Low-level task - 물체 개수, 깊이 추정
Clevr 데이터 셋으로 물체 개수 세기(Count)와 거리 추정(Dist)을 평가
I-JEPA는 깊이 추정(Dist)에서 DINO, iBOT을 큰 차이로 앞서고 물체 개수 세기(Count)에서도 augmentation 기반 방법들과 비슷하거나 더 나음
이는 Crop, scaling 같은 augmentation을 쓰지 않기 때문에 '물체가 몇 개 있었는지, 얼마나 큰 지' 같은 정보가 왜곡되지 않는다고 저자는 해석함
MAE는 픽셀 단위 복원을 목표로 하다 보니 오히려 Count에서 가장 높은 점수를 기록함
5) 데이터, 모델 스케일업
데이터를 키우면 대부분 지표가 개선되지만
모델을 키우면 의미적 과제는 더 좋아지지만 저수준 과제는 오히려 살짝 떨어짐
이는 모델을 키울 수록 더 큰 patch를 사용하며 세밀한 지역 정보 손실이 야기되기 때문임
6) 계산 효율성
GPU 시간 대비 1% ImageNet 성능을 비교한 결과, I-JEPA는 MAE-data2vec보다 훨씬 적은 반복만으로 수렴함
Augmentation 기반 iBOT과 비교하면 격차가 더 큼
7) Predictor, Target Encoder 시각화
RCDM(diffusion 기반 디코더)으로 표현 벡터를 다시 픽셀로 복원해봄
아래 이미지는 순서대로 predictor, target encoder 출력 그리고 MSN과 비교임
predictor 출력은 같은 표현을 여러 노이즈 시드로 디코딩했을 때 물체의 종류, 포즈는 샘플마다 공통적으로 유지되지만, 정확한 색, 질감, 배경은 샘플마다 달라짐(predictor가 위치의 불확실서은 인지하면서 고수준 정보만 예측하고 있다는 증거임)
target encoder 출력도 마찬가지로 물체의 고수준 정보는 유지하고 저수준 디테일, 배경은 불확실함
view-invariance 기반의 MSN과 비교에서는 디코딩 시 샘플 가느이 변동성이 매우 커 I-JEPA가 지역적 구조 정보를 더 많이 보존한다는 것을 보임
8) Ablation
마스킹 전략 비교