Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases

이 논문은 확산 모델의 시간적 귀납 편향과 활성 뉴런에서 비롯된 선입견 편향을 분석하여, 이를 해결하는 'TDPO-R' 알고리즘을 제안함으로써 보상 과최적화 문제를 효과적으로 완화합니다.

핵심

🎨 핵심 비유: "맛있는 요리를 만드는 AI 요리사"

상상해 보세요. AI 는 이제 막 요리를 배우는 신입 요리사입니다. 우리는 이 요리사에게 "더 맛있게 만들어!"라고 말하며 (이게 '보상 모델'입니다) 요리를 가르칩니다.

하지만 문제는, 요리사가 "맛있다"는 기준만 맹신하다가 정작 요리의 본질을 잃어버리는 경우가 생긴다는 것입니다.

• 과도한 최적화 (Reward Overoptimization): 요리사가 "맛있다"는 점수를 더 받기 위해, 소금을 너무 많이 넣거나, 심지어 요리에 먹을 수 없는 물건을 넣어서 점수만 높이는 식의 비정상적인 행동을 하는 것입니다. 결과물은 점수는 높지만, 실제로는 먹거나 볼 수 없는 괴물이 되어버립니다.

이 논문은 **"왜 이런 일이 일어나는지"**를 찾아내고, "어떻게 고칠지" 두 가지 관점에서 해결책을 제시합니다.

---

🔍 원인 1: "시간의 흐름을 무시한 점수 매기기" (유도 편향의 불일치)

기존 방식의 문제:
기존에는 요리사가 요리를 다 끝낸 **최종 결과물 (완성된 요리)**만 보고 점수를 매겼습니다.

• 비유: 요리사가 재료를 섞고, 가열하고, 마지막에 접시에 담는 전 과정을 보지 않고, 오직 마지막 접시만 보고 "이건 10 점, 저건 0 점"이라고 점수를 줍니다.
• 문제점: 요리사는 "어떻게" 맛있는 요리를 만들었는지 (과정) 는 배울 수 없고, "무엇"이 맛있는지 (결과) 만 외우게 됩니다. 그래서 점수만 높이는 이상한 방법 (소금 과다 등) 을 찾아내게 됩니다.

이 논문의 해결책 (TDPO): "과정마다 점수를 주자!"
이 연구는 **"시간적 유추 (Temporal Inductive Bias)"**라는 개념을 도입했습니다.

• 해결책: 요리사가 재료를 섞는 순간, 가열하는 순간, 마지막에 담는 순간 매 순간마다 점수를 매겨줍니다.
• 효과: 요리사는 "어떻게" 단계별로 맛을 내는지 자연스럽게 배우게 됩니다. 결과물만 보고 점수를 매기는 것보다 훨씬 효율적이고, 이상한 방법을 쓰지 않아도 됩니다.

---

🧠 원인 2: "오래된 기억에 갇힌 뇌" (최초 편향, Primacy Bias)

기존 방식의 문제:
AI 의 뇌 (크리틱 모델) 안에는 수많은 **뉴런 (세포)**들이 있습니다.

• 비유: 요리사가 처음 배운 레시피를 너무 강하게 기억해서, 새로운 레시피를 배우려 해도 그 오래된 기억이 방해가 되는 상황입니다.
• 발견: 연구진은 놀라운 사실을 발견했습니다.
  1. 활발한 뉴런 (Active Neurons): 지금 가장 많이 쓰이는 세포들입니다. 이 세포들이 "처음 배운 나쁜 습관 (과도한 점수 맞추기)"을 너무 강하게 고집하며 문제를 일으킵니다.
  2. ** dormant(잠자는) 뉴런:** 잘 쓰이지 않는 세포들입니다. 기존 연구들은 이 세포들을 깨워야 한다고 했지만, 이 논문은 **"잠자는 세포들은 오히려 AI 를 정상적으로 유지해주는 '안전장치' 역할을 한다"**고 발견했습니다.

이 논문의 해결책 (TDPO-R): "활발한 세포만 주기적으로 리셋하자"

• 해결책: AI 가 너무 오래된 나쁜 습관 (최초 편향) 에 갇히지 않도록, **가장 활발하게 움직이는 세포들만 주기적으로 초기화 (리셋)**해줍니다.
• 효과: 마치 요리사가 "너무 익숙해진 나쁜 버릇을 잊고, 다시 초심으로 돌아가 새로운 방법을 배우게" 하는 것과 같습니다. 잠자는 세포 (안전장치) 는 건드리지 않고, 문제되는 세포만 리셋해서 AI 가 더 똑똑하고 균형 잡힌 요리를 만들게 됩니다.

---

🏆 결론: 무엇이 달라졌나요?

이 연구는 TDPO-R이라는 새로운 방법을 제안했습니다.

1. 과정 중심 학습: 요리 (생성) 과정의 매 순간을 평가해서, AI 가 결과만 쫓지 않게 합니다.
2. 뇌 리셋 전략: AI 가 나쁜 습관 (과도한 최적화) 에 빠지지 않도록, 문제되는 뇌 세포만 주기적으로 초기화합니다.

