서론 (Introduction)

만약 여러분이 고전 컴퓨터(Classical computer)에서 1,000개의 양자 큐비트(Quantum qubits)를 시뮬레이션하고 40,000개의 파라미터(Parameters)에 대한 정확한 전체 그래디언트(Full gradients)를 계산하고자 한다면, 어느 정도의 컴퓨팅 파워가 필요할까요?

1,000-큐비트, 10-레이어 회로(Circuit)의 경우, 단일 상태 진폭(State amplitude)을 계산하는 것은 감당할 만한 계산 오버헤드(Computational overhead)를 가질 수 있습니다. 하지만 시스템의 변분 양자 고유치 계산기 (Variational Quantum Eigensolver, VQE)를 해결하고 모든 파라미터 그래디언트(Parameter gradients)를 계산하려고 시도하면, 난이도가 급격히 치솟습니다. 언뜻 보기에 이것은 엄격하게 슈퍼컴퓨팅 클러스터(Supercomputing cluster)가 필요한 작업처럼 들립니다.

이러한 복잡성 폭발의 근본적인 원인은 두 가지입니다:

1. 해밀토니안 (Hamiltonian)의 도입: VQE 기대값(Expectation value)을 계산하려면 $E(\boldsymbol\theta)=\langle\psi(\boldsymbol\theta)\vert H\vert\psi(\boldsymbol\theta)\rangle$를 평가해야 합니다. 이는 텐서 네트워크(Tensor network)를 단측 상태 수축(Single-sided state contraction)에서

우리는 벤치마크로 개방 경계 조건 (Open boundary conditions)을 가진 1D 횡장 이징 모델 (Transverse Field Ising Model, TFIM)을 사용합니다. 회로의 각 레이어는 큐비트 체인을 따라 인접한 쌍 $(0,1),(1,2),\ldots,(n-2,n-1)$에 대각 엔탱글링 게이트 (Diagonal entangling gates)를 순차적으로 적용합니다:

$$R_{ZZ}(\theta)=\exp(-i\theta Z\otimes Z/2)$$

그 후 단일 큐비트 회전 (Single-qubit rotations)이 뒤따릅니다. TFIM 해밀토니언 (Hamiltonian)은 다음과 같습니다:

$$H_{\rm TFIM}=\sum_{i=0}^{n-2}X_iX_{i+1}+\sum_{i=0}^{n-1}Z_i$$

우리는 또한 이 해밀토니언을 정확하고 압축된 행렬 곱 연산자 (Matrix Product Operator, MPO, 결합 차원 (Bond Dimension) 3)로 표현할 수 있습니다.

아래 표는 세 가지 서로 다른 수축 모드에 대한 벤치마크 결과입니다 (complex64 정밀도를 사용하는 NVIDIA RTX 6000D 기준):

이러한 접근 방식들은 서로를 보완합니다:

• 전역 트리 (Global Tree) 방식은 뛰어난 정상 상태 실행 시간 (1.09초)을 제공하지만, 5시간 이상의 가혹한 콜드 컴파일 (Cold compilation) 시간과 52GB에 달하는 막대한 최대 VRAM 사용량이라는 단점이 있습니다.
• 로컬 윈도우 (Local Window) 방식은 컴파일 시간을 약 10초로 대폭 단축하지만, 인접한 로컬 항들의 공유된 환경 (Shared environments)을 중복 계산하여 단일 GPU 실행 속도를 느리게 만듭니다.
• 공간 전이 행렬 (Spatial Transfer Matrix, STM) 방식은 반복적인 1D 공간 구조를 jax.lax.scan에 명시적으로 전달함으로써 컴파일 비용을 시스템 길이로부터 성공적으로 분리하며, 컴파일과 실행 효율성 사이의 완벽한 균형을 맞춥니다.

scatter plot for efficiency of three methods

수축 경로 지표: FLOPs, 최대 텐서, 그리고 쓰기

텐서 네트워크 수축 (tensor network contraction)에서 경로의 품질은 다음과 같은 몇 가지 리소스 지표로 정의됩니다:

