기술 요약: 정류 흐름을 위한 발산 억제 결합

문제 정의

정류 흐름 (RectFlow) 은 직선 궤적을 가진 속도장을 학습함으로써 생성 모델링을 가속화하며, 적거나 심지어 단일 수치 적분 단계 (NFE) 로도 고품질 샘플 생성을 가능하게 합니다. 이 방법은 반복적인 "재흐름 (reflow)" 절차에 의존합니다. 독립 결합으로 훈련된 기본 모델이 새로운 소스 - 타겟 쌍을 생성하면, 이를 사용하여 더 직선적인 속도장을 가진 모델을 재훈련합니다.

그러나 해당 논문은 근본적인 병목 현상인 결합 왜곡을 식별합니다. 실제로 기본 모델이 생성한 궤적은 종종 구부러지고 서로 얽힙니다. 이러한 왜곡은 학습된 속도장의 비영 (nonzero) 발산 영역과 자주 연관됩니다. 이러한 영역에서 국소적인 팽창 또는 수축은 궤적이 벌어지거나 압축되게 하여 입자들을 이상적인 종착점에서 벗어나게 합니다. 결과적으로 재훈련된 모델은 이러한 오류를 물려받으며, 추가적인 정류 라운드는 모델이 이미 구부러진 수송을 "재흐름"하고 있기 때문에 체감 효과가 감소합니다. 결합 쌍을 생성하는 데 사용되는 전방 적분기의 품질이 훈련 손실이나 모델 용량보다 RectFlow 의 효과성을 제한합니다.

방법론: DS-RectFlow

저자들은 재흐름 과정에서 훈련 목적 함수, 모델 아키텍처, 또는 추론 절차를 변경하지 않고 더 깨끗한 결합 쌍을 생성하도록 설계된 오프라인 교정 방법인 DS-RectFlow(발산 억제 정류 흐름) 를 제안합니다.

핵심 메커니즘

이 방법은 궤적의 구부러짐 (결합 왜곡) 이 속도장의 압축성 구성 요소에 의해 주도된다는 관찰에 기반합니다. 헬름홀츠 분해를 사용하여 속도장 $v$는 발산이 없는 수송 부분 $u$와 장의 모든 압축성을 운반하는 쌍극자 부분 $\nabla \phi$로 분할됩니다. 논문은 발산 $|\nabla \cdot v|$의 크기를 억제함으로써 궤적을 직선으로 유지할 수 있다고 주장합니다.

발산 억제 적분기

결합 쌍의 오프라인 생성 (즉, "재흐름" 단계) 동안 DS-RectFlow 는 표준 오일러 적분기를 발산 억제 오일러 적분기로 대체합니다:

1. 상태 교정: 각 적분 단계에서 현재 상태 $x_t$에서 직접 진행하는 대신, 추정된 발산 $|\hat{\nabla} \cdot v_\theta(x^*_t, t)|$이 최소화되는 인접한 상태 $x^*_t$를 탐색합니다.
2. 허치슨 추정: 고차원에서 정확한 발산 계산의 높은 비용을 피하기 위해, 이 방법은 $\nabla \cdot v$를 근사하기 위해 허치슨의 대각합 추정기를 사용합니다. 이는 현재 상태 주변의 $m$개의 가우스 섭동을 생성하고 추정된 발산 크기를 최소화하는 후보를 선택합니다.
3. 적분: 그런 다음 원래 수정되지 않은 속도장을 사용하여 이 교정된 상태 $x^*_t$에서 오일러 단계를 계산합니다: $x_{t+\Delta t} = x^*_t + v_\theta(x^*_t, t) \cdot \Delta t$.

구현 세부 사항

• 오프라인 전용: 교정은 재훈련을 위한 결합 데이터셋 생성 동안에만 적용됩니다.
• 추론: 최종 훈련된 모델은 추가적인 계산 오버헤드 없이 일반 오일러 적분을 사용하여 배포됩니다. 발산 억제 로직은 추론 시점에 존재하지 않습니다.
• 예산 제어: 교정은 시간 $t_{stop}$(일반적으로 0.5) 까지 적용되며, 검색 반경 매개변수 $\delta$에 의해 통제됩니다.

주요 기여

