[논문 리뷰] Directly Fine-Tuning Diffusion Models on Differentiable Rewards

Kevin Clark, Paul Vicol외 2명|arXiv (Cornell University)|2023. 09. 29.

Generative Adversarial Networks and Image Synthesis인용 수 7

한 줄 요약

이 논문은 다양한 미분 가능 보상 함수를 사용하여 확산 모델을 미분 가능 보상 함수를 통해 전체 확산 샘플링 과정을 직접 역전파하는 방법인 Direct Reward Fine-Tuning (DRaFT)을 소개한다. Gradient checkpointing과 LoRA를 활용함으로써 강화학습 기반 기준 대비 200배 이상 빠른 성능과 효율성을 달성하며, DRaFT-$K$ 및 DRaFT-LV와 같은 변종들은 학습 안정성 향상과 분산 감소를 개선한다.

ABSTRACT

We present Direct Reward Fine-Tuning (DRaFT), a simple and effective method for fine-tuning diffusion models to maximize differentiable reward functions, such as scores from human preference models. We first show that it is possible to backpropagate the reward function gradient through the full sampling procedure, and that doing so achieves strong performance on a variety of rewards, outperforming reinforcement learning-based approaches. We then propose more efficient variants of DRaFT: DRaFT-K, which truncates backpropagation to only the last K steps of sampling, and DRaFT-LV, which obtains lower-variance gradient estimates for the case when K=1. We show that our methods work well for a variety of reward functions and can be used to substantially improve the aesthetic quality of images generated by Stable Diffusion 1.4. Finally, we draw connections between our approach and prior work, providing a unifying perspective on the design space of gradient-based fine-tuning algorithms.

연구 동기 및 목표

• 예측된 확산 모델과 바람직한 행동(예: 미적 품질 또는 인간 선호도) 간의 괴리 문제를 해결하기 위해, 모델 파라미터를 기울기 기반으로 직접 최적화할 수 있도록 하는 것.
• 확산 모델 정렬을 위한 강화학습 기반 미세조정 방법의 비효율성과 샘플 복잡도 문제를 해결하는 것.
• 흑상자 보상 최적화가 필요 없이, 분석적 기울기에서 유도된 다양한 보상 함수를 효율적으로 활용하는 방법을 개발하는 것.
• 이미지 품질, 객체 검출, 악성 예제 생성 등 다양한 보상 함수를 위한 확장 가능하고 효율적이며 조합 가능한 미세조정을 가능하게 하는 것.
• 다양한 목표를 위한 다중 목적 합성 및 보간을 지원하는 통합된 기울기 기반 프레임워크를 제공하는 것.

제안 방법

• 최종 생성된 이미지의 보상에서 기울기를 도출하여 전체 전개된 확산 샘플링 과정을 거쳐 기울기를 역전파함으로써, 확산 모델의 엔드 투 엔드 미세조정을 가능하게 하는 DRaFT를 제안한다.
• 메모리 비용을 줄이기 위해 기울기 체크포인팅을 사용하고, 파라미터의 소수만을 미세조정하기 위해 LoRA를 활용하여 계산 효율성을 유지한다.
• DRaFT-$K$ 를 도입하여, 기울기 역전파를 마지막 $K$ 개의 노이즈 제거 단계로만 제한함으로써 학습 안정성 향상과 기울기 폭발 방지에 기여한다.
• DRaFT-LV 는 다수의 노이즈 샘플을 사용해 더 낮은 분산 기울기 추정치를 계산함으로써 학습 안정성과 속도를 향상시키는 DRaFT-1의 변종이다.
• LoRA 가중치 적응을 통해 Stable Diffusion 1.4 를 다양한 분석 가능 보상 함수(예: LAION Aesthetics, PickScore, HPSv2)에 대해 미세조정한다.
• 다양한 보상 목표에서 학습된 LoRA 가중치를 혼합하거나 스케일링하여 모델의 조합 및 보간을 가능하게 하여, 재학습 없이 다중 목적 생성을 지원한다.

실험 결과

연구 질문

• RQ1 전체 확산 샘플링 과정을 거쳐 기울기 기반 미세조정을 수행할 경우, 분석 가능 보상 함수에서 강화학습 대비 더 뛰어난 성능을 달성할 수 있는가?
• RQ2 기울기 역전파를 마지막 $K$ 단계로 제한하는 것(DRaFT-$K$)은 전체 역전파 대비 학습 안정성과 수렴성에 어떤 영향을 미치는가?
• RQ3 다수의 노이즈 샘플 평균화와 같은 분산 감소 기법(DRaFT-LV)은 학습 효율성과 수렴 속도를 크게 향상시킬 수 있는가?
• RQ4 DRaFT 는 이미지 압축성, 객체 제거, 악성 예제 생성과 같은 다양한 보상 함수로 일반화될 수 있는가?
• RQ5 다양한 보상 목표에서 학습된 미세조정된 모델은 LoRA 가중치 혼합을 통해 조합되거나 보간될 수 있으며, 여러 목표에서 성능을 유지할 수 있는가?

주요 결과

• LAION Aesthetics 점수를 최대화할 때, DRaFT 는 강화학습 기반 기준 대비 학습 속도가 200배 이상 빠르게 성능을 뛰어넘는다.
• DRaFT-LV 는 더 낮은 분산 기울기 추정치 덕분에 ReFL 과 같은 이전 기울기 기반 방법 대비 약 2배 빠른 학습 속도를 기록한다.
• K=2 인 DRaFT-$K$ 는 기울기 폭발 문제로 인해 실질적으로 전체 역전파보다 더 뛰어난 성능을 달성한다.
• PickScore 와 HPSv2 에 기반한 미세조정된 모델은 인간 선호도 점수에서 뚜렷한 향상을 보이며, LoRA 가중치 스케일링으로 인해 과적합이 완화된다.
• 다양한 보상에서 학습된 LoRA 가중치 간의 효과적인 보간이 가능하여, 재학습 없이도 복합적 생성이 가능해진다.
• DRaFT 는 다양한 프롬프트에서 이미지 품질 향상을 성공적으로 달성했으며, 이미지 압축성 및 악성 예제 생성과 같은 복잡한 목표도 지원한다.