Diffusion Alignment as Variational Expectation-Maximization

Each language version is independently generated for its own context, not a direct translation.

🎨 배경: 왜 이 기술이 필요한가요?

지금까지의 AI 그림 그리기 기술 (확산 모델) 은 아주 훌륭한 화가였습니다. 하지만 이 화가를 "더 예쁜 그림을 그려달라"고 주문하면, 두 가지 문제가 생겼습니다.

과도한 최적화 (Reward Over-optimization): 화가가 "예쁘게"라는 말만 듣고, 사람 얼굴을 그릴 때 눈이 너무 크거나 피부가 너무 매끄러운 등 비현실적이고 기괴한 그림만 그리게 됩니다. (너무 잘 하려다 망친 경우)
모드 붕괴 (Mode Collapse): 화가가 "가장 예쁜 그림"을 하나만 찾아서, 그걸 수천 번 복사해서 내어놓습니다. 다양성이 사라지고 모든 그림이 똑같아집니다.

기존 방법들은 이 화가를 훈련시킬 때, "정답"을 강제로 주입하거나 (역전파), 화가가 실수할 때마다 매를 들며 훈련시켰습니다 (강화학습). 하지만 이 방식들은 화가의 창의성을 죽이거나, 이상한 그림만 그리게 만드는 부작용이 있었습니다.

💡 DAV 의 해결책: "탐험가"와 "스승"의 듀오

이 논문은 이 문제를 해결하기 위해 DAV라는 새로운 훈련 방식을 제안합니다. 이 방식은 **E-Step(탐험)**과 **M-Step(정리)**라는 두 단계를 반복하는 EM 알고리즘을 기반으로 합니다.

이를 **유명한 화가 (AI 모델)**와 **그의 예술가 지망생 (탐험가)**의 관계로 비유해 볼까요?

1 단계: E-Step (탐험 - "잠깐, 화가님, 잠시 쉬세요!")

상황: 화가 (AI) 가 그릴 준비를 했지만, 아직 완벽하지 않습니다.
행동: 대신 **탐험가 (테스트 타임 서치)**가 나갑니다. 탐험가는 화가의 기본 실력을 바탕으로, "어떤 그림이 가장 반응이 좋을지" 찾아다니며 **수백 개의 시안 (후보 그림)**을 그려봅니다.
특징: 이때 탐험가는 단순히 하나만 고르는 게 아니라, 다양한 스타일의 좋은 그림들을 골고루 모읍니다. (다양성 유지)
비유: 마치 요리사가 "최고의 맛"을 찾기 위해 실험실로 가서 다양한 재료를 섞어보며 최고의 레시피를 찾아내는 과정입니다.

2 단계: M-Step (정리 - "자, 이 레시피를 배우세요!")

상황: 탐험가가 찾아온 최고의 시안들 (다양하고 좋은 그림들) 이 모였습니다.
행동: 이제 **스승 (AI 모델)**이 이 시안들을 보고 "아, 이렇게 그리면 사람들이 좋아하구나!"라고 배웁니다.
핵심: 기존 방식은 "가장 좋은 것 하나"만 따라 하게 했지만, DAV 는 다양한 좋은 것들 전체를 학습합니다. 그래서 화가는 "예쁜 그림"을 그리면서도, 다양한 스타일을 잃지 않게 됩니다.
비유: 요리사가 실험실에서 찾아온 10 가지의 맛있는 레시피를 모두 노트에 적어두고, 자신의 요리 스킬을 업그레이드하는 것입니다.

이 과정을 탐험 (E-Step) → 학습 (M-Step) → 다시 탐험 → 다시 학습으로 반복하면, AI 는 점점 더 똑똑해지면서도 창의성을 잃지 않게 됩니다.

🌍 이 기술로 무엇을 할 수 있나요?

이 논문은 이 기술이 두 가지 다른 분야에서 잘 작동함을 증명했습니다.

연속적인 데이터 (이미지 생성):
- 상황: "고양이 그림을 그려줘"라고 했을 때, 단순히 고양이만 그리는 게 아니라 더 예쁘고, 사람 취향에 맞으면서도 고양이마다 다른 표정과 자세를 가진 그림을 그립니다.
- 결과: 기존 방법들은 고양이 눈이 너무 커지거나 (과도한 최적화), 모든 고양이 그림이 똑같아지는 (다양성 붕괴) 문제가 있었지만, DAV 는 이를 해결했습니다.
이산적인 데이터 (DNA 설계):
- 상황: "기능이 좋은 DNA 서열을 만들어줘"라고 했을 때, 단순히 기능만 좋게 만드는 게 아니라 자연스러운 DNA 구조를 유지하면서 다양한 변이를 만들어냅니다.
- 결과: 생물학 실험에서 쓸 수 있는 실제 DNA 를 설계할 때, 너무 인위적이거나 위험한 구조가 나오지 않도록 도와줍니다.