1. RectFlow 한계의 진단: 이 논문은 기본 모델이 생성한 결합의 품질, 특히 궤적의 압축/팽창을 유발하는 비영 발산의 존재가 RectFlow 가 단일 단계 생성이라는 이론적 잠재력을 달성하는 것을 방해하는 주요 병목 현상임을 식별합니다.
2. 오프라인 교정 전략: 결합 생성 동안 속도장의 발산 성분을 감쇠시키는 새로운 오프라인 교정을 도입합니다. 추론 시 개선 (예: 고차 솔버 또는 가이드) 과는 달리, 이 방법은 교정 비용을 완전히 훈련 데이터 생성 단계로 이동시킵니다.
3. 제로 추론 오버헤드: 이 방법은 배포 중 표준 RectFlow 와 동일한 월클럭 비용을 갖는 모델을 산출하며, 추론 경로는 간단한 오일러 단계로 유지됩니다.
4. 복합 효과의 복원: 더 깨끗한 결합 쌍을 제공함으로써 DS-RectFlow 는 정류가 약속하지만 실제로는 드물게 달성하는 "복합" 혜택을 복원합니다. 이는 각 재흐름 라운드가 필드의 직선성을 크게 개선하는 현상으로, 표준 RectFlow 에서는 누적 왜곡으로 인해 종종 손실됩니다.

실험 결과

저자들은 2D 합성 벤치마크 및 이미지 생성 작업 (CIFAR-10, CelebA-64) 에서 DS-RectFlow 를 평가했습니다.

• 2D 벤치마크 (체크보드 및 GMM):

  • 체크보드 벤치마크에서 DS-RectFlow 는 바닐라 RectFlow 에 비해 슬라이스된 와슈터스 거리 (SWD) 를 크게 감소시킵니다. $NFE=1$에서 SWD 는 0.102(바닐라의 0.165 대비) 를 달성하여, 단일 단계로 20 단계 흐름 매칭 기준선의 성능과 효과적으로 일치합니다.
  • GMM 교차 벤치마크에서 DS-RectFlow 는 초기에 복잡한 모드 할당 결정을 더 충실히 포착함에 따라 낮은 NFE 에서 더 높은 오차를 보이지만, 빠르게 우수한 성능으로 수렴합니다. 두 번째 재흐름 라운드 ($k=2$) 까지 $NFE=1$에서 바닐라 RectFlow 대비 SWD 가 3 배 개선됩니다.
  • 시각화는 DS-RectFlow 가 "누출"된 샘플이 금지 영역으로 들어가는 것이 적은 더 선명하고 분리된 모드를 생성함을 확인시켜 줍니다.

• 이미지 생성 (CIFAR-10 및 CelebA-64):

  • CIFAR-10: DS-RectFlow 는 모든 NFE 예산에 걸쳐 일관되게 더 낮은 프리chet 시작 거리 (FID) 를 달성합니다. 오일러 -20 오프라인 솔버를 사용하면 $NFE=1$에서 FID 가 17.41 에서 12.03 으로 감소합니다 (상대적 개선 31%). 주목할 점은 $NFE=1$ DS-RectFlow 모델이 스윕 내의 모든 NFE 단계에서 바닐라 모델보다 우수하다는 것입니다.
  • CelebA-64: 유사한 경향이 관찰됩니다. $NFE=2$에서 DS-RectFlow 는 FID 4.90 을 달성하여, 모든 NFE 에서의 바닐라 모델의 최고 성능 ( $NFE=20$ 에서 5.64) 을 능가합니다.
  • 이러한 결과는 오프라인 교정이 고품질 단일 단계 생성이 가능한 모델로 직접 이전됨을 보여줍니다.

중요성 및 주장

이 논문은 결합 품질이 훈련 용량이나 추론 예산이 아닌 RectFlow 의 주된 병목 현상이라고 주장합니다. 발산 유발 왜곡을 오프라인으로 해결함으로써 DS-RectFlow 는 정류가 약속하지만 실제로는 드물게 제공하는 복합적 개선을 가능하게 합니다.

저자들은 이 접근 방식이 실용적인 이점을 제공한다고 강조합니다. 즉, 발산이 교정된 적분기의 이점 (더 직선적인 궤적, 더 나은 모드 정렬) 을 제공하면서도 표준 RectFlow 의 제로 지연 추론 프로파일을 유지합니다. 교정은 데이터 생성 동안 지불되는 일회성 비용으로, 전체 훈련 실행에 걸쳐 상각됩니다. 논문은 이 메커니즘이 발산 드리프트 성분이 샘플 품질을 저하시키는 모든 곳에서 확률적 보간 및 최적 수송 흐름 매칭과 같은 다른 흐름 기반 방법에 적용 가능할 수 있다고 제안합니다.