기술 요약: GraphGP: 비트 역전 방식을 이용한 확장 가능한 가우시안 프로세스

문제 정의

가우시안 프로세스(GP)는 공간 모델링과 불확실성 정량화를 위한 원칙적인 프레임워크를 제공하지만, 심각한 확장성 한계에 직면해 있습니다. 나이브한 구현은 $O(N^3)$의 계산 비용과 $O(N^2)$의 메모리 요구 사항을 발생시켜, 수백만 또는 수십억 개의 점이 있는 데이터셋에는 실용적이지 않습니다. 기존 방법들은 부분적인 해결책을 제시하지만, 다음과 같은 상당한 트레이드오프를 가집니다:

• 유도점(Inducing-point) 방법은 복잡도를 $O(NM^2)$로 줄이지만, 유도 변수의 수 $M$이 커져야 할 때 어려움을 겪습니다.
• **고속 푸리에 변환(FFT)**은 규칙적인 격자 위에서 $O(N \log N)$ 스케일링을 달지만, 넓은 동적 범위를 가진 임의의 점 분포에 적용하기 어렵습니다.
• **베치아 근사(Vecchia's approximation)**는 각 점을 $k \ll N$개의 선행 점들에 조건화함으로써 $O(Nk^3)$ 스케일링을 제공하지만, 임의의 점 분포에 대해 이를 효율적으로 실행하는 것은 까다로운 과제였습니다.
• **반복적 차트 정밀화(Iterative Charted Refinement, ICR)**는 단일 GPU에서 수억 개의 파라미터까지 확장 가능하지만, 점들을 계층적 격자에 제한하여 유연성이 떨어집니다.

주요 과제는 적절한 조건화 관계를 빠르게 결정하면서, 격자 기반 방식이 달성한 확장성을 희생하지 않고 임의의 점 분포에 대해 계산을 병렬화하는 방법을 찾는 것입니다.

방법론

GraphGP는 베치아 근사의 유연성과 고성능 병렬화를 결합한 GPU 알고리즘을 도입합니다. 핵심 방법론은 두 가지 주요 혁신인 새로운 점 순서 지정 전략과 맞춤형 CUDA 구현에 기반합니다.

1. 베치아 근사 (Vecchia's Approximation)

알고리즘은 $N$개 점의 결합 밀도 $p(v)$를 다음과 같은 조건부 확률의 곱으로 분해합니다:
$$p(v) \approx \prod_{i=1}^{N} p(v_i | v_{c(i)})$$
여기서 $c(i)$는 $k$개의 선행 조건화 점들의 집합입니다. 모든 이전 점들을 사용하는 대신(이는 정확하지만 비용이 많이 듭니다), GraphGP는 각 점을 이전 점들 중 가장 가까운 $k$개의 이웃에 대해서만 조건화합니다. 이는 조건화 행렬의 크기를 $k \times k$로 줄여, $O(Nk^3)$ 시간 및 $O(Nk^2)$ 메모리 복잡도를 가능하게 합니다.

2. 비트 역전 $k$-d 트리 순서 (Bit-Reversed $k$-d Tree Ordering)

베치아 근사의 효율성은 점의 생성 순서에 크게 의존합니다. 표준 $k$-d 트리 순서는 공간적으로 인접한 점들을 그룹화하여, 병렬화를 저해하는 긴 의존성 체인을 생성합니다. GraphGP는 **비트 역전 순서(bit-reversed ordering)**를 제안합니다:

• 트리 구조: 점들은 좌측 균형(left-balanced) $k$-d 트리에 의해 조직됩니다.
• 비트 역전 치환: 생성 순서를 수정하여, 각 레벨에서 모든 왼쪽 자식들이 오른쪽 자식들보다 먼저 삽ert되도록 합니다. 이는 효과적으로 공간적으로 인접한 점들을 생성 시퀀스 상에서 멀리 떨어뜨려 놓습니다.
• 이점: 이 순서는 생성 시퀀스 내의 연속적인 점들이 공간적으로 멀리 떨어져 있게 하여, 이들이 조건부 독립이 될 가능성을 높입니다. 결과적으로, 대규모 점 집합을 병렬로 생성할 수 있습니다.
• 배치 구축: 알고리즘은 이웃 관계에 의해 암시된 유향 비순환 그래프(DAG)를 분석하여 각 노드의 "깊이"(루트로부터의 최장 경로)를 계산합니다. 동일한 깊이에 있는 점들은 조건부 독립이며 병렬로 처리됩니다. 정확도를 보장하고 초기 배치가 작아지는 것을 방지하기 위해, 처음 $n_0$개의 점은 밀집 촐레스키 분해(dense Cholesky factorization)를 사용하여 생성됩니다.

3. CUDA 구현 및 최적화

JAX 기반의 베이스라인 구현이 존재하지만, GraphGP는 메모리 및 성능 병목 현상을 극복하기 위해 맞춤형 CUDA 구현을 활용합니다:

• 레지스터 기반 계산: $k \times k$ 조건화 행렬은 크기가 작고($k \lesssim 64$) 수명이 짧기 때문에, 전역 메모리에 할당하는 대신 GPU 레지스터에서 온더플라이(on-the-fly)로 계산됩니다. 이는 메모리 대역폭 제한을 피하게 해줍니다.
• 인플레이스(In-place) 연산: 구현은 인플레이스 행렬 분해 및 솔버를 사용하며, 레지스터 사용을 최소화하기 위해 행렬을 하삼각(lower-triangular) 형식으로 저장합니다.
• 미분 가능성: 베이지안 추론 및 모델 피팅을 지원하기 위해, 저자들은 촐레스키 분해 미분을 사용하여 CUDA에서 순방향 및 역방향 미분을 구현했습니다. 이는 커스텀 JAX 프리미티브로 등록되어, GraphGP가 더 큰 미분 가능한 모델 내에서 원활하게 작동할 수 있도록 합니다.
• 병렬 그래프 구축: 이웃 탐색은 맞춤형 비트 역전 순서에 적응된 병렬 정렬 기반 GPU 알고리즘을 사용합니다. 깊이 업데이트는 순차적이 아닌 모든 점에 대해 반복적으로 병렬 수행되어, 이 애플리케이션에 전형적인 얕은 그래프에 대해 더 빠른 속도를 제공합니다.

주요 기여

1. 비트 역전 $k$-d 트리 순서: 트리 구조 내에서 효율적인 이웃 탐색 능력을 유지하면서 배치 병렬성을 극대화하는 새로운 생성 순서.
2. 미분 가능한 CUDA 구현: 순수 JAX 베이스라인보다 실질적으로 훨씬 빠르고 메모리 효율적인 고성능 구현으로, 커널 파라미터 미분에 대한 자동 미분을 특징으로 함.
3. 거의 10억 개의 파라미터로의 확장성: 순서 지정 전략과 하드웨어 최적화의 결면을 통해, 단일 GPU에서 선형 시간 및 메모리 요구 사항으로 거의 $10^9$개의 파라미터까지 확장 가능함.
4. 완전한 추론 구성 요소: 순방향 생성, 역방향 적용, 로그-행렬식(log-determinant) 계산, 커널 파라미터 미분을 제공하여 완전한 베이지안 추론을 가능하게 함.

결과

• 성능: H100 GPU에서의 벤치마크 결과, CUDA 구현이 순방향 패스(forward pass)에서 JAX 베이스라인보다 10배 이상 빠름을 입증했습니다.
• 확장성: 순방향 패스는 점의 개수 $N$에 대해 선형 스케일링을 달성합니다.
• 메모리 효율성: 행렬을 레지스터에서 계산함으로써, 전역 메모리 할당 전략에 비해 메모리 풋프린트를 크게 줄였습니다.
• 정확도: 쿨백-라이블러(Kullback–Leibler) 발산을 사용한 평가에 따르면, 생성 순서가 정확도에 영향을 미치기는 하지만, 사용된 이웃의 수 $k$에 비하면 그 영향은 미미합니다. 제안된 순서는 표준 트리 또는 랜덤 순서와 비교했을 때 무시할 만한 근사 정확도 손실을 보이면서 가장 큰 병렬 배치를 달성했습니다.
• 비교: 역방향 연산과 로그-행렬식만 구현하고 GPU 이웃 쿼리가 없는 GpGpU와 달리, GraphGP는 전체 순방향 패스를 지원하며 임의의 점 분포를 다룰 수 있습니다. 격자로 점을 제한해야 하는 ICR와 달리, GraphGP는 격리되지 않은 자유로운 분포를 지원합니다.

의의

GraphGP는 불균일한 해상도를 가진 필드에 대한 추론이 필요한 성간 매질의 3차원 매핑 및 기타 물리 과학 분야에 가우시안 프로세스의 실질적인 적용을 가능하게 합니다. 격자 기반 방식의 제약과 나이브한 GP의 메모리 병목 현상을 제거함으로써, 더 낮은 메모리 요구 사항으로 고해상도 지도를 생성할 수 있게 합니다. 이 연구는 임의의 점에 대한 베치아 근사의 유연성과 구조화된 격리에서만 가능했던 거대한 확장성 사이의 간극을 메우며, 대규모 공간 데이터에 대한 차세대 베이지안 추론 도구의 문을 열어줍니다.