📊 statistics

A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models

본 논문은 기존의 점 단위 보상(pointwise-reward) 방식들을 확장하고 배치 단위 조향(batch-level steering)을 위한 원칙적인 토대를 제공하기 위해, 추론 시점에 확산 모델을 사전에 규정된 분포 수준의 보상으로 유도하는 가중 상호작용 입자 스킴(weighted interacting particle scheme)을 도출하는 이론적으로 근거가 있는 평균장 프레임워크(mean-field framework)를 제안한다.

원저자: Samuel Howard, Nikolas Nüsken

게시일 2026-08-11

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Samuel Howard, Nikolas Nüsken

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 북적이는 주방을 운영하는 숙련된 셰프라고 상상해 보세요. 인공지능의 세계에는 **확산 모델(Diffusion Models)**이라는 특별한 '생성적' 셰프들이 있습니다. 이 모델들은 마치 놀라운 재능을 가진 예술가와 같아서, 순수한 정적 노이즈(마치 채널이 맞지 않는 TV의 지지직거리는 화면 같은)라는 빈 캔버스에서 시작하여, 단계별로 그 노이즈를 정교하게 다듬어 선명하고 아름다운 그림, 실제적인 분자, 또는 단백질 구조로 만들어냅니다. 이들은 혼돈을 질서로 바꾸는 법을 배웁니다.

하지만 때때로 우리는 단순히 '아무' 그림이나 원하는 것이 아닙니다. 우리는 특정한 자질을 갖춘 그림을 원합니다. 예를 들어, 질병을 치료할 수 있도록 매우 특정한 방식으로 접히는 단백질을 원하거나, 지루한 반복을 피하기 위해 서로 모두 다른 이미지 꾸러미를 원할 수도 있습니다. 여기서 **추론 시 제어(Inference-Time Steering)**가 등장합니다. 이것은 셰프에게 처음부터 다시 교육하는 대신, 요리하는 동안에 약간의 향신료를 더하거나 재료에 부드럽게 자극을 주는 것과 같습니다. 보통 셰프들은 현재 재료가 얼마나 좋은지에 따라 개별 재료를 살짝 밀어줍니다. 하지만 만약 당신이 전체 배치(batch)가 특정 맛을 내도록, 즉 전체 묶음이 다양하거나 균형 잡히도록 '국 솥 전체'를 밀어주고 싶다면 어떻게 해야 할까요? 이것이 바로 이 논문이 다루는 까ền 질문입니다. 개별 항목이 아닌, 전체 집단이 글로벌 목표에 부합하도록 어떻게 제어할 것인가 하는 문제입니다.

저자인 사무엘 하워드(Samuel Howard)와 니콜라스 뉴스켄(Nikolas Nüsken)은 **평균장 프레임워크(Mean-Field Framework)**라는 개념을 사용하여 이 문제를 해결하는 새로운 방법을 제안합니다. 이들의 해결책을 이해하기 위해, 바다를 헤엄치는 물고기 떼를 상상해 보세요. 만약 당신이 물고기 떼 전체를 왼쪽으로 돌게 하고 싶다면, 물고기 한 마리에게 소리를 지르는 것만으로는 부족할 것입니다. 물고기들은 서로 상호작용하기 때문입니다. 한 마리의 물고기가 방향을 틀면 그것이 이웃에게 영향을 미치고, 그 이웃이 다시 또 다른 이웃에게 영향을 미치는 파동 효과가 발생합니다. 이 논문의 방법에서 '물고기'는 AI가 생성한 샘플들입니다. 이 새로운 방법은 샘플들을 고립된 개별 존재로 취급하는 대신, 모든 샘플이 서로에 대해 '알고 있는' 연결된 집단으로 취급합니다.

이 논문은 기존의 제어 방식, 즉 개별 샘플을 보상(reward)을 향해 단순히 밀어내는 방식은 군중에게 개개인에게 소리를 지르는 것과 같다고 주장합니다. 그것이 어느 정도 효과가 있을 수는 있지만, 군중이 정확히 의도한 모양을 갖추도록 보장하지는 못합니다. 저자들은 진정한 의미에서 '분포'(군중의 전체적인 형태)를 제어하려면, 샘플들이 상호작용하고 집단의 현재 상태에 따라 자신의 가중치를 조정하는 시스템이 필요하다고 보여줍니다. 그들은 샘와가 지휘자 역할을 하는 '가중 상호작용 입자 스킴(weighted interacting particle scheme)'을 개발했습니다. 이 지휘자는 샘플들이 진화함에 따라, 사용자가 원하는 정확한 분포로 자연스럽게 안착하도록 보장합니다. 그것이 실험 데이터와 일치하는 단백질이든, 가능한 모든 스타일을 아우르는 이미지 세트이든 말입니다.

