I-JEPA 논문 리뷰

Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture

---

Summary

본 논문은 비지도 표현 학습에서 기존 Joint-Embedding 아키텍처와 Generative 아키텍처의 한계를 지적하고, 이를 보완하기 위해 encoder가 표현을 직접 예측하도록 하는 Joint-Embedding Predictive 아키텍처(JEPA)를 제안한다.

제안된 방법은 모델이 더 풍부하고 고수준의 표현을 학습하도록 유도할 뿐만 아니라, 다양한 다운스트림 테스크에서 경쟁력 있는 성능을 보이며 GPU 연산 효율성 또한 향상됨을 입증한다.

---

Preliminaries

Joint-Embedding Architectures

Joint-Embedding 아키텍처는 같은 이미지를 서로 다르게 증강하여 얻은 $x$와 $y$의 표현을 정렬시키면서, 다른 이미지(negative pair)의 표현과 거리를 두는 방식의 학습 방법을 의미한다. 이러한 방법의 주된 설계는 representation collapse를 막는 것에 있으며, 대표적인 해결 접근으로는

• Contrastive loss를 통해서 negative example을 representation space에서 밀어내도록 하는 방법 (SimCSE)
• Exponential moving-average를 사용하여 reference가 되는 encoder의 모델을 안정적으로 변화시키면서 target을 학습시키는 비대칭적 방법(BYOL)
• Embedding 평균의 엔트로피를 최대화시켜 한 방향으로 수렴하는 것을 방지하는 방법

이와 같은 기법들은 풍부한 표현을 학습할 수 있다는 장점이 있으나, 특정 다운스트림 과제에는 적합하지 않거나 특정 데이터 분포에서는 성능이 저하되는 경향이 있어, 보편적(global) 해결책으로 보기는 어렵다. 또한 아무리 contrastive loss를 사용한다고 하더라도 batch training 과정에서 negative sample의 수가 많아야 representation collapse를 방지할 수 있는데, 이는 그대로 GPU 연산 및 메모리에 부담을 주기도 한다.

Generative Architecture

원본 이미지 $y$에 대해서 패치로 나눈 후 일부 패치들을 마스킹한 $x$를 만들고, $x$의 정보를 바탕으로 원본 이미지 $y$로 복구하도록 학습하는 방법이다. 이러한 방법의 경우 representation collapse에는 상당히 자유롭지만 픽셀 복원 수준이기 때문에 고수준의 표현을 학습하기에는 한계가 존재한다.

---

Architecture

I-JEPA는 앞서 preliminaries에서 제기된 두 아키텍처 문제의 절충안인 Joint-Embedding Predictive Architecture을 제시한다. 우선 target, context 데이터를 살펴보자.

위 그림을 보면 generative 아키텍처인 MAE를 연상할 수 있지만, MAE는 픽셀 단위 예측에 초점을 맞추는 반면 I-JEPA는 representation 공간에서 예측이 이루어진다는 것이 차이점이다.

Targets

입력 이미지 $y$가 주어지면 이를 $N$개의 겹치지 않는 패치 시퀀스로 분할한 뒤, 타깃 인코더 $f_{\bar{\theta}}(\cdot)$에 연결하여 $N$개의 패치 레벨 표현 집합, 즉 $s_y = \{ s_{y_1}, ... , s_{y_N} \}$을 얻는다.

이후 $s_y$에서 $M$개의 (서로 겹칠 수도 있는) 블록을 샘플링한다. 이때 각 블록은 해당 영역을 구성하는 패치들의 임베딩 벡터 집합으로 정의된다. 논문에서는 블록을 단순히 단일 고정 크기나 비율으로 만드는 것이 아니라 스케일(scale: 전체 이미지 대비 블록이 차지하는 면적 비율)과 종횡비(ratio: 블록의 가로:세로 비율)을 다양하게 만들어서 target으로 사용한다.

 

이렇게 만들어진 $i$번째 타겟 block을 $B_i$ 로 정의하고, $B_i$ 에 속하는 패치들의 임베딩 집합을 $s_y(i) = \{ s_{y_j} \}_{i \in B_i}$ 라 한다.

 

Context

Context 블록 생성 과정은 다음과 같다.

• 입력 이미지 $y$를 ViT 패치($N$개 패치)로 쪼개고 나서
• 이 패치 그리드 위에서 하나의 큰 블록 $x$를 샘플링을 하는데, 여기에서
  • Scale: 전체 이미지 대비 컨텍스트 블록 $x$가 차지하는 면적을 조정하고,
  • Aspect Ratio: 블록의 가로:세로 비율을 조정한