• FLOPs (Floating Point Operations, 부동 소수점 연산): 해당 경로에 필요한 스칼라 곱셈-덧셈 연산의 총 횟수입니다. 예를 들어, log10 FLOPs = 12.0은 대략 $10^{12}$번의 스칼라 연산을 의미합니다.
• 최대 텐서 크기 (Max Tensor Size, 폭 $w$): 수축 과정에서 생성되는 가장 큰 중간 텐서의 요소(element) 수입니다. 이는 순전파 (forward pass) 동안의 순간적인 VRAM 압박을 직접적으로 결정합니다. (모든 인덱스의 차원이 2라면, 폭 $w$는 $2^w$개의 요소에 해당합니다).
• 총 쓰기 (Total Write): 모든 중간 텐서에 걸쳐 기록되는 요소의 총합입니다. 역전파 (backpropagation) 과정에서, 그래디언트 (gradients)를 계산하기 위해 이 기록된 모든 텐서들이 잔차 (residuals)로서 메모리에 유지되어야 합니다.

기댓값 (expectations)과 그래디언트를 동시에 계산할 때는 이 세 가지 지표를 모두 주시해야 합니다. 하지만 궁극적으로는 경험적인 컴파일 시간 (compile time), 실행 시간 (execution time), 그리고 피크 GPU VRAM이 실질적인 기준 (ground truth)이 됩니다.

관점 1: 전역 수축 트리 (Global Contraction Tree) — 전체 그래프에 대한 브루트 포스 (Brute-Forcing)

완전한 bra-MPO-ket 네트워크에 대해 value_and_grad를 직접 호출함으로써, 단순한 전역 수축 (naive global contraction)의 실제 비용을 정량화할 수 있습니다.

이를 실행 가능하게 만들기 위해, 먼저 네트워크 표현 방식을 변경합니다. 각 $R_{ZZ}$ 게이트를 정확한 Rank-2 분해 형태로 작성하는 것이 매우 중요합니다:

$$R_{ZZ}(\theta)=\cos(\theta/2)I\otimes I-i\sin(\theta/2)Z\otimes Z$$

원래의 왼쪽에서 오른쪽으로 이어지는 "사다리 (ladder)" 형태의 게이트 순서를 유지하는 것 또한 매우 중요합니다. 만약 교환 가능한 (commuting) $R_{ZZ}$ 게이트들을 교차하는 짝수/홀수 본드 "벽돌벽 (brick-wall)" 구조로 재정렬하면, 제거 경로 (elimination path)가 파괴됩니다. 이 경우 텐서 네트워크의 폭은 즉시 39.585로 급증합니다. 사다리 순서를 유지하면 폭을 관리 가능한 수준인 23.585로 유지할 수 있습니다.

이 거대한 그래프에 대해 omeco 경로 탐색을 사용한 결과, log10 FLOPs = 12.0466 및 log2 write = 32.5912를 갖는 전역 수축 경로를 찾아냈습니다.

단일 RTX 6000D에서 컴파일 후 그래디언트 실행 (post-compilation gradient execution)은 단 1.09초가 소요되며, VRAM 사용량은 최대 51 GB에 달합니다. 이는 전체 그래프 VQE 그래디언트 (full-graph VQE gradients)가 기술적으로 단일 카드에서 실행될 수 있음을 증명하지만, 병목 현상은 순수하게 전역 HLO 컴파일 시간과 역전파 잔차 (backward residuals)를 저장하는 데 필요한 메모리에서 발생합니다.

omeco vs. cotengra: 두 가지 경로 탐색 프레임워크

전역 그래프의 자원 병목이 수축 경로 탐색 (contraction path search) 및 컴파일에 크게 의존하기 때문에, 최상위 도구들은 어떤 성능을 보이는지 확인해 보았습니다. 우리는 이 거대한 텐서 네트워크 (tensor network)에서 omeco와 cotengra라는 두 가지 프레임워크를 비교했습니다.

• omeco (Rust 기반)는 TreeSA 알고리즘을 사용하여 트리 공간 (tree space)에서 강도 높은 시뮬레이티드 어닐링 (simulated annealing)을 수행합니다. 이는 대규모 그래프에서 슬라이싱되지 않은 (un-sliced) 고품질 시드 트리 (seed trees)를 빠르게 찾는 데 탁월합니다.
• cotengra는 기존 경로를 기반으로 서브트리 (subtrees)를 재구성하거나 고정된 트리에서 세밀한 슬라이스 인덱스 탐색 (slice-index searches)을 수행하는 등, 수축 트리 (contraction trees)를 조작하고 미세 조정할 수 있는 더 풍부한 인터페이스를 제공합니다.

