강화학습을 활용한 확산 모델의 최적화: DRLX 프레임워크와 DDPO 알고리즘

최근 대규모 언어 모델(LLM) 분야에서는 인간의 피드백을 통한 강화학습(RLHF, Reinforcement Learning from Human Feedback)이 모델의 지시 이행 능력과 답변 품질을 높이는 핵심 기술로 자리 잡았습니다. 이러한 흐름은 텍스트 생성 모델을 넘어 이미지 생성 모델인 확산 모델(Diffusion Models)로 확장되고 있습니다. 확산 모델 역시 프롬프트에 대한 충실도를 높이거나 사용자의 미적 선호도를 반영하기 위해 강화학습을 도입하려는 시도가 이어지고 있으며, 그 중심에 DRLX (Diffusion Reinforcement Learning X) 라이브러리가 있습니다.

강화학습 기반 확산 모델의 핵심: DDPO 알고리즘

언어 모델의 RLHF가 문장의 토큰 생성을 마르코프 결정 과정(MDP)으로 정의하듯, 확산 모델 또한 노이즈를 제거하는 역과정(Reverse Process)을 MDP로 모델링할 수 있습니다. 버클리 대학교의 Levine Lab에서 제안한 DDPO (Denoising Diffusion Policy Optimization)는 가장 대표적인 알고리즘입니다. DDPO는 강화학습의 대중적인 알고리즘인 PPO(Proximal Policy Optimization)를 확산 모델의 학습 구조에 최적화하여 적용했습니다.

DDPO는 확산 모델의 각 타임스텝을 상태 전환으로 간주하고, 모델이 예측하는 노이즈를 액션(Action)으로 정의합니다. 최종적으로 생성된 이미지에 대해 보상 모델(Reward Model)이 점수를 부여하면, 이 점수를 최대화하는 방향으로 모델의 파라미터를 업데이트합니다. 이때 학습의 안정성을 위해 기존 모델과 업데이트되는 모델 사이의 변화량을 제한하는 클리핑(Clipping) 기법을 사용합니다.

보상 모델(Reward Model)의 역할

강화학습의 목표는 특정 보상을 극대화하는 것입니다. 이미지 생성에서 보상은 주로 두 가지 측면에서 정의됩니다:

미적 점수(Aesthetic Score): 이미지가 시각적으로 얼마나 아름다운지를 평가합니다.
프롬프트 일치도(Prompt Adherence): 생성된 이미지가 입력된 텍스트 설명과 얼마나 일치하는지를 평가합니다.

보상 모델은 일반적으로 인간의 선호도가 반영된 데이터셋으로 사전 학습된 모델을 사용하며, 생성된 이미지(또는 이미지와 텍스트 쌍)를 입력받아 하나의 스칼라 점수를 출력합니다.

DRLX를 활용한 학습 구현

DRLX 라이브러리를 사용하면 복잡한 강화학습 로직을 직접 구현하지 않고도 확산 모델을 튜닝할 수 있습니다. 아래는 특정 카테고리의 이미지를 생성하도록 프롬프트를 구성하고 학습하는 예시 코드입니다.

import random
from drlx.pipeline import PromptPipeline

class CategoryPromptProvider(PromptPipeline):
    """
    학습에 사용할 카테고리 기반 프롬프트를 생성하는 파이프라인
    """
    def __init__(self, category_list, total_samples=5000):
        super().__init__()
        self.categories = category_list
        self.total_samples = total_samples

    def __getitem__(self, idx):
        # 무작위로 카테고리 선택 후 프롬프트 생성
        selected_category = random.choice(self.categories)
        return f"a high quality photo of a {selected_category}"

    def __len__(self):
        return self.total_samples

다음으로 미적 점수를 산출하는 보상 모델을 정의합니다. 여기서는 CLIP 모델을 기반으로 한 스코어링 로직을 가정합니다.

from drlx.reward_modelling import RewardModel
import torch

class ImageQualityScorer(RewardModel):
    """
    이미지의 미적 가치를 평가하는 보상 모델 클래스
    """
    def __init__(self, model_path, device="cuda"):
        super().__init__()
        self.device = device
        self.evaluator = self.load_custom_model(model_path).to(device)

    def forward(self, images, prompts):
        # 생성된 이미지 배열을 입력받아 보상 점수 반환
        with torch.no_grad():
            scores = self.evaluator(images)
        return scores

설정 파일(YAML)을 통해 모델 구조, 옵티마이저, 샘플링 스텝 등을 정의한 후, 최종적으로 트레이너를 실행합니다.

from drlx.trainer.ddpo_trainer import DDPOTrainer
from drlx.configs import DRLXConfig

# 설정 로드
training_config = DRLXConfig.load_yaml("custom_config.yml")

# 데이터 및 모델 준비
prompt_pipe = CategoryPromptProvider(category_list=["tiger", "eagle", "dolphin"])
reward_engine = ImageQualityScorer(model_path="path/to/weights")

# 학습 시작
model_trainer = DDPOTrainer(training_config)
model_trainer.train(prompt_pipe, reward_engine)

보상 해킹(Reward Hacking)과 PickScore

강화학습 과정에서 발생할 수 있는 주요 문제 중 하나는 '보상 해킹'입니다. 모델이 보상 모델의 허점을 파고들어 실제 품질은 낮지만 점수만 높은 이미지를 생성하는 현상입니다. 예를 들어, 미적 점수 모델이 특정 색감에 높은 점수를 준다면 모델은 프롬프트를 무시하고 해당 색감의 노이즈 패턴만 생성하게 될 수 있습니다.

이를 방지하기 위해 프롬프트와 이미지의 정렬 상태를 함께 평가하는 PickScore와 같은 모델이 제안되었습니다. PickScore는 수많은 유저의 선호 데이터를 학습하여, 단순히 예쁜 이미지가 아니라 사용자가 의도한 내용이 정확히 담긴 이미지에 더 높은 점수를 부여합니다. DRLX에서 PickScore를 보상 모델로 활용하면 더 정교한 정렬(Alignment)이 가능해집니다.

실험적 결과 및 의의

실제로 Stable Diffusion v2.1 모델을 DDPO와 PickScore로 학습시킨 결과, 프롬프트에 명시된 세부 묘사(예: 동물의 표정이나 도구의 형태)가 기본 모델에 비해 훨씬 명확해지는 경향을 보였습니다. 또한 학습이 진행됨에 따라 평균 보상 값이 꾸준히 상승하며 모델이 인간의 선호 영역으로 수렴하는 것을 확인할 수 있습니다.

강화학습을 통한 확산 모델의 개선은 아직 초기 단계에 있으며, 향후 DPO(Direct Preference Optimization)와 같은 최신 최적화 기법과의 결합이나 오프라인 강화학습 적용 등 발전 가능성이 매우 높은 분야입니다.

태그: Diffusion Models Reinforcement Learning DDPO DRLX PPO

7월 31일 09:40에 게시됨

괴물 클럽