📊 statistics

One-Step Generative Modeling via Wasserstein Gradient Flows

본 논문은 1.29 FID 로 ImageNet 생성에서 최첨단 성능을 달성하고 다단계 확산 모델보다 약 100 배 빠른 샘플링을 실현하기 위해 와세르슈타인 기울기 흐름을 단일 신경망 추론으로 압축하는 원스텝 생성 모델링 프레임워크인 W-Flow 를 소개합니다.

원저자: Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

게시일 2026-05-13

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

로봇에게 완벽한 고양이 그림을 그리도록 가르치고 싶다고 상상해 보세요.

기존 방식 (확산 모델):
대부분의 현대 AI 화가는 돌덩이를 조각해 나가는 조각가처럼, 혹은 흐릿한 사진을 찍어 서서히 선명하게 만드는 사진가처럼 작동합니다. 이들은 무작위 잡음에서 시작해 수백 개의 미세한 단계를 거쳐 그 잡음을 점차 고양이로 변환합니다.

문제점: 마치 상점까지 곧장 걸어가는 대신 100 개의 작고 망설이는 걸음으로 이동하는 것과 같습니다. 단일 이미지를 얻기 위해 많은 시간과 에너지 (컴퓨팅 파워) 가 소요됩니다.

새로운 방식 (W-Flow):
이 논문은 W-Flow를 소개합니다. 이는 로봇이 단 하나의 거대한 도약으로 완벽한 고양이를 그리도록 가르치는 방법입니다. 100 단계를 거치는 대신, '잡음'에서 '고양이'로 가는 완벽한 지름길을 즉시 학습합니다.

작동 원리: '강' 비유

W-Flow 가 이 지름길을 어떻게 학습하는지 이해하기 위해 두 그룹의 사람들을 상상해 보세요:

군중 (목표): 완벽한 원형으로 서 있는 사람들 (실제 데이터, 예를 들어 실제 고양이 사진).
떠돌이들 (생성기): 들판에 무작위로 흩어져 서 있는 사람들 (AI 의 현재 엉망인 추측).

목표는 떠돌이들이 군중과 동일한 완벽한 원을 형성하도록 이동시키는 것이지만, AI 는 이를 한 번에 수행할 수 있는 규칙을 학습해야 합니다.

1. '에너지' 지도 (경사)
저자들은 떠돌이들과 군중 사이의 공간을 언덕이 있는 풍경으로 상상합니다. '군중'은 계곡의 가장 아래 (최저 에너지 지점) 에 있고, 떠돌이들은 언덕 위 어딘가에 있습니다.

규칙: 만약 당신이 떠돌이라면, 군중에게 도달하기 위해 가능한 한 빠르게 언덕을 굴러 내려가고 싶을 것입니다.
혁신: 이전 방법들은 어느 방향이 아래인지 파악하기 위해 '휴리스틱 (추측)'을 사용했습니다. 때로는 추측을 잘못해 길을 잃거나, 잘못된 방향으로 사람들을 너무 강하게 밀어붙이기도 했습니다.
W-Flow 의 비법: 그들은 Sinkhorn Divergence라는 수학 도구를 사용합니다. 이는 마치 전체 그룹의 움직임을 개별이 아닌 통합적으로 고려하여, 떠돌이 그룹에 속한 모든 사람에게 언덕을 내려가는 정확한 가장 가파른 경로를 계산하는 초정밀 GPS 와 같습니다.

2. '이중 배치' 안전망
떠돌이들이 어떻게 이동해야 하는지 계산할 때 까다로운 문제가 있습니다: 만약 한 사람에게 자신의 그룹에서 멀어지라고 요청하면, 그들은 실수로 스스로로부터 멀어지려 할 수 있으며, 이는 논리적으로 맞지 않습니다.

