플로우 맵: 더 빠른 샘플링을 위한 확산 모델의 적분 학습

플로우 맵은 확산 모델 샘플링을 가속화하는 새로운 접근 방식으로, 신경망이 노이즈에서 데이터로 가는 전체 경로를 한 번에 예측하도록 훈련하여 접선 방향을 반복적으로 적분하는 대신 한 번에 처리합니다. Sander Dieleman의 블로그 게시물은 이론과 분류 체계를 분석합니다.

현재 확산 샘플링 방식

표준 확산 모델은 각 노이즈 수준에서 접선 방향(잡음 제거 추정치)을 예측하고, 작은 단계를 밟아 노이즈에서 데이터로의 경로를 적분합니다. 이는 많은 반복(예: DDPM에서 50–1000단계)을 필요로 합니다. DDIM과 같은 결정론적 샘플러는 이 과정을 ODE로 취급하지만 여전히 수치 적분이 필요합니다.

플로우 맵의 차이점

플로우 맵은 노이즈와 데이터 사이의 경로 상의 임의의 점에서 다른 점으로 직접 예측하는 방법을 학습합니다. 한 점에서 접선 방향을 예측하는 대신, 네트워크는 전체 적분 결과(잡음이 있는 입력과 목표 노이즈 수준이 주어졌을 때 최종 깨끗한 샘플)를 출력합니다. 따라서 샘플링을 한 번 또는 매우 적은 단계로 수행할 수 있습니다.

훈련 접근법

Dieleman은 Boffi 등의 분류 체계를 기반으로 플로우 맵 훈련을 여러 전략으로 분류합니다:

• 일관성 모델(Consistency models): 네트워크 출력이 궤적을 따라 모든 노이즈 수준에서 동일하도록 강제합니다(자기 일관성).
• 직접 회귀(Direct regression): 네트워크가 임의의 노이즈 수준에서 잡음이 있는 샘플로부터 깨끗한 데이터를 직접 예측하도록 훈련하고, L2 손실을 최소화하여 실제 값과 비교합니다.
• 증류(Distillation): 사전 훈련된 확산 모델을 교사로 사용하여 (잡음, 깨끗함) 쌍을 생성한 후, 플로우 맵 학생이 교사의 ODE 적분을 모방하도록 훈련합니다.

선택은 ODE 솔버를 통해 역전파할지(비용이 많이 들지만 정확함) 아니면 이를 피할지(더 저렴하지만 잠재적으로 불안정함)에 따라 달라집니다.

실용적 고려 사항

Dieleman은 가우시안 확산의 경우 알려진 노이즈-데이터 매핑을 활용하여 플로우 맵을 처음부터(교사 없이) 훈련할 수 있다고 언급합니다. 실제로 플로우 맵은 다음을 제공합니다:

• 더 빠른 샘플링: 1–4단계 (증류된 확산의 10–50단계와 비교).
• 보상 기반 학습: 적분 형태 덕분에 하류 보상(예: 이미지 품질 메트릭)의 직접적인 그래디언트 기반 최적화가 가능합니다.
• 조종 가능성: 잠재 변수에 조건을 거는 것이 더 간단해집니다. 전체 생성 경로가 단일 함수이기 때문입니다.

응용 및 확장

플로우 맵은 이미지 생성(예: LCM, SDXL-Turbo 같은 일관성 모델), 비디오, 3D 데이터에 적용되고 있습니다. 게시물은 더 깊은 수학적 처리를 위해 Lai 등의 최신 논문을 참조합니다.