시각 표현 학습에서 세계 모델의 학습과 활용

Digest

기존 자기지도 시각 표현 학습은 불변(invariant) 표현과 등변(equivariant) 표현 사이에서 양자택일을 강요받았다. 대조 학습(contrastive learning)은 augmentation에 불변한 특징을 학습하여 선형 평가에서 강점을 보이지만(MoCoV3: 76.3%, Table 2), 변환에 대한 정보를 버리기 때문에 세밀한 적응에 한계가 있다. 반면 마스크 이미지 모델링(MAE 등)은 풍부한 정보를 보존하지만 선형 분리 가능성이 떨어진다(MAE: 60.2%, Table 2). 이 논문은 **JEPA(Joint-Embedding Predictive Architecture)**를 확장하여 **Image World Models(IWM)**을 제안하며, 잠재 공간에서 전역적 광도 변환(global photometric transformation)의 효과를 예측하는 세계 모델을 학습한다. 핵심 통찰은 predictor의 조건화(conditioning), 예측 난이도(prediction difficulty), 용량(capacity) 세 가지 축이 세계 모델 품질을 결정한다는 것이다. 이를 통해 불변과 등변 표현 사이를 자유롭게 보간(interpolate)할 수 있으며, 등변 IWM(IWM₁₈,₃₈₄ᴱᑫᵘⁱ)은 predictor fine-tuning만으로 ViT-B/16에서 84.4% top-1 정확도(Table 4)를 달성하여, 5배 이상 긴 학습(1600 epoch)을 수행한 MAE(83.6%)와 DINO(82.8%)를 300 epoch만에 능가한다.

메타데이터

항목	내용
제목	Learning and Leveraging World Models in Visual Representation Learning
저자	Quentin Garrido, Mahmoud Assran, Nicolas Ballas, Adrien Bardes, Laurent Najman, Yann LeCun
소속	FAIR at Meta, Université Gustave Eiffel (CNRS, LIGM), INRIA, NYU
연도	2024
발표	arXiv:2403.00504
링크	arXiv
키워드	JEPA, World Model, Self-Supervised Learning, Equivariant Representation, Vision Transformer

왜 이 연구를 하는가?

핵심 질문

자기지도 시각 표현 학습에서, 불변성(invariance)과 등변성(equivariance)을 동시에 제어할 수 있는 통합 프레임워크를 만들 수 있는가?

기존 접근법의 한계

한계	설명
대조 학습의 정보 손실	DINO, MoCoV3 등은 augmentation에 불변한 표현을 학습하여 변환 정보를 버린다. 선형 평가는 강하지만(76.3%, Table 2), fine-tuning 시 개선 폭이 제한적이다(+0.8, Table 3).
마스크 모델링의 낮은 선형 분리도	MAE는 풍부한 정보를 보존하지만 선형 평가 성능이 극히 낮다(60.2%, Table 2). 유용한 표현을 얻으려면 end-to-end fine-tuning이 필수적이다.
I-JEPA의 제한된 예측 범위	I-JEPA는 마스크된 패치만 예측하여 공간적 등변성에 국한되며, 색상·밝기 등 광도 변환에 대한 세계 모델을 학습하지 못한다.
표현 추상화 수준 고정	기존 방법들은 학습 후 표현의 추상화 수준이 고정되어, 하류 과제에 따라 유연하게 조절할 수 없다.

핵심 통찰

세계 모델로서의 predictor: JEPA의 predictor를 단순 마스크 복원기가 아닌, 잠재 공간에서 변환의 효과를 예측하는 세계 모델로 재해석할 수 있다.
세 가지 축의 발견: 조건화(conditioning) 방식, 변환의 난이도(difficulty), predictor 용량(capacity)이 세계 모델 품질을 결정하며, 이 세 축을 조절하면 불변-등변 스펙트럼을 자유롭게 탐색할 수 있다.
Predictor fine-tuning의 효율성: 학습된 세계 모델(predictor)을 하류 과제에 fine-tuning하면, encoder의 1/3 파라미터만으로도 end-to-end fine-tuning에 필적하는 성능을 달성한다.

방법 (Method)

프레임워크 개요

