[논문 리뷰] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models

Ying Fan, Olivia Watkins외 8명|arXiv (Cornell University)|2023. 05. 25.

Generative Adversarial Networks and Image Synthesis인용 수 17

한 줄 요약

DPOK는 KL 규제를 이용한 온라인 강화학습을 통해 확산 기반 텍스트-이미지 모델을 미세조정하며, 감독형 미세조정보다 텍스트-이미지 정렬 및 이미지 충실도가 더 우수하다.

ABSTRACT

Learning from human feedback has been shown to improve text-to-image models. These techniques first learn a reward function that captures what humans care about in the task and then improve the models based on the learned reward function. Even though relatively simple approaches (e.g., rejection sampling based on reward scores) have been investigated, fine-tuning text-to-image models with the reward function remains challenging. In this work, we propose using online reinforcement learning (RL) to fine-tune text-to-image models. We focus on diffusion models, defining the fine-tuning task as an RL problem, and updating the pre-trained text-to-image diffusion models using policy gradient to maximize the feedback-trained reward. Our approach, coined DPOK, integrates policy optimization with KL regularization. We conduct an analysis of KL regularization for both RL fine-tuning and supervised fine-tuning. In our experiments, we show that DPOK is generally superior to supervised fine-tuning with respect to both image-text alignment and image quality. Our code is available at https://github.com/google-research/google-research/tree/master/dpok.

연구 동기 및 목표

• 인간 피드백으로부터 학습하여 텍스트-이미지 생성 개선의 동기를 부여한다.
• 확산 모델의 미세조정을 온라인 RL 문제로 공식화하고 KL 규제를 적용한다.
• 텍스트와 이미지의 정렬에서 온라인 RL이 감독형 미세조정보다 갖는 이점을 조사한다.
• ImageReward를 보상 모델로 사용하는 Stable Diffusion에서 RL 미세조정을 평가한다.
• RL 미세조정과 KL-규제된 감독형 미세조정 간의 분석 및 비교를 제공한다.

제안 방법

• DDPM의 디노이징 프로세스를 상태 s_t와 행동 a_t를 갖는 다단계 MDP로 프레이밍한다.
• 사전 학습된 모델에 KL 규화를 두고 기대 보상을 최대화하는 온라인 정책 그래디언트를 도출한다.
• DPOK: KL 규정화된 디퓨전 정책 최적화(Diffusion Policy Optimization with KL regularization)와 그 그래디언트를 제공한다(Eq. 9).
• 온라인 미세조정을 안정시키기 위해 출력 이미지에 대한 상한 KL 규제를 제안한다(Lemma 4.2 & Eq. 8-9).
• 감독형 미세조정으로 KL 규제를 확장하고 KL-D와 KL-O 변형을 비교한다(Lemmas 4.3, 4.3).
• LoRA와 ImageReward를 사용한 Stable Diffusion v1.5에서 RL vs. SFT를 평가하기 위한 실험을 수행한다.

실험 결과

연구 질문

• RQ1 온라인 RL 미세조정이 감독형 미세조정에 비해 이미지 충실도를 희생하지 않고 텍스트-이미지 정렬을 개선할 수 있는가?
• RQ2 KL 규제가 RL 미세조정을 안정화하고 확산 모델의 일반적인 실패 모드를 완화하는가?
• RQ3 훈련 프롬프트 외의 다중 프롬프트나 데이터셋에 대해 RL 미세조정은 어떻게 성능을 보이는가?
• RQ4 RL에서의 KL 규제와 감독형 설정에서의 보상 및 미학에 대한 비교 효과는?
• RQ5 온라인 RL이 보상 기반 최적화를 통해 프리트레인된 확산 모델의 편향을 줄이는가?

주요 결과

데이터셋	원래 ImageReward	RL ImageReward	원래 미학	RL 미학
MS-CoCo	0.22	0.55	5.39	5.43
Drawbench	0.13	0.58	5.31	5.35

• 온라인 RL 미세조정은 일반적으로 더 높은 ImageReward 점수와 감독형 미세조정보다 경쟁력 있거나 더 나은 미학 점수를 달성한다.
• RL 미세조정 모델은 프롬프트와의 정렬이 개선되며 보지 못한 객체를 포함하고 SFT에서 보이는 과다 채도 같은 이슈를 피한다.
• KL 규제가 RL을 안정화하고 보상과 미학을 모두 향상시키는 데 도움이 되며, KL 없이 RL은 이미지 품질을 저하시킬 수 있다.
• 인간 평가에서 이미지-텍스트 정렬 및 이미지 품질 면에서 RL이 SFT보다 우수하다.
• 여러 프롬프트(MS-CoCo, Drawbench)에 대한 RL은 보상을 크게 높이면서도 원래 모델에 비해 미학을 유지한다.