🚀 요약: 왜 이것이 중요한가요?

기존의 AI 훈련 방식은 "정답 하나만 찾아서 그걸 외우게 하는" 방식이었다면, DAV 는 "다양한 정답들을 찾아보고, 그걸 종합해서 더 똑똑하게 만드는" 방식입니다.

창의성 유지: AI 가 똑같은 결과물만 반복하지 않게 합니다.
과도한 욕심 방지: "더 잘해"라는 주문 때문에 AI 가 망가지는 것을 막아줍니다.
범용성: 그림을 그리는 AI 이든, DNA 를 설계하는 AI 이든 모두 적용 가능합니다.

결론적으로, DAV 는 AI 가 **인간이 원하는 목표를 달성하면서도, 여전히 다양하고 자연스러운 결과물을 만들어낼 수 있게 해주는 '지혜로운 훈련 코치'**라고 할 수 있습니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 이미지 생성, 로봇 제어, 계산 생물학 등 다양한 분야에서 고품질 샘플 생성에 탁월한 성능을 보이지만, 실제 응용에서는 외부 기준 (예: 이미지의 미적 품질, DNA 서열의 생물학적 활성) 에 최적화된 샘플이 필요합니다. 이를 위해 기존에 제안된 확산 모델 파인튜닝 방법들은 주로 두 가지 접근법을 취해 왔으나, 각각 심각한 한계를 가지고 있습니다.

강화학습 (RL) 기반 방법: 역 KL 발산 (Reverse-KL) 을 최소화하는 방식으로 최적화합니다. 이는 단일 우세한 모드 (Mode) 를 찾게 되어 **모드 붕괴 (Mode Collapse)**와 **과도한 최적화 (Reward Over-optimization)**를 유발하며, 샘플의 다양성과 자연스러움이 크게 저하됩니다.
직접 역전파 (Direct Backpropagation) 방법: 학습된 보상 함수를 통해 직접 기울기를 전파합니다. 샘플 효율성은 높지만, 보상 함수의 기울기 신호가 날카롭고 취약하여 (Brittle) 역시 심각한 과도한 최적화를 초래합니다.

따라서, 보상 함수를 극대화하면서도 사전 학습된 확산 모델의 **다양성 (Diversity)**과 **자연스러움 (Naturalness)**을 유지할 수 있는 새로운 파인튜닝 프레임워크가 시급히 필요했습니다.

2. 방법론 (Methodology)

저자들은 **DAV (Diffusion Alignment as Variational Expectation-Maximization)**라는 새로운 프레임워크를 제안합니다. 이는 변분 기대치 최대화 (Variational EM) 알고리즘을 확산 모델 정렬에 적용한 것으로, **E-step (탐색)**과 **M-step (평균화/Amortization)**을 반복하는 과정으로 구성됩니다.

핵심 아이디어:

E-step (Posterior Exploration):
- 현재 모델 파라미터 ( $\theta_k$ ) 하에서 변분 사후 분포 (Variational Posterior) 를 탐색합니다.
- **테스트 시간 검색 (Test-time Search)**을 활용하여 높은 보상을 얻는 다양한 트래젝토리를 생성합니다.
- 구체적으로, 소프트 Q 함수 (Soft Q-function) 를 기반으로 **기울기 가이드 (Gradient Guidance)**와 **중요도 샘플링 (Importance Sampling)**을 결합하여, 사후 분포에 가까운 고품질 샘플을 효율적으로 탐색합니다.
- 이 단계는 모델 가중치를 변경하지 않고, 추론 시 추가 계산을 통해 최적의 샘플을 찾습니다.
M-step (Amortization via Forward-KL):
- E-step 에서 발견된 샘플들을 통해 확산 모델 파라미터 ( $\theta$ ) 를 업데이트합니다.
- 기존 RL 기반 방법들이 역 KL 발산을 최소화 (단일 모드 집중) 하는 것과 달리, DAV 는 **순방향 KL 발산 (Forward-KL)**을 최소화합니다.
- 이는 E-step 에서 발견된 **모든 다양한 모드 (Multi-modal)**를 모델이 포괄하도록 유도하여, 샘플의 다양성을 유지하면서 보상을 최적화합니다.

수식적 배경:

최적성 변수 $O=1$ 의 우도 (Likelihood) 를 최대화하는 문제로 재정의합니다.
할인 인자 (Discount factor, $\gamma$ ) 를 도입하여 확산 과정의 초기 단계에서의 불확실성을 고려하고, TWEEDIE 공식을 활용하여 소프트 Q 함수를 근사합니다.
사전 학습된 모델의 능력을 잃지 않도록 KL 정규화 항 (DAV-KL) 을 추가할 수 있습니다.

3. 주요 기여 (Key Contributions)

새로운 정렬 프레임워크 (DAV) 제안: RL 과 직접 역전파의 단점을 극복하고, 테스트 시간 검색과 변분 EM 을 결합하여 보상 최적화와 다양성 유지를 동시에 달성하는 새로운 접근법을 제시했습니다.
모드 커버링 (Mode-Covering) 최적화: 순방향 KL 발산을 최소화하는 M-step 업데이트를 통해, 기존 방법들이 겪던 모드 붕괴 문제를 해결하고 다양한 고품질 샘플을 생성하도록 유도합니다.
연속 및 이산 도메인에서의 검증:
- 연속 확산 (Continuous): 텍스트 - 이미지 합성 (Stable Diffusion v1.5) 에서 미적 품질, 비차분 가능 보상 (압축성 등) 최적화를 수행했습니다.
- 이산 확산 (Discrete): DNA 서열 설계 (Masked Diffusion) 에서 고활성 DNA 엔핸서 생성 및 자연스러움/다양성 유지에 성공했습니다.
비차분 가능 보상 (Non-differentiable Rewards) 지원: 기울기 기반 접근법이 불가능한 블랙박스 보상 함수 (예: 이미지 압축률, 생물학적 활성 예측기) 에도 적용 가능한 범용성을 입증했습니다.

4. 실험 결과 (Results)

A. 텍스트 - 이미지 합성 (Text-to-Image)

성능: DAV 는 기존 RL 기반 방법 (DDPO, TDPO) 과 직접 역전파 방법 (DRaFT) 보다 **높은 미적 점수 (Aesthetic Score)**를 기록했습니다.
다양성 유지: 기존 방법들은 보상을 높이는 과정에서 이미지 품질 (ImageReward) 이나 다양성 (LPIPS) 이 급격히 떨어지는 '과도한 최적화' 현상을 보인 반면, DAV 는 높은 보상과 함께 높은 ImageReward 와 다양성을 유지했습니다.
비차분 가능 보상: 이미지 압축성 (Compressibility) 과 같은 비차분 가능 보상에서도 DDPO 대비 더 짧은 학습 시간으로 우수한 성능을 보였습니다.

B. DNA 서열 설계 (DNA Sequence Design)

성과: DAV 는 기존 방법들 (DRAKES, VIDD, DDPO) 보다 **예측된 활성 (Pred-Activity)**이 높은 DNA 서열을 생성했습니다.
생물학적 유효성: 과도한 최적화로 인해 생물학적 유효성 (ATAC-Acc) 이 떨어지는 다른 방법들과 달리, DAV 는 높은 활성을 유지하면서도 **자연스러움 (3-mer Correlation)**과 **다양성 (Levenshtein Distance)**을 잘 보존했습니다.
테스트 시간 검색의 효과: DAV Posterior (학습된 모델 + 테스트 시간 검색) 는 가장 높은 보상 점수를 기록하면서도 다양성을 유지했습니다.

5. 의의 및 결론 (Significance & Conclusion)

균형 잡힌 최적화: DAV 는 확산 모델 정렬 분야에서 '보상 극대화'와 '다양성/자연스러움 유지'라는 상충되는 목표를 효과적으로 조화시킨 최초의 체계적인 프레임워크 중 하나입니다.
범용성: 연속 데이터 (이미지) 와 이산 데이터 (DNA) 모두에 적용 가능하며, 보상 함수의 미분 가능성 여부에 구애받지 않습니다.
미래 연구 방향: 테스트 시간 검색의 계산 비용은 여전히 과제로 남아있으나, E-step 이 모듈러하게 설계되어 향후 더 효율적인 검색 알고리즘과 결합될 경우 성능을 더욱 향상시킬 수 있는 잠재력을 가지고 있습니다.

요약하자면, DAV는 변분 EM 알고리즘의 E-step(탐색) 과 M-step(학습) 을 반복하여, 확산 모델이 다양한 고품질 샘플을 생성하면서도 외부 보상 기준에 정렬되도록 하는 혁신적인 파인튜닝 방법론입니다.