Composing diffusion priors with explicit physical context via generative Gibbs sampling

기술적 요약: 생성적 깁스 샘플링을 통한 명시적 물리적 맥락으로 확산 사전 지식 구성

문제 제기
사전 훈련된 확산 모델은 과학적 샘플링을 위한 강력한 학습된 사전 지식을 제공하지만, 종종 시스템의 자유도 중 선택된 부분 집합 (예: 단백질 골격 또는 분자 단편) 만을 기술할 뿐 전체 시스템 상태를 설명하지는 못합니다. 과학적 응용 분야에서 목표 분포는 종종 용매, 이온, 외부 장, 또는 다른 서브시스템과의 상호작용과 같은 물리적 맥락에 의존하는데, 이는 단일 생성 모델로 완전히 표현되지 않습니다. 가이드나 사후 샘플링과 같은 표준 추론 시 접근법은 모든 맥락을 생성 모델의 변수로 표현해야 합니다. 이는 표현되지 않은 자유도를 유효 자유 에너지 항으로 주변화 (marginalizing) 해야 함을 의미하며, 이는 고차원 환경에서는 종종 계산 불가능하거나, 다른 서브시스템이 이미 별도의 사전 지식이나 힘장 (force fields) 으로 잘 모델링된 경우 불필요한 중복이 됩니다. 해결된 핵심 과제는 모델을 재훈련하지 않고 명시적 시스템 수준의 물리적 맥락과 여러 부분 학습된 사전 지식을 추론 시에 구성하는 것입니다.

방법론: 물리 인식 샘플링을 위한 생성적 깁스 (GG-PA)
저자들은 학습된 부분 사전 지식과 명시적 물리적 맥락의 구성을 확장된 상태 공간에서의 결합 목표 분포에 대한 추론으로 공식화하는 훈련 없는 프레임워크인 GG-PA 를 제안합니다.

1. 확장된 상태 공간: 이 방법은 전체 시스템 상태 $s$(예: 용매를 포함한 모든 원자 좌표) 의 명시적 표현을 유지하고, 투영 연산자 $\Phi_i: S \to X_i$를 통해 $K$개의 사전 훈련된 확산 사전 지식과 결합시킵니다. 확장된 상태는 $Z = S \times \prod X_i$입니다.
2. 결합 목표 분포: 확산 시간 $t$로 인덱싱된 결합 목표 밀도 함수의 집합이 정의됩니다:
   $$\pi_t(s, \{x_i\}) \propto q_{\text{ctx}}(s, t) \prod_{i=1}^K \left[ p_i(x_i) \cdot q^{(i)}_t(\Phi_i(s) | x_i) \right]$$
   여기서 $p_i$는 사전 훈련된 사전 지식, $q^{(i)}_t$는 결합으로 작용하는 순방향 확산 커널, $q_{\text{ctx}}$는 명시적 물리적 맥락 인자 (예: 볼츠만 인자) 입니다. $t \to 0$일 때, 결합 커널은 엄격한 일관성 ($\Phi_i(s) = x_i$) 을 강제하여 사전 지식이 특정 부분 집합을 지배하고 맥락이 나머지를 지배하는 구성된 분포를 복원합니다.
3. 생성적 깁스 샘플러: 샘플링은 두 단계 사이를 번갈아 수행합니다:
   • 병렬 디노이징: 각 사전 지식 변수 $x_i$는 사전 지식 $p_i$와 현재 투영된 상태 $\Phi_i(s)$를 노이즈가 있는 관측치로 간주하여 유도된 사후 분포에서 샘플링함으로써 업데이트됩니다. 이는 사전 훈련된 역시간 샘플러를 실행하여 수행됩니다.
   • 맥락 인식 집계: 전체 시스템 상태 $s$는 현재 $x_i$값과 명시적 맥락에 조건부로 업데이트됩니다. 이 단계는 맥락과 순방향 커널의 로그 가능도에서 유도된 유효 전위 $U_{\text{eff}}$를 최소화합니다.
4. 리플리카 교환: 충실도 (작은 $t$) 와 혼합 (큰 $t$) 간의 트레이드오프를 해결하기 위해, 저자들은 확산 시간에 대한 리플리카 교환을 도입합니다. 여러 리플리카가 서로 다른 $t$값에서 실행되며, 계산 불가능한 사전 지식 밀도가 상쇄되는 계산 가능한 수용 비율을 기반으로 스왑 이동이 제안됩니다.

이론적 특성

