CREPE: Controlling Diffusion with Replica Exchange

이 논문은 재학습 없이 확산 모델의 추론을 제어하기 위해 시퀀셜 몬테 카를로 (SMC) 대안으로 레플리카 교환 기반의 CREPE 방법을 제안하며, 이 방법은 높은 샘플 다양성 유지와 온라인 정제 기능을 통해 다양한 제어 작업에서 경쟁력 있는 성능을 보입니다.

핵심

1. 문제 상황: "그림을 그리는데, AI 가 제멋대로 구는 경우"

AI 가 그림을 그릴 때, 우리는 "노란 택시"라고 주문을 하지만 AI 는 가끔 "검은색 버스"를 그리거나, "노란색이지만 너무 어색한" 그림을 그릴 수 있습니다.

기존의 해결책 (SMC 라고 부름) 은 다음과 같았습니다:

• 비유: 100 명의 화가를 동시에 불러모아 그림을 그리게 합니다.
• 과정: 그림이 조금씩 완성될 때마다, "아, 이 화가는 노란 택시를 못 그리네?"라고 판단하여 그 화가를 쫓아내고, "노란 택시를 잘 그리는 화가"를 여러 명 복사해서 대열에 넣습니다.
• 단점:
  1. 비효율적: 100 명을 모두 고용하고 관리하는 데 비용이 많이 듭니다.
  2. 다양성 부족: 결국 모두 똑같은 "노란 택시"만 그리게 되어, 그림이 다 비슷비슷해집니다 (모드 붕괴).
  3. 수정 불가: 그림이 다 그려진 후에 "아, 창문 하나 더 추가해줘"라고 하면, 처음부터 다시 그리는 수고를 해야 합니다.

2. CREPE 의 등장: "여러 번의 시도를 병렬로, 한 번에 하나씩 완성하기"

저자들은 CREPE라는 새로운 방법을 제안합니다. 이는 물리학의 '복제 교환 (Replica Exchange)'이라는 개념을 차용한 것입니다.

CREPE 의 핵심 비유: "등산 팀의 교대 작전"

想像해 보세요. 100 명의 등산객이 있습니다. 그들은 모두 같은 산 (완성된 그림) 을 오르고 싶지만, 각자 다른 난이도의 코스를 오릅니다.

1. 다양한 난이도의 코스 (온도 조절):

   • 어떤 팀은 아주 험한 길 (완성도 0%, 노이즈가 가득한 상태) 을 오릅니다.
   • 어떤 팀은 거의 정상에 가까운 길 (완성도 90%) 을 오릅니다.
   • CREPE 는 이 모든 팀을 동시에 (병렬로) 이동시킵니다.

2. 팀원 교환 (복제 교환):

   • 중간에 험한 길을 오르는 팀원 A 와, 쉬운 길을 오르는 팀원 B 가 만납니다.
   • "너는 이 험한 길이 너무 힘들어 보이니, 내가 대신 가볼게. 너는 내가 오던 쉬운 길로 가서 쉬어."라고 서로 위치를 바꿉니다.
   • 이 교환은 AI 가 "이 조합이 더 좋은 결과를 낼 것 같다"고 판단할 때만 일어납니다.

3. 한 번에 하나씩 완성 (순차적 생성):

   • 기존 방법처럼 100 명을 다 고용할 필요 없이, 한 번에 한 명씩 (또는 소수씩) 새로운 등산객을 데려와서 이 교환 과정을 반복합니다.
   • 시간이 지나면, 이 등산객들은 모두 정상 (완성된 그림) 에 도달하게 됩니다.

3. CREPE 가 기존 방법보다 좋은 점

• 다양한 그림을 보장합니다:
  • 기존 방법은 "노란 택시"만 그리는 화가들만 남게 되어 그림이 다 똑같아집니다.
  • CREPE 는 등산로가 여러 개이고 서로 교환되므로, "노란 택시", "노란 트럭", "노란 자전거" 등 다양한 변형이 자연스럽게 만들어집니다.
• 수정이 가능합니다 (온라인 수정):
  • 그림이 거의 다 그려진 상태에서 "아, 배경을 겨울로 바꿔줘"라고 하면, CREPE 는 이미 그려진 그림을 가지고 바로 수정을 시작할 수 있습니다.
  • 기존 방법은 처음부터 다시 그리는 수고를 해야 합니다.
• 효율적입니다:
  • 100 명을 동시에 고용할 필요 없이, 적은 인원으로 순차적으로 작업을 진행하면서도 좋은 결과를 냅니다.

4. 실제 적용 사례 (논문에서 보여준 것들)

이 방법은 다양한 분야에서 테스트되었습니다:

