시각 표현 학습에서 세계 모델을 학습하고 활용하기

Digest

이 논문은 JEPA의 predictor를 새롭게 해석한다. 기존에는 predictor가 사전학습 때만 쓰이고 버려지는 보조 모듈에 가까웠다. 저자들은 이 predictor가 사실상 잠재 공간의 world model이 될 수 있다고 본다. 즉, 이미지에 색상 변화, 밝기 변화, 블러 같은 변환을 주었을 때 표현이 어떻게 달라질지를 예측하는 모델로 본 것이다. 이 관점은 LeCun의 AMI(Autonomous Machine Intelligence) 구상에서 나온 “표현 공간에서 예측하는 world model” 아이디어를 시각 SSL 실험으로 구체화한 흐름에 가깝다.

배경은 자기지도 시각 표현 학습의 오래된 긴장 관계다. 대조 학습은 augmentation에 흔들리지 않는 불변 표현을 잘 만든다. 그래서 선형 평가에는 강하지만, 변환 정보를 많이 버린다. 반대로 MAE 같은 마스크 이미지 모델링은 정보를 더 많이 보존한다. 하지만 선형 분리도는 낮고, 좋은 성능을 얻으려면 큰 fine-tuning이 필요하다.

저자들은 이 문제를 Image World Models, IWM으로 푼다. IWM은 JEPA 구조를 유지하되, predictor에게 어떤 변환이 일어났는지 알려주고, target 표현이 어떻게 바뀌는지를 맞히게 한다. 이때 성능을 좌우하는 요인은 세 가지다. 첫째, predictor가 변환 정보를 받는지. 둘째, 예측할 변환이 충분히 어렵고 의미 있는지. 셋째, predictor의 용량이 충분한지.

결과적으로 등변 IWM은 300 epoch만 학습하고도 ViT-B/16에서 predictor fine-tuning으로 84.4% top-1을 달성한다. 이는 1600 epoch를 학습한 MAE의 83.6%, DINO의 **82.8%**보다 높다. ViT-L에서는 predictor만 fine-tuning한 결과가 encoder fine-tuning보다도 높았다. 이때 predictor는 encoder보다 훨씬 작은 모듈이므로, 새 모듈을 붙이는 PEFT가 아니라 사전학습 중 이미 사용한 predictor를 재활용하는 적응 경로로 볼 수 있다. 다만 이 논문은 주로 전역 광도 변환만 다루므로, 기하학적 변환이나 비디오 동역학으로 일반화되는지는 아직 열려 있다.


먼저 잡을 직관

이 논문을 읽을 때 가장 먼저 잡아야 할 직관은 다음이다.

JEPA의 predictor는 단순한 예측 헤드가 아니라, 변환을 이해하는 작은 시뮬레이터처럼 볼 수 있다.

용어 정리: predictor, 잠재 공간, world model

  • Predictor: 이미지를 직접 맞히는 모듈이 아니라, encoder가 만든 표현을 보고 다른 표현이 어떻게 생길지 예측하는 모듈이다.
  • 잠재 공간(latent space): 픽셀처럼 사람이 바로 보는 공간이 아니라, 모델이 이미지를 숫자 벡터로 압축해 놓은 내부 표현 공간이다.
  • World model: 어떤 행동이나 변환이 일어나면 다음 상태가 어떻게 바뀔지 예측하는 모델이다. 이 논문에서는 실제 세계가 아니라 이미지 변환 후 표현 변화를 예측한다.

예를 들어 원본 이미지와 색이 바뀐 이미지가 있다고 하자. 불변 표현 학습은 두 이미지를 거의 같은 표현으로 만들려고 한다. 색이 바뀌었다는 사실을 표현에서 지우는 쪽에 가깝다. 반면 등변 표현 학습은 색이 바뀌었다는 사실을 보존한다. 입력이 바뀌면 표현도 예측 가능한 방식으로 바뀌어야 한다.

등변 표현 학습 vs 불변 표현 학습

  • 불변 표현(invariant representation): 입력에 변화가 있어도 표현은 거의 그대로 유지되는 경우다. 예를 들어 사진 밝기가 바뀌어도 “고양이”라는 의미가 같게 남는 것이 목표다.
  • 등변 표현(equivariant representation): 입력이 바뀌면 표현도 그 변화에 맞춰 함께 바뀌는 경우다. 예를 들어 이미지가 이동하거나 색이 변하면, 표현도 그 변화를 추적할 수 있어야 한다.
  • 분류 과제는 불변성이 유리한 경우가 많고, 분할·검출처럼 위치나 구조가 중요한 과제는 등변성이 유리할 수 있다.

IWM(Image World Model)은 이 둘을 대립 관계로만 보지 않는다. predictor에게 변환 정보를 주지 않으면 모델은 자연스럽게 변환을 무시하는 불변 표현으로 간다. 반대로 predictor에게 변환 정보를 주고, 충분히 어려운 변환을 맞히게 하면 모델은 변환 정보를 보존하는 등변 표현으로 간다.

따라서 이 논문의 핵심은 “불변이냐 등변이냐”가 아니라, predictor 설계로 표현의 추상화 수준을 조절할 수 있다는 주장이다.


메타데이터

항목내용
제목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, PDF
키워드JEPA, World Model, Self-Supervised Learning, Equivariant Representation, Vision Transformer, Predictor Fine-tuning

논문이 던지는 질문

자기지도 시각 표현 학습에서는 보통 두 방향이 있었다.

