Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling

기술 요약: 텐서 트레인 디퓨전 (Tensor Train Diffusion)

문제 정의
본 논문은 복잡하고 고차원적인 확률 밀도 $p_{\text{target}} = \rho_{\text{target}}/Z$로부터 샘플링하는 문제의 어려움을 다룹니다. 여기서 비정규화된 밀도 $\rho_{\text{target}}$는 점별(pointwise)로 계산 가능하지만, 정규화 상수 $Z$는 계산이 불가능하며 $p_{\text{target}}$로부터의 샘플 또한 가용하지 않은 상태입니다. 이는 분자 역학 및 베이지안 통계학을 포함한 계산 과학 분야의 핵심적인 문제입니다.

디퓨전 모델은 노이징된 확률 미분 방정식(SDE)을 역전시키는 과정을 통해 이러한 작업에 강력한 프레임워크를 제공하지만, 일반적으로 데노이징 스코어 매칭(denoising score matching)을 통해 샘-플로부터 스코어 함수를 학습해야 합니다. 본 특정 설정에서는 스코어가 샘플에 대한 접근 없이 비정규화된 밀도로부터 직접 유도되어야 합니다. 시간 역전 역학(time-reversed dynamics)은 해밀턴-자코비-벨만(Hamilton-Jacobi-Bellman, HJB) 편미분 방정식(PDE)에 의해 지배됩니다. PINN(Physics-Informed Neural Networks)이나 궤적 기반 기법과 같은 기존의 고차원 PDE 해결 방식은 긴 학습 시간, 하이퍼파라미터에 대한 높은 민감도, 그리고 $\rho_{\text{target}}$ 및 그 고계 도함수의 비용이 많이 드는 평가가 필요하다는 단점이 있습니다.

방법론
저자들은 함수적 텐서 트레인(Functional Tensor Train, FTT) 표현을 활용하여 HJB 방정식을 해결하는 새로운 솔버인 **텐서 트레인 디퓨전(Tensor Train Diffusion, TTD)**을 제안합니다.

1. 수학적 정식화:

   • 문제를 전방 노이징 프로세스를 역전시키는 샘플링 SDE를 위한 최적 제어(optimal control, drift)를 찾는 것으로 프레이밍합니다.
   • 전방 프로세스의 로그 밀도 $V = -\log \vec{p}_Y$는 특정 HJB PDE를 만족합니다.
   • HJB 방정식과 후방 확률 미분 방정식(Backward Stochastic Differential Equations, BSDEs) 사이의 연결 관계를 사용하여, 저자들은 PDE 해를 구하는 것을 일련의 회귀 문제로 재정의합니다. 이를 통해 솔루션은 역방향 시간 반복 스킴(backward-in-time iterative scheme)을 통해 계산될 수 있습니다.

2. 함수적 텐서 트레인(FTT) 근사:

   • 신경망 대신, 가치 함수(value function) $V$와 그 기울기(gradient)를 FTT를 사용하여 근사합니다. 이 형식은 고차원 함수의 잠재적인 저계수(low-rank) 구조를 활용하여, 이를 저계수 텐서 코어(tensor cores)의 곱으로 표현합니다.
   • 본 방법은 특정 단변량 기저 함수(예: 르장드르 다항식, B-스플라인, 푸리에 모드)와 TT 형식을 결합한 **확장 텐서 트레인(Extended Tensor Trains, xTT)**을 사용합니다.
   • FTT 계수의 최적화는 확률적 경사 하강법(SGD)의 불안정성을 피하기 위해 **교대 최소 제곱법(Alternating Least Squares, ALS)**이라는 회귀 기반 접근 방식을 사용하여 수행됩니다.

