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

[Dongwoo-Im]

github : https://github.com/facebookresearch/ijepa

---

1. Introduction

invaraiance-based (contrastive) pre-training은 hand-crafted view-based augmentation에 의존하기에 결국 특정 down-stream task에 적합한 방식이다. (ex. 이미지 분류)

masked pre-training은 상기한 augmentaiton의 영향을 받지 않기에 다른 modality로의 확장이 용이하지만 (data2vec)
reconstruction하는 decoder의 영향으로 더 낮은 수준의 semantic을 갖는 representation을 추출하는 것으로 보인다. (MAE, MSN)

I-JEPA는 extra prior knowledge 없이도 높은 semantic 정보를 갖는 representation 추출이 가능하다. how?

• abstract prediction target : I-JEPA는 target representation을 예측한다.
• multi-block masking strategy

---

2. Background

이미지 관점에서 각 architecture 설명

• (a) JEA : invariance-based learning (ex. contrastive learning)
• (b) GA : reconstrunction-based learning (ex. masked image modeling)
  • (learnable) mask와 position token이 z에 해당
  • 즉, 이 경우에 z가 y에 비해 적은 정보를 갖기에 representation collapse가 발생하지 않는다고 한다.
• (c) JEPA : I-JEPA
  • JEA와 다른 점은 invariant representation을 추출하는 것이 아니라 additional z를 바탕으로 예측한다는 것
  • GA와 다른 점은 input space가 아닌 embedding space에 걸리는 loss function이다.
  • x-encoder와 y-encoder는 asymmetric design

---

3. Method

• context block이 주어지면 각 target block의 representation을 예측하는 구조이고, 최종적으로 target encoder를 evaluation에 사용
• context encoder, predictor, target encoder 모두 ViT
  • predictor는 narrow ViT지만, self-attention head는 encoder와 동일하게 맞춰줌
  • [CLS] token 사용 X (average pooled representation on last layer or last 4 layers)
• I-JEPA가 MAE와 다른 점은 non-generative 모델이고 representation level에서 prediction이 수행된다는 점

Targets

• target scale = (0.15, 0.2)
• target끼리는 overlap되어도 상관 없고, 저자들은 target을 4개로 설정
• target encoder의 output에 masking

data2vec의 figure로, teacher의 representation을 예측한다는 점에서 I-JEPA와 비슷한 점이 있다.

Context

• context scale = (0.85, 1.0)
• 각 prediction이 유의미할 수 있도록, context block 중 target block과 겹치는 영역은 제거한다.

Prediction

• 1개의 predictor로 4개의 target을 예측한다.

Loss

• patch-level representation에 L2 distance loss 사용
• context encoder와 predictor를 update하며, target encoder는 EMA로 update된다.

---

Evaluation

• 전체적으로 I-JEPA에서 기록한 모델 사이즈가 더 큰 경향이 있지만 연산 비용은 낮다.
  • hand-crafted augmentation이 존재하지 않음
  • 1 iter는 MAE보다 약간 더 오래 걸리지만 5배 더 빠르게 수렴한다고 주장
  • torch autocast bfloat16의 영향일수도

• left : linear evaluation on ImageNet
• right : low-shot semi-supervised learning on ImageNet-1%
  • linear probing과 partial fine-tuning 중 더 높은 성능 기재
  • similar recipe with semi-ViT)

• left : linear probing transfer on image classification
• right : linear probing transfer on low-level task (object counting, depth prediction)
  • ablation에서 model size를 키우더라도 low-level task 성능은 향상되지 않았다고 한다.

---

Predictor Visualization

predictor가 각 mask token의 위치에 적절한 latent를 예측하고 있는지 (JEPA 구조에서 z에 condition된 예측이 가능한지)
meta의 RCDM framework로 시각화한 결과입니다. (3,4,5,6 column은 서로 다른 seed 결과)

Representational Conditional Diffusion Model (RCDM)의 figure로,
pretrained SSL representation을 condition으로 하는 diffusion-based decoder를 학습하여 시각화에 활용한다.

---

Ablations

MAE와 동일하게 pixel을 target으로 할 경우, 심각한 성능 저하
저자들은 representation을 target함으로써 pixel detail에 의존적이지 않은 abstract prediction이 주효했을 것이라 추측

저자들은 target과 겹치는 context block을 제거해준 것이 주효했을 것이라 추측