용어 정리: 자기지도 학습과 augmentation

  • 자기지도 학습(self-supervised learning, SSL): 사람이 붙인 정답 라벨 없이, 데이터 자체에서 학습 문제를 만들어 표현을 배우는 방식이다.
  • 시각 표현 학습(visual representation learning): 이미지를 분류하기 전에, 이미지의 의미와 구조를 잘 담는 벡터 표현을 배우는 단계다.
  • Augmentation: 원본 이미지를 자르거나 색을 바꾸거나 흐리게 만드는 데이터 변환이다. 모델이 사소한 변화에 과하게 흔들리지 않도록 만드는 데 자주 쓴다.

하나는 불변 표현을 만드는 방향이다. DINO, MoCoV3, SimCLR 같은 대조 학습 계열이 여기에 가깝다. 같은 이미지에 서로 다른 augmentation을 적용해도 표현은 비슷해야 한다고 학습한다. 이 방식은 분류처럼 “무엇이 있는가”를 묻는 과제에 강하다.

다른 하나는 등변 표현을 만드는 방향이다. MAE나 I-JEPA처럼 입력의 일부를 가리거나 바꾼 뒤, 그 변화와 관련된 정보를 복원하거나 예측한다. 이 방식은 변환 정보를 더 많이 보존하므로 fine-tuning이나 dense prediction 과제에 유리할 수 있다.

용어 정리: 대조 학습, MAE, dense prediction

  • 대조 학습(contrastive learning): 같은 이미지에서 나온 두 뷰는 가깝게, 다른 이미지에서 나온 뷰는 멀게 만들며 표현을 학습하는 방법이다.
  • MAE(Masked Autoencoder): 이미지 일부를 가린 뒤, 가려진 부분을 복원하도록 학습하는 자기지도 학습 방법이다.
  • Dense prediction: 이미지 전체에 대해 하나의 답만 내는 것이 아니라, 픽셀이나 패치마다 답을 내는 과제다. Segmentation, depth estimation 등이 여기에 속한다.

문제는 두 방향이 서로 다른 장단점을 가진다는 점이다. 불변 표현은 선형 평가가 좋지만 세밀한 변환 정보가 사라진다. 등변 표현은 정보를 보존하지만 선형 분리도가 낮을 수 있다.

이 논문은 다음 질문을 던진다.

하나의 JEPA 기반 프레임워크 안에서 불변 표현과 등변 표현을 조절할 수 있을까? 그리고 학습된 predictor를 버리지 않고 하류 과제에 재사용할 수 있을까?


기존 접근의 한계

한계쉬운 설명
대조 학습의 정보 손실DINO나 MoCoV3는 augmentation에 흔들리지 않는 표현을 만든다. 선형 평가는 강하지만, 색상이나 밝기 변화 같은 변환 정보는 많이 버린다.
마스크 모델링의 낮은 선형 분리도MAE는 정보를 많이 보존하지만, frozen representation 위에서 선형 분류기를 학습하면 성능이 낮다. 논문 기준 MAE 300 epoch의 선형 평가는 60.2%다.
I-JEPA의 제한된 예측 범위I-JEPA는 주로 마스크된 패치를 예측한다. 공간적 등변성은 다루지만, 색상·밝기·블러 같은 광도 변환의 효과는 잘 모델링하지 못한다.
표현 추상화 수준의 고정한 번 사전학습한 표현이 특정 수준으로 고정된다. 하류 과제마다 필요한 불변성 정도가 달라도 쉽게 조절하기 어렵다.
Predictor의 폐기JEPA 계열에서 predictor는 학습 후 버려지는 경우가 많다. 하지만 이 predictor 안에는 변환을 예측하며 얻은 학습 신호가 들어 있을 수 있다.

용어 정리: frozen representation과 선형 분리도

  • Frozen representation: 사전학습된 encoder를 더 이상 학습하지 않고 고정한 상태에서 꺼낸 표현이다.
  • 선형 평가(linear evaluation): frozen representation 위에 선형 분류기 하나만 얹어 성능을 보는 평가 방식이다.
  • 선형 분리도(linear separability): 단순한 직선이나 초평면만으로 클래스가 잘 나뉘는 정도다. 값이 높으면 표현이 분류에 바로 쓰기 좋다는 뜻이다.

핵심 아이디어

1. Predictor를 world model로 본다

JEPA는 두 이미지 뷰를 같은 잠재 공간에 올려놓고, 한쪽 표현으로 다른 쪽 표현을 예측한다. 여기서 predictor는 source 표현을 보고 target 표현을 맞힌다.

용어 정리: JEPA와 view

  • JEPA(Joint-Embedding Predictive Architecture): 두 입력을 같은 임베딩 공간에 넣고, 한쪽 임베딩으로 다른 쪽 임베딩을 예측하는 구조다.
  • View: 같은 원본 이미지에서 augmentation을 거쳐 만들어진 서로 다른 입력 버전이다.
  • Source / target 표현: source는 예측의 출발점이 되는 표현, target은 predictor가 맞혀야 하는 목표 표현이다.

IWM은 이 predictor를 이미지 변환의 효과를 예측하는 world model로 본다. world model이라는 말은 보통 환경의 동역학을 예측하는 모델을 뜻한다. 이 논문에서는 환경 대신 이미지 변환을 다룬다. 즉, “이미지에 어떤 변환이 일어나면 잠재 표현이 어떻게 바뀌는가”를 예측하는 모델이다.

2. 변환 정보를 주면 등변 표현이 생긴다

Predictor에게 어떤 변환이 적용되었는지 알려주지 않으면, predictor는 변환을 무시하는 쪽으로 학습하기 쉽다. 이 경우 encoder는 불변 표현을 만든다.

