기술 요약: 리만 플로우 매칭을 통한 매니폴드 상의 학습 (Learning on the Manifold with Riemannian Flow Matching)

1. 문제 정의

최근 생성 모델링의 발전은 합성 정밀도와 수렴 속도를 개선하기 위해 고차원 시맨틱 표현 인코더(예: DiNOv2, SigLigP)를 활용하는 방향으로 이동했습니다. 그러나 표준 디퓨전 트랜스포머(DiT) 및 플로우 매칭(Flow Matching, FM) 프레임워크는 이러한 표현(representation)에 대해 직접 학습할 때 효과적으로 수렴하지 못합니다.

이전 연구인 표현 오토인코더(Representation Autoencoders, RAE)는 이러한 실패의 원인을 **용량 병목 현상(capacity bottleneck)**으로 보고, 문제를 해결하기 위해 "너비 스케일링(width scaling, 인코더의 잠재 차원에 맞춰 트랜스포머의 너비를 늘리는 것)"이 필요하다고 제안했습니다. 본 논문은 이 가설에 이의를 제기하며, 근본적인 원인이 불충분한 모델 용량이 아니라 기하학적 불일치(geometric mismatch), 즉 **기하학적 간섭(Geometric Interference)**에 있다고 주장합니다.

핵심 문제는 표현 인코더가 특징(feature)을 거의 제로에 가까운 반경 분산을 가진 초구형 매니폴드($S^{d-1}$, "단단한 껍질" 구조) 상의 하이퍼스피어(hypersphere)로 투영한다는 점입니다. 반면, 표준 유클리드 플로우 매칭은 데이터와 가우시안 노이즈 사이의 선형 확률 경로(현, chord)를 가정합니다. 이 선형 보간은 하이퍼스피어의 저밀도 내부를 가로지르며, 이로 인해 모델은 특징 공간의 정의되지 않은 영역에서 속도장(velocity field)을 학습하도록 강요받습니다. 결과적으로 모델은 시맨틱 각도 전이(angular transitions)를 학습하는 대신, 구조적으로 0인 반경 방향 오차(radial errors)를 최소화하는 데 모델 용량을 낭비하게 됩니다.

2. 방법론: 야코비 정규화를 포함한 리만 플로우 매칭 (RJF)

용량 확장을 위한 건축적 수정 없이 기하학적 간섭을 해결하기 위해, 저자들은 RJF를 제안합니다. 이는 생성 프로세스를 잠재 공간의 내재적 기하학에 정렬시키는 프레임워크입니다.

2.1 리만 플로우 매칭 (궤적 교정)

유클리드 공간에서의 선형 보간 대신, RJF는 **구면 선형 보간(SLerical Linear Interpolation, SLERP)**을 사용하여 확률 경로를 초구형 매니폴드 상에 직접 정의합니다.

• 측지선 경로 (Geodesic Paths): 데이터 포인트 $x$와 노이즈 $\epsilon$(둘 다 단위 구면에 투영됨) 사이의 궤적은 현(chord)이 아닌 측지선(대원 호, great circle arc)을 따릅로 합니다.
• 접공간 속도 (Tangent Space Velocity): 목표 속도장은 SLERP 경로를 미분하여 계산됩니다. 이는 예측된 속도 벡터가 엄격하게 매니폴드의 접공간에 존재하도록($v_t \cdot x_t = 0$) 보장하며, 반경 성분을 학습할 필요성을 제거합니다.
• 샘생 (Sampling): 추론 과정에서 모델은 지수 맵(Exponential Map) 적분기를 사용하여 잠재 상태를 업데이트함으로써, 생성 과정 전반에 걸쳐 궤적이 매니폴드 표면에 머물도록 합니다.

2.2 야코비 정규화 (곡률 교정)

저자들은 단순히 궤적을 교정하는 것만으로는 불충분하다는 점을 관찰했습니다. 왜냐하면 표준 평균 제곱 오차(MSE) 손실은 속도 오차를 균일하게 취급하기 때문입니다. 양의 곡률을 가진 초구형 위에서는 측지선들이 집중(focusing, 수렴)되므로, 속도 섭동에 의한 오차가 비선형적으로 전파됩니다.

• 야코비 장 (Jacobi Fields): 이를 고려하기 위해 저자들은 속도 섭동에 의해 발생하는 측지선 간의 간격을 정량화하는 야코비 장으로부터 유도된 가중치 인자를 도입합니다.
• 손실 재가중 (Loss Reweighting): 손실 함수는 기하학적 가중 인자 $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$에 의해 곱해집니다. 여기서 $\Omega$는 측지선 거리입니다. 이는 섭동이 완화되는 데이터 근처($t=0$)의 오차는 낮추고, 노이즈 근처($t=1$)의 정밀도를 우선시함으로써, 매니폴드의 곡률에 맞춰 학습 신호를 어닐링(annealing)합니다.