실제 결과:
기존 방법들 (DDPO, AlignProp 등) 은 점수를 높이면 높일수록 이미지가 뭉개지거나 이상해지는 현상이 심했지만, 이 논문의 방법 (TDPO-R) 은 점수도 높으면서도 이미지의 질과 다양성도 유지했습니다. 마치 "점수만 높은 괴물 요리"가 아니라, "진짜 맛있는 요리"를 만드는 데 성공한 것입니다.

💡 한 줄 요약

"AI 가 점수만 쫓다가 망가지는 것을 막기 위해, 과정마다 칭찬을 해주고, 나쁜 습관을 가진 뇌세포만 주기적으로 청소해 주니, AI 가 훨씬 더 똑똑하고 안정적인 요리를 만들게 되었습니다."

기술 요약

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

• 배경: 확산 모델 (Diffusion Models) 은 텍스트 - 이미지 생성 등에서 최첨단 성능을 보이지만, 이를 실제 응용에 맞게 인간 선호도 (Human Preferences) 와 정렬 (Alignment) 하는 과정이 필수적입니다. 이를 위해 학습된 보상 모델 (Reward Model) 을 사용하여 모델을 미세 조정 (Finetuning) 하는 방식이 주로 사용됩니다.
• 핵심 문제: 보상 과최적화 (Reward Overoptimization): 보상 모델을 지나치게 최적화할 경우, 모델은 실제 인간이 원하는 품질 (Ground-truth) 을 희생하면서 보상 점수만 극대화하는 현상이 발생합니다. 이는 생성된 이미지의 충실도 (Fidelity) 저하, 도메인 외 일반화 능력 상실, 그리고 샘플 효율성 (Sample Efficiency) 과 과최적화 사이의 트레이드오프를 야기합니다.
• 기존 연구의 한계: 기존 정렬 방법들은 확산 모델의 다단계 탈노이즈 (Multi-step Denoising) 과정의 시간적 인덕티브 바이어스 (Temporal Inductive Bias) 를 무시하고 최종 생성된 이미지만을 기준으로 보상을 계산합니다. 또한, 강화학습 (RL) 에이전트에서 초기 학습 경험에 과도하게 의존하는 프라이머시 바이어스 (Primacy Bias) 가 보상 과최적화의 원인이 될 수 있다는 점은 충분히 연구되지 않았습니다.

2. 제안 방법: TDPO-R (Methodology)

저자들은 인덕티브 바이어스 (Inductive Bias) 와 프라이머시 바이어스 (Primacy Bias) 두 가지 관점에서 문제를 해결하기 위해 TDPO-R (Temporal Diffusion Policy Optimization with critic active neuron Reset) 알고리즘을 제안합니다.

A. 시간적 인덕티브 바이어스 활용 (TDPO)

기존 방법과의 불일치를 해결하기 위해 확산 모델의 탈노이즈 과정을 시간 단계 (Timestep) 별 보상을 갖는 MDP 로 재정의합니다.

• 시간 의존적 보상 (Temporal Rewards): 최종 이미지뿐만 아니라 각 탈노이즈 단계 ($x_t$) 에서도 보상을 부여합니다.
• Temporal Critic 학습: 기존 보상 모델은 최종 깨끗한 이미지 ($x_0$) 에만 훈련되어 있으므로, 중간 노이즈 이미지 ($x_t$) 에 대한 보상을 직접 구할 수 없습니다. 이를 해결하기 위해, 최종 보상 $R(x_0, c)$ 를 기준 (Anchor) 으로 하고, Temporal Critic을 학습하여 중간 단계의 잔차 (Residual) 를 예측하는 방식을 사용합니다.
  • $T(x_t, c) \approx R(x_0, c) - R_\phi(x_t, c)$
• Timestep 별 그래디언트 업데이트: 기존 배치 단위 업데이트 대신, 각 탈노이즈 타임스텝마다 정책 ($\theta$) 과 크리티크 ($\phi$) 를 동시에 업데이트합니다. 이는 시간적 인덕티브 바이어스를 활용하여 샘플 효율성을 높이고 과최적화를 완화합니다.

B. 프라이머시 바이어스 해결 (TDPO-R)

TDPO 프레임워크 내에서 크리티크 모델의 뉴런 상태를 분석하여 프라이머시 바이어스를 식별하고 해결합니다.

• 뉴런 상태 분석:
  • 휴면 뉴런 (Dormant Neurons): 학습 중 활성화가 낮은 뉴런. 기존 RL 연구에서는 이들을 재초기화해야 한다고 보았으나, 본 연구에서는 이들이 보상 과최적화에 대한 적응적 정규화 (Adaptive Regularization) 역할을 한다는 것을 발견했습니다.
  • 활성 뉴런 (Active Neurons): 활성화가 높은 뉴런. 이들이 프라이머시 바이어스 (초기 학습 경험에 대한 과적합) 를 반영하여 과최적화를 유발합니다.