반대로 predictor에게 변환 정보를 조건으로 주면, predictor는 변환의 효과를 설명해야 한다. 그러면 encoder는 변환 정보를 완전히 지우지 않고, 표현 안에 보존하려는 압력을 받는다. 이것이 등변 표현으로 이어진다.

3. 세 가지 축이 world model 품질을 결정한다

저자들은 world model 품질을 좌우하는 요인을 세 가지로 정리한다.

의미왜 중요한가
ConditioningPredictor에게 변환 정보를 알려주는 방식변환 정보를 모르면 predictor는 변환을 무시하는 자명한 해로 갈 수 있다.
Difficulty예측해야 하는 변환의 난이도너무 쉬운 변환만 주면 world model을 배울 필요가 적다.
CapacityPredictor의 깊이와 폭복잡한 변환 관계를 모델링하려면 predictor 자체가 충분히 커야 한다.

이 세 축 중 하나라도 부족하면 predictor는 진짜 world model처럼 작동하지 못한다.


방법

프레임워크 개요

용어 정리: pre-training과 fine-tuning

  • Pre-training: 큰 데이터로 범용적인 표현을 먼저 학습하는 단계다. 여기서는 IWM이 이미지 변환을 예측하며 표현을 배운다.
  • Fine-tuning: 사전학습된 모델을 특정 과제에 맞게 추가 학습하는 단계다. 이 논문은 encoder를 고정하고 predictor만 조정하는 경로를 특히 강조한다.
  • Task head: 분류나 분할처럼 최종 과제의 답을 내기 위해 모델 뒤에 붙이는 작은 출력 모듈이다.
graph TB
    subgraph Training["IWM 학습: pre-training"]
        direction TB
        X["원본 이미지 x"] --> AugS["Source augmentation<br/>crop, color jitter,<br/>grayscale, blur,<br/>solarization, mask"]
        X --> AugT["Target augmentation<br/>crop, light color jitter"]

        AugS --> Enc["Context encoder f_θ<br/>ViT-B/16 or ViT-L/16"]
        AugT --> EMA["Target encoder f_θ̄<br/>EMA, stop-gradient"]

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

        AugS -.추출.-> TransInfo["변환 파라미터<br/>a_{x→y}"]
        TransInfo --> Cond["Conditioning<br/>feature or sequence"]

        SrcRepr --> Pred["Predictor g_φ<br/>Transformer depth 12~36"]
        Cond --> Pred

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

        PredRepr --> Loss["L2 loss<br/>‖ŝ_y − sg(s_y)‖²"]
        TgtRepr -.stop-grad.-> Loss
    end

    subgraph Inference["하류 과제 적응: predictor FT"]
        direction TB
        EncFrozen["Frozen encoder f_θ"] --> PredFT["Predictor fine-tuning<br/>encoder보다 적은 파라미터"]
        PredFT --> Task["Classification /<br/>Segmentation"]
    end

    Training ==> Inference

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

학습 절차

하나의 원본 이미지에서 source view와 target view를 만든다. 두 뷰에는 서로 다른 augmentation이 적용된다. Source view는 context encoder를 통과해 source 표현 s_x가 된다. Target view는 EMA target encoder를 통과해 target 표현 s_y가 된다.

Predictor는 두 가지 정보를 받는다. 하나는 source 표현이고, 다른 하나는 source에서 target으로 가는 변환 정보다. 이 둘을 이용해 target 표현을 예측한다. 손실은 예측 표현과 target 표현 사이의 L2 거리다. Target 쪽에는 stop-gradient를 걸어 학습 붕괴를 막는다.

용어 정리: EMA, stop-gradient, L2 loss

  • EMA(Exponential Moving Average): 학습 중인 encoder의 가중치를 조금씩 평균내어 더 안정적인 target encoder를 만드는 방식이다.
  • Stop-gradient: target 쪽으로는 gradient가 흘러가지 않게 끊는 기법이다. 두 네트워크가 서로를 따라가다 표현이 모두 같아지는 붕괴를 막는 데 쓰인다.
  • L2 loss: 두 벡터의 차이를 제곱해서 더한 손실이다. 여기서는 예측 표현과 target 표현이 가까워지도록 만든다.

Conditioning 방식

Predictor에게 변환 정보를 알려주는 방법은 두 가지다.

방식설명
Feature conditioning변환 파라미터를 벡터로 임베딩한 뒤 predictor 입력에 더한다.
Sequence conditioning변환 정보를 토큰처럼 만들어 predictor 입력 시퀀스에 붙인다.

중요한 점은 conditioning의 존재 자체다. 조건화가 없으면 I-JEPA의 MRR은 0.00에 머문다. 조건화를 추가하고 충분히 강한 변환을 주면 IWM의 MRR은 0.82에서 0.85 수준까지 올라간다.

용어 정리: conditioning과 MRR

  • Conditioning: 모델에게 “어떤 조건에서 예측해야 하는지”를 추가 정보로 알려주는 것이다. 여기서는 어떤 이미지 변환이 적용됐는지를 predictor에게 알려준다.
  • Feature conditioning: 변환 정보를 하나의 벡터 특징처럼 만들어 predictor 입력에 섞는 방식이다.
  • Sequence conditioning: 변환 정보를 별도 토큰처럼 만들어 Transformer 입력 시퀀스에 붙이는 방식이다.
  • MRR(Mean Reciprocal Rank): 정답 target 표현이 후보들 중 얼마나 높은 순위에 있는지 보는 검색 지표다. 1에 가까울수록 예측이 정확하다.

변환 난이도

