정류 흐름의 저차원 적응: 확산 및 확률적 지역화 관점

Low-Dimensional Adaptation of Rectified Flow: A Diffusion and Stochastic Localization Perspective

최근 몇 년 동안, 정류 흐름(Rectified flow, RF)은 뛰어난 생성 효율성과 최첨단 성능 덕분에 큰 인기를 얻고 있습니다. 본 논문에서는 RF가 표적 분포의 고유한 저차원 특성에 어떻게 자동으로 적응하여 샘플링 속도를 향상시키는지 조사합니다. 신중하게 설계된 시간 이산화 기법과 충분히 정확한 드리프트 추정치를 사용하면, RF 샘플러는 총 변동 거리에서의 정밀도 $\varepsilon$와 표적 분포의 고유 차원 $k$에 대해 $O(k/\varepsilon)$ (로그 인자 제외)의 반복 복잡도를 갖습니다. 또한, 노이즈 제거 확산 확률 모델(DDPM) 절차가 확률적 지역화를 통한 RF의 일종과 동등하다는 것을 새로운 연결 고리를 통해 보여줍니다. 이러한 연결을 바탕으로, 본 논문에서는 표적 분포의 저차원 특성에 더 완화된 드리프트 추정 정확도 및 특정 시간 계획 하에서 적응하는 확률적 RF 샘플러를 추가로 설계합니다. 제안된 샘플러는 새로 설계된 시간 이산화 기법을 구현하여 합성 데이터 및 텍스트-이미지 데이터 실험을 통해 성능 향상을 보여줍니다.

In recent years, Rectified flow (RF) has gained considerable popularity largely due to its generation efficiency and state-of-the-art performance. In this paper, we investigate the degree to which RF automatically adapts to the intrinsic low dimensionality of the support of the target distribution to accelerate sampling. We show that, using a carefully designed choice of the time-discretization scheme and with sufficiently accurate drift estimates, the RF sampler enjoys an iteration complexity of order $O(k/\varepsilon)$ (up to log factors), where $\varepsilon$ is the precision in total variation distance and $k$ is the intrinsic dimension of the target distribution. In addition, we show that the denoising diffusion probabilistic model (DDPM) procedure is equivalent to a stochastic version of RF by establishing a novel connection between these processes and stochastic localization. Building on this connection, we further design a stochastic RF sampler that also adapts to the low-dimensionality of the target distribution under milder requirements on the accuracy of the drift estimates, and also with a specific time schedule. We illustrate with simulations on the synthetic data and text-to-image data experiments the improved performance of the proposed samplers implementing the newly designed time-discretization schedules.