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

Abstract

이 논문은 수작업으로 설계된 데이터 증강에 의존하지 않고 의미 있는 이미지 표현을 학습할 수 있는 self-supervised 방법인 I-JEPA(Image-based Joint-Embedding Predictive Architecture)를 제안한다. I-JEPA는 하나의 context block을 기반으로 동일한 이미지 내 여러 target block의 표현을 예측하도록 학습된다. 이때 의미적 표현을 효과적으로 학습하기 위해 비교적 큰 target block과 공간적으로 분산된 context block을 사용하는 masking 전략을 사용한다. 실험 결과 Vision Transformer와 결합한 I-JEPA는 높은 확장성을 보였으며, ImageNet 사전학습 모델은 linear classification뿐 아니라 객체 개수 예측이나 깊이 추정과 같은 다양한 다운스트림 태스크에서도 좋은 성능을 보였다. 또한 픽셀 공간이 아니라 representation 공간에서 예측을 수행하기 때문에 학습 효율이 높고 대형 모델도 비교적 적은 계산량으로 학습할 수 있음을 확인하였다.

Introduction

기존 이미지 기반 self-supervised 학습은 크게 두 가지 계열로 나뉜다.

첫 번째는 invariance-based 방법으로, 동일한 이미지에서 데이터 증강을 통해 여러 view를 생성하고, 같은 이미지에서 생성된 서로 다른 view의 표현이 유사하도록 학습하는 방식이다. 이러한 접근은 높은 수준의 의미적 표현을 학습할 수 있다는 장점이 있지만, 랜덤 크롭, 색상 변형 등 수작업으로 설계된 데이터 증강에 강하게 의존한다는 문제가 있다. 이러한 증강은 특정 태스크에 편향을 만들 수 있으며, 이미지 분류와 같은 태스크에는 유리하지만 다른 비전 태스크나 다른 모달리티로 일반화하기 어렵다. 

 

두 번째 계열은 generative 방식으로, 입력 이미지의 일부를 가리거나 손상시킨 뒤 이를 복원하도록 학습하는 접근이다. 대표적으로 masked image modeling 방식이 있으며 이미지 패치나 픽셀을 복원하는 과정에서 표현을 학습한다. 이 방식은 데이터 증강에 대한 의존성이 적고 다양한 모달리티로 확장하기 쉽지만, 픽셀 복원 중심의 학습은 의미적 수준이 낮은 표현을 학습하는 경향이 있어 semantic classification 같은 다운스트림 태스크에서 성능이 떨어질 수 있다. 

 

이 논문은 이러한 두 접근의 장단점을 고려하여, 픽셀을 복원하는 대신 representation을 예측하는 방식을 사용하는 I-JEPA를 제안한다. I-JEPA는 이미지의 일부 영역을 context로 사용하고, 동일 이미지의 다른 영역(target block)의 embedding을 예측하는 self-supervised 학습 방식이다. 모델이 직접 픽셀을 생성하는 것이 아니라 추상적인 표현 공간에서 예측을 수행하도록 설계하여 불필요한 저수준 픽셀 정보를 제거하고 더 의미적인 표현을 학습하도록 유도한다. 또한 의미적 정보를 충분히 포함한 target block과 공간적으로 분산된 context block을 사용하는 multi-block masking 전략을 통해 모델이 이미지의 구조와 의미를 이해하도록 한다. 

 

실험적으로 I-JEPA는 데이터 증강 없이도 강력한 이미지 표현을 학습할 수 있으며, ImageNet linear probing, 1% semi-supervised 학습, 그리고 다양한 semantic transfer task에서 픽셀 복원 기반 방법인 MAE보다 높은 성능을 보인다. 또한 semantic task에서는 기존 view-invariance 기반 방법들과 경쟁력 있는 성능을 유지하면서도 객체 개수 예측이나 depth prediction과 같은 저수준 비전 태스크에서는 더 좋은 성능을 보인다. 마지막으로 representation 공간에서 예측을 수행하는 구조 덕분에 학습 효율이 높아 대규모 모델에서도 계산 비용이 크게 감소한다.

 

 

이건 처음에 이해가 잘안되어서 직접 그려본건데, 원래는 이미지 자체의 픽셀을 복원해야했기 때문에 텍스쳐나 색상, 패턴과 같은 저수준을 함께 고려해야했다. 하지만, I-JEPA에서는 인코딩된 representation(이미 pixel 정보가 일부 제거된 경우)을 예측하기 때문에 상대적으로 더 의미적인 표현을 많이 학습할 수 있다고 생각하면된다.

 