약한 color jitter만으로는 predictor가 배울 것이 많지 않다. 저자들은 grayscale, blur, solarization 같은 더 파괴적인 변환을 추가한다. 이렇게 해야 predictor가 실제로 “변환 전후의 표현 관계”를 학습해야 하는 상황이 만들어진다.

용어 정리: 광도 변환과 파괴적 변환

  • Color jitter: 밝기, 대비, 채도, 색조를 무작위로 바꾸는 augmentation이다.
  • Grayscale: 컬러 이미지를 흑백 이미지로 바꾸는 변환이다.
  • Blur: 이미지를 흐리게 만들어 고주파 세부 정보를 줄이는 변환이다.
  • Solarization: 픽셀 밝기가 특정 기준을 넘으면 색을 반전시키는 변환이다.
  • 이들은 물체의 위치를 바꾸기보다 픽셀값의 성질을 바꾸므로 **광도 변환(photometric transformation)**에 가깝다.

Predictor 용량

Predictor가 너무 작으면 복잡한 변환을 안정적으로 학습하지 못한다. 논문에서는 12층 predictor보다 18층 predictor가 더 안정적으로 색상 등변성을 학습했다. 12층은 5회 중 1회만 성공했지만, 18층은 5회 중 4회 성공했다.

용어 정리: capacity, depth, width

  • Capacity: 모델이 얼마나 복잡한 함수를 표현할 수 있는지를 뜻한다. 보통 파라미터 수가 많고 층이 깊을수록 capacity가 크다.
  • Depth: Transformer block이 몇 층 쌓였는지를 뜻한다.
  • Width: 각 토큰을 표현하는 임베딩 벡터의 차원 수다. 예를 들어 width 384는 각 토큰이 384차원 벡터라는 뜻이다.

표기법

논문에서 쓰는 표기는 다음처럼 읽으면 된다.

표기의미
IWM^{D,W}_{Inv}predictor 깊이 D, 임베딩 폭 W를 가진 불변 IWM
IWM^{D,W}_{Equi}predictor 깊이 D, 임베딩 폭 W를 가진 등변 IWM
IWM₁₈,₃₈₄ᴱᑫᵘⁱ18층, width 384 predictor를 쓰는 등변 IWM

불변 SSL과 등변 SSL의 차이

graph LR
    subgraph Invariant["불변 SSL: Contrastive, DINO, MoCo"]
        A1["x → Enc → z_x"]
        A2["T(x) → Enc → z_Tx"]
        A1 --> A3["z_x ≈ z_Tx<br/>변환 정보 소실"]
        A2 --> A3
    end

    subgraph Equivariant["등변 SSL: MAE, I-JEPA, IWM"]
        B1["x → Enc → z_x"]
        B2["T(x) → Enc → z_Tx"]
        B1 --> B3["Predictor(z_x, T)<br/>≈ z_Tx<br/>변환 정보 보존"]
        B2 --> B3
    end

    style Invariant fill:#ffe4e1
    style Equivariant fill:#e0ffe0

불변 SSL에서는 변환된 두 이미지의 표현이 같아지도록 학습한다. 그래서 배경색이나 밝기 같은 변화는 지워진다. 이것은 분류에는 유리할 수 있다.

등변 SSL에서는 표현이 변환에 따라 예측 가능한 방식으로 바뀌도록 학습한다. 이 방식은 segmentation, detection, fine-tuning처럼 세밀한 정보가 필요한 과제에 유리할 수 있다.

IWM의 장점은 이 둘을 완전히 다른 방법론으로 보지 않는다는 점이다. Conditioning과 predictor 용량을 조절하면 같은 JEPA 구조 안에서 불변 쪽과 등변 쪽을 오갈 수 있다.


주요 결과

ImageNet-1k 결과

아래 표는 ViT-B/16 기준 주요 결과다. 핵심은 등변 IWM이 predictor fine-tuning에서 강하다는 점이다.

방법Epoch선형 평가Attentive ProbeEncoder FTPredictor FTEnd-to-End
MoCoV330076.3%76.4%
MAE30060.2%73.5%
MAE160083.6%83.0%83.1%
DINO160082.0%82.8%
I-JEPA30070.0%75.0%83.0%79.1%80.0%
IWM₁₂,₃₈₄ᴵⁿᵛ30074.5%77.0%83.3%80.5%81.3%
IWM₁₈,₃₈₄ᴱᑫᵘⁱ30067.5%75.1%82.9%84.4%83.3%

출처: Table 2, 3, 4

불변 IWM은 선형 평가에서 강하다. 이는 대조 학습과 비슷한 장점이다. 반대로 등변 IWM은 predictor fine-tuning에서 가장 강하다. 이는 predictor가 변환 관계를 학습했고, 그 지식을 하류 과제에 재사용할 수 있음을 시사한다.

용어 정리: top-1, probe, FT

  • Top-1 accuracy: 모델이 가장 높은 확률로 고른 답이 정답인 비율이다.
  • Attentive probe: frozen encoder 위에 attention 기반의 작은 평가 모듈을 붙여 표현 품질을 보는 방법이다. 선형 평가보다 더 강한 probe다.
  • Encoder FT: encoder까지 하류 과제에 맞게 fine-tuning하는 방식이다.
  • Predictor FT: encoder는 고정하고 predictor 중심으로 fine-tuning하는 방식이다.
  • End-to-End: encoder부터 task head까지 전체 모델을 함께 학습하는 방식이다.

ViT-L/16 스케일링

방법Encoder FTPredictor FTEnd-to-End
I-JEPA84.1%79.9%
IWM₁₈,₃₈₄ᴵⁿᵛ84.3%81.5%
IWM₃₆,₅₁₂ᴱᑫᵘⁱ83.7%85.0%85.4%