연구진은 이 아이디어를 두 가지 방식으로 테스트했습니다. 첫째, 정답을 정확히 확인할 수 있는 단순한 저차원 수학 문제에서 시뮬레이션을 실행했습니다. 이 테스트에서 그들의 새로운 방법은 목표 분포를 성공적으로 맞춘 반면, 기존의 '강압적인' 방법들은 목표에서 약간 벗어난 형태를 만들어냈습니다. 둘-째로, 그들은 이 방법을 생물학의 실제 영역으로 가져가 단백질 구조 생성을 유도하는 데 사용했습니다. 그들은 모델이 HIV-1 프로테아제(protease) 및 아데닐레이트 키나아제(adenylate kinase)의 실제 거동과 일치하도록 유도하려고 시도했습니다. 이러한 고차원적이고 복잡한 과업에서, 그들의 방법은 표준적인 접근 방식보다 원하는 실험 데이터를 맞추는 데 있어 더 신뢰할 수 있고 정확하다는 것을 입증했습니다.

본질적으로, 이 논문은 만약 당신이 AI가 생성한 항목들의 '전체 집단'을 제어하고 싶다면, 그들을 외로운 개인으로 취급하는 것을 멈추고 서로 대화하는 하나의 팀으로 취급해야 한다고 제안합니다. '사회적 상호작용'의 층위와 영리한 교정 시스템(Feynman-Kac reweighting이라 불리는)을 추가함으로써, 저자들은 AI 모델을 복잡하고 글로벌한 목표로 안내하는 더 원칙적이고 수학적으로 견고한 방법을 제공합니다. 이 방법은 이러한 상호작용을 실행하기 위해 더 많은 컴퓨팅 자원을 필요로 하지만, 시뮬레이션과 단백질 실험 결과는 이 방법이 AI 생성을 더욱 정밀하고 제어 가능하게 만드는 강력한 도구임을 시사합니다.

기술 요약: 확산 모델의 추론 시 분포 제어를 위한 평균장 프레임워크 (A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models)

문제 정의

확산 모델은 생성 과정을 특정 목적을 충족하도록 추론 시점에 유도하는 제어 가능한 샘플러로서 점점 더 많이 활용되고 있습니다. 기존 방법들은 점별 보상(pointwise rewards, 개별 샘플에 대해 정의된 $r: \mathbb{R}^d \to \mathbb{R}$ )을 효과적으로 처리하지만, 많은 핵심적인 응용 분야에서는 분포적 보상(distributional rewards, 생성된 측도에 대해 정의된 $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$ )을 요구합니다. 이러한 분포적 목적 함수에는 실험적 관측값에 모집단을 정렬하거나, 다양성을 장려하거나, 모드(mode)의 균형을 맞추는 것 등이 포함됩니다.

분포적 유도를 위한 현재의 접근 방식들은 보상의 그래디언트 항을 추론 역학에 포함시키지만(예: 반발 퍼텐셜 또는 앙상블 가능도), 이러한 휴리스틱은 샘플링된 분포에 대한 이론적 보장을 결-여하고 있습니다. 단순히 그래디언트 항을 더하는 것만으로는 입자들이 다음의 KL 정규화된 목적 함수를 최대화하는 기울어진 측도(tilted measure) $\mu^*$ 를 추적하도록 보장할 수 없습니다:
$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$
점별 사례의 경우 기울어진 타겟이 명시적인 지수 기울기( $\mu^* \propto p_1 e^r$ )인 것과 달리, 분포적 사례의 경우 타겟이 측도 자체에 의존하는 기울임 퍼텐셜(tilting potential)을 갖는 암시적(implicit) 타겟을 생성합니다. 현재 이 암시적 분포를 타겟팅하기 위해 역학 및 가중치를 교정할 수 있는 이론적으로 근거가 있는 프레임워크는 존재하지 않습니다.

