JAX 기반 저랭크 행렬 업데이트로 양자 몬테카를 가속화

초록

lrux는 JAX 위에 구현된 라이브러리로, 행렬의 저랭크 변화에 대해 행렬식과 Pfaffian을 O(n²k) 비용으로 갱신한다. GPU에서 JIT·벡터화·자동미분을 활용해 대규모 양자 몬테카를로(QMC) 시뮬레이션의 핵심 병목을 1000배 가량 가속한다.

상세 분석

본 논문은 양자 몬테카를로(QMC) 계산에서 가장 빈번히 발생하는 행렬식·Pfaffian 평가를 저랭크 업데이트 기법으로 가속하는 방법을 제시한다. 행렬식의 경우, 매 단계에서 Aₜ = Aₜ₋₁ + vₜuₜᵀ (vₜ, uₜ는 n×k 행렬, k≪n) 로 표현하고, 행렬식 보조정리(det(A+UVᵀ)=det(A)·det(I+VᵀA⁻¹U))를 이용해 비율 rₜ = det(Aₜ)/det(Aₜ₋₁)=det(I+uₜᵀAₜ₋₁⁻¹vₜ) 를 O(n²k) 로 계산한다. 또한 Sherman‑Morrison‑Woodbury 식을 적용해 역행렬 Aₜ⁻¹를 동일 복잡도로 갱신한다. 이때 메모리 복잡도는 O(n²)이며, 연속적인 업데이트를 위해 A⁻¹를 유지한다.

연산량이 메모리 대역폭에 제한되는 경우를 위해 ‘지연 업데이트(delayed update)’ 전략을 도입한다. 각 단계에서 aₜ = Aₜ₋₁⁻¹vₜ, bₜ = (Aₜ₋₁⁻¹)ᵀuₜ·Rₜ⁻¹을 저장하고, τ 단계까지 누적된 aₜbₜᵀ를 실제 A⁻¹에 적용하기 전에 한 번에 수행한다. 이렇게 하면 매 단계의 n·k·k 행렬곱을 회피해 메모리 트래픽을 크게 감소시킨다. τ는 하드웨어와 행렬 크기에 따라 적절히 조정되며, 일반적으로 τ≈n/(10k) 로 설정한다.

Pfaffian 업데이트는 스키워 대칭 행렬 A에 대해 Aₜ = Aₜ₋₁ – uₜJ uₜᵀ (J는 2k×2k 스키워 항등행렬) 로 표현한다. Pfaffian의 행렬식 유사 성질 pf(A+UVᵀ)=pf(A)·pf(I+VᵀA⁻¹U)·pf(J) 를 이용해 비율 rₜ = pf(Aₜ)/pf(Aₜ₋₁)=pf(Rₜ)·pf(J) 로 계산한다. 여기서 Rₜ = J + uₜᵀAₜ₋₁⁻¹uₜ 이며, 역시 O(n²k) 복잡도로 갱신한다. 역행렬 업데이트는 Woodbury 식의 스키워 버전을 사용해 Aₜ⁻¹ = Aₜ₋₁⁻¹ + (Aₜ₋₁⁻¹uₜ)Rₜ⁻¹(Aₜ₋₁⁻¹uₜ)ᵀ 로 수행한다.

lrux는 JAX의 JIT, vmap, grad와 완벽히 호환되며, 실수·복소수 모두 지원한다. 사용자는 lrux.det_lru·lrux.pf_lru 등을 JIT으로 컴파일해 GPU에서 밀집 BLAS 연산을 최적화한다. 코드 예시에서는 rank‑1 업데이트와 연속 업데이트, 지연 업데이트를 각각 JIT‑wrapped 함수로 구현하고, 자동 미분을 통해 파라미터에 대한 그래디언트를 손쉽게 얻을 수 있다.

성능 평가에서는 n=4096 정도의 대형 행렬에 대해 기존 NumPy·SciPy 구현 대비 GPU에서 최대 1000배 속도 향상을 보고한다. 특히 k가 1~5 정도의 작은 값일 때 효율이 극대화되며, k가 커질수록 O(n³)와의 격차가 줄어든다. 수치 안정성을 위해 double precision 사용을 권고하고, Rₜ·det(Rₜ)·pf(J) 가 0에 가까워지는 경우 예외 처리를 제안한다.

전체적으로 lrux는 QMC 워크플로우에서 행렬식·Pfaffian 평가를 저랭크 구조에 맞춰 재구성함으로써, GPU 가속과 자동 미분을 동시에 활용할 수 있는 실용적인 툴킷을 제공한다. 향후 k‑adaptive 전략, 멀티‑GPU 스케일링, 그리고 더 복잡한 스키워 구조(예: 블록 스키워) 지원이 기대된다.