출처: Table 6

ViT-L에서는 predictor fine-tuning이 encoder fine-tuning보다 높게 나온다. 이 결과는 이 논문의 중요한 근거다. 다만 “진짜 world knowledge 재사용 때문인지”와 “추가 파라미터 효과인지”는 완전히 분리되어 증명된 것은 아니다.

World model 품질

MRR은 예측된 표현이 실제 target 표현을 얼마나 잘 찾아내는지 보는 지표다. 1에 가까울수록 좋다.

Predictor 구성Jitter Only+ Destructive+ Strong Jitter
I-JEPA (12,384)0.000.000.00
IWM (12,384)0.110.090.25
IWM (18,384)0.250.790.85

출처: Table 1

이 표에서 보이는 메시지는 명확하다. Conditioning만으로는 부족하다. 변환이 충분히 어렵고, predictor도 충분히 커야 한다.

ADE20k 의미론적 분할

방법EncoderPredictorEnd-to-End
I-JEPA44.245.445.1
IWM₁₂,₃₈₄ᴵⁿᵛ45.645.746.5
IWM₁₈,₃₈₄ᴱᑫᵘⁱ44.246.847.0

출처: Table 5, mIoU

분할 과제에서도 등변 IWM의 predictor가 강하다. 분할은 위치와 구조 정보가 중요하므로, 변환 정보를 보존하는 표현이 도움이 되었을 가능성이 있다.

용어 정리: ADE20k, semantic segmentation, mIoU

  • ADE20k: 실내외 장면 이미지에 픽셀 단위 라벨이 붙어 있는 대표적인 semantic segmentation 데이터셋이다.
  • Semantic segmentation: 이미지의 각 픽셀이 어떤 클래스에 속하는지 분류하는 과제다.
  • mIoU(mean Intersection over Union): 예측 영역과 정답 영역이 얼마나 겹치는지 클래스별로 계산한 뒤 평균낸 지표다. 값이 높을수록 분할 품질이 좋다.

다중 과제 일반화

단일 predictor를 ImageNet, iNat18, SUN397, Places205 네 데이터셋에 동시에 fine-tuning해도 평균 성능이 거의 유지된다. 논문에서는 73.5%에서 73.4%로만 변했다. 이는 predictor가 특정 데이터셋에만 맞는 얕은 패턴이 아니라, 비교적 범용적인 변환 지식을 담았을 가능성을 보여준다.

용어 정리: 데이터셋 이름 읽기

  • ImageNet-1k: 1,000개 물체 클래스를 가진 대규모 이미지 분류 데이터셋이다.
  • iNat18: 자연 생물 종 분류를 위한 iNaturalist 2018 데이터셋이다.
  • SUN397 / Places205: 장면이나 장소 범주를 분류하는 데이터셋이다. 물체보다 “어떤 장소인가”를 묻는 성격이 강하다.

이론적 의의

1. 대조 학습과 마스크 모델링을 하나의 틀로 본다

기존에는 대조 학습과 마스크 모델링이 서로 다른 패러다임처럼 다뤄졌다. 이 논문은 둘을 world model 관점에서 다시 묶는다.

불변 표현은 변환 정보를 버리는 쪽에 가깝다. 등변 표현은 변환 정보를 보존하고 설명하는 쪽에 가깝다. IWM은 이 둘이 완전히 다른 것이 아니라, predictor가 어떤 정보를 받고 무엇을 예측하느냐에 따라 달라지는 스펙트럼이라고 주장한다.

2. Predictor를 버리지 않는다

이 논문의 가장 실용적인 메시지는 predictor 재사용이다. JEPA류의 predictor는 보통 사전학습 손실을 만들기 위한 부품으로만 쓰였다. 하지만 IWM에서는 predictor가 변환의 효과를 배운 모듈이므로, 하류 과제에서도 쓸 수 있다.

이 관점은 PEFT와도 연결된다. LoRA나 Adapter처럼 새 모듈을 붙이는 대신, 이미 학습 과정에서 사용한 predictor를 다시 활용한다.

논문이 강조하는 효율성은 여기서 나온다. ViT-B 설정에서는 predictor가 encoder 파라미터의 약 1/3 수준인데도 predictor fine-tuning이 end-to-end fine-tuning에 가까운 성능을 낸다. ViT-L 설정에서는 predictor FT가 encoder FT보다 높게 나오므로, “버려지던 predictor가 실제로 하류 적응에 쓸모 있는가”라는 질문에 대한 강한 실험적 근거가 된다.

용어 정리: PEFT, LoRA, Adapter

  • PEFT(Parameter-Efficient Fine-Tuning): 거대한 모델 전체를 다시 학습하지 않고, 일부 파라미터만 학습해 적은 비용으로 적응시키는 방법들의 묶음이다.
  • LoRA: 기존 가중치는 고정하고, 작은 저랭크 행렬만 추가로 학습하는 PEFT 방법이다.
  • Adapter: 모델 중간에 작은 모듈을 끼워 넣고 그 모듈만 학습하는 PEFT 방법이다.

3. 표현의 추상화 수준을 조절할 수 있다

모든 하류 과제가 같은 표현을 원하는 것은 아니다. 이미지 분류는 색상이나 위치 변화에 둔감한 표현이 좋을 수 있다. 반면 분할이나 검출은 위치와 변환 정보를 더 많이 보존해야 한다.

IWM은 conditioning, 변환 난이도, predictor 용량을 조절해 표현의 불변성과 등변성 사이를 탐색할 수 있다는 가능성을 제시한다.