방법론

저자들은 분포적 제어를 공식화하고 해결하기 위해 **평균장 프레임워크(Mean-Field Framework)**를 제안합니다.

1. 이론적 정식화

이 문제는 각 시간 단계 $t \in [0,1]$ 에서 다음과 같이 암시적으로 정의되는 기울어진 측도 $\mu^*_t$ 를 타겟팅하는 것으로 프레임화됩니다:
$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$
여기서 $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ 는 시간 의존적 보상 $R_t$ 의 첫 번째 변분(functional derivative)입니다. 이는 퍼텐셜이 측도에 의존하고, 다시 그 측도가 퍼텐셜에 의존하는 자기 일관성 관계(self-consistency relation)를 생성합니다.

2. 가중 McKean-Vlasov 역학

이 경로를 추적하기 위해 저자들은 가중 McKean-Vlasov 시스템을 도출합니다. 표준 그래디언트 유도와 달리, 유도 항 $\nabla_x \Psi_t(X_t; \mu_t)$ 는 진화하는 경험적 측도 $\mu_t$ 에 의존하기 때문에 입자 간의 상호작용을 유발합니다.

입자 $X_t$ 와 그 로그 가중치 $A_t$ 에 대한 이상적인 역학은 다음과 같습니다:
$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$
$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$
여기서 $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ 는 결정적인 교정 항입니다. 점별 사례에서 $\dot{\Psi}_t$ 가 단순히 $\partial_t r$ 인 것과 달리, 여기서는 $\mu_t$ 의 시간 진화와 관련이 있습니다. 저자들은 $\dot{\Psi}_t$ 가 보상 범함수의 2차 변분을 포함하는 암시적 적분 방정식을 만족함을 보여줍니다.

3. 유한 입자 알고리즘

이를 구현하기 위해 저자들은 가중 상호작용 입자 스킴(Algorithm 1)을 제안합니다:

위치 업데이트: 입자는 기본 드리프트, 스코어, 그리고 현재 경험적 측도 $\hat{\mu}_t$ 에 대한 보상의 그래디언트를 사용하여 업데이트됩니다.
로그 가중치 업데이트: 가중치는 유도된 역학과 타겟 기울어진 경로 사이의 불일치를 교정하기 위해 업데이트됩니다. 이를 위해서는 $\dot{\Psi}_t$ 를 계산해야 합니다.
$\dot{\Psi}_t$ 계산하기: 저자들은 암시적 시간 미분 항을 계산하기 위한 두 가지 방법을 제공합니다:
1. 암시적 선형 시스템: 적분 방정식으로부터 유도된 선형 시스템을 해결함 (Proposition 3.3).
2. 고정점 반복법(Fixed-Point Iteration): 가중치 업데이트를 근사하는 경량 피카르(Picard) 반복 스킴.
재표본 추출(Resampling): 가중치 퇴화를 방기하기 위해 표준 순차 몬테카를로(Sequential Monte Carlo, SMC) 재표본 추출이 사용됩니다.

이 프레임워크는 $R(\mu) = \int r(x) d\mu(x)$ 인 경우 표준 점별 보상 유도를 특수 사례로 복구합니다.

주요 기여

분포적 유도의 정식화: 저자들은 분포 수준의 유도를, 기울임이 측도 자체에 의존하는 기울어진 측도를 타겟팅하는 것으로 정식화하여, 점별 보상-기울어진 타겟에 대한 분포적 유사성을 확립했습니다.
가중 McKean-Vlasov 역학의 도출: 그들은 가중 상호작용 입자들이 원하는 기울어진 경로 $\mu^*_t$ 를 정확히 추적하도록 보장하는 교정적 Feynman-Kac 재가중화가 포함된 원칙적인 역학계를 도출했습니다.
이론적 복구 및 토대 마련: 이 프레임워크는 표준 점별 유도를 특수 사례로 복구하며, 다양성 샘플링이나 실험적 가이던스와 같은 기존의 배치 수준 유도 휴리스틱에 대한 이론적 토대를 제공합니다.
수렴 보장: 논문은 유한 입자 시스템이 입자 수 $N \to \infty$ 일 때 이상적인 역학으로 수렴함을 증명합니다.

