기술 요약: 단일 단계 디퓨전 설계를 위한 숏컷 경로(Shortcutting Flow Paths) 연구

문제 정의

디퓨전 기반 생성 모델은 딥 생성 모델링의 지배적인 패러다임이 되었으나, 일반적으로 수십 또는 수백 번의 신경 함수 평가(NFE)를 필요로 하는 느린 추론 속도라는 문제를 안고 있다. 컨시스턴시 모델(Consistency Models)과 증류(Distillation) 기술이 단일 단계 생성을 가능하게 했지만, 이는 대개 비용이 많이 드는 2단계 학습 과정(사전 학습된 디퓨전 모델 이후의 증류)에 의존한다. 최근에는 직접적인 매핑을 학습하기 위해 처음부터 훈련되는 "숏컷 모델"(예: Consistency Training, Shortcut Diffusion, Inductive Moment Matching)이 등장했으나, 이러한 모델들의 설계 공간은 여전히 불분명하다. 이론적 유도, 실제 구현, 그리고 특정 구성 요소(시간 샘플러, 플로우 경로, 손실 메트릭 등)의 선택은 서로 밀접하게 결합되어 있어, 체계적으로 개선점을 식별하거나 특정 모듈의 필요성을 이해하기 어렵다.

방법론

저자들은 대표적인 숏컷 모델(이산 시간 및 연속 시간 변형 모두 포함)의 구성 요소를 분리하기 위한 공통 설계 프레임워크를 제안한다. 이 프레임워크는 이러한 방법들을 2단계 플로우 맵 타겟을 1단계 파라미터화된 예측으로 근사한다는 원칙 아래 통합한다.

1. 통합 설계 프레임워크

본 논문은 플로우 맵 솔버(flow map solvers)를 사용하여 숏컷 모델을 정식화한다. 핵심 목표는 컨시스턴시 속성을 통해 구축된 2단계 타겟을 근사하는 1단계 플로우 맵 $X^\theta_{t,r}(x_t)$를 예측하도록 신경망 $F_\theta$를 훈련하는 것이다:
$$X^\theta_{t,r}(x_t) \approx \text{sg}(\hat{X}_{s,r} \circ \hat{X}_{t,s}(x_t))$$
여기서 $r \le s \le t$이며, $\text{sg}$는 스톱 그래디언트(stop-gradient) 연산자를 나타내고, 타겟은 조건부 속도 또는 중간 네트워크 예측을 사용하여 구축된다. 이 프레임워크는 기존 방법들(CT, SCD, IMM, sCT, MeanFlow)을 다음의 구체적인 선택 사항에 따라 분류한다:

• 플로우 경로 (Flow Paths): 선형(Linear) vs 코사인(Cosine) 경로.
• 시간 샘플러 (Time Samplers): 이산적(고정된 간격) vs 연속적(무한소 간격).
• 네트워크 파라미터화 (Network Parameterization): 순간 속도($v_\theta$) vs 평균 속도($u_\theta$).
• 손실 메트릭 (Loss Metrics): LPIPS, 제곱 $L_2$, 또는 커널 기반 메트릭.

2. 이론적 정당성

저자들은 데이터 분포와 생성된 분포 사이의 Wasserstein-2 거리(Wasserstein-2 distance)에 대한 이론적 경계치를 제공한다. 그들은 이산 시간(DTSC) 및 연속 시간(CTSC) 숏컷 모델의 훈련 목적 함수가 생성 오차에 대한 유효한 상계(upper bound) 역할을 한다는 것을 증명한다. 결정적으로, 저자들은 데이터 조건부 설정에서 처음부터 훈련되는 숏컷 모델의 경우 **선형 경로(linear paths)**가 코사인 경로에 비해 더 낮은 볼록 수송 비용(convex transport costs)과 감소된 궤적 곡률을 제공하여 최적임을 입증한다.

3. 제안된 개선 사항

분리된 프레임워크를 바탕으로, 본 논문은 연속 시간 숏컷 모델(특히 MeanFlow 변형)의 훈련을 향상시키기 위한 세 가지 핵심 기술적 개선 사항을 도입한다:

