기술 요약: 이미지 생성을 위한 구형 흐름 매칭을 위한 잠재 공간 기하학 정렬

문제 제기

이미지 생성을 위한 잠재 흐름 매칭 (및 확산 모델) 은 일반적으로 가우시안 노이즈를 유클리드 공간의 직선 경로를 따라 변분 오토인코더 (VAE) 잠재 공간으로 이동시키는 방식으로 작동합니다. 그러나 이 접근법은 고차원 VAE 잠재 공간의 두 가지 중요한 구조적 속성을 간과합니다:

1. 측도의 집중 (Concentration of Measure): 가우시안 노이즈 사전 분포와 인코딩된 데이터 잠재 공간 모두 공간의 부피를 채우는 것이 아니라 얇은 구형 껍질에 집중되어 있습니다. 이러한 껍질 위의 두 점을 연결하는 직선 유클리드 선 (현) 은 두 분포 모두 점유하지 않는 내부 반지름을 통과하게 되어, 모델이 거의 채워지지 않은 잠재 공간 영역을 통한 전이를 학습하도록 강제합니다.
2. 방향적 우세 (Directional Dominance): 디코더의 출력은 잠재 토큰의 길이 (반경 성분) 보다 방향 (각 성분) 에 훨씬 더 민감합니다. 동일한 클래스의 이웃 토큰들로부터 반경과 방향을 서로 바꾸는 실험 결과, 방향을 유지하면서 반경을 변경하는 것은 지각적 변화가 미미한 반면, 방향을 변경하는 것은 의미론적 변화를 크게 초래함이 입증되었습니다.

표준 선형 흐름 매칭은 이러한 사실을 무시하여 훈련 속도 목표의 상당 부분 (토크나이저에 따라 최대 50-90%) 을 반경 운동에 할당합니다. 디코더가 반경에 덜 민감하므로, 이러한 감독은 비효율적이며, 근본적인 기하학을 존중하는 방법에 비해 수렴 속도가 느리고 이미지 품질 (FID) 이 저하됩니다.

방법론

저자들은 확산 아키텍처를 수정하거나 보조 인코더를 필요로 하지 않으면서 잠재 공간 기하학을 데이터의 구형 구조와 정렬시키는 구형 흐름 매칭 (Spherical Flow Matching) 을 제안합니다. 이 접근법은 네 가지 핵심 구성 요소로 이루어집니다:

1. 토큰 단위 구형 투영:
   이 방법은 사전 훈련된 VAE 의 각 잠재 토큰 $z$를 고정 반지름 초구면 $S^{d-1}(\sqrt{d})$ 위로 투영합니다. 이는 고정된 인코더와 디코더 사이에 삽입된 $L_2$ 정규화 단계를 통해 달성됩니다:
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   이를 통해 모든 토큰이 표준 가우시안 사전 분포의 집중 반지름과 일치하는 반지름을 가진 구 위에 위치하게 됩니다.

2. 구형 사전 분포:
   원시 가우시안 노이즈 대신 노이즈 끝점 $z_0$는 등방성 가우시안 노이즈 $\epsilon \sim \mathcal{N}(0, I)$를 샘플링하여 구 위에 투영함으로써 구성됩니다:
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   이는 투영된 데이터의 기하학과 일치하는 구 위의 균일 분포를 생성합니다.

3. 디코더 파인튜닝:
   기하학적 제약을 수용하기 위해 디코더는 몇 에포크 (실험에서 5 에포크) 동안 파인튜닝되며, 인코더와 투영 레이어는 고정된 상태로 유지됩니다. 재구성 목적 함수는 표준 픽셀 손실 (L1, LPIPS, 적대적) 을 유지하지만, 투영된 잠재 공간이 결정적이므로 KL 항은 제거됩니다.

4. Slerp 이동 및 접선 속도:
   노이즈와 데이터 끝점 사이의 이동 경로는 구 위의 측지선 호를 따르는 구형 선형 보간 (slerp) 으로 대체됩니다.

   • 측지선 경로: 경로 $z_t$는 모든 $t \in [0, 1]$에 대해 구 위에 머뭅니다.
   • 속도 목표: 조건부 속도 $u_t$는 측지선의 시간 미분으로, 접선 공간 $T_{z_t}S^{d-1}$ 위로 투영됩니다. 이는 속도가 구성상 순수한 각도 성분을 가지며 반경 성분이 0 임을 보장합니다.
   • 훈련 및 샘플링: 모델은 이 접선 속도를 예측하도록 훈련됩니다. 추론 중에는 샘플러가 지수 사상 (exponential map) 을 사용하여 속도를 적분하여, 각 단계에서 반경 투영 보정이 필요하지 않도록 궤적이 구 위에 머무르도록 보장합니다.

