Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

이 논문은 보상 과최적화 문제를 해결하고 다양성을 유지하면서 타겟 보상을 극대화하기 위해, 재파라미터화된 정책 경사와 소프트 Q-함수 추정을 기반으로 한 새로운 KL-정규화 강화학습 방법인 SQDF 를 제안합니다.

핵심

🎨 SQDF: "현명한 예술가"를 만드는 새로운 방법

이 논문은 **확산 모델 (Diffusion Model)**이라는 최신 AI 기술이 어떻게 더 똑똑하고 창의적으로 변할 수 있는지에 대한 이야기를 담고 있습니다. 특히, AI 가 원하는 대로 그림을 그리게 하되, 그 과정에서 AI 가 "미친 듯이" 변해버리는 것을 막는 새로운 방법인 SQDF를 소개합니다.

이 복잡한 내용을 요리사와 요리 레시피에 비유해서 쉽게 설명해 드릴게요.

---

1. 문제: "맛있는 요리"를 만들려다 "재앙"이 된 경우

지금까지의 AI 그림 그리기 기술 (확산 모델) 은 아주 훌륭한 기본 실력을 가지고 있습니다. 마치 기본 레시피를 완벽하게 익힌 요리사처럼 말이죠. 하지만 우리는 이 요리사에게 "더 맛있게", "더 화려하게" 만들어달라고 주문합니다 (이를 '보상'이라고 부릅니다).

그런데 기존 방법들은 이 주문을 받자마자 과도하게 반응했습니다.

• 상황: "더 맛있게 해줘!"라고 말하자, 요리사는 설탕을 산더미처럼 붓고, 소금 통을 비워버립니다.
• 결과: 맛은 '최고'로 측정되지만, 실제로는 먹을 수 없는 괴물 요리가 됩니다. (이걸 논문에서는 '과도한 최적화 (Over-optimization)'라고 부릅니다.)
• 문제점: AI 가 원하는 스타일만 쫓다가, 그림의 다양성이 사라지고 (모든 그림이 똑같이 생김), 원래의 자연스러운 느낌이 사라져 버립니다.

2. 해결책: SQDF (소프트 Q 기반 확산 미세 조정)

이 논문은 SQDF라는 새로운 방법을 제안합니다. 이 방법은 요리사에게 "맛만 쫓지 말고, 원래의 요리 스타일도 잊지 말라"고 가르치는 현명한 코치 역할을 합니다.

SQDF 는 세 가지 핵심 비법 (혁신) 을 사용합니다.

🍳 비법 1: "시간의 중요성"을 아는 할인율 (Discount Factor)

• 비유: 요리를 할 때, **불을 켠 직후 (초기 단계)**에 소금을 얼마나 넣느냐보다, 마무리 단계에 어떤 재료를 넣느냐가 맛을 결정합니다.
• 기존 방법: 요리 시작부터 끝까지 모든 단계에 똑같은 점수를 매겨서, 초기 단계의 작은 실수가 전체를 망치거나, 반대로 초기 단계에 너무 많은 에너지를 쏟게 만들었습니다.
• SQDF: "초기 단계는 크게 신경 쓰지 말고, 마무리 단계에 집중해라"라고 가르칩니다. 이렇게 하면 AI 는 불필요한 초기 노이즈에 휘둘리지 않고, 최종 결과물의 질을 높이는 데 집중하게 됩니다.

🎨 비법 2: "예측의 달인" (Consistency Model)

• 비유: 요리사가 "이 재료를 넣으면 어떻게 될까?"라고 상상할 때, 기존 방법은 어렴풋한 느낌으로 예측해서 자주 틀렸습니다. (소금 한 스푼 넣었는데 짠맛이 안 날 수도 있죠.)
• SQDF: 더 정확한 예측을 위해 **전문 예견가 (Consistency Model)**를 고용합니다. 이 예견가는 "지금 이 상태의 재료를 섞으면, 최종 요리가 이렇게 될 거야!"라고 정확하게 예측해 줍니다.
• 효과: AI 는 정확한 예측을 바탕으로 학습하므로, 엉뚱한 방향으로 나가지 않고 효율적으로 실력을 늘립니다.

📚 비법 3: "실패와 성공의 기록장" (Replay Buffer)

• 비유: 요리사가 매일 새로운 재료만 쓰다가, 가장 맛있었던 요리나 다양한 스타일의 요리를 잊어버릴 수 있습니다.
• SQDF: 과거에 만든 **가장 맛있는 요리 (높은 보상)**와 가장 다양한 요리들을 모아둔 **기록장 (리플레이 버퍼)**을 사용합니다.
• 효과: AI 는 과거의 성공 사례와 다양한 실험들을 반복해서 학습하며, "맛"과 "다양성" 사이의 균형을 잘 맞추게 됩니다.

---

3. SQDF 의 성과: "자연스러움"과 "고급스러움"의 동시 달성

