LoRA를 활용한 효율적인 Stable Diffusion Fine-Tuning

Stable Diffusion 모델의 성능을 높이기 위해 Fine-Tuning은 필수적이지만, 이 과정은 큰 비용과 시간이 소요되는 단점이 있습니다. 이에 대한 해결책으로 Microsoft에서 개발한 LoRA(Low-Rank Adaptation) 기술이 주목받고 있습니다. LoRA는 원래 대규모 언어 모델의 효율적인 Fine-Tuning을 위해 개발되었으나, 최근에는 이미지 생성 모델인 Stable Diffusion에도 적용되고 있습니다. 이 글에서는 "Stable Diffusion Fine-Tuning"과 "Stable Diffusion LoRA"에 대해 알아보고, LoRA가 어떻게 모델 성능을 향상시키면서도 효율성을 높이는지 설명하겠습니다.

---

LoRA란 무엇인가?

LoRA는 대규모 사전 학습 모델을 대상으로 새로운 작업에 맞춰 효율적으로 Fine-Tuning을 수행하는 방법입니다. 이 기술의 핵심은 사전 학습된 모델의 가중치를 고정하고, 필요한 부분에만 훈련 가능한 매트릭스를 추가하여 업데이트하는 방식에 있습니다. 이를 통해 모든 가중치를 업데이트하는 기존 Fine-Tuning 방식보다 훨씬 적은 양의 메모리와 계산 자원을 사용합니다. 이러한 방식은 GPT-3와 같은 수십억 개의 파라미터를 가진 모델에서도 효과적으로 적용될 수 있습니다.

---

LoRA가 Stable Diffusion에 적용되는 방식

LoRA는 본래 언어 모델에 주로 사용되었으나, Stable Diffusion과 같은 이미지 생성 모델에도 응용될 수 있습니다. Stable Diffusion의 경우 텍스트와 이미지 간의 관계를 구축하는 교차-어텐션(cross-attention) 레이어에 LoRA를 적용할 수 있습니다.

 

교차-어텐션 레이어는 텍스트 프롬프트와 이미지 표현을 연결하는 핵심 요소로, 여기에 LoRA의 트레이닝 가능한 매트릭스를 주입함으로써 기존 모델의 가중치를 고정한 채 추가적인 Fine-Tuning을 수행할 수 있습니다.

 

Stable Diffusion의 교차-어텐션 레이어에 LoRA를 주입하면 다음과 같은 장점이 있습니다:

1. 훈련 속도가 크게 향상됩니다.
2. 메모리 사용량이 감소하여, 11GB VRAM의 GPU에서도 훈련이 가능합니다.
3. 모델 크기가 대폭 줄어듭니다. LoRA로 훈련된 가중치 파일은 약 3MB에 불과해, 원본 모델의 크기보다 약 1,000배 작습니다.

---

Stable Diffusion LoRA Fine-Tuning의 장점

LoRA를 통해 Stable Diffusion의 Fine-Tuning을 수행하면, Dreambooth 및 Textual Inversion과 같은 경량화된 Fine-Tuning 방법보다 더욱 효율적인 모델 훈련이 가능합니다.

 

기존 방식들은 전체 모델을 업데이트하기 때문에 큰 용량과 메모리, 시간이 소요되는 반면, LoRA는 필요한 매트릭스만 업데이트하기 때문에 훈련이 훨씬 빠릅니다.

 

LoRA의 주요 장점은 다음과 같습니다:

• 빠른 훈련 속도: 필요한 파라미터만 업데이트하므로, 전체 모델을 업데이트하는 것보다 훈련 시간이 훨씬 짧습니다.
• 적은 컴퓨팅 자원: VRAM이 적은 GPU에서도 Fine-Tuning이 가능해졌습니다.
• 저용량 가중치 파일: 훈련된 가중치 파일의 크기가 작아, 다른 사용자가 Fine-Tuning된 모델을 쉽게 다운로드하고 사용할 수 있습니다.

예를 들어, Hugging Face의 diffusers 라이브러리에서 제공하는 스크립트를 활용하면, 11GB RAM을 사용하는 GPU에서도 Stable Diffusion을 Fine-Tuning할 수 있습니다.

 

 또한, Lambda Labs의 포켓몬 데이터셋을 사용하여 모델을 Fine-Tuning할 때, W&B 대시보드를 통해 약 5시간의 훈련 시간으로 훌륭한 결과를 얻을 수 있었습니다.

---

LoRA를 활용한 Dreambooth 적용

Dreambooth는 Stable Diffusion 모델에 새로운 개념을 학습시키는 기술로, 일반적으로 5~10개의 이미지로도 학습이 가능합니다. LoRA는 Dreambooth와 호환이 가능하며, 이 과정을 더욱 효율적으로 만들어 줍니다.

 

