[논문 리뷰] Practical and Asymptotically Exact Conditional Sampling in Diffusion Models

Luhuan Wu, Brian L. Trippe외 3명|UvA-DARE (University of Amsterdam)|2023. 06. 30.

Machine Learning in Materials Science인용 수 13

한 줄 요약

이 논문은 Twisted Diffusion Sampler (TDS)를 소개한다. SMC 기반의 방법으로 확산 모델에서 조건부 샘플링을 점근적으로 정확하게 수행하며, twisting을 사용하여 조건부를 구현하고 Task-specific한 훈련 없이도 이를 가능하게 하며, 인페이팅, 클래스-조건 MNIST, 및 단백질 모티프-스캐폴딩에 걸쳐 경험적 이점을 보여준다.

ABSTRACT

Diffusion models have been successful on a range of conditional generation tasks including molecular design and text-to-image generation. However, these achievements have primarily depended on task-specific conditional training or error-prone heuristic approximations. Ideally, a conditional generation method should provide exact samples for a broad range of conditional distributions without requiring task-specific training. To this end, we introduce the Twisted Diffusion Sampler, or TDS. TDS is a sequential Monte Carlo (SMC) algorithm that targets the conditional distributions of diffusion models through simulating a set of weighted particles. The main idea is to use twisting, an SMC technique that enjoys good computational efficiency, to incorporate heuristic approximations without compromising asymptotic exactness. We first find in simulation and in conditional image generation tasks that TDS provides a computational statistical trade-off, yielding more accurate approximations with many particles but with empirical improvements over heuristics with as few as two particles. We then turn to motif-scaffolding, a core task in protein design, using a TDS extension to Riemannian diffusion models. On benchmark test cases, TDS allows flexible conditioning criteria and often outperforms the state of the art.

연구 동기 및 목표

• 무작위 확산 모델에서 조건부 샘플링을 태스크 특화 훈련이나 무거운 휴리스틱 없이 동기부여한다.
• twisting을 사용하여 p_theta(x^0|y)를 타깃으로 하는 실용적인 순차 몬테카를로(SMC) 접근법을 개발한다.
• 더 넓은 조건부 작업을 위해 인페이팅과 리만 확산 모델로 TDS를 확장한다.
• MNIST 인페인팅 및 클래스 조건, 단백질 모티프-스캐폴딩 문제에서 휴리스틱 방법 대비 TDS의 실험적 이점을 입증한다.

제안 방법

• 확장된 확산 시퀀스와 SMC를 통해 p_theta(x^0|y)에서의 조건부 샘플링을 형식화한다.
• hat{x}_theta(x^t)인 denoising 추정치로 근사한 p_theta(y|x^t)에 기반한 twisting 함수들을 도입한다.
• 최적 SMC 선택을 근사하면서도 점근적 정확성을 유지하도록 twisted 제안분포와 가중치를 정의한다.
• 인페이팅 및 로피니언(Riemannian) 확산 모델에 대한 확장을 제공하며 자유도 조건부를 포함한다.
• 정규성 조건하에서 입자 근사가 p_theta(x^0|y)에 수렴한다는 (비공식적) 증명을 제시한다.

Figure 1: Errors of conditional mean estimations with 2 SEM error bars averaged over 25 replicates. TDS applies to all three tasks and provides increasing accuracy with more particles.

실험 결과

연구 질문

• RQ1 TDS가 태스크 특성 조건화 학습 없이도 p_theta(x^0|y)의 점근적으로 정확한 샘플을 제공할 수 있는가?
• RQ2 TDS가 정확도와 입자 효율성 측면에서 휴리스틱 방법 및 조건부 학습과 비교하여 어떤 차이를 보이는가?
• RQ3 TDS를 인페이팅, 자유도 조건부, 및 리만 확산 모델로 확장할 수 있는가?

주요 결과

• TDS는 조건 정보가 있는 작업에 대해 확산 모델로부터 점근적으로 정확한 조건부 샘플을 생성한다.
• 실험적으로, 적은 수의 입자와 함께하는 TDS가 MNIST 분류-조건 생성 및 인페인팅 설정에서 휴리스틱 접근법을 자주 능가한다.
• TDS는 단백질 모티프-스캐폴딩에 대한 유연한 조건부를 보여주며 이 문제들에서 최첨단 조건부 모델을 능가할 수 있다.
• TDS는 입자의 수가 증가함에 따라 정확도가 향상되는 우호적인 계산-통계적 트레이드오프를 보인다(O(1/K) 추정에서의 속도).
• 인페이팅 및 리만 매니폴드로의 확장은 점근적 정확성을 유지하고 적용 범위를 확장한다.

(a) Approximate conditional samples of each method with 64 particles given class $y=7$ .