이 새로운 방법 (SQDF) 으로 훈련된 AI 요리사들은 어떤 결과를 보여줄까요?

1. 원하는 맛을 정확히 냈습니다: "더 맛있게"라는 주문을 잘 들어주어 점수가 매우 높습니다.
2. 괴물 요리를 만들지 않았습니다: 설탕을 너무 많이 넣거나 모양이 이상해지는 '과도한 최적화'가 일어나지 않았습니다.
3. 다양성을 유지했습니다: 모든 요리가 똑같은 맛이 아니라, 각각 개성이 살아있는 다양한 요리를 만들어냅니다.

4. 결론: 왜 이 논문이 중요한가요?

기존의 AI 학습 방법은 "점수만 높이면 된다"는 생각에 AI 를 미치게 만들었습니다. 하지만 SQDF는 AI 가 원래 가진 자연스러운 능력 (다양성, 자연스러움) 을 잃지 않으면서도, 우리가 원하는 목표 (고퀄리티, 특정 스타일) 를 달성할 수 있게 해줍니다.

마치 유능한 요리사에게 "맛을 더 살려줘"라고 말하면서도, "너의 원래 스타일과 창의성은 유지해"라고 조언해주는 현명한 멘토와 같은 역할을 하는 것입니다.

이 기술은 앞으로 우리가 AI 에게 그림을 그리게 하거나, 새로운 분자를 설계할 때, 원하는 대로 만들면서도 AI 가 망가지지 않게 하는 핵심 열쇠가 될 것입니다.

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 확산 모델 (Diffusion Models) 은 텍스트 - 이미지 합성, 비디오 생성 등 다양한 분야에서 고품질 샘플 생성 능력을 입증했습니다. 그러나 생성된 샘플을 하류 작업의 목표 (예: 미적 품질, 인간 선호도, 분자 활성도) 에 맞추기 위해 보상 (Reward) 을 최적화하는 파인튜닝이 필요합니다.
• 핵심 문제: 기존 보상 최적화 기반 파인튜닝 방법 (RL 기반 또는 직접 역전파 기반) 은 보상 과최적화 (Reward Over-optimization) 현상에 시달립니다. 이는 높은 보상 점수를 얻지만, **의미론적 붕괴 (Semantic Collapse, 프롬프트와 무관한 이미지 생성)**와 **다양성 붕괴 (Diversity Collapse, 유사한 패턴만 반복)**를 초래하여 실제 활용성을 떨어뜨립니다.
• 기존 방법의 한계:
  • KL 발산 정규화를 도입한 방법들은 과최적화를 완화하려 하지만, 불안정한 가치 함수 (Value Function) 학습이나 높은 분산을 가진 몬테카를로 추정기에 의존해야 하는 문제가 있습니다.
  • 직접 역전파 (Direct Backpropagation) 방식은 보상 기울기를 활용하지만, 과최적화 방지를 위한 안정적인 메커니즘이 부족합니다.

2. 제안 방법: SQDF (Methodology)

저자들은 **Soft Q 기반 확산 파인튜닝 (Soft Q-based Diffusion Finetuning, SQDF)**을 제안합니다. 이는 KL 정규화 강화학습 (RL) 프레임워크 내에서 학습이 불필요한 (Training-free) 소프트 Q 함수의 재파라미터화 정책 기울기를 활용하는 새로운 접근법입니다.

핵심 메커니즘

1. 소프트 Q 함수의 근사 (Approximation of Soft Q-Function):
   • 기존 RL 방법처럼 별도의 Q-네트워크를 학습시키지 않고, Tweedie's formula를 기반으로 한 단일 단계 후사 평균 (Posterior Mean) 근사 ($\hat{x}_0$) 를 사용합니다.
   • 이 근사된 $\hat{x}_0$에 보상 함수 $r(\cdot)$를 적용하여 소프트 Q 함수를 직접 추정합니다. 이는 가치 함수 학습의 불안정성을 제거합니다.
2. 재파라미터화 정책 기울기 (Reparameterized Policy Gradient):
   • 추정된 Q 함수의 기울기를 보상 기울기로 직접 활용합니다.
   • 샘플링 과정 ($x_{t-1} = \mu_\theta(x_t, t) + \sigma_t \epsilon$) 을 재파라미터화하여, 확률 변수인 $x_{t-1}$을 통해 보상 기울기를 모델 파라미터 $\theta$에 직접 전파 (Backpropagation) 할 수 있게 합니다. 이는 낮은 분산과 높은 샘플 효율성을 보장합니다.
