[논문 리뷰] V-JEPA 요약, 코드, 구현

논문을 상세히 번역하고 한단어씩 해석해주는 포스팅은 많다.

나는 논문을 누구나 알아듣도록 쉽고 간결하게 전달하고자 한다.

 

V-JEPA

V-JEPA: Latent Video Prediction for Visual Representation Learning

BARDES, Adrien, et al. V-JEPA: Latent Video Prediction for Visual Representation Learning. 2023.

(ICLR 2024에 reject 되어 수정 중. 24/04/30 현재는 오픈리뷰 논문만 존재. 오피셜 논문 공개 시 수정 예정.)

 

논문 원문 링크

 

메타의 움직임이 심상치 않다.

V-JEPA가 오픈리뷰와 ICLR 논문 심사를 거치고 있는 중이다.

그런데 여기 CV의 대가 얀 르쿤 교수님이 참여하셨다.

얀 르쿤 교수님이 이 모델에 대한 자신감이 넘치신다.

 

언어 분야의 서비스 최강자 구글, MS가 NLP 기반모델을 완성시킨 것처럼,

비전 분야의 서비스 최강자 메타가 CV 기반모델 을 완성시킬까.

 

얀 르쿤 "사람처럼 추론하고 계획하는 AI 에이전트 개발 중"(기사 원문)

 

저자의 의도

LLM에 masked modeling 방법을 사용해 효율적으로 학습시켜보자.

비디오에서 SSL으로 시공간적 마스킹 영역을 예측하는 V-JEPA 모델을 제시한다.

(여기서 SSL은 self-supervised learning으로 자기지도학습을 말한다.)

 

레이턴트 비디오 예측 전략은 다운스트림에서 추가로 모델을 최적화할 필요가 없다.
비디오 학습이 이미지 학습 보다 모션을 이해시키는데 유리하다.

 

기존 문제점

인지 이론을 활용한 SSL은 NLP에 큰 발전을 기여했다.

SSL은 생물이 학습하는 방식을 모방한 학습 방식으로 포텐셜이 뛰어나다.

비전 분야 중 이미지 분야에서는 MAE에서 I-JEPA로 발전하고 있다.
비전 분야 중 비디오에 어떻게 적용해야할지 아직 특출난 방법이 없다.

 

해결 아이디어

JEPA 기법과 비디오 데이터셋을 활용해보자.

비디오를 시공간 마스크로 가리고 이 부분을 예측하며 비주얼 인코더를 학습한다.

 

1. I-JEPA

I-JEPA의 논문을 비디오로 진행하는 것이 핵심이다.

내가 리뷰한 포스팅인데 읽어도 되고 안읽어도 된다.

[논문 리뷰] I-JEPA 요약, 코드, 구현

 

I-JEPA를 2줄로 요약하면 다음과 같다.

1개의 context으로 4개의 target을 추상적인 레이턴트 스페이스에서 예측한다.

모델에서 중요한 모듈은 3개다. (context encoder, target encoder, predictor)

 

저자들은 이 아키텍처를 비디오와 마스킹에 적용해 V-JEPA를 구상했다.

 

2. Architecture

Architecture

정말로 I-JEPA와 아키텍처 컨셉이 완전히 똑같다.

backbone으로는 ViT를 사용한다. (ViT-L or ViT-H)

이 ViT는 스탠다드 트랜스포머 블럭에 joint space-time attention을 적용한 것이다.

 

3. Input

Input

ViT는 1D 시퀀스를 인풋으로 받는다.

비디오는 2D 이미지가 시간에 따라 변하는 3D 자료다.

3D 비디오를 1D 토큰 시퀀스로 바꿔야 한다.

 

sequence 관점에서 먼저 보자. (= 위 그림)

3D Conv filter에 비디오를 통과시킨다.

필터는 d개, 사이즈는 (2, 16, 16), 시간 stride 2, 공간 stride 16 이다.

이 필터를 통과한 결과 텐서의 모양은 (8, 14, 14, d) 이다.
여기 3D sin-cos 포지션 임베딩을 더해준다.
flatten 사용해서 1D로 펼쳐준다.
최종 시퀀스 모양은 (1568, d) 이다.

 

그럼 이 sequence를 만들 비디오는 어떤걸 넣을지 보자. (= 위 그림에 넣을 인풋)

각 인풋 비디오에서 stride 4로 두고 랜덤하게 16프레임의 클립을 추출한다.
16*4 = 64 이니까, 총 64프레임으로, 약 2초 정도 되는 길이의 클립이다.
spatial 사이즈(비디오 가로세로 크기)를 224로 통일한다.
최종 비디오 모양은 (16, 224, 224, 3) 이다.

 

4. 3D multi-block masking

3D multi-block masking

블럭 마스킹 전략을 간단하게 3D로 확장한다.

인풋으로 연속적인 비디오(툭툭 끊기면 안됨)가 주어지면 10개의 블럭으로 샘플링한다.

블럭은 곧 타겟을 말하며, 영상의 일부분(공간적 일부분) 이다.

