효율적인 Stable Diffusion 파인튜닝을 위한 LoRA 활용

Hugging Face는 Low-Rank Adaptation(LoRA)을 diffusers 라이브러리에 통합하여, 훨씬 낮은 VRAM 요구량과 3 MB 정도의 작은 모델 가중치로 Stable Diffusion 파인튜닝을 더 빠르게 수행할 수 있게 했습니다. 이 구현을 통해 사용자는 전체 모델 파인튜닝에 드는 막대한 연산 비용 없이 대규모 diffusion 모델을 특정 스타일이나 개념에 맞게 조정할 수 있습니다.

Stable Diffusion에서의 LoRA 아키텍처 및 적용

Low-Rank Adaptation(LoRA)은 Microsoft 연구원들이 GPT-3과 같은 대형 언어 모델(LLM)을 위해 처음 개발한 기술입니다. 사전 학습된 모델 가중치를 고정하고, 트랜스포머 블록에 학습 가능한 랭크 분해 행렬을 삽입하는 방식으로 동작합니다. 이 접근법은 대부분의 모델 가중치에 대해 그래디언트를 계산하지 않기 때문에 학습 가능한 파라미터 수와 GPU 메모리 요구량을 감소시킵니다.

Stable Diffusion에서는 LoRA가 교차 주의(cross-attention) 레이어에 적용됩니다. 이 레이어들은 이미지 표현과 이를 설명하는 텍스트 프롬프트 간의 관계를 설정하는 역할을 합니다. 이러한 특정 레이어를 대상으로 함으로써 LoRA는 전체 모델 파인튜닝에 버금가는 품질을 유지하면서도 더 빠르고 연산 효율적인 파인튜닝을 달성합니다.

LoRA 파인튜닝의 주요 장점

  • 컴퓨팅 요구량 감소: 파인튜닝을 일반 소비자용 하드웨어에서도 수행할 수 있습니다. 예를 들어, NVIDIA 2080 Ti 11 GB VRAM을 사용해 전체 파인튜닝된 모델을 만들 수 있습니다.
  • 빠른 학습: 학습 가능한 파라미터가 감소함에 따라 학습 주기가 크게 빨라집니다.
  • 극도로 작은 가중치 파일: 원본 모델이 고정된 상태이므로, 삽입된 레이어의 가중치만 저장하면 됩니다. 이로 인해 가중치 파일 크기가 약 3 MB(일부 경우 3.29 MB로 명시) 정도가 되며, 이는 원래 UNet 모델보다 대략 1,000배 작습니다.

구현 및 추론 워크플로우

파인튜닝 과정

Hugging Face는 train_text_to_image_lora.py라는 전용 LoRA 파인튜닝 스크립트를 제공하며, 최소 11 GB GPU RAM으로도 실행할 수 있습니다. 주목할 만한 기술적 세부 사항은 LoRA가 일반 파인튜닝보다 높은 학습률을 필요로 하는 경우가 많다는 점이며, 예를 들어 전체 파인튜닝에서 일반적으로 사용되는 ~1e-6에 비해 1e-4 학습률을 사용합니다.

추론 및 로딩

추론 과정은 수정되지 않은 Stable Diffusion 기본 모델 위에 LoRA 가중치를 로드하도록 설계되었습니다. 워크플로우는 다음과 같습니다:

  1. 기본 모델 식별: Hub API를 사용하여 학습 시 사용된 기본 모델(예: CompVis/stable-diffusion-v1-4 또는 runwayml/stable-diffusion-v1-5)을 확인합니다.
  2. 파이프라인 초기화: 식별된 기본 모델을 사용해 표준 StableDiffusionPipeline을 로드합니다.
  3. 가중치 주입: pipe.unet.load_attn_procs(model_path)를 사용해 Hub에서 LoRA 가중치를 직접 일반 모델 가중치에 로드합니다.

Dreambooth와의 호환성

LoRA는 Stable Diffusion에 새로운 특정 개념을 학습시키는 방법인 Dreambooth와 호환됩니다. 이 두 기술을 결합하면 다음과 같은 여러 이점을 제공합니다:

  • 표준 Dreambooth에 비해 더 빠른 학습 속도
  • 낮은 데이터 요구량, 일반적으로 대상의 이미지 5~10장만 필요합니다.
  • 선택적 텍스트 인코더 조정을 통해 대상에 대한 충실도를 높일 수 있습니다.

다른 파인튜닝 방법과의 비교

Stable Diffusion을 적용하는 다른 방법도 존재하지만, LoRA는 사용 사례에 따라 뚜렷한 장점을 제공합니다:

  • Textual Inversion: LoRA와 유사하게 작고 공유하기 쉬운 가중치를 생성하지만, 일반적으로 단일 대상이나 소수의 개념에만 제한됩니다.
  • LoRA: 일반 목적 파인튜닝에 적합하며, 모델을 완전히 새로운 도메인이나 데이터셋에 맞게 조정할 수 있습니다.
  • Pivotal Tuning: Textual Inversion으로 토큰 임베딩을 얻은 뒤, LoRA를 사용해 해당 임베딩을 학습하는 하이브리드 접근법입니다.

Sources