Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion

"Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion" 논문에 대한 상세한 기술적 요약입니다.

문제 제기

소수 샘플링 (Minority sampling) 은 데이터 매니폴드 상의 저밀도 인스턴스를 생성하는 것을 목표로 하며, 이는 의료 진단, 이상 탐지, 창의적 AI 와 같은 응용 분야에서 중요합니다. 그러나 기존 접근법들은 모델 자체가 학습한 생성적 사전 (generative prior) (즉, 훈련 데이터의 분포) 에 상대적으로 "소수"를 정의합니다. 이는 근본적인 한계를 야기합니다: "희귀"하다고 식별된 샘플들은 종종 실제 세계의 진정한 비정형적 인스턴스라기보다는 특정 훈련 세트의 산물이거나 모델의 귀납적 편향의 산물에 불과합니다. 예를 들어, 표준 데이터셋으로 훈련된 모델은 훈련 데이터에서 과소 대표되어 있을 뿐인 흰색 배경 위의 개를 단순히 "희귀"하다고 간주할 수 있지만, 실제 세계에서는 진정으로 희귀한 스텔스 항공기는 모델의 사전이 좁다면 그렇게 포착되지 않을 수 있습니다. 이 논문은 현재 방법론들이 진정한 의미적 희귀성을 포착하는 데 필수적인 실제 세계 사전 (real-world priors) 과 정렬되지 못한다고 주장합니다.

방법론: JEPA 가이드

저자들은 JEPA 가이드 (JEPA Guidance) 를 제안하는데, 이는 소수 샘플링을 생성기 중심의 관점에서 세계 중심 (world-centric) 관점으로 전환하는 프레임워크입니다. 이는 광범위하고 의미적으로 풍부한 표현을 인코딩하는 것으로 알려진 세계 모델 (예: DINOv2) 인 Joint-Embedding Predictive Architecture(JEPA) 의 암시적 밀도를 사용하여 확산 모델을 안내함으로써 달성됩니다.

핵심 메커니즘:

1. 세계 중심 정의: 확산 모델의 밀도 $p_\theta(x)$ 하에서 확률을 최소화하는 대신, 이 방법은 JEPA 유도 밀도 $p_\phi(x)$ 하에서 확률을 최소화하는 것을 목표로 합니다. 소수 샘플은 JEPA 표현 공간의 저밀도 영역에 위치하는 것으로 정의됩니다.
2. JEPA-SCORE: 이 방법은 사전 훈련된 JEPA 인코더 $f_\phi$ 의 야코비안 (Jacobian) 에서 유도된 밀도 측정치인 JEPA-SCORE를 활용합니다. 구체적으로, 야코비안 $J_f(x)$ 의 특이값 (singular values) 로그의 합입니다. 낮은 JEPA-SCORE 값은 세계 사전에서의 낮은 밀도 (높은 비정형성) 를 나타냅니다.
3. 무작위 SVD 를 통한 근사: 모든 확산 단계에서 전체 야코비안과 그 특이값 분해 (SVD) 를 계산하는 것은 계산적으로 불가능합니다. 저자들은 압축된 야코비안의 상위-$k$개 특이값을 효율적으로 근사하기 위해 무작위 SVD (Randomized SVD) 를 사용합니다. 그들은 근사 오차에 대한 이론적 상한을 제공하며, 이를 무작위 SVD 추정 오차와 절단 오차로 분해합니다.
4. 효율성을 위한 Envelope Theorem: 입력 $x_t$에 의존하는 무작위 SVD 의 내부 최적화 루프를 통한 역전파를 피하기 위해, 저자들은 Envelope Theorem을 적용합니다. 이를 통해 최적 투영 행렬 $Q^*$ 를 기울기 계산 중 상수로 취급할 수 있어, 올바른 1 차 기울기를 유지하면서 메모리 사용량과 계산 오버헤드를 크게 줄일 수 있습니다.
5. 지연 가이드 (Deferred Guidance): 확산 모델의 노이즈가 있는 중간 잠재 변수와 JEPA 인코더가 기대하는 깨끗한 입력 사이의 도메인 격차를 해결하기 위해, 저자들은 지연 가이드를 도입합니다. JEPA 가이드는 중간 시간 단계 $\tau T$(예: $t=0.8T$) 이후에만 적용되어, JEPA 인코더가 효과적으로 평가할 수 있을 만큼 충분히 명확한 탈노이즈 추정치 $\hat{x}_{0|t}$가 보장되도록 합니다. 이는 또한 조건부 생성 (예: 텍스트-이미지) 으로 확장할 수 있게 하는데, 그 이유는 세계 중심 가이드가 시작되기 전에 조건이 초기에 설정되기 때문입니다.