해결책: 논문은 'Two-Batch (이중 배치)' 전략이라는 교묘한 트릭을 사용합니다. 떠돌이들을 두 개의 별도 줄로 나누어 상상해 보세요. A 줄이 어떻게 이동해야 하는지 B 줄을 기준으로 계산하고, 그 반대로도 계산합니다. 이는 그들이 자신의 반사에 혼동되는 것을 방지하고, 멈추지 않고 목표 방향으로 부드럽게 이동하도록 보장합니다.

3. 여정 압축
여기가 마법의 단계입니다:

먼저, AI 는 이 '언덕을 굴러 내려가는' 과정을 여러 번 시뮬레이션합니다 (떠돌이들이 서서히 원을 형성하는 영화를 보는 것과 같습니다).
그런 다음, 신경망 (생성기) 이 그 전체 영화를 기억하도록 훈련시킵니다.
결과: 훈련이 완료되면, 네트워크는 더 이상 영화를 볼 필요가 없습니다. 시작과 끝을 알고 있으므로 '무작위 잡음'에서 '완벽한 고양이'로 바로 점프할 수 있습니다.

이것이 중요한 이유는 무엇인가?

속도: 이 논문은 이 방법이 기존 다단계 방식보다 약 100 배 빠르다고 주장합니다. 기존 방식이 이미지를 만드는 데 10 초가 걸렸다면, 이 방식은 10 분의 1 초도 걸리지 않습니다.
품질: 단 한 단계임에도 불구하고 이미지는 놀라울 정도로 고품질입니다. 유명한 ImageNet 테스트에서 그들은 1.29의 점수 (FID) 를 기록했는데, 이는 단일 단계 생성기로서 새로운 기록입니다. 이는 사진이 실제 사진과 거의 구별할 수 없음을 의미합니다.
안정성: 추측이 아닌 견고한 원리 (Wasserstein Gradient Flows) 에 기반한 수학이기 때문에, AI 가 '붕괴' (한 가지 유형의 고양이만 그리고 나머지는 무시하는 현상) 할 가능성이 낮습니다. 데이터의 모든 다른 '모드'를 더 잘 포괄합니다.

요약

W-Flow를 수학 문제를 풀도록 학생을 가르치는 것으로 생각해 보세요.

기존 방법: 선생님이 학생에게 단계별로 해답을 보여주고, 학생이 그 단계를 반복해서 연습하여 수행할 수 있게 될 때까지 기다립니다.
W-Flow: 선생님이 학생에게 해답의 논리 (경사 흐름) 를 보여주고, 그 논리를 연습하게 한 다음, 즉시 최종 답안을 쓰도록 요구합니다. 학생은 그 '지름길'을 너무 잘 학습하여 더 이상 풀이 과정을 보여줄 필요가 없습니다.

이 논문은 이 특정 수학 나침반 (Sinkhorn divergence) 을 사용하여 훈련을 안내함으로써, 매우 빠르고 매우 정확한 생성기를 구축할 수 있음을 증명합니다.

기술 요약: Wasserstein 기울기 흐름을 통한 1 단계 생성 모델링

문제 정의

현대적인 생성 모델, 특히 확산 모델 (diffusion models) 과 흐름 기반 방법 (flow-based methods) 은 최첨단 샘플 품질을 달성했으나, 다단계 반복적 샘플링에 대한 의존성으로 인해 높은 계산 비용과 지연 시간을 겪습니다. 1 단계 또는 소수 단계 증류 (distillation) 방법이 존재하지만, 이들은 종종 반복적 생성을 위해 원래 설계된 경로의 학습 신호를 계승하거나 수렴 보장이 없는 휴리스틱 동역학 (예: Drifting Models) 에 의존하여 모드 붕괴 (mode collapse) 에 취약합니다. 핵심적인 과제는 생성기가 참조 분포에서 목표 데이터 분포로 직접적인 1 단계 매핑을 학습하면서도 안정적인 수렴과 전역 모드 커버리지를 보장하는 학습 동역학을 설계하는 것입니다.

방법론: W-Flow