4. LeCun의 AMI 구상을 시각 SSL로 좁혀 실험한다

이 논문은 A Path Towards Autonomous Machine Intelligence에서 제시된 큰 구상과 연결된다. AMI에서 LeCun은 world model이 픽셀을 그대로 생성하기보다 추상 표현 공간에서 미래 상태를 예측해야 한다고 주장했다. I-JEPA는 이 아이디어를 이미지 마스크 예측으로 구현했고, IWM은 한 걸음 더 나아가 색상, 밝기, 블러 같은 광도 변환이 표현에 미치는 효과를 예측한다.

따라서 IWM은 “자율 지능을 위한 일반 world model”을 바로 만든 논문이라기보다는, JEPA식 world model이 시각 표현 학습 안에서 실제로 어떤 형태의 학습 신호와 하류 적응 이점을 줄 수 있는지 검증한 논문으로 읽는 편이 정확하다.


한계와 주의할 점

한계설명
광도 변환 중심논문이 주로 다루는 변환은 color jitter, grayscale, blur, solarization이다. 회전, affine, 3D, 시간 변화 같은 더 복잡한 변환에는 아직 확장되지 않았다.
코드와 체크포인트 공개 제한arXiv 버전 기준 공식 코드와 체크포인트 공개가 명확하지 않다. 재현성 측면에서 감점 요인이다.
주요 결과의 분산 보고 부족일부 ablation은 반복 실험이 있지만, 핵심 ImageNet 성능은 단일 실행 중심으로 보인다.
World model 재사용의 귀인 문제Predictor FT의 성능 향상이 진짜 변환 지식 재사용 때문인지, 추가 용량이나 좋은 초기화 때문인지는 더 엄밀한 분리가 필요하다.
MRR의 해석 한계MRR은 예측 표현이 target 표현을 잘 검색하는지 보는 프록시 지표다. 이것만으로 모델이 물리적 의미의 world knowledge를 배웠다고 말하기는 어렵다.
소형 모델에서 이점 축소 가능성Predictor FT의 장점은 ViT-L에서 특히 두드러진다. 작은 모델에서도 일관되게 유리한지는 더 확인해야 한다.

용어 정리: affine, ablation, 귀인

  • Affine 변환: 이동, 회전, 확대·축소, 기울이기처럼 이미지의 기하학적 형태를 바꾸는 변환이다.
  • Ablation study: 모델 구성 요소를 하나씩 빼거나 바꿔 보면서 어떤 요소가 성능에 영향을 주는지 확인하는 실험이다.
  • 귀인(attribution): 성능 향상의 원인이 무엇인지 따지는 문제다. 여기서는 predictor가 배운 지식 때문인지, 단순히 파라미터가 더 많아서인지 분리해야 한다.

재현성 및 신뢰도 평가

항목등급비고
코드 공개⚠️arXiv 본문 기준 공식 저장소 링크가 명확하지 않다. FAIR의 선행 I-JEPA 코드를 바탕으로 일부 재현은 가능할 수 있다.
데이터 공개ImageNet-1k, ADE20k, iNat18, SUN397, Places205 모두 공개 데이터셋이다.
하이퍼파라미터Optimizer, EMA 계수, 증강 강도, predictor depth와 width가 본문과 부록에 기재되어 있다.
실험 환경⚠️ViT-B/L 학습에 필요한 GPU 수와 총 시간은 요약 수준이다. 정확한 재현 비용 추정은 어렵다.
통계적 신뢰도⚠️일부 ablation은 반복 성공 빈도를 보고한다. 주요 성능 수치의 분산 보고는 제한적이다.
종합 등급C구조와 하이퍼파라미터 설명은 충분하지만, 코드 미확인과 단일 실행 중심 결과가 약점이다.

주장별 신뢰도

#주장근거신뢰도
1Conditioning, difficulty, capacity가 world model 품질을 결정한다.Table 1의 MRR ablation. 0.00에서 0.85까지 상승한다.🟢
2등변 IWM은 predictor FT로 ViT-B 84.4%, ViT-L 85.0%를 달성한다.Table 4, Table 6🟢
3Predictor FT가 encoder FT와 동등하거나 더 좋다.ViT-L에서 predictor FT가 encoder FT보다 1.3%p 높다. 다만 귀인은 간접적이다.🟡
4IWM은 불변과 등변 사이의 스펙트럼을 조절할 수 있다.불변 IWM과 등변 IWM 결과 비교. 중간 지점의 연속 곡선은 제한적이다.🟡
5단일 predictor가 여러 데이터셋에 일반화된다.Table 7에서 평균 성능이 73.5%에서 73.4%로 유지된다.🟢
6기하학적 변환이나 시간 동역학으로 확장 가능하다.본문에서 직접 실증하지 않았다.🔴

읽기 난이도

난이도: ⭐⭐

JEPA, I-JEPA, 대조 학습, MAE에 대한 배경지식이 있으면 읽기 쉽다. 수식은 L2 예측 손실 수준이라 부담이 크지 않다. 다만 불변성과 등변성의 차이를 먼저 이해해야 논문의 메시지가 선명해진다.


관련 연구 비교

