기술 요약: x-Prediction Is All You Need

문제 정의

확산(Diffusion) 및 플로우 매칭(Flow matching) 모델은 생성 작업에서 최첨단 성능을 확립했으나, 샘플당 일반적으로 수십에서 수백 번의 신경 함수 평가(NFE)를 필요로 하는 높은 추론 비용 문제를 겪고 있습니다. 궤적 정류(Rectified Flow), 증류(Progressive, Consistency-based, Bootstrapping), 그리고 고급 ODE 솔버와 같은 기존의 가속 전략들은 일반적으로 모델의 훈련 레시피, 아키텍처 또는 궤적 자체에 대한 수정을 요구합니다. 결과적으로, 이러한 방법들은 재학습 비용 없이 기존의 방대한 사전 학습된 체크포인트 라이브러리(예: SDXL, SD3.5M, FLUX)와 호환되지 않는 경우가 많습니다. 현재 이 분야에는 모델을 "있는 그대로" 사용하는 학습 불필요한 가속 전략이 부족합니다.

방법론: 절단 점프 샘플링 (Truncated Jump Sampling, TJS)

본 논문은 종단 디코더 가능성(Endpoint Decodability) 개념에 기반한 학습 불필요 추론 전략인 **절단 점프 샘플링(TJS)**을 제안합니다.

핵심 통찰: 종단 디코더 가능성 (Endpoint Decodability)

저자들은 임의의 비퇴화 아핀 확률 경로(diffusion 및 flow matching 모델의 표준 구조인 $x_t = \alpha_t x_0 + \sigma_t \epsilon$)에 대해, 임의의 중간 상태 $x_t$와 그 경로 속도 $u_t$로부터 클린 샘플 $x_0$를 폐쇄형 디코더(closed-form decoder)를 통해 유일하게 복구할 수 있음을 공식화합니다.

• 대수적 복구: 위치와 속도의 결합된 선형 시스템이 주어졌을 때, $x_0$는 $x_0 = (\sigma_t u_t - \dot{\sigma}_t x_t) / \Delta_t$ (여기서 $\Delta_t$는 경로 행렬식)로 복구됩니다.
• 최적성: 표준 $\ell_2$ 훈련 목적 함수 하에서, 이 대수적 디코더는 최소 평균 제곱 오차(MMSE) 추정치인 $E[x_0 | x_t]$에 대응합니다. 이는 표준적인 사전 학습된 모델들이 비록 노이즈나 속도를 예측하도록 훈련되었더라도, 모든 타임스텝에서 종단점 $x_0$를 예측하도록 암묵적으로 학습한다는 것을 의미합니다.
• 곡률 독립성: 결정적으로, 이 종단점 예측의 오차는 모델의 추정 정확도와 상태 $x_t$의 기약 불확실성에만 의존하며, 궤적의 곡률($\|\ddot{x}_t\|$)에는 의존하지 않습니다. 이는 정확한 소수 단계 생성을 위해 직선 궤적(Rectified Flow와 같은)이 필수적이라는 기존의 가설에 도전합니다.

알고리즘

TJS는 이 통찰력을 활용하여 ODE 적분을 조기에 중단하고 $t=0$까지 계속 진행하는 대신, 즉시 디코딩된 $x_0$를 출력함으로써 작동합니다.

1. 부분 적분: 모델은 전체 단계의 일 fraction $\gamma$만큼 ODE를 적분합니다 ($k^* = \lceil \gamma K \rceil$).
2. 조기 종료: 단계 $k^*$에서, 알고리즘은 학습된 속도 필드와 폐쇄형 디코더를 사용하여 종단점 추정치 $\hat{x}_0$를 계산합니다.
3. 종료: 프로세스를 중단하고 $\hat{x}_0$를 최종 샘플로 반환합니다.

이 접근 방식은 재학습, 증류, 아키텍처 변경이 전혀 필요하지 않습니다. 이는 모든 ODE 솔버, 노이즈 스케줄, 분류기 없는 가이드(CFG) 스케일 또는 기존 체크포인트와 호환됩니다.

주요 기여