주요 기여

• 세계 중심 소수 샘플링: 생성 모델의 좁은 사전이 아닌 (JEPAs 를 통한) 실제 세계 사전에 기반하여 소수 샘플을 재정의하는 패러다임 전환.
• JEPA 가이드 프레임워크: JEPA 유도 사전 하의 저밀도 영역으로 궤적을 유도하는 원칙적인 확산 샘플링 방법.
• 효율적인 계산: 원칙적인 근사 전략 (무작위 SVD) 개발과 Envelope Theorem 의 적용을 통해 JEPA 가이드를 계산적으로 실용적으로 만들었으며, 이론적 오차 한계를 동반함.
• 훈련 없는 통합: 이 접근법은 완전히 훈련이 필요 없으며, 사전 훈련된 확산 모델과 JEPA 모델만 필요로 하고, 하위 모델을 수정하지 않고 추론 시 통합할 수 있습니다.

실험 결과

저자들은 CelebA, ImageNet, Stable Diffusion 1.5, SDXL-Lightning 을 사용하여 무조건부, 클래스 조건부, 텍스트-이미지 생성 작업 전반에 걸쳐 JEPA 가이드를 평가했습니다.

• 비정형성 (Atypicality): Sehwag 등 (2022), MG, SGMS, MinorityPrompt 와 같은 베이스라인에 비해 JEPA 가이드는 일관되게 가장 낮은 JEPA-SCORE(세계 사전 하에서 더 높은 비정형성 표시) 를 달성했습니다.
• 충실도 및 유효성: 희귀 영역을 목표로 함에도 불구하고, 이 방법은 높은 충실도를 유지합니다. 텍스트-이미지 작업에서 CLIPScore, PickScore, ImageReward 와 같은 텍스트 정렬 점수에서 경쟁력 있는 성과를 거두면서도, 다른 방법론들이 생성하지 못하는 의미적으로 구별되는 샘플 (예: 고령 군인, 아프리카계 미국인 야구 선수) 을 생성합니다.
• 정성적 비교: 시각화 결과, 생성기 중심 방법론들은 종종 비정형적인 맥락(예: 크레인 근처의 독수리) 을 가진 샘플을 생성하는 반면, JEPA 가이드는 비정형적인 의미나 행동(예: 짝짓기 중인 독수리, 비정상적인 동물 형태) 을 가진 샘플을 생성하여 실제 세계의 희귀성과 더 잘 부합합니다.
• 하류 유틸리티: 분류기 훈련을 위한 데이터 증강 작업에서, JEPA 가이드된 소수 샘플로 증강한 경우 다른 방법론들보다 적은 샘플로 더 우수하거나 동등한 성능을 보여주어, 이러한 샘플이 더 정보적인 훈련 신호를 제공함을 시사합니다.

의의 및 주장

이 논문은 JEPA 가이드가 기존 소수 샘플러들이 생성 모델의 자체 분포에 갇히는 한계에 대한 강력한 해결책을 제공한다고 주장합니다. 세계 모델 (JEPAs) 을 활용함으로써, 이 방법은 특정 데이터셋에 상대적으로 희귀한 것이 아니라 실제 세계에서 진정으로 희귀한 샘플을 생성합니다.

저자들은 이 접근법이 다음과 같음을 강조합니다:

1. 소수의 개념을 생성적 사전에서 분리 (Decouples) 합니다.
2. 생성 모델링과 세계 모델링 사이의 격차를 연결 (Bridges) 하여, 일반 인공지능 (AGI) 추구에서 핵심 구성 요소가 됩니다.
3. 견고성 테스트, 공정성 개선, 창의적 콘텐츠 생성과 같이 다양하고 비정형적인 데이터가 필요한 응용 분야를 위한 실용적이고 훈련이 필요 없는 도구 (Provides) 를 제공합니다.

저자들은 또한 잠재적인 이중 사용 우려를 인정합니다: 가이드 방향을 반대로 하면 세계 사전의 기존 편향을 강화할 수 있으므로 책임 있는 사용을 장려합니다. 그들은 이 방법이 표준 샘플러에 비해 계산 오버헤드를 도입하지만, 일부 복잡한 베이스라인 (예: MinorityPrompt) 보다 여전히 빠르며, 세계 중심 비정형성의 상당한 이득에 대해 합리적인 절충안을 제공한다고 지적합니다.