• 분자 구조 (약물 개발): 분자의 모양을 더 안정적으로 찾아내는 데 사용되었습니다.
• 이미지 생성: "노란 택시"를 그릴 때, 배경을 어둡게 하거나 특정 색상을 강조하는 등 세부 지시사항을 더 정확하게 반영했습니다.
• 미로 찾기: 짧은 경로 조각들을 이어 붙여 긴 미로 길을 만들 때, 벽에 부딪히지 않고 목표 지점에 도달하는 경로를 찾아냈습니다.
• 텍스트 생성: 감정을 조절하는 글쓰기 (예: "기쁜 감정으로 글을 써줘") 에서, AI 가 원래 가졌던 편향을 줄이고 더 자연스러운 글을 만들어냈습니다.

5. 결론: 왜 이것이 중요한가요?

CREPE 는 AI 가 그림을 그리거나 글을 쓸 때, 사용자의 의도를 더 정확하게 이해하고, 더 다양한 결과물을 내며, 중간에 수정도 쉽게 할 수 있게 해주는 '스마트한 조율사' 역할을 합니다.

기존의 "많은 사람을 모아놓고 고르기" 방식에서, "적은 사람으로 서로 도와가며 최적의 길을 찾기" 방식으로 패러다임을 바꾼 혁신적인 방법이라고 볼 수 있습니다.

기술 요약

1. 문제 정의 (Problem)

확산 모델은 이미지, 비디오, 텍스트 등 다양한 모달리티에서 고품질 샘플을 생성하는 데 혁신을 가져왔습니다. 그러나 학습된 모델을 재학습 (Retraining) 하지 않고도 새로운 제약 조건 (예: 특정 클래스 조건, 보상 함수, 모델 조합 등) 을 만족하도록 생성 과정을 제어하는 추론 시간 제어가 중요한 과제로 대두되었습니다.

기존의 주요 접근 방식은 다음과 같은 한계를 가집니다:

• 휴리스틱 가이드 (Heuristic Guidance): Classifier-free guidance 등의 방법은 근사적이며 편향 (Bias) 이 발생할 수 있습니다.
• 순차 몬테 카를로 (SMC): 편향을 보정하기 위해 SMC 를 사용하는 방법들이 제안되었으나, 다음과 같은 치명적인 단점이 있습니다:
  1. 높은 메모리 비용: 전체 탈노이즈 (Denoising) 경로에 걸쳐 많은 수의 파티클 (Particles) 을 병렬로 유지해야 합니다.
  2. 샘플 다양성 저하: 작은 배치 크기에서는 모드 붕괴 (Mode Collapse) 가 발생하기 쉽고, 샘플 다양성이 낮아집니다.
  3. 온라인 수정 불가: 샘플링이 완료된 후에는 생성된 샘플을 추가로 정제하거나 새로운 제약을 추가할 수 없어 처음부터 다시 생성해야 합니다.

2. 방법론 (Methodology: CREPE)

저자들은 병렬 템퍼링 (Parallel Tempering, PT) 또는 레플리카 교환 (Replica Exchange) 알고리즘을 확산 모델의 추론 시간 제어에 적용하여 CREPE를 제안합니다. 이는 SMC 와 계산적으로 이중 (Computationally Dual) 인 구조를 가집니다.

핵심 아이디어

• SMC vs CREPE:
  • SMC: 시간 축을 따라 병렬로 파티클을 이동시키고, 중요도 재샘플링 (Importance Resampling) 을 통해 제약을 만족하는 파티클을 선택합니다.
  • CREPE: 여러 다른 확산 시간 단계 (Diffusion steps) 에서 병렬로 체인 (Chains) 을 실행하고, 시간 축을 따라 순차적으로 파티클을 생성합니다.
• 알고리즘 흐름 (Algorithm 1):
  1. 어닐링 경로 (Annealing Path): 목표 분포 $\pi_0$와 추론이 용이한 기준 분포 $\pi_1$ (예: 가우시안 노이즈) 사이를 연결하는 경로 $(\pi_t)_{t \in [0,1]}$를 정의합니다.
  2. 커뮤니케이션 단계 (Communication Step): 인접한 시간 단계 $t$와 $t'$에 있는 샘플들 사이에서 **가속화된 병렬 템퍼링 (Accelerated PT, APT)**을 사용하여 스왑 (Swap) 을 시도합니다.
     • 제안 과정 (Proposal) 으로 전진 및 후진 확산 경로를 시뮬레이션합니다.
     • **라돈 - 니코딤 추정기 (Radon-Nikodym Estimator, RNE)**를 사용하여 제안된 경로의 확률 비율을 계산하고, 메트로폴리스 - 헤이스팅스 (Metropolis-Hastings) 기준에 따라 스왑을 수락하거나 거절합니다.
  3. 로컬 탐색 (Local Exploration): 커뮤니케이션 후, 목표 분포 $\pi_t$의 스코어 함수를 사용하여 로컬 MCMC 업데이트 (예: Langevin step) 를 수행하여 샘플을 정제합니다.
  4. 온라인 정제: 생성된 샘플은 순차적으로 업데이트되므로, 생성 도중 새로운 제약 조건이 추가되거나 품질이 부족할 경우 즉시 정제 (Refinement) 가 가능합니다.

