Variational inference and density estimation with non-negative tensor of hierarchical tucker format

기술 요약: 비음수 계층적 턱커(Hierarchical Tucker) 형식을 이용한 변분 추론 및 밀도 추정

문제 정의
본 논문은 $[n]^d$ 상의 텐서 $P: [n]^d \to \mathbb{R}_{\geq 0}$로 표현되는 고차원 이산 확률 분포를 다루기 쉬운 저계수(low-rank) 형식으로 압축하는 과제를 다룹니다. 이러한 텐서를 직접 저장하는 것은 $O(n^d)$의 복잡도를 요구하며, 이는 중간 정도의 차원(예: 수십 개의 변수)에서도 실행 불가능해집니다. 목표는 $P$를 효율적인 저장, 샘플링 및 가능도(likelihood) 평가가 가능한 파라미터 가족 내에서 근사하는 것입니다. 본 연구는 두 가지 접근 모델을 고려합니다:

1. 변분 추론 (Variational Inference, VI): 상수 차이만큼의 미정규화된 값 $P(i_1, \dots, i_d)$를 제공하는 오라클을 통해 $P$에 접근합니다.
2. 밀도 추정 (Density Estimation, DE): 분포에서 추출된 샘플 집합을 통해 $P$에 접근합니다.

텐서 네트워크 모델인 텐서 트레인(Tensor-Train, TT) 형식은 "면 법칙(area law)"을 만족하는 분포에는 효과적이지만, 면 법칙이 성립하지 않는 2D 격자 모델이나 복잡한 상호작용 구조를 가진 모델에서는 어려움을 겪습니다. 또한, 기존의 텐서 네트워크는 확률 분포에 필수적인 비음수성(non-negativity)을 본질적으로 보장하지 못하며, 이는 쿨백-라이블러 발산(KL divergence) 계산이나 샘플링과 같은 후속 작업에서 "부호 문제(sign problems)"를 야기할 수 있습니다.

방법론
저자들은 $P$를 비음수 계층적 턱커(Non-Negative Hierarchical Tucker, NHT) 형식으로 압축하기 위한 2단계 엔드 투 엔드 알고리즘을 제안합니다.

• 1단계: 계층적 턱커 보간 (Hierical Tucker Interpolation, 참조 구축)

  • 목표: 부호가 있는 계층적 턱커(HT) 근사 $\tilde{P} \approx P$를 구축합니다.
  • VI 케이스: 보간적 분해(interpolative decomposition, CUR 유사 분해)에 기반한 기존 보간법을 사용하여 $P$의 항목들을 쿼리하고 $\tilde{P}$를 재구성합니다.
  • DE 케이스: 샘 samples로부터 $\tilde{P}$를 구축하기 위해 기존의 밀도 추정 알고리즘을 사용합니다.
  • 출력: 부호가 있는 성분을 가진 참조 텐서 $\tilde{P}$입니다. 이 단계에서는 비음수성을 강제하지 않습니다.

• 2단계: 비음수 계층적 턱커 피팅 (Non-Negative Hierarchical Tucker Fitting)

  • 목표: 부호가 있는 참조 $\tilde{P}$에 대하여 비음수 안사(ansatz) $P_G$ (비음수 텐서 인자 $G = \{G_q\}$로 파라미터화됨)를 피팅합니다.
  • 최적화: 문제를 $\tilde{P}$와 $P_G$ 사이의 제곱 프로베니우스 노름(Frobenius norm) 및 비음수성을 강제하기 위한 로그 배리어 페널티로 구성된 손실 함수 $\ell(G)$를 최소화하는 것으로 정식화합니다:
    $$\ell(G) = \|P_G - \tilde{P}\|_F^2 + \sum_q \mu_q \phi_q(G_q)$$
    여기서 $\phi_q$는 $G_q$의 항목들에 대한 로그 배리어입니다.
  • 알고리즘: 교대 최소화(alternating minimization) 접근법이 사용됩니다. 각 텐서 성분 $G_q$에 대해, 알고리즘은 2차 뉴턴 단계(second-order Newton step)를 사용하여 변수를 업데이트합니다.
    • 그래디언트/헤시안(Gradient/Hessian): 크로네커 곱(Kronecker products)을 활용한 이진 트리 구조상의 메시지 패싱을 통해 계산됩니다.
    • 효율성: 조밀한 헤시안을 역행렬 계산하는 데 드는 $O(\rho^9)$의 비용을 피하기 위해, 중간 노드에는 공액 경사법(Conjugate Gradient, CG)을 사용하고 리프 노드에는 디커플드 솔브(decoupled solves)를 사용합니다.
  • 웜 초기화 (Warm Initialization): 핵심적인 구성 요소는 "게이지 웜 초기화(gauge warm initialization)" 전략(알고리즘 2.3)입니다. 단순한 임계값 처리 대신, 이 방법은 텐서 네트워크의 게이지 자유도(gauge freedom)를 활용합니다. 이는 참조 성분 $F_q$를 "leave-one-out" 방식으로 가장 잘 맞추는 직교 게이지 행렬 $Q$와 양의 성분 $G_q$를 찾기 위해 디커플드 스윕(leaves to root)을 수행합니다. 이는 수렴성을 크게 향상시키고 지역 최솟값(local minima)을 피하게 해줍니다.