graph TB
    subgraph Training["IWM 학습 과정"]
        direction TB
        X["원본 이미지 x"] --> AugS["Source Augmentation<br/>(crop + color jitter + grayscale<br/>+ blur + solarization + mask)"]
        X --> AugT["Target Augmentation<br/>(crop + color jitter)"]

        AugS --> Enc["Encoder f_θ<br/>(ViT-B/16)"]
        AugT --> EMA["EMA Encoder f_θ̄<br/>(stop-gradient)"]

        Enc --> SrcRepr["Source 표현 s_x"]
        EMA --> TgtRepr["Target 표현 s_y"]

        AugS --> TransInfo["변환 정보 추출<br/>(Δcolor, Δbrightness, ...)"]
        TransInfo --> Cond["Conditioning<br/>(feature 또는 sequence)"]

        SrcRepr --> Pred["Predictor g_φ<br/>(Transformer, depth=12~18)"]
        Cond --> Pred

        Pred --> PredRepr["예측 표현 ŝ_y"]

        PredRepr --> Loss["L2 Loss<br/>||ŝ_y - s_y||²"]
        TgtRepr --> Loss
    end

    subgraph Inference["하류 과제 적응"]
        direction TB
        EncFrozen["Frozen Encoder"] --> PredFT["Predictor Fine-tuning<br/>(1/3 파라미터)"]
        PredFT --> Task["분류 / 세분화"]
    end

    Training --> Inference

    style Training fill:#f0f8ff,stroke:#4682b4
    style Inference fill:#f0fff0,stroke:#228b22

핵심 구성요소

1. JEPA에서 IWM으로의 확장

JEPA(Joint-Embedding Predictive Architecture)는 원래 마스크된 이미지 패치를 잠재 공간에서 예측하는 구조다. IWM은 이를 확장하여 마스킹뿐 아니라 전역적 광도 변환(color jitter, grayscale, blur, solarization)의 효과를 잠재 공간에서 예측한다. 손실 함수는 예측된 표현과 EMA 타겟 표현 사이의 제곱 L2 거리다.

2. 조건화(Conditioning)

Predictor가 어떤 변환이 적용되었는지 아는 것이 핵심이다. 두 가지 방식이 있다:

Feature conditioning: 변환 파라미터를 벡터로 인코딩하여 predictor에 추가 입력으로 제공
Sequence conditioning: 변환 정보를 토큰 시퀀스로 변환하여 입력 시퀀스에 연결

조건화 없이는 predictor가 변환을 무시하고 불변 표현을 학습한다(MRR 0.00, Table 1). 조건화를 추가하면 등변 표현이 학습된다(MRR 0.82~0.85, Table 1).

3. 예측 난이도(Prediction Difficulty)

약한 augmentation(color jitter만)으로는 세계 모델이 충분히 학습되지 않는다. 파괴적 변환(grayscale, blur, solarization)을 추가하고 color jitter 강도를 높여야 predictor가 복잡한 변환 관계를 학습한다. 강한 변환 + 조건화 시 MRR이 0.11에서 0.85로 극적으로 상승한다(Table 1).

4. Predictor 용량(Capacity)

12층 predictor는 5회 중 1회만 색상 등변성을 달성하는 반면, 18층 predictor는 5회 중 4회 성공한다. 복잡한 변환을 모델링하려면 충분한 용량의 predictor가 필수적이다.

5. 표기법: IWMˣ·ʸ_Z로 표기하며, X=predictor 깊이, Y=임베딩 차원, Z=Inv(불변) 또는 Equi(등변)이다.

발견 (Findings)

주요 결과 — 선형 평가 및 Fine-tuning (ImageNet-1k, ViT-B/16)

방법	Epoch	선형 평가	Attentive Probe	Encoder FT	Predictor FT	End-to-End
MoCoV3	300	76.3%	76.4%	—	—	—
MAE	300	60.2%	73.5%	—	—	—
MAE	1600	—	—	83.6%	83.0%	83.1%
DINO	1600	—	—	82.0%	—	82.8%
I-JEPA	300	70.0%	75.0%	83.0%	79.1%	80.0%
IWM₁₂,₃₈₄ᴵⁿᵛ	300	74.5%	77.0%	83.3%	80.5%	81.3%
IWM₁₈,₃₈₄ᴱᑫᵘⁱ	300	67.5%	75.1%	82.9%	81.5%	83.3%

출처: Table 2, 3, 4

주요 결과 — ViT-L/16 스케일링

방법	Encoder FT	Predictor FT	End-to-End
I-JEPA	84.1%	79.9%	—
IWM₁₈,₃₈₄ᴵⁿᵛ	84.3%	81.5%	—
IWM₃₆,₅₁₂ᴱᑫᵘⁱ	83.7%	85.0%	85.4%

출처: Table 6

세계 모델 품질 (Mean Reciprocal Rank)

Predictor 구성	Jitter Only	+ Destructive	+ Strong Jitter
I-JEPA (12,384)	0.00	0.00	0.00
IWM (12,384)	0.11	0.09	0.25
IWM (18,384)	0.25	0.79	0.85

출처: Table 1

의미론적 분할 (ADE20k, mIoU)