3. KL 정규화:
   • 사전 학습된 모델 ($p'$) 에서의 KL 발산을 패널티로 추가하여, 파인튜닝된 모델이 원래 데이터 분포 (자연스러움과 다양성) 에서 크게 벗어나지 않도록 제약합니다.

세 가지 주요 혁신 (Key Innovations)

1. 할인 계수 (Discount Factor, $\gamma$) 도입:
   • 확산 과정의 초기 단계 (고잡음 상태) 는 최종 샘플 품질에 미치는 영향이 적습니다. 따라서 시간 단계 $t$에 따라 보상을 $\gamma^{t-1}$만큼 할인하여, 초기 단계의 부정확한 근사 오차가 학습에 미치는 영향을 줄이고 신용 부여 (Credit Assignment) 를 개선합니다.
2. 일관성 모델 (Consistency Model) 통합:
   • 초기 단계의 Tweedie's formula 근사는 잡음이 많을 때 정확도가 낮습니다. 이를 해결하기 위해 **일관성 모델 (Consistency Model)**을 사용하여 $x_0$를 더 정확하게 예측하고, 이를 통해 소프트 Q 함수 추정의 신뢰도를 높입니다.
3. 오프-폴리시 리플레이 버퍼 (Off-policy Replay Buffer):
   • 희귀하지만 높은 보상을 받는 샘플과 다양한 샘플을 버퍼에 저장하여 재사용합니다. 이는 모드 커버리지 (Mode Coverage) 를 향상시키고, 보상과 다양성 간의 트레이드오프를 관리하는 데 기여합니다.

3. 주요 실험 결과 (Results)

저자들은 텍스트 - 이미지 파인튜닝과 온라인 블랙박스 최적화 두 가지 시나리오에서 SQDF 를 평가했습니다.

• 텍스트 - 이미지 파인튜닝 (Stable Diffusion 1.5 & XL):

  • 목표: LAION 미적 점수 (Aesthetic Score) 와 인간 선호도 점수 (HPSv2) 최적화.
  • 결과: 기존 방법 (DDPO, DRaFT, ReFL 등) 은 보상을 높이면 다양성과 정렬 (Alignment) 이 급격히 떨어지는 반면, SQDF 는 높은 보상 점수를 유지하면서도 다양성과 프롬프트 정렬을 효과적으로 보존했습니다.
  • KL 정규화 비교: 단순히 기존 방법에 KL 항을 추가한 방법들보다 SQDF 가 더 나은 Pareto 최적성 (보상 vs 다양성/정렬) 을 보였습니다.

• 온라인 블랙박스 최적화:

  • 설정: 제한된 쿼리 예산 내에서 블랙박스 오라클 (미적 점수) 을 최적화하면서 자연스러움과 다양성을 유지해야 하는 과제.
  • 결과: SEIKO (기존 KL 정규화 직접 역전파 방법) 와 PPO+KL 대비 **SQDF 는 더 높은 표본 효율성 (Sample Efficiency)**을 보이며, 오라클 보상 점수, 정렬, 다양성 모든 지표에서 우월한 성능을 달성했습니다. 특히 부정확한 보상 모델 (Surrogate) 하에서도 강건함을 입증했습니다.

• 성분 분석 (Ablation Study):

  • 할인 계수 $\gamma$를 제거하면 초기 학습이 느려지고 정렬/다양성 점수가 하락합니다.
  • 일관성 모델 (CM) 을 제거하면 수렴 속도가 느려지고 목표 보상이 감소합니다.
  • 리플레이 버퍼를 제거하면 다양성 점수가 떨어집니다.

4. 기여 및 의의 (Contributions & Significance)

1. 안정적인 RL 기반 파인튜닝 프레임워크: 가치 함수 학습 없이도 보상 기울기를 직접 활용할 수 있는 학습 불필요 (Training-free) 소프트 Q 함수 근사를 통해, 확산 모델의 RL 파인튜닝 안정성을 획기적으로 개선했습니다.
2. 과최적화 해결: 보상 과최적화로 인한 의미론적/다양성 붕괴를 효과적으로 억제하면서도, 하류 작업의 목표를 극대화하는 방법을 제시했습니다.
3. 실용성: 텍스트 - 이미지 생성뿐만 아니라 분자 설계 등 다양한 생성 모델 적용 분야에서, 제한된 데이터/쿼리 예산 하에서도 고품질 생성을 가능하게 하는 고효율 최적화 알고리즘으로서의 가치를 입증했습니다.
4. 개방성: 코드와 모델을 공개하여 (GitHub 링크 제공), 연구 커뮤니티의 재현과 추가 연구를 장려했습니다.

5. 결론

이 논문은 확산 모델의 파인튜닝 과정에서 발생하는 "보상 과최적화" 문제를 해결하기 위해, 재파라미터화 정책 기울기와 학습 불필요 소프트 Q 함수를 결합한 SQDF를 제안했습니다. 할인 계수, 일관성 모델, 오프-폴리시 버퍼라는 세 가지 기법을 통해 학습 안정성과 효율성을 동시에 확보한 SQDF 는 기존 방법들보다 우수한 보상 최적화 성능과 함께 자연스러움 및 다양성을 유지하는 새로운 기준을 제시했습니다.