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

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

저자: Mahmoud Assran, Quentin Duval, Ishan Misra
발행년도: 2023년
인용수: 450회
논문 링크: https://arxiv.org/abs/2301.08243
arXiv ID: 2301.08243

---

이미지 자기지도학습의 새로운 방향: I-JEPA가 제시하는 표현 예측의 가능성

문제 정의 (Problem Definition)

컴퓨터 비전 분야에서 자기지도학습(self-supervised learning)은 라벨이 없는 대량의 이미지 데이터로부터 유의미한 표현을 학습하는 핵심 기술이다. 최근 몇 년간 대조학습(contrastive learning)이나 마스크된 이미지 모델링(masked image modeling) 같은 방법들이 큰 성공을 거두었지만, 여전히 근본적인 한계를 안고 있었다.

기존 방법들의 핵심 딜레마는 "무엇을 학습 목표로 삼을 것인가"였다. 대조학습은 같은 이미지의 서로 다른 증강(augmentation)끼리는 가깝게, 다른 이미지끼리는 멀게 만드는 방식으로 학습했다. 하지만 이는 증강 방법에 따라 성능이 크게 좌우되고, 본질적으로 이미지의 의미적 구조를 이해하기보다는 불변성(invariance)을 학습하는 데 치중했다.

반면 MAE(Masked Autoencoder) 같은 마스크 예측 방법은 가려진 패치의 픽셀값을 직접 복원하는 방식을 택했다. 이는 직관적이지만 픽셀 수준의 디테일에 과도하게 집중하게 만들어, 고수준의 의미적 특징을 학습하는 데는 비효율적이었다. 예를 들어, 나뭇잎의 정확한 텍스처를 복원하는 것보다 "이것이 나무의 일부"라는 개념을 학습하는 것이 더 중요할 수 있다.

실제 환경에서 이러한 한계는 더욱 명확해진다. 조명 변화, 시점 변화, 부분적 가림 등이 있는 실제 이미지에서는 픽셀 수준의 완벽한 복원이 불가능하고 무의미할 때가 많다. 우리에게 필요한 것은 이미지의 의미적 구조를 이해하고 예측할 수 있는 표현 학습 방법이었다.

기존 방법의 한계 (Motivation)

자기지도학습의 대표적인 접근법들을 살펴보면 각각의 구조적 한계가 명확히 드러난다.

대조학습 방법들(SimCLR, MoCo 등)은 negative sample에 대한 의존성이 컸다. 같은 이미지의 두 view를 가깝게 만들고 다른 이미지들과는 멀게 만드는 방식은 직관적이지만, 대량의 negative sample이 필요했고 배치 크기에 민감했다. 또한 어떤 증강을 사용할지에 따라 학습되는 특징이 크게 달라지는 문제가 있었다. 색상 변화에 강건한 특징을 학습하려면 color jittering을 써야 하고, 기하학적 변화에 강건하려면 rotation이나 flip을 써야 했다.

마스크 기반 방법들(MAE, BEiT 등)은 다른 극단을 택했다. 이미지의 일부를 가리고 복원하는 방식은 언어 모델의 성공에서 영감을 받았지만, 이미지와 텍스트의 본질적 차이를 간과했다. 텍스트의 마스크된 단어는 명확한 정답이 있지만, 이미지의 마스크된 패치는 무수히 많은 가능한 복원 결과를 가질 수 있다. 이로 인해 모델은 고수준 의미보다는 저수준 패턴 복원에 집중하게 되었다.

저자들이 지적한 기존 방법의 한계는 타당하지만, 각 방법의 장점도 간과하면 안 된다. 예를 들어 MAE는 비록 픽셀 복원에 집중하지만, 그 과정에서 학습된 표현이 다양한 downstream task에서 좋은 성능을 보였다. 문제는 "픽셀 복원 자체"가 아니라 "어떤 abstraction level에서 복원할 것인가"였을 것이다.

이러한 한계들은 결국 하나의 질문으로 귀결된다. "이미지의 일부로부터 다른 부분을 예측할 때, 무엇을 예측 대상으로 삼아야 하는가?" 픽셀값도, 클래스 레이블도 아닌 새로운 예측 대상이 필요했고, 이것이 I-JEPA가 태어난 배경이다.

제안 방법의 핵심 아이디어 (Key Idea)

I-JEPA의 핵심은 한 문장으로 요약할 수 있다: "픽셀이 아닌 추상적 표현(representation) 공간에서 이미지 패치 간의 관계를 예측하자."

