Hugging Face가 Dreambooth 파인튜닝 실험을 통해 학습률·스텝·텍스트 인코더 설정 조합으로 Stable Diffusion의 오버피팅 문제 해결

Context

Dreambooth는 Stable Diffusion에 새로운 개념을 학습시키는 파인튜닝 기법이지만, 적절한 하이퍼파라미터 선택이 어렵고 빠른 오버피팅이 발생하는 문제가 있었다. 특히 얼굴 학습 시 노이즈 증가, 품질 저하, 프롬프트 해석 능력 감소 현상이 나타났다.

Technical Solution

• 학습률과 스텝 수 조정: 저학습률(1e-6 ~ 2e-6)과 단계적 스텝 증가를 통해 오버피팅과 언더피팅 사이의 최적점 탐색
• 객체 vs 얼굴별 파라미터 차등 설정: 객체는 배치 4, 400스텝, 학습률 2e-6 / 얼굴은 배치 2, 800-1200스텝, 학습률 1e-6 적용
• 사전 보존(Prior Preservation) 활성화: 얼굴 학습 시 사전 보존을 통해 원본 개념 유지 및 오버피팅 방지
• 텍스트 인코더 함께 파인튜닝: 기존 UNet만 학습하던 방식에서 텍스트 인코더도 함께 학습하도록 변경
• 추론 스케줄러 변경: PNDM/LMSDiscrete 대신 DDIM 스케줄러 사용 및 추론 스텝을 100으로 증가
• 메모리 최적화 기법 조합: 8-bit Adam, fp16 학습, 그래디언트 누적을 통해 16GB GPU에서 동작 가능하도록 구현

Impact

텍스트 인코더 파인튜닝 적용 시 얼굴 생성 이미지의 현실성이 증대되었고, 오버피팅 발생률이 감소했으며, 복잡한 프롬프트 해석 능력이 향상되었다.

Key Takeaway

Dreambooth 파인튜닝은 단순히 하이퍼파라미터를 높이는 것이 아니라 대상(객체/얼굴)별로 학습률과 스텝 수를 맞춤 조정하고, 텍스트 인코더 학습과 추론 설정을 함께 최적화할 때 최고 품질의 결과를 얻을 수 있다.