TRL을 통한 DDPO를 이용한 Stable Diffusion 파인튜닝
Hugging Face는 DDPOTrainer를 통해 trl 라이브러리에 Denoising Diffusion Policy Optimization(DDPO)를 통합하여 강화 학습(RL)을 사용하여 Stable Diffusion 모델을 인간 선호도 및 미적 기준에 더 잘 맞도록 파인튜닝할 수 있게 했습니다.
DDPO와 Reward-Weighted Regression 비교
DDPO는 계산 오버헤드를 줄이고 근사 오차를 제거하여 diffusion 모델에 대한 이전 RL 방법(예: Reward-weighted regression(RWR))보다 개선되었습니다. RWR은 디노이징 손실 함수를 재사용하고 최종 샘플의 보상에 따라 가중치를 적용하여 중간 단계를 실제로 무시하는 반면, DDPO는 전체 디노이징 프로세스를 다단계 마르코프 결정 프로세스(MDP)로 간주합니다.
By framing the process as an MDP and using a fixed sampler, DDPO allows the agent policy to be an isotropic Gaussian. This enables the calculation of the exact likelihood of each denoising step rather than relying on the approximate likelihood of the final sample, leading to better performance and the ability to handle more complex objectives.
DDPO 알고리즘 및 RLHF 워크플로
DDPO는 최적화 문제를 해결하기 위해 정책 그래디언트 방법인 근접 정책 최적화(PPO)를 활용합니다. DDPO 구현에서의 주요 커스터마이징은 PPO 알고리즘의 궤적 수집 부분입니다.
Reinforcement Learning from Human Feedback(RLHF) 워크플로에 통합될 때, diffusion 모델을 정렬하는 프로세스는 세 단계로 간소화됩니다:
- 사전 학습된 모델: 사전 학습된 Diffusion Model(예: Stable Diffusion)으로 시작합니다.
- 보상 모델: 선호 데이터를 수집하고 신호 역할을 하는 보상 모델을 훈련시킵니다.
- DDPO 파인튜닝: 보상 모델을 신호로 사용하여 DDPO로 모델을 파인튜닝합니다.
이미지 미적 향상의 경우, Hugging Face는 Aesthetic Visual Analysis(AVA) 데이터셋에서 훈련된 훈련 가능한 회귀 헤드가 있는 동결된 CLIP 모델을 사용하여 보상 신호자로 작동하게 했습니다.
구현 및 학습 지침
DDPO를 사용하여 Stable Diffusion을 학습하려면 상당한 하드웨어 리소스가 필요하며, NVIDIA A100 GPU는 메모리 부족(OOM) 오류를 피하기 위한 최소 요구 사항입니다. 구현은 trl 라이브러리의 DDPOTrainer 및 DDPOConfig 클래스를 통해 처리됩니다.
권장 하이퍼파라미터
단일 GPU 학습에 대해 다음과 같은 하이퍼파라미터가 권장됩니다:
| 매개변수 | 권장 값 |
|---|---|
num_epochs |
200 |
train_batch_size |
3 |
sample_batch_size |
6 |
gradient_accumulation_steps |
1 |
sample_num_steps |
50 |
sample_num_batches_per_epoch |
4 |
per_prompt_stat_tracking |
True |
per_prompt_stat_tracking_buffer_size |
32 |
mixed_precision |
True |
train_learning_rate |
3e-4 |
주요 발견 및 배운 점
Hugging Face의 실험 결과에 따르면, DDPO로 튜닝된 모델은 훈련 프롬프트 세트가 최소 크기라도 다양한 프롬프트에 걸쳐 잘 일반화됩니다.
LoRA와 전체 파인튜닝 비교
- LoRA: 권장되며 여러 테스트에서 안정성이 입증되었습니다.
- 전체 파인튜닝 (non-LoRA): LoRA보다 더 복잡한 이미지를 생성할 수 있지만 안정화하기 훨씬 더 어렵습니다. non-LoRA 실행의 경우, 학습률을 약
1e-5정도로 낮게 설정하고mixed_precision을None으로 설정하는 것이 좋습니다.
현재 제한 사항
trl 라이브러리의 현재 DDPOTrainer 구현은 바닐라 Stable Diffusion 모델의 파인튜닝에만 제한됩니다. LoRA가 주요 초점이며 잘 작동하지만, 전체 학습도 가능하지만 더 정확한 하이퍼파라미터 튜닝이 필요합니다.