저자들은 W-Flow를 제안합니다. 이는 학습 동역학의 설계와 생성기의 파라미터화를 분리하는 프레임워크입니다. 핵심 아이디어는 **Wasserstein 기울기 흐름 (Wasserstein Gradient Flows, WGF)**을 사용하여 모델 분포가 목표 분포로 진화하도록 규정한 후, 정적 신경망을 훈련하여 이 다단계 진화를 단일 추론 단계로 압축하는 것입니다.

1. Wasserstein 기울기 흐름을 통한 학습 동역학

W-Flow 는 직접적으로 노이즈 제거 단계 시퀀스나 속도장을 학습하는 대신, 학습 궤적을 확률 측도 공간에서의 연속 시간 진화로 정의합니다.

에너지 함수 (Energy Functional): 진화는 에너지 함수 $F(q) = D(q \| p)$ 의 가장 가파른 하강에 의해 주도됩니다. 여기서 $q$ 는 현재 모델 분포이고 $p$ 는 목표 데이터 분포입니다.
속도장 (Velocity Field): 연속 시간 속도장 $V_t$ 는 에너지 함수의 1 차 변분에서 유도됩니다: $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$ .
이산화 (Discretization): 연속 흐름은 명시적 오일러 방식 (explicit Euler scheme) 을 사용하여 이산화됩니다. 각 학습 단계 $k$ 에서, 모델 분포 $q^{(k)}$ 는 로컬 운송 맵 $g^{(k)}(x) = x + \eta V^{(k)}(x)$ 를 통해 업데이트되며, 여기서 $\eta$ 는 단계 크기입니다.

2. Sinkhorn 발산을 통한 구체화

속도장이 고차원에서 계산적으로 다루기 쉽고 기하학적으로 잘 행동하도록 보장하기 위해, 저자들은 Sinkhorn 발산( $S_\epsilon$ )을 사용하여 에너지 함수를 구체화합니다. 이는 엔트로피 정규화 최적 운송 (OT) 지표입니다.

속도장 공식화: 유도된 속도장은 두 개의 바리센터 투영 (barycentric projections) 의 차이로 주어집니다:
$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$
여기서 $T^\epsilon_{q,p}$ 는 생성된 분포 $q$ 에서 목표 $p$ 로 질량을 운송하고, $T^\epsilon_{q,q}$ 는 현재 분포와의 일관성을 강제하는 자체 운송 항을 나타냅니다.
전역 조정 (Global Coordination): 국소 커널 상호작용 (예: MMD 또는 Drifting Models) 과 달리, Sinkhorn 속도장은 전역 OT 계획에서 유도되므로, 입자 이동이 전체 배치에 걸쳐 조정되어 모드 붕괴를 방지합니다.
이중 배치 추정기 (Two-Batch Estimator): 입자가 0 비용으로 자기 자신과 매칭되는 자기 매칭 아티팩트를 방지하기 위해, 저자들은 이중 배치 전략을 사용합니다. 단일 배치의 비용 행렬 대각선을 마스킹하는 대신, 현재 분포에서 추출한 독립적인 두 번째 배치 샘플을 사용하여 자체 운송 항 $T^\epsilon_{q,q}$ 를 계산합니다.

3. 생성기 훈련

생성기 $f_\theta$ 는 WGF 로 정의된 이산 입자 업데이트를 모방하도록 훈련됩니다.

목적 함수: 손실 함수는 생성기 출력과 Sinkhorn 업데이트에서 유도된 "목표" 입자 위치 사이의 제곱 거리를 최소화합니다:
$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$
여기서 $\text{sg}(\cdot)$ 는 stop-gradient 연산자를 나타내며, 업데이트된 입자 위치를 고정된 회귀 목표로 취급합니다.
클래스리스 가이드 (Classifier-Free Guidance, CFG): 프레임워크는 속도장을 직접 수정함으로써 CFG 를 지원합니다. 저자들은 분포를 선형적으로 혼합하는 대신, 조건부 및 무조건부 속도장 사이의 차이에 비례하는 항을 추가하는 **속도 가이드 (velocity guidance)**를 제안합니다. 이 접근법은 KL 발산을 사용할 때 표준 지수적으로 기울어진 목표 분포를 복원하고, Sinkhorn 발산에 대해 기하학적으로 일관된 확장을 제공하는 것으로 입증되었습니다.