1. 이론적 공식화: 본 논문은 종단 디코더 가능성(Theorem 5)을 아핀 경로의 보편적 속성으로 확립합니다. 저자들은 유도된 디코더가 MMSE 최적 추정치임을 증명하고(Theorem 6), 오차 분해가 궤적 곡률과 독립적임을 증명합니다(Theorem 9).
2. 궤적 직선화에 대한 도전: 저자들은 직선 궤적이 정확한 종단점 예측을 위해 충분조건이지만 필요조건은 아님을 증명합니다(Proposition 12). 이는 추론 가속을 궤적 정류의 필요성과 분리하며, Euler 적분이 실패하는 곡선 경로에서도 TJS가 작동하는 이론적 근거를 제공합니다.
3. 학습 불필요 알고리즘 (TJS): 모델이나 훈련 데이터를 수정하지 않고 다양한 모델 패밀리에서 NFE를 20~70% 감소시키는 실용적인 샘플러입니다.
4. 포괄적인 실증적 검증: 6개의 모델 패밀리(SDXL, SD3.5M, Z-Image-Turbo, ImageNet-256, CIFAR-10, MNIST)에 대한 광범한 실험을 통해 TJS가 전체 ODE 궤적과 유사하거나 더 우수한 품질을 달성함을 입증했습니다.

실험 결과

저자들은 클래스 조건부 및 텍스트-이미지 벤치마크 전반에 걸쳐 TJS를 평가했습니다:

• NFE 감소: TJS는 거의 대등한 품질을 유지하면서 NFE를 20~70% 감소시킵니다.
  • SDXL/SD3.5M: 약 37~40%의 NFE 절감(30단계 중 $k^* \approx 18$에서 중단) 시 Full ODE 품질의 95%(ImageReward 기준)를 달성합니다.
  • Z-Image-Turbo: 증류된 10단계 모델에서 TJS는 70%의 추가 절감(총 3 NFE, 즉 $k^*=2$에서 중단)을 달 통해 95%의 품질 유지를 보여주며, 종단 디코더 가능성이 증류된 궤적에서도 지속됨을 입증합니다.
  • 클래스 조건부: MNIST 및 CIFAR-10에서 TJS는 때때로 Full ODE보다 뛰어난 성능을 보이는데("오버슈트" 현상), 이는 Euler 적분의 마지막 단계에서 축적되는 이산화 오차를 피하기 때문입니다.
• 메트릭 계층 구조: 실험은 시맨틱 콘텐츠(CLIP, PickScore)는 조기에 해결되는 반면(종종 $k^* \approx 12$까지), 미세한 미적 디테일(ImageReward, AES)은 더 깊은 적분이 필요하다는 수렴 계층 구조를 드러냅니다.
• 강건성: TJS의 성능은 ODE 솔버(DDIM, DPM++, UniPC 등), 노이즈 스케줄(Beta, Karras 등), CFG 스케일의 선택에 관계없이 불변하며, 이는 기존 추론 구성 요소들과 직교함을 확인시켜 줍니다.

의의 및 주장

본 논문은 적당한 수준의 소수 단계 가속은 궤적 직선화, 증류 또는 훈련 재설계를 요구하지 않는다고 주장합니다. 그 의의는 종단 디코더 가능성을 아핀 경로 상의 생성 모델이 가진 고유한 구조적 속성으로 인식했다는 데 있습니다.

• 실질적 영향: TJS는 재학습이나 증류가 불가능한 경우가 많은 커뮤니티의 기존 미세 조정된 체크포인트 라이브러리에 대해 즉각적이고 비용이 들지 않는 가속 경로를 제공합니다.
• 이론적 전환: 궤적을 직선화하는 것이 종단점 예측을 위한 필수 조건이 아님을 증명함으로써, 이 연구는 추론 가속을 위한 Rectified Flow 및 Consistency Models의 근본적인 동기에 도전합니다.
• 보완적 성격: 저자들은 TJS를 증류와 직교하는 것으로 규정합니다. 증류가 전체 궤적 길이를 압축한다면, TJS는 남은 궤적의 불필요한 꼬리 부분을 제거하여, Z-Image-Turbo와 같은 증류된 모델에 TJS를 적용하는 식의 가산적 가속(additive acceleration)을 가능하게 합니다.

결론적으로, 이 논문은 이러한 속성을 인식하는 데는 "비용이 들지 않으며", 이를 무시하는 것은 성능을 낭비하는 것이라고 말하며, 생성된 샘플의 품질을 희생하지 않으면서 적분 깊이와 추론 속도 사이의 트레이드오프를 수행하는 원칙적인 방법을 제시합니다.