3. 주요 기여

1. 수렴 실패에 대한 기하학적 분석: 본 논문은 표준 디퓨전이 표현 인코더에서 겪는 병목 현상이 기하학적 간섭임을 식별했습니다. 이는 모델 용량의 부족이 아니라, 유클리드 목적 함수가 모델로 하여금 하이퍼스피어의 매니폴드 외부 내부를 통해 학습하도록 강제하는 데서 기인함을 입증했습니다.
2. RJF 프레임워크: 리만 플로 매칭(궤적 정렬을 위함)과 야코비 정규화(곡률 인지 오차 전파를 위함)를 결합한 기하학적으로 일관된 훈련 프레임워크를 제안했습니다.
3. 너비 스케일링 없는 효율성: 표준 DiT 아키텍처(예: 131M 파라미터의 DiT-B)가 이전 연구에서 제안된 비용이 많이 드는 너비 스케일링 없이도 고차원 시맨틱 잠재 변수에 대해 효과적으로 수렴할 수 있음을 입증했습니다.

4. 실험 결과

실험은 다양한 모델 규모(DiT-B, DiT-L, DiT-XL)와 표현 인코더(DiNOv2-B, SigLIP, MAE)를 사용하여 256×256 해상도의 ImageNet-1K에서 수행되었습니다.

• 수렴 및 FID 성능:
  • DiT-B (131M): DiNOv2-B 특징에 대한 표준 유클리드 플로우 매칭은 수렴에 실패했습니다(FID ~24.21). RJF는 80 에포크에서 6.77의 FID를, 200 에포크에서 가이던스 적용 시 3.37의 FID를 달성했습니다.
  • DiT-XL (677M): RJF는 가이던스 없이 80 에포크 만에 3.62의 FID를 달로 달성하여, 유클리드 베이스라인(FID 4.28) 및 VAE 기반의 LightingDiT(FID 4.29)를 능가했습니다.
  • 스케일링: 이 방법은 모델 크기에 따라 일관된 개선을 보였으며, 최첨단 Inception Score(IS 186.2)와 정밀도(Precision 0.82)를 달성했습니다.
• 절제 연구 (Ablation Studies):
  • 반경 vs 각도 (Radial vs. Angular): 반경 방향 손실을 마스킹했을 때 미미한 이득만 있었는데(FID 21.99), 이는 선형 경로(현, chord) 자체가 주요 문제임을 확인시켜 줍니다.
  • 궤적 (Trajectory): 리만 플로우 매칭(SLERP)만 적용해도 성능이 급격히 향려되었습니다(FID 7.06), 이는 궤적 불일치 가설을 검증합니다.
  • 정규화 (Regularization): 야코비 정규화를 추가함으로써 결과를 더욱 정교화했습니다(FID 6.77).
• 일반화: 이 방법은 표준 방식이 완전히 실패했던 하이퍼스피어 제약이 있는 다른 인코더들(SigLIP FID 10.39, MAE FID 19.82)에서도 성공적으로 수렴했습니다.

5. 의의 및 주장

본 논문은 표현 인코더에서 디퓨전 모델을 학습시키는 어려움이 아키텍처의 문제가 아니라 근본적으로 기하학적인 문제라고 주장합니다. 잠재 공간의 초구형 위상(topology)을 존중함으로써, 표준 디퓨전 트랜스포머는 높은 충실도의 시맨틱을 효율적으로 학습할 수 있습니다.

저자들은 RJF가 표준 DiT 아키텍처의 잠재력을 끌어올린다고 단언하며, 이를 통해 훨씬 더 크거나 너비가 확장된 모델 또는 복잡한 VAE 기반 파이프라인과 대등하거나 이를 능가하는 성능을 낼 수 있다고 주장합니다. 이 연구는 효율적인 생성 모델링이 브루트 포스(brute-force) 방식의 용량 스케일링이 아닌 기하학적 정렬(정확한 궤적 및 곡률 인지 손실)을 필요로 한다는 점을 확립했습니다. 이 접근 방식은 DiT-B 모델이 이전 방법들이 실패했던 지점에서 FID 3.37을 달성한 것에서 알 수 있듯이, 더 적은 파라미터와 빠른 수렴 속도로 고충실도 합성을 가능하게 합니다.