우리는 100-큐비트, 10-레이어 TFIM 그래프 (6,460개 텐서, 8,441개 인덱스)를 대상으로 통제된 탐색 비교를 수행했습니다:

데이터에 따르면 omeco TreeSA는 더 적은 시간 동안 훨씬 더 많은 어닐링 단계 (annealing steps)를 실행했으며, 더 낮은 복잡도의 트리를 반환하여 이 슬라이싱되지 않은 (un-sliced) 시나리오에서 대안들보다 뛰어난 성능을 보였습니다. 그러나 깊은 슬라이싱 (deep slicing)이 필요한 시나리오의 경우, 현재로서는 두 프레임워크 모두 이상적이지 않습니다.

슬라이싱 (Slicing)의 효율성 한계

VRAM이 부족할 때, 메모리 사용량 (memory footprints)을 제어하기 위해 수축 트리 (contraction tree)를 슬라이싱하는 것은 피할 수 없는 선택입니다. 하지만 슬라이싱되지 않은 경로를 위한 충분한 VRAM이 있다면, 여러 GPU에 분산시키기 위해 강제로 슬라이싱을 적용하는 것이 효율성을 선형적으로 확장시키지는 않습니다.

근본적으로, 텐서 네트워크 (Tensor Network)를 슬라이싱 (Slicing)한다는 것은 선택된 인덱스들의 수축 (Contraction)을 경로의 맨 마지막으로 강제로 지연시키는 것을 의미합니다. 최적화된 비슬라이싱 (Un-sliced) 경로에서는 이러한 인덱스들이 초기에 제거되어, 중간 텐서 (Intermediate Tensors)의 크기가 급격히 커지는 것을 방지합니다. 슬라이싱은 이러한 최적의 제거 순서를 방해합니다. 이는 모든 하위 작업 (Sub-task)이 최종 합산 (Summation) 전에 공유될 수 있었던 중간 결과들을 중복해서 계산하도록 강제합니다.

1,000-큐비트 글로벌 경로를 8개, 32개, 128개의 작업으로 강제 슬라이싱했을 때, 슬라이스당 log2 write는 각각 32.5912에서 32.5632, 32.5491, 32.5365로 거의 감소하지 않았습니다. 이는 역전파 (Backward Pass)의 주요 메모리 구조를 유의미하게 해체하는 데 실패했을 뿐만 아니라, 최적의 수축 순서에서 크게 벗어났기 때문에 중복 계산이 급증했기 때문입니다. 총 log10 FLOPs는 12.0466에서 12.9420, 13.5391, 14.1355로 급증했습니다.

핵심 요약 (Takeaway): VRAM에 들어갈 수 있다면, 비슬라이싱 (Un-sliced) 경로를 사용하는 것이 일반적으로 하드웨어 활용도를 극대화합니다.

framework of the three methods

관점 2: 로컬 인과적 슬라이딩 윈도우 (Local Causal Sliding Window) — 로컬 물리학의 활용

회로 깊이(Circuit depth)가 $L=10$인 경우, 2-큐비트 게이트 (2-qubit gates)의 가환성 (Commutativity)을 활용하여 사다리 구조 (Ladder structure)를 벽돌 벽 (Brick-wall) 구조와 동일하게 변환할 수 있습니다. 결과적으로, 모든 로컬 TFIM 항은 역전파 (Backpropagation) 과정 동안 정확한 역방향 인과적 원뿔 (Backward causal cone)을 가집니다. 대각 엔탱글링 게이트 (Diagonal entangling gates)는 가환하므로, 이 인과적 윈도우 (Causal window)의 너비는 엄격하게 $2L+2$이며, $L=10$일 때 단 22개 사이트 (Sites)에 불과합니다. 글로벌 트리 탐색 (Global tree search)에서는 단점이었던 벽돌 벽 구조가 여기서는 거대한 장점이 됩니다.