• 활성 뉴런 리셋 전략:
  • 휴면 뉴런을 재초기화하면 과최적화가 심화되지만, 크리티크 모델의 활성 뉴런을 주기적으로 재초기화 (Reset) 하면 과최적화가 완화됩니다.
  • 이는 모델이 새로운 정규화 패턴을 학습하도록 유도하면서도, 휴면 뉴런이 제공하는 중요한 정규화 메커니즘은 보존하는 전략입니다.

3. 주요 기여 (Key Contributions)

1. 새로운 관점의 발견: 확산 모델 정렬에서 보상 과최적화의 근본 원인을 인덕티브 바이어스 불일치와 프라이머시 바이어스의 관점에서 최초로 체계적으로 분석했습니다.
2. TDPO 프레임워크 제안: 확산 모델의 시간적 인덕티브 바이어스를 활용하여, 시간 단계별 보상을 기반으로 한 RL 정렬 프레임워크 (TDPO) 를 설계했습니다. 이는 과최적화를 줄이면서도 샘플 효율성을 동시에 개선합니다.
3. TDPO-R 알고리즘 개발: 크리티크 모델의 활성 뉴런이 프라이머시 바이어스를 반영한다는 사실을 발견하고, 이를 주기적으로 리셋하는 전략 (TDPO-R) 을 도입하여 과최적화를 추가적으로 완화했습니다.
4. 정량적 평가 지표 개발: 보상 과최적화를 평가하기 위해 교차 보상 일반화 (Cross-reward Generalization) 지표를 도입했습니다. 이는 특정 보상 함수로 학습된 모델이 다른 도메인의 보상 함수에서도 얼마나 잘 수행하는지를 측정하여 과최적화 정도를 수치화합니다.

4. 실험 결과 (Results)

• 실험 설정: Stable Diffusion v1.4 를 기반으로 Aesthetic Score, PickScore, HPSv2 등 다양한 보상 함수로 학습 및 평가 수행.
• 샘플 효율성 (Sample Efficiency): TDPO 및 TDPO-R 은 기존 DDPO (Denoising Diffusion Policy Optimization) 및 AlignProp 대비 적은 샘플 수 (Reward Queries) 로 더 높은 보상 점수를 달성했습니다. 이는 Timestep 별 업데이트 전략이 학습 효율을 높였음을 의미합니다.
• 과최적화 완화 (Mitigation of Overoptimization):
  • 교차 보상 일반화: 특정 보상 (예: Aesthetic Score) 으로 학습한 모델을 다른 보상 (예: ImageReward, PickScore) 으로 평가했을 때, TDPO-R 이 DDPO 나 AlignProp 보다 훨씬 우수한 일반화 성능을 보였습니다. 이는 과최적화가 효과적으로 억제되었음을 의미합니다.
  • 시각적 품질: 고보상 영역에서 생성된 이미지를 비교했을 때, 기존 방법들은 스타일, 배경, 조명 등에서 과적합 (Saturation) 현상을 보인 반면, TDPO-R 은 더 높은 충실도 (Fidelity) 와 다양성을 유지했습니다.
• 뉴런 리셋 전략의 효과: 크리티크의 활성 뉴런만 리셋한 TDPO-R 이 가장 좋은 성능을 보였으며, 휴면 뉴런을 리셋하거나 모든 뉴런을 리셋하는 경우 오히려 성능이 저하되거나 과최적화가 심화됨을 확인했습니다.

5. 의의 및 결론 (Significance)

• 이론적 통찰: 확산 모델의 다단계 생성 과정을 단순히 최종 결과물이 아닌, 시간적 흐름을 가진 과정으로 접근해야 함을 증명했습니다. 또한, 신경망 내 뉴런의 상태 (활성/휴면) 가 과최적화 현상과 어떻게 연결되는지에 대한 새로운 통찰을 제공했습니다.
• 실용적 가치: 제안된 TDPO-R 은 보상 과최적화 문제를 해결하면서도 학습 효율성을 높여, 실제 산업 환경에서 인간 선호도와 정렬된 고품질 생성 모델을 개발하는 데 기여할 수 있습니다.
• 미래 연구 방향: 본 연구는 확산 모델의 다중 보상 학습 (Multi-reward Learning) 영역에서의 과최적화 연구에 대한 새로운 길을 열었으며, 다른 강화학습 분야에서의 바이어스 해결에도 적용 가능한 가능성을 제시합니다.

요약하자면, 이 논문은 확산 모델의 정렬 과정에서 발생하는 보상 과최적화 문제를 시간적 구조의 불일치와 신경망의 초기 학습 편향이라는 두 가지 핵심 원인을 규명하고, 이를 해결하기 위해 시간적 보상을 도입한 RL 프레임워크와 활성 뉴런 리셋 전략을 결합한 혁신적인 솔루션을 제시했습니다.