• 컨텍스트 블록을 가져온다. 논문에서는 가로:세로 비율은 1:1, 즉 정사각형으로 두고, scale은 0.85~1.0을 사용한다.
• 이는 타깃 블록의 생성 과정과 독립적으로 진행되기 때문에, 자칫하면 컨텍스트 블록과 이미지가 겹치는 부분이 존재하여 모델이 정답을 쉽게 맞추어버리는 문제점이 존재한다.
• 따라서 컨텍스트 블록에서 타깃 블록과 겹치는 영역은 제거하여, 컨텍스트는 타깃을 직접적으로 포함하지 않는 부분 정보의 벡터 집합이 된다.

Prediction & Loss

컨텍스트 인코더의 출력 $s_x$가 주어지면 $M$개의 타깃 표현에 대한 예측을 수행한다. 여기서 predictor $g_{\phi}(\cdot, \cdot)$은 $s_x$와 출력을 예측하려는 각 패치에 대한 마스크 토큰 $\{m_j\}_{j \in B_i}$ 를 입력으로 사용한다.

즉 출력되는 $i$번째 블록의 패치에 대한 예측 값은 위와 같은 형태로 출력되는 것이다.

 

결국 $M$개의 타겟 블록에 대한 예측을 수행하고, 매번 예측하려는 타겟 블록 위치에 해당하는 마스크 토큰을 조정하며 예측을 한다.

 

컨텍스트 인코더는 예측된 패치 레벨 표현과 타겟 패치 레벨 표현 사이의 L2 평균 거리를 loss로 하여 업데이트된다.

타겟 인코더는 EMA를 통해서 업데이트되는데, 이는 컨텍스트 인코더의 안정적인 학습을 위해서 필수적이다.

 

결론적으로, I-JEPA는 단순히 한 이미지의 증강 뷰를 매칭하는 방식이 아니라, (1) 이미지의 일부 영역(컨텍스트)을 기반으로 가려진 영역(타깃)을 예측하는데, 이때 픽셀 수준이 아닌 표현 단위에서의 예측을 수행하기 때문에 collapse를 피하면서 의미적 수준이 높은 표현을 학습할 수 있으며, (2) negative pair를 사용하지 않고 positive pair 기반의 예측 과제만 진행하기 때문에 GPU 연산 및 메모리 부담이 적고, (3) EMA를 활용하여 target encoder를 안정적으로 업데이트함으로써 학습 안정성을 보장하는 추가적인 장치를 제공한다.

---

Experiments

아키텍처 설명의 결론 부분이 실제 실험과 얼마나 부합하는지 확인해보자.

I-JEPA가 epoch 수를 압도적으로 줄이고도 더 좋은 성능을 내는 것을 확인할 수 있다. MAE의 경우 픽셀 레벨 예측이라 저수준 정보뿐 아니라 필요 없는 부분까지 전부 다 학습해야 해서 많은 반복이 필요하며, BYOL과 SimCLR의 경우 이미지 페어 간 학습이기 때문에 오랜 반복이 필요할 수밖에 없다. 특히 SimCLR은 배치 크기도 커야 하기 때문에 비효율적이다.

왼쪽은 image classification 전이 성능(high-level classification)에 관한 표고, 오른쪽은 object counting, depth prediction 관련 성능(low-level geometry)에 관한 표다.

다른 방법론들은 둘 중 하나에서만 우수한 성능을 보이는 한편, I-JEPA는 두 가지 측면에서 모두 강하다는 것을 확인할 수 있다. 즉 본 논문에서 제기한 특정 방법론들의 다운스트림 테스크에서의 약점에 대한 대안으로써 충분하다는 것을 확인할 수 있다.

 

Scalability에서도 우수한 것을 확인할 수 있다.

적은 데이터(IN1k)에서도 성능이 DINO나 iBOT에 비교하여도 우수한 것을 확인할 수 있으며, 데이터셋을 확장하는 경우 성능이 더욱 확실하게 상승하는 것 또한 확인할 수 있다. (사실 모델 스케일링에 대해서 우수하다는 근거는 잘 모르겠다...ㅎㅎ)

더 적은 epoch과 GPU 훈련 시간으로도 좋은 성능을 가져오는 것을 확인할 수 있다.

 

이번에는 Ablation study를 확인해보자.

Multi-block은 target을 여러 개 블록 (scale: 0.15-0.2), context를 큰 블록에서 타깃과 겹치는 부분을 제거한 것을 의미한다. Rasterized는 격자식 분할, block은 타깃을 하나의 큰 블록으로 삼는 것이며, random은 말 그대로 랜덤으로 타깃을 정하고, 나머지를 context로 한 것이다.

 

Rasterized의 경우 경계가 인위적이고 너무 단순하여 표현이 비자연적이라 학습 성과가 낮으며, block는 예측 단위가 너무 크다는 문제점, random의 경우 spatial 구조가 깨져서 위치 정보 및 문맥 반영이 어려운 것을 확인할 수 있다.

픽셀 단위 예측 테스크보다 representation 기반 예측 테스크가 더욱 성능도 좋으며 효율적인 것을 확인할 수 있다.