IWM (본 논문)I-JEPA (2023)MAE (2022)DINO (2021)SimCLR (2020)
대표 연구Garrido et al.Assran et al.He et al.Caron et al.Chen et al.
핵심 접근잠재 공간에서 광도 변환 예측잠재 공간에서 마스크 패치 예측픽셀 공간 마스크 재구성EMA teacher 기반 self-distillationInfoNCE 대조 손실
문제 정의불변과 등변의 통합 제어, predictor 재사용공간적 마스크 완성가려진 픽셀 복원view-invariant 표현 학습augmentation-invariant 표현
데이터ImageNet-1k, ADE20k 등ImageNet-1kImageNet-1kImageNet-1kImageNet-1k
핵심 메트릭ViT-B 300ep: Linear 74.5%, Pred FT 84.4%ViT-B 300ep: Linear 70.0%, E2E 80.0%ViT-B 1600ep: E2E 83.6%ViT-B 1600ep: E2E 82.8%ResNet-50: Linear 69.3%
표현 유형불변과 등변 선택 가능공간적 등변재구성 기반 등변불변불변
확장성Predictor FT로 PEFT 가능ViT-H까지 스케일 검증ViT-H까지 검증ViT-B까지ResNet 위주
Predictor 재사용✅ 핵심 기여❌ 학습 후 폐기❌ Decoder 폐기N/AN/A
한계광도 변환 중심, 코드 미확인광도 변환 학습 없음선형 평가 약함긴 학습 필요큰 batch와 증강에 의존
코드 공개⚠️ 미확인✅ facebookresearch/ijepa✅ facebookresearch/mae✅ facebookresearch/dino✅ google-research

관련 연구

  • A Path Towards Autonomous Machine Intelligence (LeCun, 2022): JEPA와 world model을 큰 인지 아키텍처의 핵심 부품으로 제안한 position paper다. IWM은 이 비전을 이미지 변환 예측 문제로 좁혀 실험한 후속 흐름으로 볼 수 있다.
  • I-JEPA (Assran et al., 2023): IWM의 직접적 선행 연구다. 마스크된 패치만 예측하는 JEPA 구현체이며, IWM은 예측 대상을 광도 변환까지 넓힌다.
  • V-JEPA (Bardes et al., 2024): 비디오 도메인의 JEPA다. IWM을 시간적 변환과 동역학으로 확장할 때 연결된다.
  • MAE (He et al., 2022): 픽셀 공간 마스크 재구성의 대표 방법이다. IWM은 픽셀 대신 잠재 공간에서 예측한다는 점에서 MAE와 대비된다.
  • DINO (Caron et al., 2021): Self-distillation 기반 불변 표현 학습 방법이다. IWM의 불변 버전과 비교된다.
  • MoCoV3 (Chen, Xie, and He, 2021): 모멘텀 대조 학습을 ViT에 확장한 방법이다. 이 논문에서는 선형 평가의 강한 기준선 역할을 한다.
  • A Simple Framework for Contrastive Learning of Visual Representation (Chen et al., 2020): SimCLR 논문이다. 강한 augmentation과 InfoNCE 손실로 augmentation-invariant 표현을 학습하는 대표적 대조 학습 기준선이다.
  • data2vec (Baevski et al., 2022): 여러 모달리티에서 잠재 공간 예측을 사용하는 방법이다. 픽셀이나 토큰을 직접 복원하기보다 표현을 예측한다는 점에서 JEPA 계열과 개념적으로 가깝다.
  • VICReg (Bardes, Ponce, and LeCun, 2022): 분산, 불변성, 공분산 정규화를 통해 표현 붕괴를 막는 방법이다. LeCun 계열의 non-contrastive SSL 흐름과 연결된다.
  • Platonic Representation Hypothesis: 서로 다른 모델이 유사한 표현으로 수렴한다는 가설이다. IWM의 불변-등변 스펙트럼과 연결해 생각할 수 있다.

원자적 인사이트

Predictor의 conditioning이 불변성과 등변성을 결정한다

출처: Learning and Leveraging World Models in Visual Representation Learning (Garrido et al., 2024)

JEPA 계열의 predictor에게 변환 정보 a_{x→y}를 주지 않으면, 손실을 줄이는 쉬운 방법은 변환을 무시하는 것이다. 이때 encoder는 augmentation에 둔감한 불변 표현을 만들게 된다.

반대로 predictor에게 변환 정보를 주면, predictor는 그 변환을 설명해야 한다. 그러면 encoder는 변환 정보를 완전히 지우기 어렵다. 즉, 불변성과 등변성은 encoder 구조만의 문제가 아니라 predictor가 어떤 정보에 접근할 수 있는가의 문제다.

활용 가능성: 분류처럼 불변성이 중요한 과제와, 검출·분할처럼 등변성이 중요한 과제에 맞춰 사전학습 조건화 정책을 조절할 수 있다.

Predictor fine-tuning은 값싼 하류 적응 경로다

출처: Learning and Leveraging World Models in Visual Representation Learning (Garrido et al., 2024)

Encoder를 얼려 두고 predictor만 fine-tuning해도 강한 성능이 나온다. ViT-B에서는 end-to-end fine-tuning에 필적하고, ViT-L에서는 predictor FT가 encoder FT보다 높다.

이 결과는 predictor가 단순 보조 모듈이 아니라, 하류 과제에서 재사용 가능한 변환 지식을 담고 있을 수 있음을 보여준다. 특히 ViT-B에서 predictor는 encoder의 약 1/3 파라미터 수준이므로, 전체 encoder를 업데이트하지 않고도 강한 성능을 얻는 파라미터 효율적 경로가 된다. 물론 이 해석은 아직 간접적이다. 추가 용량 효과나 좋은 초기화 효과와 분리하는 실험이 더 필요하다.

활용 가능성: 대규모 vision foundation model을 저비용으로 적응시키는 PEFT 경로로 볼 수 있다.

불변성과 등변성은 연속 스펙트럼이다