Dreambooth를 사용하여 새로운 객체나 개념을 모델에 추가하고자 할 때, LoRA를 함께 사용하면 추가적인 하드웨어 설정 없이도 높은 품질의 모델을 빠르게 얻을 수 있습니다.

---

다른 Fine-Tuning 방법과 LoRA의 비교

Stable Diffusion에서 Fine-Tuning을 수행하는 방식으로 Dreambooth와 Textual Inversion이 널리 알려져 있습니다.

 

Dreambooth는 특정 개념을 모델에 학습시키는 데 매우 유용하며, 소수의 이미지를 가지고도 주어진 개념을 모델에 추가할 수 있다는 장점이 있습니다. 예를 들어, 사용자가 특정 인물이나 사물을 모델이 잘 이해하도록 훈련하고 싶을 때 Dreambooth를 사용하면 적합합니다. 반면, Dreambooth는 개별 개념을 학습시키는 데 특화되어 있어, 다양한 주제를 일반화하는 데는 한계가 있을 수 있습니다.

 

Dreambooth에 대한 자세한 설명은 아래 포스팅을 참고해주세요.

  

Stable Diffusion과 파인 튜닝 방법: 완벽 가이드

 

DreamBooth를 활용한 이미지 파인튜닝: 코드와 구현 방법

 

Textual Inversion은 더 적은 양의 데이터로 특정 주제를 학습하는 방식으로 인기를 끌고 있습니다. 주로 하나의 특정 이미지나 아주 제한적인 주제를 학습시키는 데 사용되며, 일반화된 개념보다는 개별 객체나 스타일을 추가하는 데 적합합니다. 다만, Textual Inversion은 모델에 새로운 개념을 깊이 학습시키기보다는, 제한적인 상황에서 텍스트와 이미지 간의 관계를 세밀하게 조정하는 역할을 합니다.

 

반면, LoRA는 Dreambooth나 Textual Inversion보다 더욱 유연하게 다양한 주제에 적응할 수 있도록 설계되었습니다. LoRA는 사전 학습된 모델의 주요 가중치를 유지하면서, 필요한 부분에 트레이닝 가능한 파라미터를 추가로 삽입하여 다양한 개념과 스타일을 학습할 수 있게 합니다. 이를 통해 더 많은 파라미터를 업데이트할 수 있어, Dreambooth가 특정 개념에 한정되는 반면, LoRA는 일반적인 주제나 스타일에도 쉽게 적용할 수 있습니다. 또한, LoRA는 전체 모델을 재훈련하는 방식이 아니기 때문에 GPU 메모리와 자원을 아끼면서도 효율적으로 학습이 가능하다는 장점이 있습니다.

 

따라서, Dreambooth는 특정 개념의 깊이 있는 학습에 강점이 있는 반면, LoRA는 다양한 주제와 스타일에 대해 더 유연하고 경제적인 Fine-Tuning을 제공하여, 더 넓은 범위의 데이터셋이나 주제에 대한 일반화가 필요한 경우 적합한 솔루션이 됩니다.

 

---

LoRA Fine-Tuning 적용 예제

Stable Diffusion에서 LoRA Fine-Tuning을 직접 적용하려면 다음과 같은 스크립트를 사용할 수 있습니다:

 

train_text_to_image_lora.py 는 이 링크를 참고해주세요.

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/pokemon"
export HUB_MODEL_ID="pokemon-lora"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="Totoro" \
  --seed=1337

 

이 스크립트를 사용하면, Hugging Face Hub에 모델을 푸시할 수 있어 다른 사용자와 Fine-Tuning 모델을 쉽게 공유할 수 있습니다. 또한 W&B를 사용하여 훈련 과정의 모니터링이 가능합니다.

---

LoRA Fine-Tuning 결과물 공유의 장점

LoRA의 큰 장점 중 하나는 Fine-Tuning된 결과물을 쉽게 공유할 수 있다는 점입니다. 기존의 Dreambooth 모델들은 전체 모델을 공유해야 하므로 용량이 컸지만, LoRA를 사용하면 3MB의 파일만으로도 Fine-Tuning된 결과를 공유할 수 있습니다. 이는 많은 사용자에게 큰 이점을 제공합니다.

---

결론

LoRA는 Stable Diffusion과 같은 대규모 모델을 효율적으로 Fine-Tuning할 수 있는 혁신적인 방법입니다. 교차-어텐션 레이어에 최적화된 매트릭스를 주입하여 모델의 훈련 속도를 높이고 메모리 사용량을 줄여, 기존의 Fine-Tuning보다 훨씬 더 실용적이고 효율적인 솔루션을 제공합니다.

 

Dreambooth, Textual Inversion과 함께 활용할 수 있어, 다양한 데이터를 학습시키고자 하는 사용자들에게 이상적인 선택지입니다.