[논문 리뷰] Diffusion Models for Constrained Domains

Nic Fishman, Leo Klarner외 3명|arXiv (Cornell University)|2023. 04. 11.

Reservoir Engineering and Simulation Methods인용 수 5

한 줄 요약

이 논문은 기하적 지오데식선이 정의되지 않는 부등식 제약 조건으로 정의된 제약 조건이 있는 다양체에 대해 두 가지 새로운 확산 모델을 소개한다. 기존 리만 기반 확산 방법은 이러한 상황에서 실패한다. 본 논문은 로그-바리어 기반 노이징 과정과 반사 브라운 운동 접근법을 제안하여 볼록 다면체와 대칭 양의 정부호 행렬에서 스코어 기반 생성 모델링을 가능하게 하였으며, 단백질 기저 뼈대와 로봇 팔의 구조적 변형 모델링에서 실험적으로 성공하였다.

ABSTRACT

Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.

연구 동기 및 목표

• 기하적 지오데식선이 전역적으로 정의되어야 하는 기존 리만 기반 확산 모델의 한계를 해결하며, 이는 부등식 제약 조건으로 정의된 다양체에서는 지오데식선이 정의되지 않기 때문이다.
• 볼록 다면체, 대칭 양의 정부호 행렬, 그리고 로봇 공학 및 단백질 설계 분야의 다양체와 같은 제약 조건이 있는 도메인에서의 생성 모델링을 가능하게 한다.
• 시간 역전 및 스코어 매칭 프레임워크와 호환되는 원칙적인 노이징 과정—로그-바리어와 반사 브라운 운동—을 개발한다.
• 기하적 지오데식선이 완비되지 않은 다양체에 대해서도 확산 모델의 이론적 및 실용적 프레임워크를 확장하여 실제 과학 분야의 적용 가능성을 넓힌다.
• 합성 및 실제 응용 작업을 통해 고차원 제약 조건이 있는 다양체에서 제안된 모델의 확장성과 유연성을 입증한다.

제안 방법

• 부등식 제약 조건을 존중하고 잘 정의된 확산 경로를 보장하는 리만 기하학적 구조를 유도하는 로그-바리어 메트릭 기반의 전방 노이징 과정을 설계한다.
• 제약 조건이 있는 영역의 경계에서 궤적을 반사하는 반사 브라운 운동 기반의 노이징 과정을 도입하여 다양체의 기하학적 구조를 유지한다.
• 비유클리드적이고 제약 조건이 있는 설정으로 스코어 매칭 및 스티븐 스코어 추정을 확장하여, 두 노이징 과정에 대해 시간 역전 SDE를 유도한다.
• 제약 조건이 있는 다양체에서 노이즈 제거 목표를 사용하여 신경망을 훈련시켜 뒤집힌 드리프트(스티븐 스코어)를 근사함으로써 스코어 기반 생성 모델을 구축한다.
• 표본의 정확도와 다양성 간의 트레이드오프를 제어하기 위해 온도 매개변수를 사용한 저온 샘플링을 적용한다.
• 실제 샘플과 생성된 샘플 간의 분포 유사도 평가를 위해 최대 평균 차이(MMD)의 U통계량 근사를 사용한다.

Figure 1 : The behaviour of different types of noising processes considered in this work defined on the unit interval. $\mathbf{B}_{t}$ : Euclidean (unconstrained) Brownian motion. $\hat{\mathbf{B}}_{t}$ : Log barrier forward noising process. $\bar{\mathbf{B}}_{t}$ : Reflected Brownian motion. All s

실험 결과

연구 질문

• RQ1 기하적 지오데식선이 모든 시간에 대해 정의되지 않는 부등식 제약 조건으로 정의된 다양체에 대해 확산 모델을 확장할 수 있는가?
• RQ2 기하학적 구조를 유지하면서 시간 역전이 가능한 원칙적인 노이징 과정을 제약 조건이 있는 다양체에 대해 어떻게 설계할 수 있는가?
• RQ3 로그-바리어 메트릭과 반사 브라운 운동이 기존의 지오데식선 기반 노이징과 비교해 타당한 대안이 될 수 있는가?
• RQ4 제안된 모델은 바이르코프 다면체와 단백질 구조적 변형 공간과 같은 고차원 제약 조건이 있는 영역에서 어떻게 성능을 발휘하는가?
• RQ5 로봇 공학 및 계산 생물학 분야의 실제 응용에 대해 모델이 얼마나 일반화되는가?

주요 결과

• 제안된 로그-바리어 및 반사 브라운 운동 기반의 노이징 과정은 기존 리만 기반 확산 모델이 실패하는 제약 조건이 있는 다양체에서 확산 모델링을 성공적으로 가능하게 하였다.
• 초입체, 단-simplex, 바이르코프 다면체에서의 실험 결과, MMD 값이 낮게 나타나 실제 데이터와 생성된 데이터 간의 분포 일치도가 높음을 확인하였다.
• 크기 $ n=4 $ 인 바이르코프 다면체에서, 높은 정확도로 유효한 이중 확률 행렬을 생성하였으며, 쌍별 및 마진널 샘플 시각화를 통해 이를 확인하였다.
• 모델은 고차원 환경으로까지 일반화되며, $ n=10 $ 단백질 기저 뼈대의 구조적 변형과 로봇 팔의 구성에 대해 확장성을 입증하였다.
• 온도 매개변수 $ heta_0 \in \{1,2,3,5,10\} $ 를 사용한 저온 샘플링을 통해 표본의 다양성과 정확도를 제어할 수 있었으며, 높은 $ heta_0 $ 는 더 날카우며 농축된 표본을 생성하였다.
• 단순체에서의 시각화 결과, 스코어 네트워크가 예측한 드리프트가 시간에 따라 적절히 변화하는 것으로 나타나 안정적인 훈련과 의미 있는 스코어 함수 학습이 이루어졌음을 확인하였다.

Figure 2 : A convex polytope defined by six constraints $\{f_{i}\}_{i\in\mathcal{I}}$ , along with the log barrier potential, and ‘straight trajectories’ under the log-barrier metric and under the Euclidean metric with and without reflection at the boundary.