주요 기여

본 논문은 표준 잠재 흐름 매칭에서의 "반경 - 껍질 불일치 (radial-shell mismatch)"를 식별하고 정량화하며, 이를 해결하기 위한 기하학적 개입을 도입합니다:

1. 불일치 식별: 저자들은 선형 경로가 디코더가 민감하지 않은 반경 운동을 예측하도록 모델을 강제함으로써 용량을 낭비하고 수렴을 지연시킨다는 것을 입증합니다.
2. 토큰 단위 투영: 인코더를 다시 훈련하거나 보조 표현 인코더를 사용할 필요 없이 사전 훈련된 VAE 잠재 공간을 고정 반지름 구로 제한하는 방법을 도입합니다. 이는 오직 디코더 파인튜닝만으로 달성됩니다.
3. Slerp 기반 흐름 매칭: 접선 속도 목표를 가진 slerp 측지선을 따라 흐름 매칭을 훈련하여 전체 궤적이 구형 다양체를 존중하도록 합니다.
4. 실증적 검증: 이 방법은 아키텍처 변경 없이 여러 토크나이저 계열 (FLUX.2, VA-VAE, REPA-E FLUX.1) 과 모델 규모 (SiT-B 에서 SiT-XL) 에서 성능을 향상시키는 것으로 나타났습니다.

결과

클래스 조건부 ImageNet-256 생성에 대한 실험은 일관된 개선을 보여줍니다:

• FID 개선: 구형-slerp 방법은 모든 테스트된 구성에서 바닐라 선형 흐름 매칭보다 낮은 FID 점수를 달성합니다. 예를 들어, FLUX.2 를 사용하는 SiT-B/2 에서 FID 는 26.35(선형) 에서 20.55(구형-slerp) 로 감소합니다.
• 훈련 효율성: 이 방법은 바닐라 선형 기준선에 비해 약 2.2$\times$ 적은 훈련 단계에서 목표 FID 30 에 도달합니다.
• 일반화: 다양한 가이드 스케일 (CFG) 과 모델 크기 (B 및 XL) 에서 개선이 관찰되어, 이 이점이 특정 하이퍼파라미터 튜닝이 아닌 잠재 공간 기하학 자체에 내재된 것임을 나타냅니다.
• Ablation 연구:
  • 바닐라 및 구형 파이프라인 간 디코더를 교체하면 양방향 모두에서 성능이 저하되어, 학습된 흐름이 특정 잠재 공간 기하학과 밀접하게 결합되어 있음을 확인합니다.
  • 바닐라 잠재 공간에서 "껍질 분해" 경로 (반경과 방향을 별도로 보간) 를 사용하는 것은 동일한 이점을 제공하지 않아, 반경 자유도를 완전히 제거하는 것이 중요함을 확인합니다.
  • 파인튜닝된 구형 디코더의 재구성 품질 (rFID) 은 경쟁력을 유지하여, 기하학적 제어가 VAE 의 이미지 재구성 능력을 심각하게 훼손하지 않음을 보여줍니다.

의의

본 논문은 잠재 지지 기하학 (latent support geometry) 을 생성 모델링을 위한 중요한 설계 축으로 확립합니다. 잠재 분포의 형태가 생성기가 이동해야 하는 공간을 정의하며, 이를 무시하고 (유클리드 선형성을 가정함으로써) 기하학을 간과하면 비효율성이 초래된다고 주장합니다.

이 연구의 의의는 이미지 생성 품질과 훈련 속도를 향상시키면서도 다음을 수반하지 않는다는 점에 있습니다:

• 확산 아키텍처 변경 (예: SiT 백본).
• 보조 인코더 추가 (고정된 DINOv2 특징에 대한 리만 흐름 매칭과 달리).
• VAE 인코더를 처음부터 다시 훈련.

기존 사전 훈련된 잠재 공간을 단순히 구 위에 투영하고 흐름 경로를 결과적인 측지선과 정렬시킴으로써, 이 방법은 최소한의 계산 오버헤드로 최첨단 결과를 달성합니다. 저자들은 향후 연구가 더 강력한 압력 하에서 이러한 기하학적 제약을 탐구하고, 본질적으로 이러한 구형 구조와 일치하는 토크나이저를 설계할 것을 제안합니다.