방법	Encoder	Predictor	End-to-End
I-JEPA	44.2	45.4	45.1
IWM₁₂,₃₈₄ᴵⁿᵛ	45.6	45.7	46.5
IWM₁₈,₃₈₄ᴱᑫᵘⁱ	44.2	46.8	47.0

출처: Table 5

핵심 발견

불변-등변 스펙트럼의 존재: IWMᴵⁿᵛ는 대조 학습과 유사하게 선형 평가에서 강하고(74.5%), IWMᴱᑫᵘⁱ는 마스크 모델링과 유사하게 fine-tuning에서 강하다(84.4%, Table 4). 동일 프레임워크에서 조건화와 용량 조절만으로 이 스펙트럼을 탐색할 수 있다.

Predictor fine-tuning의 효율성: 등변 IWM은 predictor만 fine-tuning해도 encoder의 1/3 파라미터로 end-to-end에 필적하는 성능을 달성한다. ViT-L에서 predictor FT(85.0%)가 encoder FT(83.7%)를 1.3%p 능가한다(Table 6).

다중 과제 일반화: 단일 predictor를 4개 데이터셋(ImageNet, iNat18, SUN397, Places205)에 동시 fine-tuning해도 평균 성능이 유지된다(73.5% vs 73.4%, Table 7). 이는 세계 모델이 범용적 변환 지식을 습득했음을 시사한다.

이론적 의의

세계 모델 관점의 표현 학습 통합

이 연구는 대조 학습과 마스크 모델링이라는 자기지도 학습의 양대 패러다임을 세계 모델이라는 상위 프레임워크로 통합한다. 불변 표현은 “변환 정보를 버리는 세계 모델”이고, 등변 표현은 “변환 정보를 보존하는 세계 모델”이라는 관점에서 두 접근법이 동일 스펙트럼의 양 극단임을 보여준다. 이는 LeCun의 JEPA 비전을 시각 도메인에서 실증적으로 검증한 것이다.

Predictor as World Model의 새로운 패러다임

기존에 predictor는 학습 보조 도구로 여겨져 추론 시 버려졌다. 이 연구는 predictor 자체가 잠재 공간의 물리 법칙을 학습한 세계 모델이며, 하류 과제에 직접 활용 가능함을 보여준다. 이는 파라미터 효율적 적응(parameter-efficient adaptation)의 새로운 경로를 열어준다.

표현의 추상화 수준 제어 가능성

조건화, 변환 난이도, predictor 용량이라는 세 축을 통해 표현의 추상화 수준을 연속적으로 제어할 수 있다. 이는 “하나의 사전학습, 다양한 하류 과제”라는 foundation model의 이상에 더 가까이 다가가는 메커니즘을 제시한다.

핵심 용어 정리

용어	정의
JEPA (Joint-Embedding Predictive Architecture)	두 뷰의 임베딩을 공유 잠재 공간에 매핑하고, 한쪽에서 다른 쪽을 예측하는 자기지도 학습 아키텍처. LeCun이 제안한 비전의 핵심 구조.
IWM (Image World Models)	JEPA를 확장하여 마스킹뿐 아니라 전역적 광도 변환의 효과를 잠재 공간에서 예측하는 세계 모델. 이 논문의 핵심 제안.
불변 표현 (Invariant Representation)	데이터 augmentation(색상 변환, 크롭 등)에 대해 변하지 않는 표현. 변환 정보를 버려 선형 분류에 유리하다.
등변 표현 (Equivariant Representation)	입력에 변환을 적용하면 표현도 예측 가능한 방식으로 변하는 표현. 변환 정보를 보존하여 fine-tuning에 유리하다.
EMA (Exponential Moving Average)	학습 중 encoder 가중치의 지수 이동 평균을 타겟 네트워크로 사용하는 기법. 학습 안정성을 제공한다.
Predictor	Source 표현과 변환 조건으로부터 target 표현을 예측하는 Transformer 네트워크. IWM에서는 세계 모델 역할을 한다.
Conditioning	Predictor에게 어떤 변환이 적용되었는지 알려주는 방법. Feature conditioning(벡터 추가)과 sequence conditioning(토큰 시퀀스 연결) 두 방식이 있다.
MRR (Mean Reciprocal Rank)	세계 모델의 품질을 측정하는 지표. 예측된 표현이 실제 타겟 표현을 얼마나 잘 식별하는지 평가한다. 1.0이 완벽한 예측.
Attentive Probing	Frozen encoder 위에 cross-attention 기반 프로브를 학습하여 표현 품질을 평가하는 방법. 선형 평가보다 표현의 풍부함을 더 잘 포착한다.
Photometric Transformation	이미지의 색상, 밝기, 대비 등을 변경하는 변환. IWM이 예측하는 대상이다.

Juhyeon's Blog

탐색기

Learning and Leveraging World Models in Visual Representation Learning