[논문 리뷰] Scaling Riemannian Diffusion Models

Aaron Lou, Minkai Xu외 1명|arXiv (Cornell University)|2023. 10. 30.

Machine Learning in Healthcare인용 수 4

한 줄 요약

이 논문은 고차원 다양체에서의 정확한 열핵 계산을 가능하게 하는 확장 가능한 리만 확산 모델을 제안한다. 이는 이전에 처리할 수 없었던 스코어 매칭에서의 근사치 문제를 해결한다. 열핵을 정확히 계산하고 확률 흐름 상미방정식을 통해 정확한 샘플링을 가능하게 함으로써, $SU(3)$ 격자(128차원)와 초구형 대비 임베딩(127차원)을 포함한 고차원 다양체에서 최신 기술 수준의 성능을 달성한다. 이는 수백 차원에서의 실제 작업에 대해 미분방정식 기반 생성 모델을 처음으로 성공적으로 적용한 것이다.

ABSTRACT

Riemannian diffusion models draw inspiration from standard Euclidean space diffusion models to learn distributions on general manifolds. Unfortunately, the additional geometric complexity renders the diffusion transition term inexpressible in closed form, so prior methods resort to imprecise approximations of the score matching training objective that degrade performance and preclude applications in high dimensions. In this work, we reexamine these approximations and propose several practical improvements. Our key observation is that most relevant manifolds are symmetric spaces, which are much more amenable to computation. By leveraging and combining various ansätze, we can quickly compute relevant quantities to high precision. On low dimensional datasets, our correction produces a noticeable improvement, allowing diffusion to compete with other methods. Additionally, we show that our method enables us to scale to high dimensional tasks on nontrivial manifolds. In particular, we model QCD densities on $SU(n)$ lattices and contrastively learned embeddings on high dimensional hyperspheres.

연구 동기 및 목표

• 리만 확산 모델에서 복잡한 기하학적 구조로 인해 발생하는 노이즈 제거 스코어 매칭의 비처리 가능성 문제를 해결하기 위해.
• 물리학 및 표현 학습에 관련된 고차원 다양체에서 확장 가능하고 정확한 리만 확산 모델의 훈련을 가능하게 하기 위해.
• 성능 저하와 고차원으로의 확장성 제한을 초래하는 이전의 스코어 매칭 근사치를 극복하기 위해.
• 미분방정식 기반 생성 모델이 $SU(n)$ 격자와 초구형 임베딩과 같은 실제 고차원 다양체에 적용 가능한지 보여주기 위해.
• 최종 은닉층 특징에 의존하지 않고 진정한 초구형 분포를 모델링함으로써 대비 학습에서 모델의 해석 가능성과 성능을 향상시키기 위해.

제안 방법

• 리만 대칭 공간을 활용하여 열핵의 효율적이고 정확한 계산을 가능하게 하여 노이즈 제거 스코어 매칭에 필수적인 요소를 확보한다.
• 기존의 열핵 기법을 리만 스코어 매칭의 구체적 맥락에 맞게 적응시켜 신속하고 안정적인 훈련을 가능하게 한다.
• 최대 토러스에서 표준 상미방정식 해법기를 사용해 확률 흐름 상미방정식을 풀어 열핵에서 정확한 샘플링 방법을 개발한다.
• 최대 토러스를 활용해 계산 복잡도를 감소시키고 대칭 공간에서의 시뮬레이션 없는 효율적 샘플링을 가능하게 한다.
• 정확한 열핵 추정기를 노이즈 제거 스코어 매칭 목표에 통합하여 훈련 안정성과 성능을 향상시킨다.
• 이 방법을 $4\times4$ $SU(3)$ 격자에서 윌슨 작용 학습 및 $S^{127}$에서의 대비 학습 임베딩 모델링과 같은 고차원 작업에 적용한다.

(a) $S^{2}$

실험 결과

연구 질문

• RQ1 리만 확산 모델은 $SU(n)$ 격자와 고차원 초구와 같은 고차원 다양체로 확장될 수 있는가?
• RQ2 대칭 공간에서 정확한 열핵 계산은 스코어 매칭에서 정확하지 못한 근사치가 필요 없도록 할 수 있는가?
• RQ3 확률 흐름 상미방정식을 통한 정확한 샘플링은 다양체에서 모델 성능과 훈련 안정성을 향상시키는가?
• RQ4 리만 확산 모델은 대비 학습된 초구형 임베딩에서 충실한 분포를 회복할 수 있는가? 이는 해석 가능성 향상과 OOD 탐지 향상에 기여하는가?
• RQ5 미분방정식 기반 접근법을 사용해 비트라이비얼 다양체에서 고차원 생성 모델링에서 최고 성능을 달성할 수 있는가?

주요 결과

• 4\times4 $SU(3)$ 격자에서 윌슨 작용을 학습할 때, 제안된 방법은 효과적 표본 크기가 0.62에 도달한다. 이는 고차원 128차원 다양체이다.
• 127차원 초구에서의 대비 학습 임베딩에서, 리만 확산 모델은 최종 은닉층 특징을 사용한 최신 기술 수준의 방법과 비교해 유사하거나 더 뛰어난 OOD 탐지 성능을 달성한다.
• 평탄한 토러스에서 복잡한 칸트리크 분포를 성공적으로 학습하였으며, RFM와 플로우 매칭은 복잡도 증가에 따라 성능이 저하되는 것과는 달리 이에 뛰어선다.
• 정확한 열핵 추정기가 없을 경우, 노이즈 제거 스코어 매칭은 열악한 결과를 낳으며, 이는 정확한 기하학적 계산의 필요성을 입증한다.
• 이 방법은 수백 차원에서의 실제 고차원 작업에 대해 미분방정식 기반 생성 모델을 처음으로 성공적으로 적용할 수 있게 한다.
• 대비 학습에서 붕괴된 투영 헤드 표현을 복구하여, 진정한 초구형 분포를 모델링함으로써 모델의 해석 가능성을 향상시킨다.

(b) $S^{127}$