Background

self-supervised learning의 목적은 입력들 사이의 관계를 학습하여 호환되는 입력 쌍은 낮은 energy를, 호환되지 않는 입력 쌍은 높은 energy를 갖도록 학습하는 것으로 볼 수 있다. 이 관점에서 기존 방법들은 크게 세 가지 구조로 분류된다.

 

첫 번째는 Joint-Embedding Architecture(JEA)이다. 이 방식은 동일한 이미지에서 생성된 두 view를 encoder에 입력하여 두 임베딩이 서로 유사하도록 학습한다. 일반적으로 데이터 증강을 통해 view를 생성하며, contrastive loss나 redundancy reduction 등을 사용해 representation collapse를 방지한다.

 

두 번째는 Generative Architecture이다. 이 접근은 입력의 일부를 마스킹하거나 손상시킨 뒤 이를 복원하도록 학습하는 방식으로, 대표적으로 masked image modeling이나 autoencoder 기반 방법이 있다. 이 구조에서는 decoder가 손실된 정보를 복원하며, reconstruction loss를 통해 표현을 학습한다.

 

세 번째는 Joint-Embedding Predictive Architecture(JEPA)이다. JEPA는 generative방식과 유사하게 입력의 일부 정보를 이용해 다른 부분을 예측하지만, 픽셀을 직접 복원하는 대신 embedding 공간에서 target representation을 예측한다는 점이 차별점이다. 즉 입력 x로부터 다른 신호 y의 representation을 예측하도록 학습하며, 예측 과정에서 추가적인 정보 z가 조건으로 사용된다.

 

I-JEPA는 이러한 JEPA 구조를 이미지 학습에 적용한 모델로, 데이터 증강에 의존하지 않고 representation 공간에서 예측을 수행함으로써 보다 의미적인 이미지 표현을 학습하는 것을 목표로 한다.

 

 

**이것도 이해가 잘안돼서 지피티에게 물어보며 정리했다..

representation collapse는 모델이 서로 다른 입력에도 불구하고 모든 데이터를 거의 동일한 embedding으로 매핑하는 현상이다. (a) Joint-Embedding 구조에서는 이러한 문제가 쉽게 발생할 수 있다. 두 입력 x와 y의 임베딩을 sx, sy라고 할 때, 두 임베딩의 거리를 최소화하는 목적 함수 D(sx,sy)를 사용하면 모든 입력을 동일한 상수 벡터로 출력하는 경우에도 손실이 최소가 될 수 있다. 이런 경우 모델이 입력과 무관하게 같은 표현을 출력하는 trivial solution으로 수렴하는 representation collapse가 발생할 수 있다. (b) Generative 구조에서는 이와 다르게, 만약 x-encoder가 collapse되어 항상 같은 표현을 출력한다면 디코더는 동일한 입력으로부터 서로 다른 이미지를 복원할 수 없기 때문에 reconstruction loss를 최소화할 수 없다. 따라서 일반적으로 collapse가 자연스럽게 방지된다. 그러나 추가 입력인 z가 충분한 정보를 포함하고 있다면, 디코더가 z만을 이용해 y를 복원할 수 있게 된다. 이 경우 encoder의 출력이 상수여도 reconstruction이 가능해지므로 representation collapse가 다시 발생할 수 있다.

 

Method

 

핵심 목표는 이미지의 일부 영역(context)으로부터 다른 영역(target)의 표현을 예측하도록 학습하는 것이다. 이 구조는 context encoder, target encoder, predictor의 세 가지 모듈로 구성되며, 전체적으로 Vision Transformer(ViT)를 기반으로 구현된다. 

 

먼저 입력 이미지는 여러 개의 패치로 분할되고, target encoder를 통해 각 패치의 representation이 계산된다. 이 representation 중에서 무작위로 여러 개의 target block(빨강, 파랑, 노랑 박스)을 선택하여 학습 목표로 사용한다. 중요한 점은 target이 픽셀 값이 아니라 encoder가 생성한 representation이라는 것이다. 계속 언급했듯이 이를 통해 모델이 불필요한 저수준 픽셀 정보를 예측하는 대신 더 의미적인 특징을 학습하도록 유도한다.

 

다음으로 context를 구성한다. 이미지에서 비교적 큰 영역을 context block으로 선택하는데, 만약 context에 target 영역이 그대로 포함되면 모델이 이미 본 정보를 그대로 예측하는 trivial한 문제가 된다. 이를 방지하기 위해 context와 target block이 겹치는 영역은 제거한다. 이후 남은 context 패치들만 context encoder에 입력되어 context representation이 생성된다.

 

 