3. 알고리즘 특징:

   • 반복적 정교화(Iterative Refinement): 알고리즘은 제어 정책을 반복적으로 업데이트하는 외부 루프를 채택합니다. 초기 추측(예: 랑주뱅 역학)에서 시작하여, 훈련 분포를 타겟 분포에 더 잘 정렬하도록 샘플링 궤적을 정교화함으로써 새로운 모드(mode)를 효과적으로 발견합니다.
   • 적응성(Adaptivity): 프레임워크는 다음을 위한 적응형 전략을 포함합니다:
     • 랭크 선택(Rank Selection): 특이값 감소(singular value decay)에 따라 TT 랭크를 동적으로 조정하여 정확도와 계산 비용 사이의 균형을 맞춥니다.
     • 기저 차수(Basis Degrees): 서로 다른 시간 단계에서의 솔루션 매끄러움(smoothness)에 따라 기저 함수의 다항식 차수를 적응시킵니다.
     • 정규화(Regularization): 도함수 학습을 안정화하기 위해 정규화 크기를 자동으로 튜닝합니다.
   • 도메인 처리: 방법은 움직이는 근사 도메인을 처리하며, 훈련 도메인 외부 영역으로 학습된 정책을 아핀 선형 확장(affine linear extensions)을 통해 확장하여 샘플링 중 안정성을 보장합니다.

주요 기여

• 새로운 솔버: 디퓨전 기반 샘플링, HJB 이론, 그리고 BSDE를 통합하여 계산 비용과 정확도 사이의 유리한 절충안을 달로 달성하는 HJB 방정식용 솔버인 TTD를 도입했습니다.
• 강건성 및 유연성: 다양한 기저 함수(르장드르, B-스플라인, 푸리에)를 통합하고, 특히 주요 훈련 도메인 외부의 궤적을 처리함에 있어 기존 텐서 트레인 기반 PDE 솔버보다 개선된 강건성을 보여줌으로써 방법의 유연성을 입증했습니다.
• 성능: TTD가 다양한 차원의 복잡한 다중 모드 분포로부터 샘플링할 수 있음을 보여주며, 속도와 정확도 측면 모두에서 기존 디퓨전 기반 샘플러를 능가함을 증명했습니다.

실험 결과
저자들은 여러 도전적인 벤치마크에서 TTD를 평가합니다:

• 고차원 다중 웰(Multiwell) 문제: $d=10$ 및 $d=50$ 차원의 다중 모드 문제에서 TTD는 고충실도 샘플링을 달성합니다. 이 방법은 디퓨전 기반 샘플링의 흔한 문제인 모드 붕괴(mode collapse)를 피합니다. 실행 시간은 시간 단계 수에 선형적으로 스케일링되며, 신경망 기반 대안(예: DIS)보다 현저히 빠릅니다.
• 비등방성 가우시안 타겟: 실험은 알고리즘이 타겟 분포가 표준 정규 분포로 수렴함에 따라 랭크를 적응적으로 줄이는 능력을 확인시켜 주며, 이는 랭크 적응 메커니즘을 검증합니다.
• $\phi^4$ 스칼라 장 이론: 이 방법은 표준 MCMC 방법이 모드 붕괴로 인해 실패하는 경우가 많은 시스템인, 지수적으로 많은 메타스테이블 모드와 강한 최근접 이웃 상관관계를 가진 이산화된 긴즈버그-랜다우(Ginzburg-Landau) 모델로부터 성공적으로 샘플링합니다.
• 비교: TTD는 유효 샘플 크기(Effective Sample Size, ESS), 로그-분산 다이버전스(log-variance divergence), 그리고 추정된 정규화 상수의 정확도 측면에서 최신 기술 기반 베이스라인(PIS, DDS, DIS 및 다양한 MCMC 변형 포함)을 지속적으로 능가하며, 종종 더 적은 시간 내에 더 나은 결과를 달ach합니다.

의의 및 주장
본 논문은 TTD가 고차원 HJB 방정식을 해결하기 위한 신경망 기술에 대한 빠르고 정확하며 안정적인 대안을 제공한다고 주장합니다. 저계수 텐서 구조를 활용함으로써, 이 방법은 전통적인 격자 기반 방법의 "차원의 저주"와 SGD 기반 신경망 솔버의 높은 훈련 비용 및 하이퍼파라미터 민감도를 우회합니다. 저자들은 TTD를 HJB 방정식을 해결하기 위한 강력한 기초로 위치시키며, 더 적은 타겟 밀도 평가를 요구하고 고차원 샘플링 문제에 대해 향상된 효율성을 제공한다고 강조합니다. 저자들은 매우 복잡한 분포에서 모든 모드를 식별하는 것에 대한 보장과 특정 타겟 밀도에 대한 FTT의 표현력 한계를 인정하면서도, 테스트된 영역에서 이 방법의 우수한 성능을 강조합니다.