Hamiltonian (해밀토니안)의 로컬 Pauli 항들을 반복하기 위해 lax.scan을 사용하고, 단일 항들을 jax.checkpoint로 감싸면, 컴파일 규모는 오직 윈도우 크기와 회로 깊이에만 의존하게 되며, 전체 시스템 사이트(system sites)의 총 개수와는 완전히 분리됩니다:

$$T_{\rm window}=O(n C_{\rm local}(L)),\qquad M_{\rm window}=O(M_{\rm local}(L))$$

이 접근 방식은 빠른 컴파일 시간(~10초)과 8.86 GB라는 가벼운 VRAM 점유율을 자랑합니다. 만약 8개의 GPU(GPU당 125개의 로컬 항)에 분산하여 실행할 경우, 정상 상태(steady-state)의 그래디언트(gradient) 시간은 18.71초로 단축됩니다. 하지만, 중첩되는 인과적 윈도우(causal windows)는 막대한 계산 중복을 초래합니다.

게다가, 이 방법은 대각 게이트(diagonal gates)의 가환성(commutativity)에 크게 의존합니다. 만약 이를 비가환(non-commuting) 2-큐비트 게이트로 교체한다면, 사다리 구조(ladder structure) 하에서의 인과적 원뿔(causal cone)이 $O(n)$으로 확장되어 이 방법은 무용지물이 됩니다.

관점 3: 공간 전이 행렬 (Spatial Transfer Matrix) — 손실 없는 정확한 스캐닝

공간 전이 행렬(spatial transfer matrix) 전략은 공간 큐비트 차원을 따라 시스템을 왼쪽 끝 블록, 반복되는 중간 블록들, 그리고 오른쪽 끝 블록으로 나눕니다. 중간 블록은 왼쪽으로부터 완전한 경계 텐서(boundary tensor)를 전달받고, 내부의 bra-MPO-ket 서브 네트워크를 수축(contract)시킨 후, 다음 블록으로 출력합니다.

재귀 관계식은 다음과 같습니다:

$$B_{k+1}=\mathcal{T}_k(\boldsymbol\theta_k)B_k$$

깊이 $L=10$에서, 공간적 절단(spatial cut)은 10개의 ket 본드(ket bonds), 10개의 bra 본드(bra bonds), 그리고 1개의 MPO 본드(차원 3)를 가로지릅니다. 정확한 경계 크기는 다음과 같습니다:

$$D_{\rm boundary}=3\times2^{2L}=3\times4^L = 3\times4^{10}$$

이는 약 24 MiB의 메모리 트래픽에 해당하며, 왜 이 방법이 1,000-큐비트 규모에서 선형 스캐닝(linear scanning)을 손쉽게 수행할 수 있는지를 완벽하게 설명해 줍니다.

3-사이트 블록 분할을 사용할 경우, 이 모드는 콜드 컴파일(cold compilation)에 단 6.367초가 소요되며, 스텝당 1.054초로 실행되고, VRAM 피크는 8.62 GB입니다.

각 스캔 스텝의 공간적 너비(블록 크기 $b$)는 조절 가능한 시공간 트레이드오프(time-space tradeoff) 역할을 합니다:

• **더 작은 블록 (Smaller blocks)**은 내부 수축 그래프 (internal contraction graph)를 축소하여 컴파일 압력 (compile pressure)과 역전파 (backprop)를 위해 저장되는 잔차 (residuals)를 줄이지만, 경계 텐서 (boundary tensor)를 통과하는 데 필요한 스캔 스텝 (scan steps)의 수를 증가시킵니다.
• **더 큰 블록 (Larger blocks)**은 스캔 스텝을 줄이지만, 최대 텐서 크기 (maximum tensor size), 총 쓰기 (total write), 컴파일 시간 (compile time), 그리고 피크 VRAM (peak VRAM)을 팽창시킵니다.

우리의 테스트 결과, 3-사이트 블록 (3-site block)이 최적의 처리량 (throughput)을 제공했습니다. 5-사이트 블록 (5-site block)은 실행 시간을 1.133s로 약간 증가시키는 대가로 피크 VRAM을 5.44 GB까지 더 낮추었습니다.

직관에 반하는 가속화: 대역폭을 위한 연산량의 트레이드오프 (Trading Compute for Bandwidth)