이후 predictor는 context encoder의 출력 sx를 입력으로 받고, 추가적으로 target의 위치 정보를 나타내는 mask token을 조건 변수 z로 함께 입력받는다. predictor는 이 정보들을 이용해 해당 위치의 target representation을 예측한다. 모델은 여러 개의 target block에 대해 서로 다른 위치 정보를 가진 mask token(z)을 조건으로 하여 동일한 예측 과정을 반복적으로 수행한다.

 

학습은 예측된 representation과 실제 target representation 사이의 L2 거리 손실을 최소화하도록 진행된다. 이때 context encoder와 predictor는 gradient 기반으로 학습되고, target encoder의 파라미터는 context encoder 파라미터의 exponential moving average(EMA)로 업데이트된다. 이러한 teacher–student 구조는 학습 안정성과 representation collapse 방지에 중요한 역할을 한다. 

 

Related Work

먼저 초기 연구에서는 이미지의 일부를 손상시키거나 제거한 뒤 이를 복원하도록 학습하는 방식이 사용되었다. 예를 들어 denoising autoencoder나 context encoder는 입력 이미지의 일부 정보를 제거한 후 주변 정보를 이용해 이를 복원하도록 학습하면서 표현을 학습한다. 또한 이미지 색상 복원과 같은 과제도 self-supervised 표현 학습의 한 방법으로 활용되었다. 

 

최근에는 Vision Transformer 기반의 masked image modeling 접근이 등장하였다. MAE는 입력 이미지에서 일부 패치를 가린 뒤 이를 픽셀 공간에서 복원하는 방식으로 효율적인 학습 구조를 제안했으며, BEiT는 픽셀 대신 토큰화된 표현을 예측하는 방식을 사용한다. SimMIM은 픽셀 대신 특징 공간(HOG 기반)을 예측하는 방식도 탐색하였다. 이러한 방법들은 주로 입력 공간에서 reconstruction을 수행한다는 공통점이 있다.

 

I-JEPA와 가장 유사한 연구로는 data2vec와 Context Autoencoder가 있다. data2vec은 target encoder가 생성한 representation을 예측하도록 학습하며, 다양한 모달리티에 적용 가능하다. Context Autoencoder는 reconstruction loss와 representation alignment를 함께 사용해 가려진 패치의 표현을 예측하도록 한다. I-JEPA는 이러한 접근과 비교했을 때 계산 효율이 높고 더 의미적인 표현을 학습할 수 있다는 점을 강조한다. 

 

또한 joint-embedding 기반 self-supervised 방법으로는 DINO, MSN, iBOT 등이 있다. 이 방법들은 여러 개의 데이터 증강 view를 사용하여 동일 이미지의 표현이 유사하도록 학습하지만, 여러 view를 처리해야 하기 때문에 계산 비용이 증가하고 확장성이 제한될 수 있다. 반면 I-JEPA는 하나의 이미지 view만 사용하면서 representation을 예측하는 방식을 통해 이러한 한계를 완화한다. 

 

Image Classification

ImageNet-1K 데이터셋을 사용하여 linear probing(classifier만 학습)과 일부 fine-tuning(encoder + classifier 모두 학습) 설정에서 성능을 평가한다. 

 

ImageNet-1K linear evaluation 결과에서 I-JEPA는 MAE, CAE, data2vec과 같은 데이터 증강을 사용하지 않는 self-supervised 방법들보다 적은 계산으로 더 높은 성능을 보였다. 특히 모델 크기를 확장할수록 성능이 향상되며, 큰 모델에서는 데이터 증강을 사용하는 방법들과도 경쟁력 있는 결과를 보인다.

 

또한 ImageNet의 1% 라벨만 사용(이미지는 전부 사용, 라벨만 1%만 사용)하는 low-shot 학습 설정에서도 I-JEPA는 MAE보다 높은 성능을 보였으며, 더 적은 학습 epoch로도 data2vec과 유사한 수준의 성능을 달성한다. 해상도를 높여 학습한 대형 모델의 경우 기존의 DINO, iBOT 같은 view-invariance 기반 방법들과도 경쟁력 있는 성능을 보인다. 

 

마지막으로 CIFAR100, Places205, iNaturalist2018과 같은 다양한 downstream classification task에서 linear probing을 수행한 결과, I-JEPA는 MAE와 data2vec보다 높은 성능을 보이며 데이터 증강을 사용하는 방법들과의 성능 격차도 줄이는 것으로 나타났다.

 