실험 결과

저자들은 저차원 및 고차원 환경 모두에서 프레임워크를 검증했습니다:

저차원 (1D 혼합 가우시안):
- 태스크: 제곱 최대 평균 불일치(MMD) 보상을 사용하여 이봉 가우시안 혼합 분포를 타겟 분포의 다른 모드 가중치로 기울임.
- 결과: 표준 그래디언트 유도는 입자를 타겟 쪽으로 밀어내지만, 기울어진 분포 $\mu^*$ 의 정밀한 형태를 맞추는 데 실패합니다. 반면, 평균장 방법은 $\lambda$ 가 목적 함수 파라미터와 일치할 때 최적성 격차 $J(\mu) - J(\mu^*)$ 를 최소화하며 $\mu^*$ 를 정확하게 타겟팅합니다. 이는 다양한 노이즈 스케줄에 걸쳐 유지됩니다.
고차원 (단백질 구조):
- 태스크 1 (HIV-1 Protease): 실험적으로 관찰된 잔기 거리 분포(DEER 데이터)와 일치하도록 Boltz-2 생성 앙상블을 유도. 베이스 모델은 "닫힌(closed)" 상태로 편향되어 있었으며, 목표는 "세미 오픈(semi-open)" 상태의 비율을 회복하는 것이었습니다.
  - 결로: 평균장 유도는 경험적 잔기 거리 분포를 정확하게 맞춘 반e, 그래디언트 전용 유도는 약간 다른 형태와 더 강한 피크를 가진 분포를 생성했습니다.
- 태스크 2 (Adenylate Kinase): 분자 역학(MD) 시뮬레이션에서 유도된 2D 결합각 분포를 향해 유도.
  - 결과: 평균장 유도는 그래디언트 전용 방법보다 타겟 분포를 향한 더 강력한 이동을 생성했습니다.
- 태스크 3 (X-선 결정학): 이질적인 구조를 보이는 전자 밀도 맵(4OLE 구조)에 맞도록 Protenix 생성을 유도.
  - 결과: 동일한 유도 강도에 대해, 평균장 유도는 그래디end-only 방법보다 코사인 정렬도가 더 높았고 대안 구조(Conformation B)에 더 많은 샘플을 배치했습니다. 이는 그래디end-only 방법이 실제 역문제 타겟에 비해 과소 유도(under-steer)됨을 시사합니다.

의의 및 주장

이 논문은 분포적 제어를 위한 원칙적인 이론적 토대를 제공함으로써 점별 및 분포적 추론 제어 사이의 간극을 메운다고 주장합니다.

정확성: 주요 주장은 제안된 절차가 무한한 입자의 극한에서 규정된 보상-기울어진 분포 $\mu^*$ 를 정확하게 타겟팅하며, 휴리스틱한 그래디언트 유도에 내재된 불일치를 교정한다는 것입니다.
해석 가능성: 프레임워크는 기존의 배치 수준 유도 방법들이 효과적으로 그래디언트 단계를 수행하고 있지만, 진정한 사후 분포/기울어진 측도를 샘플링하는 데 필요한 교정적 재가중치가 결여되어 있음을 명확히 합니다.
실용성: 이 방법이 $\dot{\Psi}_t$ 를 해결하기 위한 추가적인 계산 비용을 발생시키지만, 저자들은 고차원 설정(예: 단백질)에서의 네트워크 평가에 비해 이 비용이 미미하며 고정점 솔버가 효율적이라고 언급합니다.
한로사항: 저자들은 재표본 추출에 대한 의존이 이미지와 같은 양식(modality)에서 한계가 될 수 있음(잠재적으로 중복 샘플 생성)을 인정하며, 이 방법이 미분 가능한 보상 함수를 필요로 한다고 밝힙니다. 또한, 분산 감소 기법(예: 푸아송 방정식 솔버)과의 결동을 통해 재표본 추출 문제를 완화할 수 있다고 제안합니다.

이 연구는 기존의 모든 방법을 대체하려는 것이 아니라, 왜 특정 휴리스틱이 작동하는지를 설명하고 올바른 분포를 타겟팅하도록 보장하는 메커니즘을 제공하는 엄격한 근거를 구축하는 데 목적이 있습니다.