이 블럭들은 종횡비가 다양하며, 이를 모두 커버하는 1개의 공간적 마스크를 구한다.
이 공간적 마스크를 전체 시간적 차원으로 확장하여 최종 마스크가 된다.

V-JEPA에서는 2가지 마스크(short-range, long-range)를 사용한다.
short-range mask(8개 랜덤 타겟 블럭, 전체 이미지의 0.15배 크기)
long-range mask(2개 랜덤 타겟 블럭, 전체 이미지의 0.7배 크기)
두 마스크 모두 종횡비는 랜덤하게 결정(0.75~1.5) 된다.

이 마스크의 최종 확정은 모든 마스크의 면적을 겹쳤을때 전체 이미지의 90%는 차지해야 한다는 규칙을 따른다.

이런 저런 알고리즘이 적용된 마스크로 90%나 가리고 모두 예측해야 한다는 말이다.
다시 말해 context 인코더에 인풋으로 들어가는 것은 10% 밖에 안된다.

 

5. Patch-level loss

Patch-level loss

loss를 설명하기 위해 우리가 싫어하는 수학이 등장한다.

그래도 최대한 쉽게 설명할테니 따라와보자.

 

시퀀스부터 보자. (Fig 3의 가장 앞부분)

먼저 비디오 클립과 비디오 마스크를 샘플링 한다.

비디오 클립은 1D 토큰 시퀀스이고 길이 L은 1568 이다.
마스크의 시퀀스 길이는 M, 언마스크의 시퀀스 길이는 N이다.

따라서 N = L-M 이다.

(위 그림에서 d는 배치사이즈라고 생각하면 된다.)

 

다음은 시퀀스가 인코더로 들어가 context representation가 되는 과정이다.

비디오 클립에 마스킹을 하고(x_N), context 인코더에 넣는다.

인코더(E)의 결과(z)라고 할 때 z_N = E(x_N) 으로 표기한다.

 

다음은 z_N 이 predictor로 들어간다.

context 인코더의 아웃풋(z_N)을 predictor에 넣는다.
가려야 하기 때문에 마스크 토큰(m)을 같이 넣는다.
예측 모델(P)의 결과(s)라고 할 때 s_M = P(z_N, m_M) 으로 표기한다.

 

다음은 s_M 과 비교할 target representation을 보자.

마스킹 안된 전체 비디오 클립(x_L)를 target 인코더에 넣는다.
인코더(E)의 결과(s)라고 할 때 s_L = E(x_L) 으로 표기한다.

 

이제 예측한 값도 있고 타겟 값도 있으니 Loss를 구할 수 있다.

predictor의 s_M과 target 인코더의 s_L으로 L1 loss를 계산한다.

Equation 1 과 같이 표기할 수 있다.

equation 1

 

6. Multi-mask prediction

Multi-mask prediction

I-JEPA나 다른 모델과 마찬가지로 타겟 블럭을 여러개 예측한다.
1번의 iter에서 타겟 representation을 1번만 계산하므로 효율적이다.

 

7. Collapse prevention mechanism

EMA

representation collapse라는 개념이 있다.
인코더가 입력과는 관계없이 계속 일정한 값을 출력하는 현상이다.
이를 피하기 위해 target 인코더의 파라미터 업데이트를 EMA 방식을 사용한다.

기존의 방식인 SGD나 Adam의 gradient 방식이 아니라,

context 인코더의 exponential moving average(EMA)를 사용한다.
이는 I-JEPA에서도 사용했던 방법이다.

 

결과 분석

1. Experimental setting

VideoGLUE benchmark에 속한 다양한 task 진행했다.
Kinetics-400, action recognition classification
Something-Something-v2, motion classification
AVA, action localization

 

이미지 관련 classification 평가도 진행했다.
ImageNet, object recognition classification
Places205, scene classification
iNaturalist 2021, fine-grained recognition

 

2. Eval on video tasks

Eval on video tasks

frozen evaluation protocol을 사용했다.

fine-tuning할 때 모든 블럭을 업데이트 하는 것이 아니라 헤드만 업데이트 한다.

해당 모델의 representation 추출 능력만을 보기 위함이다.

(이미지에서 linear probing 처럼)

V-JEPA의 ViT-H/16 384 모델이 기존 비디오 best 모델(InternVideo) 보다 성능이 좋다.
V-JEPA가 기존 이미지 best 모델(DINOv2)에 근접했다.
이미지 모델들은 기본적으로 데이터셋이 매우 크기 때문에 이정도도 대단한 성과다.

 

3. Eval on image tasks

Eval on image tasks

attentive probe 방식을 사용했다.

앞서 비디오에서 말한 것처럼 해당 모델의 representation 추출 능력만을 보기 위함이다.

이전 논문들은 linear probing으로 전체 레이어를 고정(frozen) 했다.

attentive probing은 일부 레이어만 고정하고 그 중 제일 좋은 결과를 쓴다.

마찬가지로 비디오 best 모델들에 비해서는 높은 성능이 나왔다.

 

코드 및 구현

오피셜 깃허브 코드

(해당 문서는 24.04.30. 현재 수정 중인 문서다.)

 

 

관련 논문 리스트 (스크롤 내려서 Paper List 참고)