[논문 리뷰] Aligning Text-to-Image Diffusion Models with Reward Backpropagation

Mihir Prabhudesai, Anirudh Goyal외 2명|arXiv (Cornell University)|2023. 10. 05.

Generative Adversarial Networks and Image Synthesis인용 수 5

한 줄 요약

이 논문은 텍스트-이미지 확산 모델을 노이즈 제거 과정을 통해 미분 가능한 보상 기울기의 전방향 백프로파게이션을 통해 엔드 투 엔드 백프로파게이션으로 미세조정하는 AlignProp을 제안한다. 저랭크 어댑터와 기울기 체크포인팅을 사용함으로써, 강화학습 기반 모델 대비 더 적은 학습 스텝 수로 더 높은 보상 성능을 달성하여 이미지-텍스트 정렬, 미적 요소, 제어 가능성과 같은 다양한 목표를 효율적으로 최적화할 수 있다.

ABSTRACT

Text-to-image diffusion models have recently emerged at the forefront of image generation, powered by very large-scale unsupervised or weakly supervised text-to-image training datasets. Due to their unsupervised training, controlling their behavior in downstream tasks, such as maximizing human-perceived image quality, image-text alignment, or ethical image generation, is difficult. Recent works finetune diffusion models to downstream reward functions using vanilla reinforcement learning, notorious for the high variance of the gradient estimators. In this paper, we propose AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient through the denoising process. While naive implementation of such backpropagation would require prohibitive memory resources for storing the partial derivatives of modern text-to-image models, AlignProp finetunes low-rank adapter weight modules and uses gradient checkpointing, to render its memory usage viable. We test AlignProp in finetuning diffusion models to various objectives, such as image-text semantic alignment, aesthetics, compressibility and controllability of the number of objects present, as well as their combinations. We show AlignProp achieves higher rewards in fewer training steps than alternatives, while being conceptually simpler, making it a straightforward choice for optimizing diffusion models for differentiable reward functions of interest. Code and Visualization results are available at https://align-prop.github.io/.

연구 동기 및 목표

• 이미지 품질, 공정성, 텍스트-이미지 정렬과 같은 하류 보상 함수에 맞추어 사전 학습된 텍스트-이미지 확산 모델을 정렬하는 데 도전하는 것.
• 확산 모델의 긴 노이즈 제거 체인을 거쳐 엔드 투 엔드 백프로파게이션을 수행할 경우 발생하는 금방 높은 메모리 비용을 해결하는 것.
• 강화학습의 고분산 기울기 방식을 피하고, 보상 기울기를 직접적으로 미분 가능한 방식으로 최적화할 수 있는 방법을 개발하는 것.
• LoRA 어댑터의 파라미터 평균화를 통해 여러 보상 함수 간의 보간을 가능하게 하여 다중목표 정렬을 달성하는 것.

제안 방법

• 노이즈와 텍스트 프롬프트를 이미지로 매핑하는 가역적 순환 정책으로 노이즈 제거 과정을 간주함으로써, 보상에서 모델 가중치로의 기울기 흐름을 가능하게 한다.
• 저랭크 어댑터(LoRA) 모듈을 사용하여 파라미터의 소수의 부분만 미세조정함으로써, 전체 미세조정 대비 메모리 사용을 극적으로 줄인다.
• 기울기 체크포인팅을 적용하여 기울기를 사전에 저장하는 대신 필요에 따라 부분 도함수를 계산함으로써, 타르비트 수준의 메모리 요구량을 피한다.
• 고정된 컨텍스트 길이 K를 가진 절단된 시간을 통한 백프로파게이션(BPTT)을 적용하여 기울기 품질과 메모리 효율성 간의 균형을 맞춘다.
• 다른 보상 함수에 대해 별도로 훈련된 모델의 LoRA 가중치를 평균화하여 모델 보간을 가능하게 하여, 다중목표 생성을 지원한다.
• CLIP 기반 이미지-텍스트 정렬, 미적 점수, JPEG 압축 가능도와 같은 미분 가능한 보상 함수를 사용하여 최적화를 이끈다.

실험 결과

연구 질문

• RQ1 확산 샘플링 과정을 거쳐 엔드 투 엔드 백프로파게이션을 수행하는 것이 대규모 텍스트-이미지 모델에 대해 실용적으로 메모리 효율적으로 수행될 수 있는가?
• RQ2 샘플 효율성과 최종 보상 성능 측면에서 AlignProp은 강화학습 기반 모델 대비 어떻게 비교되는가?
• RQ3 별도로 훈련된 모델의 LoRA 어댑터를 조합하여 여러 보상 함수 간의 효과적인 보간이 가능한가?
• RQ4 저랭크 어댑터와 기울기 체크포인팅의 사용이 성능 유지를 하면서도 메모리 및 계산 비용을 줄이는 데 기여하는가?

주요 결과

• AlignProp는 ReFL 및 DDPO와 같은 기반 모델 대비 유의미하게 높은 평균 보상을 달성하며, HPSv2 벤치마크에서 K=1일 때 2.91 대비 3.30의 평균 보상을 기록한다.
• 고성능 보상에 도달하기 위해 필요한 학습 스텝 수를 줄여 강화학습 방법 대비 뛰어난 샘플 효율성을 입증한다.
• 미적 요소와 압축 가능성을 각각 훈련한 모델의 LoRA 가중치를 보간하면 혼합 계수 α=0.5에서 가장 높은 총합 보상을 달성하는 하이브리드 모델을 얻을 수 있다.
• 절단 실험 분석 결과, LoRA 미세조정이 전체 UNet 미세조정보다 성능이 뛰어나며, EDICT 레이어는 메모리 절감 효과가 있지만 성능 저하와 추론 시간 증가를 유발함을 확인했다.
• K=10일 때 K=1보다 더 높은 성능을 기록하여, 더 긴 컨텍스트가 기울기 품질을 향상시키며 과도한 메모리 비용 없이도 유의미한 이점을 제공함을 시사한다.
• 파라미터 보간을 통해 이미지-텍스트 정렬과 미적 품질을 동시에 향상시키는 복잡한 다중목표 보상에 효과적으로 정렬이 가능함을 입증했다.