A Few-Step Generative Model on Cumulative Flow Maps

기술적 요약: 누적 흐름 지도 기반의 소수 단계 생성 모델

문제 제기
확산 모델 및 흐름 매칭을 포함한 생성 모델들은 일반적으로 현재 상태 $x(t)$에 조건부인 국소 상태 업데이트를 예측하는 순간 흐름 지도(또는 순간 동역학) 를 학습하도록 공식화됩니다. 데이터를 생성하기 위해 이러한 모델들은 단순한 소스 분포에서 복잡한 데이터 분포로 샘플을 이동시키는 유한 시간 수송인 누적 흐름 지도를 근사하기 위해 많은 작은 시간 단계에 걸친 반복적 수치 적분이 필요합니다. 이 다단계 샘플링 과정은 상당한 계산 비용을 초래합니다.

최근 몇 단계 및 단일 단계 생성에 대한 연구들이 있었으나, 기존 접근법들은 다음과 같은 한계에 직면해 있습니다:

1. 증류 기반 방법들은 종종 복잡한 교사 - 학생 훈련 파이프라인을 요구합니다.
2. 일관성 모델과 평균 흐름(Mean Flow, 평균 속도를 학습함) 은 주로 특정 공식화 (예: $u$-예측 흐름 매칭) 를 위해 설계되었으며, DDIM, EDM, 또는 $x_1$-예측 흐름 매칭과 같은 컴퓨터 그래픽스에서 널리 사용되는 다른 프레임워크로 자연스럽게 일반화되지 않습니다.
3. 장거리 누적 지도를 직접 학습하는 것은 어렵습니다. 왜냐하면 표준 훈련 목적을 위한 조건부 대응물이 누적 장에는 존재하지 않기 때문에 근본적인 훈련 장벽이 발생하기 때문입니다.

방법론: 누적 흐름 지도 (CFM)
저자들은 명시적인 누적장 매개변수화를 통해 국소 순간 업데이트와 유한 시간 수송을 연결하는 통합 프레임워크인 **누적 흐름 지도 (Cumulative Flow Maps, CFM)**를 제안합니다.

• 통합 추상화: 이 방법은 $m_t(x)$를 순간장 (예: 속도) 으로 하는 추상 함수 $F[m_t(x), x, t, t+h]$를 사용하여 순간 흐름 지도 $\psi_{t \to t+h}$를 공식화합니다. 이 추상화는 다양한 매개변수화 ($u$-FM, $x_1$-FM, DDIM, EDM) 를 통합합니다.
• 누적 매개변수화: CFM 는 순간 지도를 구성하는 극한으로 누적 흐름 지도 $\psi_{t \to r}$를 정의함으로써 장거리 수송으로 확장합니다. 모델은 $\psi_{t \to r}(x) = F[m_{t \to r}(x), x, t, r]$이 되도록 누적 매개변수화 장 $m_{t \to r}(x)$를 학습합니다. 이를 통해 모델은 시간 $t$에서 임의의 미래 시간 $r$로 샘플링을 단일 또는 소수 단계로 진행할 수 있습니다.
• 장 방정식을 통한 훈련: 주요 과제는 자기 일관성 있게 존재하는 조건부 누적 장 $m_{t \to r}(x | X_1)$이 없다는 점입니다. 이를 극복하기 위해 저자들은 장 방정식 기반 재공식화(정리 3) 를 유도합니다. 그들은 순간 장과 그 도함수를 통해 누적 장을 표현합니다:
  $$m_{t \to r}(x) = G(t, r)m_t(x) + H(t, r)E[\dots]$$
  이 공식화는 학습 가능한 조건부 순간 장 $m_t(x | X_1)$을 사용하여 누적 장을 감독하는 대리 손실 함수 (식 7) 를 구성할 수 있게 합니다. 이 손실은 야코비안 - 벡터 곱 (JVP) 또는 이산 근사를 통해 계산되는 도함수 항 ($\partial_t m$ 및 $\partial_x m$) 을 포함합니다.
• 구현: 이 접근법은 최소한의 아키텍처 변경만 요구합니다. 목표 시간 $r$에 대한 추가 임베더로 시간 임베딩을 보강하고 평균화된 임베딩 $(emb_t + emb_r)/2$를 사용합니다. 증류나 용량 증가는 필요하지 않습니다.