주요 기여

1. 고차원 분포를 위한 NHT 형식: 본 논문은 텐서 트레인(TT) 형식이 구조적 제약(예: 2D 격자)으로 인해 실패할 수 있는 고차원 이산 분포에 더 적합한 비음수 계층적 턱커(NHT) 형식의 사용을 도입합니다.
2. 2차 피팅 절차: 저자들은 NHT 피팅 단계를 위한 맞춤형 2차 최적화 방법을 설계했습니다. 여기에는 고차원 설정에서 지역 최솟값에 도달하는 어려움을 해결하는 새로운 웜 스타트 전략이 포함되며, 이는 최종 정확도를 종종 한 자릿수(order of magnitude) 이상 개선합니다.
3. 계산 복잡도: 제안된 방법론의 두 단계 모두 $O(d)$의 계산 복잡도를 가지므로, 고차원 설정으로의 확장이 가능합니다.
4. 엔드 투 엔드 알고리즘: 본 연구는 부호가 있는 참조를 구축하는 것과 물리적으로 유효한 비음수 모델을 피팅하는 것 사이의 간극을 메우며, 변분 추론 및 밀도 추정 시나리오 모두에 대한 완전한 파이프라인을 제공합니다.

실험 결과
수치 실험이 두 가지 접근 모델에 대해 수행되었습니다:

• 변분 추론 (Variational Inference): 2D 긴츠부르크-란다우(Ginzburg-Landau) 모델($d=32$) 및 헤비 테일(heavy-tailed) 모델($d=30$)에 대해 테스트되었습니다. 적응형 배리어 프리컨디셔닝 뉴턴 반복법은 멀티플리케이티브 업데이트(multiplicative update) 벤치마크보다 월등히 우수한 성능을 보였으며, 동일한 실제 실행 시간 내에 더 낮은 상대적 제곱 프로베니우스 손실을 달로 달성했습니다.
• 밀도 추정 (Density Estimation): 2D 이징(Ising) 모델($d=32$) 및 2D 횡장 이징(Transverse-Field Ising) 모델($d=64$)에 대해 테스트되었습니다. 방법론은 참값(ground-truth) 분포를 성공적으로 근사하였으며, 음의 로그 가능도(NLL) 값이 실제 값에 근접함을 보였습니다.
• 절제 연구 (Ablation Studies):
  • NHT vs. NTT: 2D 긴츠부르크-란다우 모델에서 NHT 안사스는 NTT 안사스가 두 배의 파라미터를 가졌을 때보다 더 나은 피팅 정확도를 달성했습니다.
  • 초기화: 제안된 게이지 웜 초기화(알고리즘 2.3)는 멀티플리케이티브 업데이트 알고리즘(알고리즘 3.1)보다 일관되게 더 나은 피팅 결과를 산출했습니다.

의의 및 주장
본 논문은 제안된 NHT 방법론이 구조적 제약(예: 2D 격자)으로 인해 전통적인 텐서 트레인 모델이 실패하는 고차원 이산 분포를 압축하는 강력한 솔루션을 제공한다고 주장합니다. 안사스를 통해 비음수성을 보장하고 정교한 2차 피팅 절차와 특화된 웜 스타트를 활용함으로써, 이 방법은 비제한적 텐서 네트워크와 관련된 부호 문제 없이 정확한 밀도 추정과 변분 추론을 가능하게 합니다. 저자들은 이 방법이 완전 이진 트리(complete binary trees)에서 입증되었으나, 코드는 불완전한 트리(incomplete trees)도 지원하며, 향ما는 기능적 계층적 텐서 안사스로 확장하거나 모멘트 보존 제약을 통합할 수 있다고 언급합니다.