빠른 확산 샘플링을 위한 미분 가능 솔버 탐색

Differentiable Solver Search for Fast Diffusion Sampling

May 27, 2025

저자: Shuai Wang, Zexian Li, Qipeng zhang, Tianhui Song, Xubin Li, Tiezheng Ge, Bo Zheng, Limin Wang

cs.AI

초록

디퓨전 모델은 뛰어난 생성 품질을 보여주지만, 이를 위해서는 수많은 함수 평가가 필요하다는 단점이 있습니다. 최근에는 제한된 샘플링 단계에서 역디퓨전 문제를 해결하기 위한 상당한 계산 요구를 완화하기 위해 고급 ODE 기반 솔버가 개발되었습니다. 그러나 Adams류의 다단계 방법에서 크게 영감을 받은 이러한 솔버들은 시간(t) 관련 라그랑주 보간에만 의존하고 있습니다. 우리는 시간 관련 라그랑주 보간이 디퓨전 모델에 최적이 아니라는 것을 보여주고, 시간 단계와 솔버 계수로 구성된 간결한 탐색 공간을 제시합니다. 이 분석을 바탕으로, 우리는 더 최적의 솔버를 식별하기 위한 새로운 미분 가능한 솔버 탐색 알고리즘을 제안합니다. 탐색된 솔버를 장착한 Rectified-Flow 모델들, 예를 들어 SiT-XL/2와 FlowDCN-XL/2는 ImageNet256에서 단 10단계만으로 각각 2.40과 2.35의 FID 점수를 달성합니다. 동시에 DDPM 모델인 DiT-XL/2도 단 10단계만으로 2.33의 FID 점수에 도달합니다. 특히, 우리가 탐색한 솔버는 기존 솔버들을 상당한 차이로 능가합니다. 더욱이, 우리의 솔버는 다양한 모델 아키텍처, 해상도 및 모델 크기에 걸쳐 일반성을 보여줍니다.

English

Diffusion models have demonstrated remarkable generation quality but at the cost of numerous function evaluations. Recently, advanced ODE-based solvers have been developed to mitigate the substantial computational demands of reverse-diffusion solving under limited sampling steps. However, these solvers, heavily inspired by Adams-like multistep methods, rely solely on t-related Lagrange interpolation. We show that t-related Lagrange interpolation is suboptimal for diffusion model and reveal a compact search space comprised of time steps and solver coefficients. Building on our analysis, we propose a novel differentiable solver search algorithm to identify more optimal solver. Equipped with the searched solver, rectified-flow models, e.g., SiT-XL/2 and FlowDCN-XL/2, achieve FID scores of 2.40 and 2.35, respectively, on ImageNet256 with only 10 steps. Meanwhile, DDPM model, DiT-XL/2, reaches a FID score of 2.33 with only 10 steps. Notably, our searched solver outperforms traditional solvers by a significant margin. Moreover, our searched solver demonstrates generality across various model architectures, resolutions, and model sizes.