• 점근적 정확성: 분해 가능한 시스템의 경우, 주변 목표는 $t \to 0$일 때 실제 물리적 분포를 복원합니다.
• 유한 시간 정확성: 상호작용이 2 차 (선형 - 가우시안) 인 설정에서, 맥락 일정이 특정 모멘트 매칭 조건 (가우시안 디컨볼루션) 을 만족하도록 매개변수화되면, 이 방법은 유한 $t$에서도 정확성을 유지합니다. 이는 최대 허용 확산 시간 $t_{\text{max}}$에 대한 중요한 경계를 제공합니다.
• 분할 깁스와의 연관성: 이 프레임워크는 선형 역문제에 대한 분할 깁스 샘플러를 일반화하여, 표준 구현에 존재하는 편향을 피하는 공분산 보정을 제공합니다.

실험 결과
저자들은 복잡도가 증가하는 세 가지 시스템에 대해 GG-PA 를 평가했습니다:

1. 결합된 더블 웰 시스템: 유한 시간 정확성과 리플리카 교환의 효능을 검증하기 위해 사용된 2 차 2 차원 시스템입니다. GG-PA 는 환경에 의해 유도된 비대칭성을 성공적으로 복원했습니다. 리플리카 교환은 고정 $t$ 샘플링 및 분자 역학 (MD) 에 비해 뻣뻣하고 낮은 $t$ 영역에서 혼합을 크게 가속화했습니다.
2. $\phi^4$ 격자 모델: 훈련 분포에 없는 다체 집단 행동을 구성하는 것을 테스트하는 2 차원 긴즈버그 - 랜다우 모델입니다. 이 모델은 로컬 온 - 사이트 더블 웰 인자만 훈련되었습니다. GG-PA 는 평형 상전이, 자발적 대칭성 깨짐, 그리고 임계 지수를 성공적으로 재현했습니다. 리플리카 교환은 임계점 근처에서 수 차례의 속도 향상을 제공했습니다.
3. 알라닌 디펩타이드 시스템: 비 2 차 상호작용을 포함하는 원자 수준 모델입니다.
   • AD–Na+: GG-PA 는 이온 배위로 인해 유도된 카르보닐 산소 거리 분포의 변화를 포착하여, 진공에서 훈련된 사전 지식을 직접 사용한 경우보다 우수한 성능을 보였습니다.
   • AD 이량체: 수소 결합 이량체를 모델링하기 위해 단량체 사전 지식의 두 사본이 구성되었습니다. GG-PA-RE 는 상호작용의 비 2 차적 성격과 정확한 유한-$t$ 보장이 없음에도 불구하고, 질적인 대칭성 깨진 조직 (반평행 대 평행 토폴로지) 과 조건부 비틀림 상태 선호도를 복원했습니다.

주요 기여

1. 공식화: 계산 불가능한 주변화를 우회하는 명시적 전체 시스템 상태에 대한 추론으로서의 부분 확산 사전 지식 구성에 대한 새로운 공식화.
2. 알고리즘 및 이론: 점근적 정확성, 2 차 상호작용에 대한 유한 시간 정확성에 대한 증명, 그리고 분할 깁스 샘플러를 위한 공분산 보정을 포함한 GG-PA 샘플러의 유도.
3. 실용적 증명: 2 차 상호작용이 있거나 없는 시스템에서 모듈식 다중 사전 지식 구성의 수치적 증명으로, 재훈련 없이 맥락에 의해 유도된 변화와 발현된 집단 행동을 복원할 수 있음을 보여줌.

의의 및 주장
이 논문은 GG-PA 를 사전 훈련된 생성적 사전 지식과 명시적 물리적 맥락을 결합하기 위한 실용적인 접근법으로 제시합니다. 저자들은 이 모듈식 패러다임이 학습된 사전 지식과 명시적 물리학을 가장 적절한 곳에 적용할 수 있게 하여, 시스템 환경이 변경될 때 모놀리식 모델을 재훈련할 필요가 없다고 주장합니다. 이 방법은 힘장이나 별도의 사전 지식으로 쉽게 처리할 수 있는 고차원 환경 자유도를 가진 과학적 시스템에 특히 유용합니다. 저자들은 유한-$t$ 정확성을 위한 2 차 구조에 대한 의존성과 여러 리플리카를 유지하는 계산 비용과 같은 한계를 인정하지만, 표준 사후 샘플링이나 가이드 기술로는 어려운 복잡하고 맥락 의존적인 샘플링 작업을 처리할 수 있는 이 방법의 능력을 강조합니다.