주요 기여

원칙적인 1 단계 프레임워크: W-Flow 는 Drifting Models 과 같은 휴리스틱 접근법에서 결여된 명확한 변분 해석을 제공하며, Wasserstein 기울기 흐름에 학습 동역학을 근거로 둔 1 단계 생성 모델링을 위한 새로운 패러다임을 확립합니다.
Sinkhorn 기반 동역학: Sinkhorn 발산을 통한 WGF 의 구체화는 효율적이고 최적 운송 기반의 업데이트 규칙을 제공하며, 커널 기반 또는 국소 상호작용 방법에 비해 전역 분포 차이를 포착하고 모드 커버리지를 개선합니다.
이론적 수렴: 저자들은 적절한 정규성 가정 하에서 W-Flow 의 유한 표본 입자 동역학이 연속 시간 분포 동역학으로 수렴함을 증명하여 학습 과정에 대한 이론적 근거를 제공합니다.
이중 배치 추정기: 단일 배치 추정기 및 대각선 마스킹 휴리스틱에 내재된 자기 매칭 편향을 제거하는 새로운 이중 배치 전략이 도입되었습니다.

실험 결과

저자들은 ImageNet 256×256과 FFHQ에서 W-Flow 를 평가했습니다.

최첨단 성능: W-Flow 는 ImageNet 에서 1 단계 클래스 조건부 생성에 대해 새로운 최첨단 FID 인 1.29(XL 규모)와 1.35(L 규모)를 달성하여, 이전 1 단계 방법들 (예: Drifting Models, MeanFlow) 을 크게 능가하고 단일 함수 평가 (NFE) 만 사용함에도 불구하고 많은 다단계 확산 모델 (예: LightningDiT, SiT) 을 능가합니다.
샘플링 효율성: 유사한 FID 점수를 가진 다단계 확산 모델과 비교할 때, W-Flow 는 약 100 배 빠른 샘플링을 제공합니다 (예: 단일 H100 GPU 에서 SiT-XL/2 대비 초당 0.93 장 대비 77.88 장).
모드 커버리지: 불균형 분포 (예: 95% 노년층 얼굴, 5% 아동 얼굴) 를 사용한 실험에서 W-Flow 는 소수 모드를 성공적으로 포착하는 반면, Drifting Models 는 소수 클래스를 완전히 무시하는 모드 붕괴를 보입니다.
수렴 속도: W-Flow 는 훈련 중 훨씬 빠르게 수렴하여, 1280 에포크 동안 훈련된 Drifting Model 의 성능을 단 384 에포크 만에 달성합니다.
도메인 전이: 이 방법은 도메인 전이 작업 (예: FFHQ 의 나이 변환) 에서 강력한 능력을 보여주며, 휴리스틱 베이스라인에 비해 더 짧고 일관된 운송 궤적을 생성합니다.

중요성과 주장

이 논문은 Wasserstein 기울기 흐름이 빠르고 고품질의 생성 모델링을 위한 원칙적이고 효과적인 기반을 제공한다고 주장합니다. 학습 동역학의 설계와 생성기 파라미터화를 분리함으로써, W-Flow 는 복잡한 다단계 반복 과정과 비교 가능한 성능을 달성하도록 훈련된 단일 단계 생성기가 가능함을 보여줍니다. 저자들은 최적 운송 동역학이 제공하는 전역 조정이 훈련을 안정화하고 포괄적인 모드 커버리지를 보장하는 핵심 요소이며, 기존 1 단계 접근법의 중요한 한계를 해결한다고 강조합니다. 이러한 결과는 W-Flow 가 실시간 창의성 도구 및 시뮬레이션과 같이 저지연 고품질 생성이 필요한 응용 분야를 위한 확장 가능한 대안이 될 수 있음을 시사합니다.