RNE: plug-and-play diffusion inference-time control and energy-based training

이 논문은 경로 분포 간의 밀도 비율 개념을 바탕으로 Radon-Nikodym Estimator(RNE) 를 제안하여, 확산 모델의 주변 밀도 추정, 추론 시 제어, 그리고 에너지 기반 학습을 단일 프레임워크로 통합하고 다양한 모달리티에서 효과적인 성능을 입증합니다.

핵심

1. 배경: 확산 모델은 어떻게 작동할까? (안개 속의 그림 그리기)

확산 모델은 마치 안개 낀 날에 그림을 그리는 과정과 비슷합니다.

• 시작: 완전한 안개 (잡음) 에서 시작합니다.
• 과정: 안개가 서서히 걷히면서 (소음 제거), 그 안에 숨겨진 그림 (데이터) 이 모습을 드러냅니다.
• 현재의 한계: 우리는 이 '안개를 걷어내는 과정 (소음 제거)'은 잘 알고 있습니다. 하지만, **과정 중간중간 그림이 얼마나 선명한지 (확률 밀도)**를 정확히 알기는 매우 어렵습니다.

이 '중간 상태의 선명함'을 알면, 그림을 그리는 동안 원하는 대로 방향을 틀거나 (제어), 여러 개의 모델을 합쳐서 더 멋진 그림을 만들 수 있습니다. 하지만 기존에는 이걸 계산하는 데 너무 많은 계산량이 들거나, 아예 불가능했습니다.

2. RNE 의 핵심 아이디어: "거꾸로 가는 길도 같다" (시간의 역행)

이 논문은 아주 재미있는 통찰을 제시합니다.

"안개를 걷어내는 과정 (정방향) 과, 반대로 안개를 다시 끼우는 과정 (역방향) 은 사실 동일한 길을 걷는 것과 같다."

• 비유: 산 정상에서 아래로 내려가는 길 (정방향) 과, 아래에서 위로 올라가는 길 (역방향) 이 있습니다. 보통은 내려가는 길만 잘 알고 있죠. 하지만 이 두 길은 동일한 지형을 공유합니다.
• RNE 의 역할: RNE 는 이 '동일한 지형'이라는 사실을 이용해, 내려가는 길의 지도 (전환 커널) 만으로도, 전체 산의 높이 (확률 밀도) 를 계산할 수 있게 해줍니다.

이걸 **RNE(라돈 - 니코딤 추정기)**라고 부릅니다. 이름은 어렵지만, 역할은 **"거꾸로 가는 길과 앞으론 가는 길 사이의 관계를 이용해, 우리가 몰랐던 정보 (확률) 를 찾아내는 자"**입니다.

3. RNE 가 해결하는 세 가지 문제

이 도구를 사용하면 세 가지 놀라운 일이 일어납니다.

① 추론 시 제어 (Inference-time Control): "나만의 길로 가자"

• 상황: 이미 훈련된 모델이 있는데, 갑자기 "더 따뜻한 색감으로 그려줘" 또는 "이 특정 모양을 포함해줘"라고 명령을 내리고 싶을 때입니다.
• 기존 방법: 모델을 처음부터 다시 훈련해야 하거나, 임의의 규칙을 적용해서 결과가 엉망이 될 수 있습니다.
• RNE 해결: RNE 는 생성 과정 중간중간에 "이제부터는 저쪽 길로 가자"라고 방향을 살짝 틀 수 있게 해줍니다.
  • 비유: 여행 중인데, "오늘은 산 정상으로 가자"고 했을 때, 지도 (RNE) 를 보고 가장 효율적인 길로 우회할 수 있게 해주는 내비게이션 같은 역할입니다.
  • 효과: 모델을 다시 훈련하지 않아도, 원하는 속성 (온도, 보상, 여러 모델 합치기) 을 가진 이미지를 생성할 수 있습니다.

② 에너지 기반 훈련 (Energy-based Training): "정확한 지도 만들기"

• 상황: 확산 모델을 '에너지' (불확실성) 로 표현하고 싶을 때, 기존 방법은 지도가 왜곡되는 경우가 많았습니다.
• RNE 해결: RNE 는 이 왜곡을 잡아주는 규제제 (Regularizer) 역할을 합니다.
  • 비유: 지도를 그릴 때, "이곳은 높은 산이야"라고 잘못 표시했다면, RNE 는 "아니야, 여기는 평지야"라고 수정해 주는 교정 펜입니다.
  • 효과: 훨씬 더 정확하고 신뢰할 수 있는 모델을 만들 수 있습니다.

③ 모든 분야에 적용 가능 (Modality-agnostic): "범용 도구"

• 특징: 이 도구는 이미지뿐만 아니라 텍스트, 분자 구조, 심지어 이산적인 (Discrete) 데이터에도 적용됩니다.
• 비유: 이 도구는 만능 열쇠와 같습니다. 자물쇠 (데이터 종류) 가 무엇이든, 열 수 있는 열쇠입니다.

4. 요약: 왜 이것이 중요한가?