• 플ue-in Velocity (플러그인 속도): 미니 배치 샘플링으로 인해 분산이 높은 조건부 속도 $v_{t|0}$에만 의존하는 대신, 저자들은 플러그인 속도 $v^*_t$를 사용할 것을 제안한다. 이는 미니 배치 내의 조건부 속도들의 가중 평균으로 계산되며, 효과적으로 주변 속도(marginal velocity)를 근사한다. 이는 계산 비용을 거의 늘리지 않으면서 속도 감독(velocity supervision)의 분산을 $O(1/B)$(여기서 $B$는 배치 크기)만큼 줄여 훈련을 안정화한다.
• 클래스 일관적 미니 배치 (Class-Consistent Mini-batching): 플러그인 속도를 사용하면서도 Classifier-Free Guidance(CFG)의 효능을 유지하기 위해, 저자들은 미니 배치 내의 모든 샘로 동일한 클래스에 속하도록 샘플링할 것을 제안한다. 이는 평균화 과정 중에 클래스별 신호가 희석되는 것을 방지한다.
• 점진적 시간 샘플러 (Gradual Time Sampler): 훈련 초기 단계에서는 터미널 타임 $r$을 0으로 고정하여(디노이징 작업 모사) 표준 연속 시간 샘플링 전략으로 점진적으로 전환하는 커리큘럼 학습 전략을 도입한다. 이는 초기 수렴을 돕는다.

주요 결과

제안된 개선 사항은 SiT 및 U-Net 아키텍처를 사용하여 ImageNet-256$\times$256 및 CIFAR-10에서 평가되었다.

• 최첨단 성능 (State-of-the-Art Performance): ESC(Explicit & Easier Shortcut)라고 명명된 결과 모델은 Classifier-Free Guidance 하에서 ImageNet-256$\times$256에 대해 FID50k 2.85를 달성하며 단일 단계 생성을 수행한다. 훈련을 연장(480 에포크)하면 2.53에 도달한다.
• 비교: 이러한 결과는 처음부터 훈련된 숏컷 모델 중 새로운 SOTA를 기록하며, 기존의 단일 단계 방법인 MeanFlow(3.43) 및 SCD(10.60)를 크게 능가한다. 이 성능은 MeanFlow의 2단계 생성 충실도(2.93)마저 넘어선다.
• 스케일링 이점: 제안된 기술의 성능 향상은 작은 모델(SiT-B/2)보다 큰 모델(SiT-XL/2)에서 더 두드러지게 나타나며, 이는 분산 감소가 숏컷 모델을 확장하는 데 매우 중요하다는 것을 시사한다.
• 효율성: 플러그인 속도 계산은 반복 시간(iteration time)을 약 0.7% 증가시키는 미미한 오버헤드만을 가지면서도 상당한 안정성을 제공한다.

의의 및 주장

본 논문은 주요 기여가 숏컷 모델의 구성 요소 수준 혁신에 대한 장벽을 낮추는 것이라고 주장한다. 공통 설계 프레임워크와 이론적 정당성을 제공함으로써 저자들은 다음을 달성했다:

1. 설계 공간의 분리: 이전에 "필수적"이라고 여겨졌던 많은 모듈이 실제로는 독립적인 선택 사항임을 보여줌으로써 체계적인 탐색과 조합을 가능하게 했다.
2. 선형 경로 검증: 처음부터 훈련하는 숏컷 모델에 대해 선형 경로의 사용을 이론적으로 정당화하여, 특정 맥락에서 코사인 경로를 선호하던 기존 관행에 도전했다.
3. 원칙적인 탐색 가능: 제안된 개선 사항(플러그인 속도, 클래스 일관적 배치)은 사전 학습, 증류, 또는 복잡한 커리큘럼 학습 없이도 고충실도 단일 단계 생성을 위한 경로를 제공한다.

저자들은 본 연구가 결합된 이론적 유도와 실용적 트릭이라는 "블랙박스"를 넘어, 단일 단계 디퓨전 모델을 설계하기 위한 보다 체계적인 접근 방식을 촉진한다고 강조한다. 또한, 현재 CFG 하이퍼파라미터 튜닝의 유연성 부족과 트랜스포머 기반 아키텍처에서의 Jacobian-vector product(JVP) 계산 비용과 같은 한계를 인정하며, 이를 향후 연구 방향으로 제시한다.