이를 일상적인 비유로 설명하면 이렇다. 퍼즐을 맞출 때 우리는 빠진 조각의 정확한 픽셀을 상상하지 않는다. 대신 "여기는 하늘이 와야 하고, 구름의 일부가 연결될 것이다"라는 식으로 추상적인 개념 수준에서 예측한다. I-JEPA도 마찬가지로 마스크된 패치의 정확한 픽셀값이 아닌, 그 패치가 가져야 할 의미적 표현을 예측한다.

기존 방법들과의 차이를 명확히 하면:

• MAE: 마스크된 패치의 RGB 픽셀값을 직접 예측
• 대조학습: 같은 이미지의 다른 view들을 표현 공간에서 가깝게 배치
• I-JEPA: 보이는 패치들의 표현으로부터 마스크된 패치의 표현을 예측

이러한 접근의 장점은 명확하다. 픽셀 수준의 노이즈나 변화에 덜 민감하면서도, 이미지의 구조적 관계를 학습할 수 있다. 또한 예측 대상이 고정된 차원의 벡터이므로 학습이 안정적이다.

하지만 이 아이디어가 완전히 새로운 것인지는 의문이다. Context Encoder나 일부 generative model들도 latent space에서의 예측을 시도했었다. I-JEPA의 진짜 기여는 아이디어 자체보다는 이를 효과적으로 구현한 아키텍처와 학습 방법에 있을 것이다.

아키텍처 설명 (Architecture)

I-JEPA의 전체 구조는 세 개의 주요 컴포넌트로 구성된다.

1. Context Encoder는 마스크되지 않은 패치들을 입력받아 표현을 생성한다. 이는 일반적인 Vision Transformer 구조를 따르며, 보이는 패치들만을 처리하여 계산 효율성을 높였다.

2. Target Encoder는 전체 이미지를 입력받아 모든 패치의 표현을 생성한다. 핵심은 이 인코더가 직접 학습되지 않고 Context Encoder의 exponential moving average(EMA)로 업데이트된다는 점이다. 이는 학습 안정성을 크게 향상시켰다.

3. Predictor는 Context Encoder의 출력과 마스크된 패치의 위치 정보를 받아, Target Encoder가 생성한 해당 패치의 표현을 예측한다. 가벼운 transformer 구조로 되어 있어 context 정보를 효과적으로 활용할 수 있다.

# Simplified I-JEPA forward pass
def forward(image):
    # 1. 마스킹 전략 적용
    visible_patches, masked_positions = apply_masking(image)

    # 2. Context encoding
    context_repr = context_encoder(visible_patches)

    # 3. Target encoding (no gradient)
    with torch.no_grad():
        all_patch_repr = target_encoder(image)
        target_repr = all_patch_repr[masked_positions]

    # 4. Prediction
    predicted_repr = predictor(context_repr, masked_positions)

    # 5. Loss computation
    loss = smooth_l1_loss(predicted_repr, target_repr)

    return loss

이 구조에서 특히 주목할 점은 비대칭성이다. Context Encoder는 일부 패치만 보고 예측해야 하므로 어려운 과제를 수행하는 반면, Target Encoder는 전체 이미지를 보고 표현을 생성한다. 이러한 비대칭성이 의미 있는 표현 학습을 가능하게 한다.

그런데 Target Encoder가 전체 이미지를 본다는 것이 일종의 "정보 누출"은 아닐까? 마스크된 영역의 정보가 다른 패치들의 표현에 영향을 줄 수 있고, 이것이 예측을 쉽게 만들 가능성도 있다. 저자들이 이 부분에 대한 ablation을 충분히 다루지 않은 것은 아쉽다.

접근 방법의 특징 및 설계 의도 (Design Choices)

I-JEPA의 설계에는 몇 가지 중요한 선택들이 있었다.

첫째, 마스킹 전략이다. 단순히 랜덤하게 패치를 마스크하는 대신, 공간적으로 인접한 패치들을 블록 단위로 마스크했다. 이는 예측 과제를 더 어렵게 만들어 모델이 진정한 의미적 이해를 하도록 유도했다. 개별 패치를 예측하는 것은 주변 텍스처를 단순 복사하는 것으로도 가능하지만, 큰 블록을 예측하려면 더 높은 수준의 이해가 필요하다.

둘째, EMA 업데이트 방식이다. Target Encoder를 별도로 학습하는 대신 Context Encoder의 가중치를 지수이동평균으로 업데이트했다. 이는 BYOL이나 MoCo에서 영감을 받은 것으로 보이며, 학습 안정성과 표현의 일관성을 보장했다.