Local Prediction Tasks

 

해당 섹션에서는 I-JEPA가 semantic classification뿐 아니라 local visual features도 잘 학습하는지를 평가한다. 이를 위해 객체 개수 예측(object counting)과 깊이 추정(depth prediction)과 같은 task에서 linear probing 방식으로 성능을 측정하였다. 실험 결과 I-JEPA는 view-invariance 기반 방법인 DINO나 iBOT보다 더 높은 성능을 보였으며, 이는 사전학습 과정에서 지역적 시각 특징까지 효과적으로 학습함을 보여준다.

 

Scalability

 

Scalability에서는 I-JEPA의 학습 효율성과 확장성을 분석한다. 실험 결과 I-JEPA는 기존 방법들보다 적은 계산량으로 높은 성능을 달성하며 MAE와 같은 reconstruction 기반 방법보다 더 적은 학습 epoch로 수렴한다. 또한 view-invariance 기반 방법들은 여러 개의 데이터 증강 view를 처리해야 하지만, I-JEPA는 하나의 이미지 view만 사용하기 때문에 학습 속도가 더 빠르다. 이러한 특성 덕분에 대형 모델에서도 계산 비용이 낮고 효율적으로 확장될 수 있다. 추가적으로 모델 크기를 확장하는 실험에서 ViT-H/14보다 더 큰 ViT-G/16 모델을 사용하면 이미지 분류와 같은 semantic task에서는 성능이 향상되지만, 입력 패치 크기가 커지기 때문에 객체 개수 예측이나 깊이 추정과 같은 local prediction task에서는 성능 향상이 제한적이거나 오히려 감소하는 경향이 나타났다.

 

Predictor Visualizations

 

Predictor Visualizations에서는 predictor가 실제로 어떤 정보를 학습했는지 시각적으로 분석한다. 이를 위해 사전학습된 context encoder와 predictor를 고정한 뒤, predictor의 출력 representation을 픽셀 공간으로 변환하는 디코더를 추가로 학습하여 예측 결과를 시각화한다. 동일한 predictor representation으로부터 여러 이미지를 생성해 비교한 결과 여러 샘플에서 공통적으로 나타나는 객체의 구조나 자세와 같은 정보는 representation에 포함되어 있고, 세부 질감이나 배경과 같은 저수준 정보는 다양하게 변하는 것을 확인할 수 있다. 이는 I-JEPA가 픽셀 수준의 세부 정보를 복원하기보다는 객체의 형태나 위치와 같은 고수준 semantic 정보를 중심으로 표현을 학습한다는 것을 보여준다.

 

Ablations

해당 섹션에서는 I-JEPA의 주요 설계 요소가 성능에 미치는 영향을 분석하기 위해 여러 ablation 실험을 수행한다.

 

먼저 예측을 픽셀 공간에서 수행하는 경우와 representation 공간에서 수행하는 경우를 비교한 결과, representation 공간에서 예측할 때 훨씬 더 높은 성능을 보였다. 이는 target encoder가 생성한 추상적인 표현을 예측하는 것이 의미적 표현 학습에 중요함을 보여준다.

 

또한 다양한 masking 전략을 비교한 결과, 여러 개의 target block을 예측하는 multi-block masking 전략이 단일 블록이나 랜덤 마스킹보다 더 좋은 성능을 보였으며, 이는 모델이 이미지의 의미적 구조를 학습하는 데 도움이 된다는 것을 확인하였다.

 

Conclusion

I-JEPA는 수작업으로 설계된 데이터 증강에 의존하지 않고도 의미적인 이미지 표현을 학습할 수 있는 self-supervised 방법을 제안한다. 이 방법은 이미지의 일부 영역으로부터 다른 영역의 정보를 픽셀 공간이 아닌 representation 공간에서 예측하도록 학습함으로써, 픽셀 수준의 세부 정보보다 의미적 구조에 집중한 표현을 학습하도록 한다. 

 

실험 결과 I-JEPA는 기존 self-supervised 방법들과 비교해 높은 성능을 보이며 representation 공간에서의 예측 방식 덕분에 학습 속도와 계산 효율 측면에서도 장점을 가진다. 또한 데이터 증강 기반 방법과 달리 보다 단순한 구조로 다양한 태스크에 적용 가능한 일반적인 표현 학습 방법이 될 가능성을 제시한다.