A scalable sample-based quantum diagonalization solver for general Pauli-sum Hamiltonians

기술 요약: 일반 파울리 합(Pauli-sum) 해밀토니언을 위한 확장 가능한 샘플 기반 양자 대각화 솔버

문제 정의
샘플 기반 양자 대각화(SQD)는 근미래의 양자 컴퓨터를 사용하여 양자 다체계의 바닥 상태 문제를 해결하기 위한 유망한 프레임워크입니다. SQD에서는 양자 장치로부터 샘플링된 양자 상태들로 정의된 작은 부분 공간(subspace)으로 고차원 해밀토니언을 투영합니다. 문제는 이러한 부분 공간의 규모입니다. 양자 컴퓨터가 기저 상태(basis states)를 제공하지만, 클래식 후처리 과정에서 부분 공간의 차원은 수백억 개(예: $10^{10}$)까지 확장될 수 있습니다.

표준적인 대각화 방법은 이 영역에서 실패하는데, 그 이유는 행렬이 극도로 희소(sparse)하지 않는 한 축소된 해밀토니언 행렬을 메모리에 저장하는 것이 불가능하기 때문입니다. 따라서 행렬을 명시적으로 형성하지 않고도 행렬-벡터 곱을 수행할 수 있는 "행렬 프리(matrix-free)" 접근 방식이 필요합니다. 또한, 이러한 계산은 현대의 고성능 컴퓨팅(HPC) 시설, 특히 점점 더 GPU 중심으로 변모하는 환경에서 효율적으로 실행될 수 있어야 합니다. 기존의 오픈 소스 구현들은 특정 해밀토니언 유형(예: 화학 또는 하드코어 보존)에 특화되어 있거나, 과도한 GPU 메모리 점유를 요구하는 고유값 솔버(예: LOBPCG)에 의존하여 범용 파울리 합 해밀토니언에 대한 확장성을 제한합니다.

방법론
저자는 단일 워크스테이션 및 다중 GPU HPC 환경에서 실행되도록 설계된 JAX API 기반의 확장 가능한 SQD 솔버를 제시합니다. 이 솔버는 두 가지 주요 구성 요소로 이루어져 있습니다:

1. 메모리 효율적인 고유값 솔버:

   • 솔버는 축소된 해밀토니언의 최저 고유쌍(eigenpair)을 찾기 위해 **국소 최적 켤레 기울기법(Locally Optimal Conjugate Gradient, LOC)**을 채택합니다.
   • 안정성을 확보하기 위해 여러 개의 큰 벡터를 유지하는 더 일반적인 LOBPCG 방법과 달리, 이 구현은 오직 최저 고유쌍을 찾는 데 집중합니다.
   • 메모리를 공격적으로 재사용함으로써 GPU 풋프린트를 최소화하며, 이를 통해 100 GB의 메모리를 가진 단일 GPU에서 최대 $10^9$ 차원의 고유값 문제를 해결할 수 있습니다.
   • 더 큰 문제를 위해, 솔버는 JAX의 내장 병렬화 기능을 활용하여 병렬 노드 간에 분산된 벡터를 지원합니다.

2. SQD 전용 행렬-벡터 연산:

   • 핵심 연산은 $\beta = P H P \alpha$를 계산하며, 여기서 $P$는 샘플링된 부분 공간으로의 투영이고 $H$는 일반 파울리 합 해밀토니언입니다.
   • 알고리즘은 해밀토니언을 파울리 문자열(Pauli strings)로 분해하고, 파울리 연산자를 처리하기 위해 **심플렉틱 표현(symplectic representation)**을 활용합니다.
   • 핵심 단계인 "X 소스 인덱스 검색"(알고리즘 1)은 파울리 $X$ 연산자의 작용 하에서 기저 상태 사이의 매핑을 식별하며, 이는 GPU 아키텍처에 최적화된 이진 인덱스의 결합 배열에 대한 안정적인 정렬(stable sorting)을 사용하여 구현됩니다.
   • 성능 최적화를 위해, 알고리즘은 해밀토니언 항들을 그들의 $X$ 시그니처별로 그룹화하여 비용이 많이 드는 검색 및 "take" 연산의 횟수를 줄입니다.
   • 구현에는 실행 속도를 위해 정적 데이터(소스 인덱스 및 대각 성분)에 대한 캐싱 메커니즘이 포함되어 있습니다.