기존의 방법들은 각자 다른 문제를 해결하기 위해 따로따로 개발된 '특수 도구'들이었습니다. 하지만 RNE는 이 모든 도구들을 하나로 통합한 범용 플랫폼을 제공합니다.

• Plug-and-Play (플러그 앤 플레이): 복잡한 수식을 다시 유도할 필요 없이, 기존 모델에 바로 꽂아 (Plug) 쓸 수 있습니다.
• 유연성: 원하는 대로 생성 과정을 조절할 수 있어, AI 가 더 창의적이고 정확한 작업을 할 수 있게 됩니다.

한 줄 요약:

"RNE 는 확산 모델이 안개 속을 걷는 동안, 정확한 지도를 보며 원하는 방향으로 길을 찾을 수 있게 해주는 '만능 나침반'입니다."

이 기술은 AI 가 더 정교하게 그림을 그리거나, 신약 개발, 재료 과학 등 다양한 분야에서 더 정확한 예측을 할 수 있는 기반을 마련해 줍니다.

기술 요약

RNE: 플러그 앤 플레이 확산 추론 시간 제어 및 에너지 기반 학습 (RNE: Plug-and-Play Diffusion Inference-Time Control and Energy-Based Training)

이 논문은 확산 모델 (Diffusion Models) 의 생성 과정에서 **한계 밀도 (Marginal Densities)**에 대한 접근성을 해결하고, 이를 통해 추론 시간 제어 (Inference-time Control) 와 에너지 기반 학습을 통합하는 새로운 프레임워크인 **라돈 - 니코딤 추정기 (Radon-Nikodym Estimator, RNE)**를 제안합니다.

1. 문제 정의 (Problem)

확산 모델은 노이즈를 점차 제거하여 데이터를 생성하는 과정 (시간 역행) 을 기반으로 합니다. 그러나 기존 모델은 주로 **전환 커널 (Transition Kernels, 즉 탈노이즈 또는 노이즈 커널)**에만 접근이 가능합니다. 많은 응용 분야 (예: 추론 시간 제어, 모델 결합, 에너지 기반 학습) 에서는 생성 궤적 (Generation Trajectory) 을 따라가는 **한계 밀도 (Marginal Densities, $p_t$)**를 알거나 추정해야 합니다.

기존의 밀도 추정 방법들은 다음과 같은 한계가 있었습니다:

• 확률 흐름 ODE (PF-ODE) 접근: 순간 변수 변환 공식을 사용하지만, 매 단계에서 스코어 네트워크의 발산 (Divergence) 을 계산해야 하므로 계산 비용이 매우 큽니다.
• 비일관된 방법론: Twisted Diffusion Sampler, Feynman-Kac Steering, Itô 밀도 추정기 등 다양한 방법들이 존재하지만, 서로 간의 연결 고리가 불분명하고 특수한 설계에 의존하여 적용 범위가 제한적입니다.
• 편향 (Bias): 휴리스틱한 가이드 (Guidance) 방법은 편향을 유발할 수 있습니다.

2. 방법론 (Methodology)

2.1 핵심 아이디어: 라돈 - 니코딘 도함수 (Radon-Nikodym Derivative)

저자들은 확산 과정과 그 시간 역행 (Time-Reversal) 과정이 경로 공간 (Path Space) 에서 동일한 확률 측도 (Probability Measure) 를 유도한다는 사실에 주목했습니다. 즉, 두 과정 사이의 라돈 - 니코딘 도함수 (밀도 비율) 는 항상 1입니다.

이 성질을 이산 시간 (Discrete Time) 으로 확장하여, **전환 커널 (Transition Kernels)**과 한계 밀도 (Marginal Densities) 사이의 관계를 유도했습니다.

2.2 RNE (Radon-Nikodym Estimator)