출처: Learning and Leveraging World Models in Visual Representation Learning (Garrido et al., 2024)

대조 학습과 마스크 모델링은 보통 서로 다른 학습 패러다임으로 소개된다. 하지만 IWM은 둘을 같은 틀 안에서 설명한다.

불변 표현은 변환을 설명하지 않는 쪽에 가깝다. 등변 표현은 변환을 설명하는 쪽에 가깝다. Conditioning, 변환 난이도, predictor 용량을 조절하면 이 스펙트럼의 어느 지점을 목표로 할지 선택할 수 있다. 다만 이 논문은 중간 지점들을 촘촘하게 그린 연속 곡선을 제시하기보다는, 불변 IWM과 등변 IWM의 대비를 통해 스펙트럼 가설을 뒷받침한다.

활용 가능성: 하류 과제별로 최적의 “불변도”를 찾는 AutoSSL 방식으로 확장할 수 있다.

World model 품질은 세 요소의 결합 함수다

출처: Learning and Leveraging World Models in Visual Representation Learning (Garrido et al., 2024)

이 논문의 ablation은 world model 품질이 세 요소에 의해 결정된다는 점을 보여준다.

요소없을 때 생기는 문제
Conditioning변환 정보를 몰라서 불변 표현으로 퇴화한다.
Difficulty예측이 너무 쉬워서 비자명한 world model을 배울 필요가 없다.
Capacity복잡한 변환을 표현할 능력이 부족하다.

셋 중 하나라도 부족하면 predictor는 진짜 world model로 작동하기 어렵다.

활용 가능성: IWM을 3D, video, robotics 같은 새 변환군으로 확장할 때 체크리스트로 쓸 수 있다.

MRR은 world model 품질의 프록시다

출처: Learning and Leveraging World Models in Visual Representation Learning (Garrido et al., 2024)

MRR이 높다는 것은 predictor가 source 표현과 변환 조건을 이용해 올바른 target 표현을 후보들 중 높은 순위로 찾아낸다는 뜻이다. 따라서 IWM의 predictor가 단순히 평균 표현을 내는 것이 아니라 변환 후 표현을 구분할 수 있다는 근거가 된다.

하지만 MRR은 어디까지나 잠재 표현 검색 성능이다. 이것이 곧 모델이 3D 물리, 객체 지속성, 시간적 인과성 같은 넓은 의미의 world knowledge를 배웠다는 뜻은 아니다. 그래서 이 논문의 “world model”은 강한 철학적 의미보다, 이미지 변환의 표현상 효과를 예측하는 제한된 world model로 읽는 것이 안전하다.


핵심 용어 정리

용어정의
JEPAJoint-Embedding Predictive Architecture. 두 뷰의 임베딩을 같은 잠재 공간에 놓고, 한쪽에서 다른 쪽을 예측하는 자기지도 학습 구조다.
IWMImage World Models. JEPA를 확장해 마스킹뿐 아니라 색상, 밝기, 블러 같은 광도 변환의 효과를 잠재 공간에서 예측한다.
World Model환경이나 데이터 도메인의 변화 규칙을 예측하는 모델이다. 이 논문에서는 이미지 변환이 표현에 미치는 효과를 예측한다.
불변 표현Augmentation이 달라져도 거의 변하지 않는 표현이다. 분류에는 유리하지만 변환 정보가 사라질 수 있다.
등변 표현입력이 변하면 표현도 예측 가능한 방식으로 변하는 표현이다. 변환 정보를 보존한다.
Contrastive Learning긍정쌍은 가깝게, 부정쌍은 멀게 만드는 학습 방식이다. SimCLR과 MoCo가 대표적이다.
Masked Image Modeling이미지 일부를 가리고 복원하거나 예측하는 자기지도 학습이다. MAE와 BEiT가 대표적이다.
EMAExponential Moving Average. 학습 중 encoder 가중치의 이동 평균을 target 네트워크로 사용하는 방식이다.
Stop-gradientTarget 경로로 gradient가 흐르지 않게 막는 기법이다. Joint-embedding 구조에서 붕괴를 막는 데 중요하다.
PredictorSource 표현과 변환 조건을 받아 target 표현을 예측하는 Transformer 모듈이다. IWM에서는 world model 역할을 한다.
ConditioningPredictor에게 어떤 변환이 적용되었는지 알려주는 방법이다. Feature conditioning과 sequence conditioning이 있다.
MRRMean Reciprocal Rank. 예측된 표현이 실제 target을 얼마나 잘 찾아내는지 평가하는 지표다. 1에 가까울수록 좋다.
Attentive ProbingFrozen encoder 위에 cross-attention 기반 probe를 학습해 표현 품질을 평가하는 방법이다.
Linear EvaluationFrozen 표현 위에 선형 분류기만 학습하는 평가 방식이다. 표현의 선형 분리 가능성을 본다.
Photometric Transformation색상, 밝기, 대비, 블러, solarization처럼 픽셀값을 바꾸는 변환이다.
Predictor Fine-tuningEncoder는 고정하고 predictor와 task head만 하류 과제에 맞게 학습하는 방식이다.

태그

paper #2024 self-supervised-learning world-model JEPA vision-transformer representation-learning equivariance


BibTeX

@article{Garrido2024IWM,
  title   = {Learning and Leveraging World Models in Visual Representation Learning},
  author  = {Garrido, Quentin and Assran, Mahmoud and Ballas, Nicolas and Bardes, Adrien and Najman, Laurent and LeCun, Yann},
  journal = {arXiv preprint arXiv:2403.00504},
  year    = {2024},
  url     = {https://arxiv.org/abs/2403.00504}
}