주요 기여

1. 순간 - 누적 흐름 지도 추상화: 이 논문은 순간 지도를 구성하여 얻어진 유한 시간 수송으로서 누적 흐름 지도를 공식화하여, 다단계 및 소수 단계 생성을 단일 수학적 프레임워크 하에서 통합합니다.
2. $u$-예측을 넘어선 일반화: CFM 은 평균 흐름의 개념을 $u$-및 $x_1$-흐름 매칭, EDM, DDIM 을 포함한 광범위한 생성 공식화로 일반화합니다. 이는 기하 분포 모델링 및 픽셀 공간 이미지 생성과 같이 이전 방법 (예: 평균 흐름) 이 적용되지 않는 환경에서 소수 단계 생성을 가능하게 합니다.
3. 모델 중립적 훈련: 장 방정식 기반 목적 함수를 유도함으로써, 이 방법은 아키텍처 변경이나 증류 없이 누적 흐름 지도를 학습할 수 있게 하여 기존 모델의 샘플링 단계를 대폭 줄입니다.

실험 결과
저자들은 CFM 을 다섯 가지 다양한 그래픽 작업에서 평가하여 생성 품질을 유지하거나 향상시키면서 10 배~200 배의 속도 향상을 입증했습니다:

• 이미지 생성 (CelebA-HQ): CFM-DDIM 을 사용하여 128 단계 기준선과 비교 가능한 FID 점수로 1 단계 및 4 단계 생성을 달성했으며, 일관성 증류 및 훈련 방법보다 우수한 성능을 보였습니다.
• 기하 분포 모델링: GeoDist 작업에서 CFM-EDM 은 60 단계 기준선 대비 체임버 거리 (Chamfer Distance) 저하 없이 6 배~10 배의 속도 향상을 달성했습니다. 주목할 점은 $u$-예측 방법들은 이 분야에서 소수 단계 생성을 지원하지 못하는 반면, CFM-EDM 은 성공했다는 것입니다.
• 결합 위치 예측 (PDT): RigNet 데이터셋에 적용된 CFM-DDIM 은 결합 예측 정확도 (CD-J2J, IoU, 정밀도, 재현율) 를 원래 1000 단계 PDT 와 비교 가능하게 유지하면서 추론을 1000 단계에서 5 단계로 줄여 200 배의 속도 향상을 이루었습니다.
• 이미지 조건부 스케치 생성: ControlSketch 데이터셋에서 CFM 은 50 단계 SwiftSketch 기준선과 동일한 충실도 (MS-SSIM, DreamSim) 로 1 또는 4 단계에서 벡터화된 스케치를 생성하여 50 배의 속도 향상을 보였습니다.
• 3D SDF 생성: 희소 조건부 SDF 재구성 (64 점) 의 경우, CFM-$x_1$-FM 은 Functional Diffusion 대비 재구성 품질을 동등하게 유지하면서 6~16 배의 속도 향상을 달성했습니다.

의의 및 주장
이 논문은 CFM 이 이미지, 점 구름, 암시적 장 등 다양한 데이터 표현과 생성 공식화를 포함하는 컴퓨터 그래픽스 응용에 특히 적합한 일반적인 수학적 프레임워크를 제공한다고 주장합니다. 그 주요 의의는 다음과 같습니다:

• 효율성과 아키텍처의 분리: 증류나 아키텍처 개편 없이 훈련 목적 함수와 시간 임베딩만 수정하여 최대 200 배에 달하는 극적인 추론 가속을 달성합니다.
• 통합 프레임워크: 순간 동역학과 장거리 수송 간의 격차를 해소하여, 이전의 "평균 속도" 접근법이 실패했던 모델 전반에 걸쳐 소수 단계 생성을 위한 원칙적인 해결책을 제공합니다.
• 실용성: 이 방법은 모델 중립적이며, 사전 훈련된 다단계 모델을 가속화하는 데 적용하거나 처음부터 새로운 소수 단계 모델을 훈련하는 데 사용할 수 있습니다.

저자들은 CFM 이 광범위한 모델에 적용 가능하지만, 현재 평가는 다섯 가지 응용과 네 가지 대표적 공식화로 제한되어 있으며, ImageNet 과 같은 대규모 데이터셋으로의 확장은 향후 과제로 남아 있다고 명시합니다.