전진 과정 (Forward) 의 전환 커널과 역전 과정 (Backward) 의 전환 커널을 이용하여 한계 밀도의 비율을 추정하는 식을 정의합니다.
$$\frac{p_\tau(Y_\tau)}{p_{\tau'}(Y_{\tau'})} = R^\nu_\mu(Y_{[\tau, \tau']})$$
여기서 $R^\nu_\mu$는 이산화된 가우스 커널들의 곱으로 계산 가능하며, 이는 **밀도 추정 (Density Estimation)**에 직접적으로 활용됩니다.

2.3 RNC (Radon-Nikodym Corrector) - 추론 시간 제어

RNE 를 기반으로 순차 몬테 카를로 (Sequential Monte Carlo, SMC) 알고리즘의 중요도 샘플링 (Importance Sampling) 가중치를 계산하는 RNC를 제안합니다.

• 플러그 앤 플레이 (Plug-and-Play): 새로운 작업 (Annealing, Reward Tilting, Model Composition 등) 을 수행할 때마다 가중치 공식을 다시 유도할 필요가 없습니다. RNE 식을 적용하여 가중치를 계산하면 됩니다.
• 유연한 설계: 제안된 프레임워크는 다양한 샘플링 과정과 목표 과정을 선택할 수 있게 하여 분산 (Variance) 을 줄이고 성능을 최적화할 수 있습니다.
• 불완전한 모델 대응: 실제 훈련된 모델은 완벽하지 않을 수 있으므로, 보상 기울기 (Reward-tilting) 작업의 경우 정확한 SMC 가중치를 보장하는 변형을 제공합니다.

2.4 안정화: 참조 과정 (Reference Process)

직접적인 이산화는 누적 오차와 불안정성을 초래할 수 있습니다. 이를 해결하기 위해 **분석적 참조 과정 (Analytical Reference Process, 예: 가우스 분포)**을 도입하여 RNE 식을 재구성합니다. 이는 전이 커널의 분산을 정렬시켜 수치적 안정성을 높이고 수렴 속도를 보장합니다.

2.5 에너지 기반 학습 정규화

RNE 는 에너지 기반 확산 모델 (Energy-based Diffusion Models) 의 훈련에도 적용됩니다. 표준 탈노이즈 스코어 매칭 (DSM) 목적 함수에 RNE 기반의 **정규화 항 (Regularisation)**을 추가하여, 학습된 에너지 함수의 정확도를 높이고 발산 (Divergence) 계산 없이 효율적으로 학습을 수행합니다.

3. 주요 기여 (Key Contributions)

1. 통합 프레임워크: 밀도 추정, SMC 기반 추론 시간 제어, 에너지 기반 훈련을 단일 관점 (RNE) 으로 통합했습니다. 이는 기존에 분리되어 있던 Twisted Diffusion Sampler, Feynman-Kac 방법론 등을 포괄합니다.
2. 플러그 앤 플레이 추론 제어: 재학습 없이 다양한 제약 조건 (Annealing, Reward Tilting, Model Composition) 하에서 샘플을 생성할 수 있는 범용적인 SMC 가중치 계산법을 제공합니다.
3. 효율적인 에너지 학습: 발산 계산 없이 에너지 기반 모델의 훈련을 개선하는 간단하고 효율적인 정규화 기법을 제안했습니다.
4. 모달리티 무관성: 연속 확산 모델뿐만 아니라 이산 확산 모델 (CTMC, Continuous-Time Markov Chains) 에도 적용 가능합니다.

4. 실험 결과 (Results)

4.1 추론 시간 제어 (Inference-time Control)

• Annealing (ALDP, Lennard-Jones 시스템): RNC 는 기존 Feynman-Kac Corrector (FKC) 보다 더 유연한 파라미터 설정을 통해 더 높은 샘플 품질과 다양성을 보여주었습니다. 특히 SMC 를 사용한 확장성 (Scaling) 이 우수했습니다.
• 모델 결합 (Multi-target SBDD): 두 개의 단백질 타겟에 대한 분자 생성 (Structure-based Drug Design) 에서 RNC 는 FKC 와 유사하거나 더 나은 성능을 보이며, 더 많은 리간드가 기준보다 좋은 도킹 스코어를 갖도록 했습니다.
• 경로 연결 (Trajectory Stitching): 미로 탐색 (Maze Navigation) 과 같은 복잡한 제어 작업에서 RNC 는 기존 가이드 방법보다 성공률을 100% 로 높였습니다.

4.2 에너지 기반 모델 훈련

• 밀도 및 에너지 추정: 2D 및 100D 가우스 혼합 모델 (GMM) 에서 RNE 정규화를 적용한 모델은 표준 DSM 보다 훨씬 정확한 에너지 함수를 학습했습니다.
• 자유 에너지 추정: 알라닌 디펩타이드 (ALDP) 의 용해 자유 에너지를 열역학적 적분 (Thermodynamic Integration) 으로 추정할 때, RNE 정규화를 적용한 모델이 기준값 (Reference Value) 에 훨씬 근접한 결과를 도출했습니다.

4.3 이산 확산 (CTMC)

• 이미지 생성: MaskGIT 모델을 기반으로 한 이산 확산 모델에 RNE 를 적용하여 텍스트 프롬프트와 이미지 생성 간의 정렬 (Alignment) 을 개선하고, ImageReward 점수를 높이는 데 성공했습니다.

5. 의의 및 결론 (Significance)

이 논문은 RNE를 통해 확산 모델의 이론적 기반을 확장하고, 실용적인 응용 분야에서의 제약을 해결했습니다.

• 이론적 통합: 다양한 밀도 추정 및 제어 방법론을 하나의 수학적 프레임워크로 통합하여 이해를 심화시켰습니다.
• 실용성: 재학습 없이도 복잡한 제약 조건 하에서 고품질 샘플을 생성할 수 있는 플러그 앤 플레이 솔루션을 제공하여, 의료 (신약 개발), 재료 과학, 로봇 공학 등 다양한 분야에서의 확산 모델 활용도를 높였습니다.
• 효율성: 발산 계산과 같은 고비용 연산을 피하면서도 높은 정확도를 보장하는 효율적인 알고리즘을 제시했습니다.

결론적으로, RNE 는 확산 모델의 밀도 추정, 제어, 학습이라는 세 가지 핵심 과제를 해결하는 강력한 도구로 자리 잡았으며, 향후 생성 모델 연구의 새로운 표준이 될 가능성이 큽니다.