적용 가능한 작업

CREPE 는 다음과 같은 다양한 제어 작업에 적용 가능합니다:

• Tempering: 온도를 조절하여 볼츠만 분포에서 샘플링.
• Reward-tilting: 보상 함수를 기반으로 사후 확률 (Posterior) 샘플링.
• Model Composition: 여러 확산 모델을 조합하여 새로운 분포 생성.
• CFG Debiasing: Classifier-free guidance 의 편향을 보정.

3. 주요 기여 (Key Contributions)

1. 새로운 프레임워크 제안: 사전 학습된 확산 모델에서 명시적인 목표 밀도 (Target Density) 없이도 병렬 템퍼링을 직접 적용할 수 있는 CREPE 알고리즘을 개발했습니다.
2. 이론적 유도: 가우시안 확산 모델과 이산적 마스크 확산 (Discrete Mask Diffusion) 모두에 대해 Tempering, Reward-tilting, Model Composition, CFG Debiasing 등 다양한 시나리오에 대한 PT 스왑 비율 (Swap Rates) 을 유도했습니다.
3. SMC 대비 우월한 성능: SMC 기반 방법론에 비해 샘플 다양성이 훨씬 높고, 온라인 정제 및 **조기 종료 (Early Termination)**가 가능하다는 것을 실험을 통해 입증했습니다. 또한, 동일한 네트워크 평가 횟수 (NFEs) 기준에서 SMC 와 유사한 효율성을 보였습니다.

4. 실험 결과 (Results)

저자들은 분자, 이미지, 궤적, 텍스트 등 다양한 도메인에서 CREPE 를 평가했습니다.

• 분자 샘플링 (Boltzmann Sampling): 알라닌 디/테트라/헥사 펩타이드에 대한 온도 어닐링 실험에서 CREPE 는 SMC 대비 에너지 및 거리 분포의 TVD (Total Variation Distance) 와 TICA 기반 MMD 에서 더 낮은 편향과 더 높은 다양성을 보여주었습니다. 특히 SMC 가 작은 배치에서 모드 붕괴를 겪는 반면, CREPE 는 다양한 모드를 잘 포착했습니다.
• 이미지 생성 (ImageNet):
  • CFG Debiasing: 이미지넷-64/512 에서 CREPE 는 FKC (SMC 기반 CFG 보정) 보다 더 높은 FID 점수와 더 다양한 시각적 샘플을 생성했습니다. 특히 샘플 수가 증가할수록 CREPE 의 성능이 우세해졌습니다.
  • Reward-tilting: ImageReward 를 활용한 프롬프트 기반 제어 실험에서, CREPE 는 Burn-in 기간 이후 프롬프트와 잘 일치하면서도 다양한 이미지를 생성했습니다.
• 모델 조합 (Maze Navigation): 짧은 궤적 확산 모델을 조합하여 미로 탐색 경로를 생성하는 작업에서, CREPE 는 조건부 모델 학습 없이도 보상 함수를 통해 긴 궤적을 성공적으로 연결했습니다. 또한, 중간 지점 (Intermediate point) 과 같은 새로운 제약을 생성 도중 추가했을 때 이를 빠르게 반영하여 경로를 수정하는 온라인 정제 능력을 입증했습니다.
• 텍스트 생성 (CTMC): 이산 확산 모델을 이용한 감정 제어 텍스트 생성에서, CREPE 는 SMC 대비 Perplexity 를 최대 5 배까지 낮추면서 감정 정확도를 유지했습니다.

5. 의의 및 결론 (Significance)

• 다양성 유지: CREPE 는 SMC 의 재샘플링으로 인한 다양성 손실 문제를 해결하며, MCMC 기반의 순차적 업데이트를 통해 자연스러운 다양성을 유지합니다.
• 유연성: 생성 과정 중 새로운 제약을 추가하거나 샘플을 정제할 수 있어, "Anytime" 알고리즘으로서의 가치가 높습니다.
• 확장성: 가우시안 확산뿐만 아니라 이산적 확산 (텍스트 등) 모델에도 적용 가능하여 범용적인 추론 시간 제어 솔루션을 제공합니다.

한계점:

• 최적의 성능에 도달하기 위해 Burn-in 기간이 필요하며, 이 기간 동안의 샘플은 버려져야 할 수 있습니다.
• 통신 단계의 수용 확률 계산 시 근사 오차가 발생할 수 있으며, 이는 반복적으로 누적될 수 있습니다.

결론적으로, CREPE 는 확산 모델의 추론 시간 제어 분야에서 SMC 의 대안으로 등장하여, 더 높은 품질과 다양성을 가진 샘플 생성과 유연한 제어 능력을 제공합니다.