셋째, 손실 함수의 선택이다. L2 loss 대신 smooth L1 loss를 사용한 것은 outlier에 더 강건한 학습을 위한 것으로 보인다. 표현 공간에서의 예측은 픽셀 예측보다 불확실성이 크므로 이러한 선택이 합리적이다.

predictor의 경량화도 흥미로운 선택이다. 복잡한 predictor는 overfitting의 위험이 있고, 너무 단순하면 context 정보를 제대로 활용할 수 없다. 저자들은 작은 transformer 구조로 균형을 맞췄다.

하지만 이러한 설계 선택들이 모두 충분히 검증되었는지는 의문이다. 특히 마스킹 비율(75%)이나 블록 크기 같은 하이퍼파라미터들이 어떻게 결정되었는지, 다른 값들에 대한 충분한 실험이 있었는지 불분명하다. 또한 왜 하필 transformer predictor인지, CNN이나 MLP와의 비교는 없다는 점도 아쉽다.

실험 결과 분석

실험 결과는 I-JEPA의 강점과 한계를 동시에 보여준다.

ImageNet-1K linear evaluation에서 ViT-H/16 모델이 87.0%의 정확도를 달성했다. 이는 MAE의 86.9%와 비슷한 수준이지만, 학습 시간은 크게 단축되었다. 2400 epoch 대신 600 epoch만으로도 경쟁력 있는 성능을 달성한 것은 주목할 만하다.

더 인상적인 것은 semantic segmentation과 depth estimation 같은 dense prediction task에서의 성능이다. ADE20K 데이터셋에서 MAE보다 높은 mIoU를 기록했고, 특히 low-shot 시나리오에서 강점을 보였다. 이는 I-JEPA가 픽셀 수준의 디테일보다는 의미적 구조를 더 잘 학습했음을 시사한다.

그러나 object detection에서는 MAE에 비해 약간 뒤처지는 결과를 보였다. 이는 bounding box regression 같은 작업에서는 여전히 세밀한 위치 정보가 중요하며, 표현 수준의 예측만으로는 부족할 수 있음을 보여준다.

실험 설계에서 아쉬운 점은 다른 joint embedding 방법들과의 직접 비교가 부족하다는 것이다. BYOL이나 SwAV 같은 방법들도 픽셀 복원 없이 표현을 학습하는데, 이들과의 공정한 비교가 있었다면 I-JEPA의 기여를 더 명확히 알 수 있었을 것이다. 또한 다양한 도메인(의료 영상, 위성 영상 등)에서의 실험이 없어 일반화 가능성을 판단하기 어렵다.

총평: 개인적 해석과 후속 연구 방향

I-JEPA는 자기지도학습의 "예측 대상"에 대한 중요한 질문을 던졌다. 픽셀도 레이블도 아닌, 학습된 표현 자체를 예측 대상으로 삼는다는 아이디어는 단순하지만 강력하다. 이는 인간이 시각 정보를 처리하는 방식과도 유사하며, 향후 멀티모달 학습으로의 확장 가능성도 열어둔다.

실무 적용 관점에서 I-JEPA의 장점은 명확하다. 학습 효율성이 높고, 특히 semantic understanding이 중요한 태스크에서 강점을 보인다. 하지만 세밀한 위치 정보가 중요한 태스크에서는 여전히 픽셀 수준 방법들이 유리할 수 있다.

이 연구의 진정한 한계는 "왜 이 방법이 잘 작동하는가"에 대한 이론적 이해가 부족하다는 점이다. Target Encoder의 표현이 정말 의미적으로 의미 있는지, 아니면 단순히 Context Encoder가 예측하기 쉬운 무언가를 학습하는지 명확하지 않다.

만약 내가 이 연구를 이어간다면 다음과 같은 방향을 탐구하고 싶다:

1. Cross-modal predictive learning: 이미지의 일부로 텍스트 설명을 예측하거나, 반대로 텍스트로 이미지 표현을 예측하는 방식으로 확장

2. Hierarchical prediction: 단일 스케일이 아닌 다양한 추상화 수준에서 예측을 수행하여 더 풍부한 표현 학습

3. Theoretical analysis: Information theory 관점에서 왜 representation prediction이 pixel prediction보다 효과적인지 분석

4. Dynamic masking: 학습 진행에 따라 마스킹 전략을 적응적으로 변경하여 curriculum learning 효과 달성

I-JEPA는 완벽한 해답이라기보다는 올바른 방향으로의 한 걸음이다. 자기지도학습이 단순히 pretext task를 잘 수행하는 것이 아니라, 진정한 의미의 이해로 나아가기 위해서는 이러한 시도들이 계속되어야 할 것이다.