[논문 리뷰] Gradient Guidance for Diffusion Models: An Optimization Perspective

Yingqing Guo, Hui Yuan외 3명|arXiv (Cornell University)|2024. 04. 23.

Spacecraft Dynamics and Control인용 수 4

한 줄 요약

이 논문은 지속적인 예측 손실을 통해 정규화된 최적화 문제로 간주되는 지속적 샘플링을 프레임화하는 새로운 기울기 가이던스 방법을 제안한다. 이는 전이 학습된 점수 함수를 통해 잠재 데이터 구조를 유지한다. 적응형 미세조정 하에서 전역 최적해로 수렴함을 증명하며, 오목 목적 함수에 대해 $\mathcal{O}(1/K)$ 수준의 수렴 속도를 달성하면서도 샘플 품질을 유지한다.

ABSTRACT

Diffusion models have demonstrated empirical successes in various applications and can be adapted to task-specific needs via guidance. This paper studies a form of gradient guidance for adapting a pre-trained diffusion model towards optimizing user-specified objectives. We establish a mathematical framework for guided diffusion to systematically study its optimization theory and algorithmic design. Our theoretical analysis spots a strong link between guided diffusion models and optimization: gradient-guided diffusion models are essentially sampling solutions to a regularized optimization problem, where the regularization is imposed by the pre-training data. As for guidance design, directly bringing in the gradient of an external objective function as guidance would jeopardize the structure in generated samples. We investigate a modified form of gradient guidance based on a forward prediction loss, which leverages the information in pre-trained score functions and provably preserves the latent structure. We further consider an iteratively fine-tuned version of gradient-guided diffusion where guidance and score network are both updated with newly generated samples. This process mimics a first-order optimization iteration in expectation, for which we proved O(1/K) convergence rate to the global optimum when the objective function is concave. Our code will be released at https://github.com/yukang123/GGDMOptim.git.

연구 동기 및 목표

• 기울기 가미된 확산 모델에서 이론적 이해의 부족, 특히 구조 유지 및 수렴 보장 측면에서의 문제를 해결하기 위해.
• 단순한 기울기 가이던스가 샘플 품질을 떨어뜨리는 문제를 해결하기 위해, 전망 손실 기반의 구조화된 가이던스 신호를 도입하기 위해.
• 가이던스된 확산을 목적 함수 최대화와 사전 훈련된 데이터 분포 사전 지식 간의 균형을 이루는 최적화 과정으로 공식화하기 위해.
• 자기 생성 샘플을 사용해 점수 네트워크를 반복적으로 개선하는 적응형 미세조정 메커니즘을 개발하기 위해.
• 오목 목적 함수 최적화에서 가이던스된 확산의 수렴 속도와 이론적 보장을 설정하기 위해.

제안 방법

• 전이 학습된 점수 함수를 활용해 잠재 구조를 유지하기 위해, $\texttt{G}_{\text{loss}}(x_t,t) := \beta(t)\left(y - g^\top\mathbb{E}[x_0|x_t]\right)g$ 로 정의된 전방 예측 손실 기반의 기울기 가이던스 신호를 도입한다.
• 실제 구현에서 조건부 기대값을 근사하기 위해 Tweedie의 공식에 기반한 전망 추정자 $\hat{\mathbb{E}}[x_0|x_t]$ 를 사용한다.
• 알고리즘 1: 정규화된 최적화 문제를 해결하는 비적응형 가이던스 샘플링 과정으로, 전이 학습된 점수 함수가 사전 지식으로 작용한다.
• 알고리즘 2: 생성된 샘플에서 기울기를 쿼리하고 새로운 데이터로 점수 네트워크를 업데이트하는 반복적 미세조정 루프를 제안한다.
• 각 시간 단계에서 가이던스 신호를 삽입하여 확산 궤적을 목적 함수가 높은 영역으로 이끄는 점수 기반 샘플링 과정을 구현한다.
• 이론적 분석 결과, 목적 함수가 오목하고 미분 가능할 경우 알고리즘 2는 $\mathcal{O}(1/K)$ 수준의 수렴 속도로 전역 최적해로 수렴함을 보였다.

실험 결과

연구 질문

• RQ1 왜 단순한 기울기 가이던스가 실무에서 실패하는가? 샘플 품질 악화의 원인은 무엇인가?
• RQ2 기울기 가이던스는 어떻게 설계되어야 전이 학습된 데이터의 잠재 저차원 구조를 유지할 수 있는가?
• RQ3 적응형 미세조정 하에서 가이던스된 확산 모델이 최적해로 수렴함을 이론적으로 보장할 수 있는가?
• RQ4 기울기 가미된 확산과 1차 최적화 방법 간의 관계는 무엇인가?
• RQ5 목적 함수가 오목할 경우, 가이던스 샘플링 과정의 수렴 속도는 무엇인가?

주요 결과

• 실험 결과에서 전방 예측 손실 기반의 제안된 기울기 가이던스($\texttt{G}_{\text{loss}}$)가 단순한 기울기 가이던스보다 잠재 부분공간의 구조를 훨씬 더 잘 유지함을 확인했다. 이는 오프서포트 오차가 낮아진 것으로 나타났다.
• 비적응형 가이던스를 사용하는 알고리즘 1은 전이 학습된 점수 함수의 정규화 효과로 인해 하위 최적해로 수렴함을 확인했으며, 이는 이론적 분석과 일치한다.
• 자기 생성 샘플을 사용해 반복적으로 미세조정하는 알고리즘 2는 오목 목적 함수에 대해 전역 최적해로 수렴하며, $\mathcal{O}(1/K)$ 수준의 증명 가능한 수렴 속도를 달성한다.
• 실험 결과, 미세조정으로 인한 분포 이탈 상황에서도 $\texttt{G}_{\text{loss}}$ 가 단순한 기울기 가이던스보다 더 높은 구조적 정밀도를 유지함을 보였다.
• 알고리즘 2에 의해 생성된 샘플의 분포가 사전 훈련된 윤곽을 기준으로 기존 데이터 분포를 초월함을 확인했으며, 이는 목적 함수가 높은 영역을 효과적으로 탐색했음을 시사한다.
• 이론적 프레임워크는 기울기 가이던스된 확산이 전이 학습된 점수 함수가 구조적 사전 지식으로 작용하는 정규화된 최적화 문제를 해결하는 것과 동일하다고 규명했다.