주요 기여

• 범용 구현: 특정 해밀토니언 유형에 국한되었던 이전 라이브러리들과 달리, 이 솔버는 큐비트 수에 제한 없이 일반 파울리 합 해밀토니언을 처리합니다.
• GPU 네이티브 확장성: 코드는 Python과 JAX를 사용하여 작성되었으므로, 코드 변경을 최소화하면서 단일 GPU 및 분산 다중 GPU HPC 클러스터에서 원활하게 실행될 수 있습니다.
• 메모리 최적화: 커스텀 고유값 솔버는 GPU 메모리 사용량을 최소화하도록 명시적으로 설계되어, 현재 하드웨어의 표준 블록 사전 조건화(block-preconditioned) 방법으로는 감당할 수 없는 차원($D \approx 10^9$)을 처리할 수 있게 합니다.
• 오픈 소스 가용성: 솔버는 rqutils라는 오픈 소스 라이브러리로 공개되며, 행렬-벡터 연산이 JAX JIT 컴파일 함수로 구현된 모든 헐미션(Hermitian) 연산자의 고유쌍을 찾는 일반적인 유틸리티를 제공합니다.

결과
논문은 두 가지 주요 실험을 통해 솔버를 검증합니다:

1. 메모리 풋프린트 검증: 솔버는 단일 NVIDIA H100 GPU(96 GB RAM)에서 30-사이트 횡장 이징 모델(transverse-field Ising model, TFIM)의 바닥 상태와 첫 번째 들뜬 상태를 성공적으로 계산했습니다. 이는 엄격한 메모리 제약 조건 내에서 비사소한(non-trivial) 해밀토니언을 처리할 수 있는 솔버의 능력을 입증했습니다.
2. 강력한 스케일링 성능: 솔버는 단일 GPU 메모리 용량을 초과하는 32-사이트 TFIM을 대상으로 8, 16, 32, 64, 128 GPU 노드에서 테스트되었습니다. 결과는 **우수한 강력한 스케일링(strong scaling)**을 보여주었으며, 모든 인스턴스에서 동일한 에너지 값을 유지하면서 노드 수가 증가함에 따라 솔루션 도달 시간(time-to-solution)이 효과적으로 감소했습니다.

의의 및 한계
본 논문은 이 솔버를 "양자 중심 슈퍼컴퓨팅(quantum-centric supercomputing)"을 위한 핵심 도구로 자리매김하며, 이를 통해 결함 허용(fault-tolerant) 이전의 양자 장치로부터 양자 이득을 추출하기 위한 클래식 후처리를 가능하게 합니다. 양자의 역할(부분 공간 차원 제안)과 클래식의 역할(대각화 및 오류 수정)을 분리함으로써, 이 프레임워크는 현재의 양자 하드웨어 한계를 극복하기 위해 거대한 클래식 컴퓨팅 파워를 활용합니다.

그러나 저자는 현재의 한계점을 명시적으로 언급합니다: X 소스 인덱스 검색은 현재 JAX에 의해 배열 크기가 $2^{31}-1$로 제한되는 글로벌 안정 정렬(global stable sorting)에 의존합니다. 결과적으로, 솔버는 현재 최대 부분 공간 차원 $N = 2^{30}-1$을 지원합니다. 저자는 이것이 현재의 많은 사용 사례를 충족하지만, 최첨단 SQD 응용 분야에서 예상되는 "수백억 개"의 차원을 지원하기 위해 향후 블록 단위 정렬(block-wise sorting)로 업그레이드할 계획임을 밝히고 있습니다. 이 논문은 결함 허용 문제를 해결하겠다고 주장하는 것이 아니라, 효율적인 클래식 스케일링을 통해 근미래 양자 컴퓨터의 효용을 극대화하는 것을 목표로 합니다.