Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces

Motivation

• Known: Gibbs distribution으로부터 샘플링하기 위한 learning-based diffusion sampler는 유한차원 공간에서 빠르게 발전해왔으며, 특히 adjoint-based diffusion sampler(Adjoint Sampling 등)는 stochastic optimal control(SOC)과 Adjoint Matching(AM) objective를 통해 simulation-free하고 scalable한 학습을 가능하게 했다.
• Gap: 기존 adjoint sampler들은 모두 유한차원 상태공간(x∈R^d)에서만 작동하여, 계산화학에서 중요한 transition path sampling과 같이 궤적(trajectory) 자체가 무한차원 함수공간에 정의되는 Gibbs-type 분포 샘플링에는 적용할 수 없었다. 기존 TPS용 diffusion sampler들도 local energy gradient 기반의 reference dynamics에 의존해 scalable AM objective를 구성하지 못하는 한계가 있었다.
• Why: Transition path sampling은 단백질 접힘, 효소 촉매, 신약 설계, 핵화 등 다양한 응용에서 반응 경로와 속도를 이해하는 데 핵심적이며, 이를 무한차원 함수공간에서 원리적으로 다룰 수 있는 scalable sampler는 계산화학과 생성모델링 이론 양쪽에 실질적 파급력을 갖는다.
• Approach: 저자들은 Gibbs-type 분포를 무한차원 정식화로 재구성하고, infinite-dimensional Ornstein–Uhlenbeck semigroup을 reference dynamics로 삼아 non-equilibrium sampling(NES)을 무한차원 SOC 문제로 캐스팅한 뒤, Stochastic Maximum Principle을 이용해 Adjoint Matching을 Hilbert space로 확장하는 FAS를 제안한다.

Achievement

1. 무한차원 diffusion sampler 제안: infinite-dimensional Ornstein–Uhlenbeck semigroup을 reference dynamics로 사용하여 NES를 infinite-dimensional SOC 문제로 정식화했다.
2. 함수공간 Adjoint Matching 이론 확립: Stochastic Maximum Principle을 기반으로 AS의 matching-type objective를 Hilbert space로 일반화하여 scalable하고 simulation-free한 학습 알고리즘 FAS를 도출했다.
3. 경계조건을 만족하는 정확한 transition 보장: Dirichlet boundary condition을 강제하는 basis를 선택함으로써 추가 projection이나 penalty term 없이 완벽한 metastable states 간 transition을 달성했다.
4. 실험적 우수성 입증: synthetic potential과 Alanine Dipeptide, Chignolin과 같은 실제 분자 시스템에서 TPS 성능이 기존 유한차원 방법 대비 우수하며 perfect transition-hit rate를 유지함을 보였다.

How

• 무한차원 Hilbert space H 위의 Gaussian measure, Q-Wiener process, Cameron-Martin space 등 필요한 확률론적 개념을 정리하고 infinite-dimensional SDE의 mild solution과 invariant law를 정의한다.
• reference dynamics로 infinite-dimensional Ornstein–Uhlenbeck C0-semigroup을 채택하고, 이를 바탕으로 NES를 Tzen and Raginsky(2019), Zhang and Chen(2022)의 유한차원 정식화를 함수공간으로 일반화한 형태로 캐스팅한다.
• Stochastic Maximum Principle(Bensoussan, 1983; Du and Meng, 2013)을 적용해 Domingo-Enrich et al.(2025)의 Adjoint Matching을 함수공간으로 확장하고, on-policy sample과 simulation-free target에 기반한 matching objective를 도출한다.
• TPS 적용을 위해 Dirichlet boundary condition을 만족하는 basis를 선택해 경계 제약을 자연스럽게 만족시키는 함수 표현을 구성한다.
• synthetic potential 및 Alanine Dipeptide, Chignolin 분자 시스템에서 FAS의 TPS 성능을 평가한다.

Originality

• Adjoint Sampling을 유한차원에서 무한차원 Hilbert space로 확장한 최초의 시도로 보이며, SOC 이론 중 Stochastic Maximum Principle을 활용해 함수공간에서의 Adjoint Matching 이론을 정립했다.
• Q-Wiener process와 Cameron-Martin space 등 무한차원 확률론 도구를 diffusion sampler 설계에 체계적으로 결합한 점이 독창적이다.
• Dirichlet boundary condition을 만족하는 basis 선택을 통해 별도의 projection이나 penalty term 없이 경계조건을 내재적으로 충족시키는 설계는 기존 TPS 방법론과 차별화된다.

Limitation & Further Study

• 발췌된 본문에서는 이론적 정식화와 방법론 소개에 집중되어 있어, 실제 실험 설정(baseline 비교, 정량적 지표, ablation)에 대한 구체적 세부사항이 부족하여 실증적 우수성 판단이 제한적이다.
• 무한차원 basis 선택(Dirichlet boundary 등)이 특정 문제 구조에 의존적일 가능성이 있어, 보다 일반적인 경계조건이나 복잡한 topology를 갖는 시스템으로의 확장성에 대한 추가 검증이 필요하다.
• 이산화(discretization) 오차나 basis truncation에 따른 근사 오차에 대한 이론적 분석이 명확히 제시되지 않아 후속 연구에서 다뤄질 필요가 있다.
• Chignolin과 같은 큰 분자 시스템에서의 계산 비용 및 스케일링 특성에 대한 논의가 추가로 필요하다.

Evaluation

총평: 유한차원에 머물러 있던 adjoint 기반 diffusion sampler를 무한차원 함수공간으로 확장한 이론적으로 탄탄하고 실용적 파급력이 큰 연구로, TPS라는 중요한 응용 문제에서 인상적인 성능 개선을 보여준다.