기술 요약: LightGP

문제 정의

가우시안 프로세스 (GP) 추론은 보정된 불확실성을 갖는 회귀를 위한 원칙적인 베이지안 프레임워크이지만, 정확한 추론의 경우 $O(N^3)$의 계산 비용과 $O(N^2)$의 메모리 요구 사항은 역사적으로 이를 소규모 데이터셋에 국한시켰습니다. 확장 가능한 근사법 (예: 유도점 방법, 켤레 기울기법, 구조화된 커널 보간법) 은 이러한 한계를 수십만 개 포인트까지 확장했지만, GPyTorch 와 GPflow 와 같은 지배적인 구현체들은 무거운 딥러닝 프레임워크 (PyTorch, TensorFlow) 에 의존합니다. 이러한 프레임워크들은 종종 2GB 를 초과하는 상당한 종속성 풋프린트, 연산당 무시할 수 없는 Python 디스패치 오버헤드, 그리고 모바일 또는 임베디드 환경 배포에 대한 제한된 지원을 부과합니다. 또한, 기존 라이브러리들은 프레임워크의 불필요한 부피 없이 Apple Metal 이나 NVIDIA CUDA 와 같은 특정 하드웨어 가속기를 활용할 수 있는 네이티브이며 종속성이 없는 C++ 구현을 종종欠缺하고 있습니다.

방법론

LightGP 는 이러한 한계를 해결하도록 설계된 종속성이 없는 C++17 라이브러리입니다. 이는 Python 바인딩을 제공하지만 PyTorch 나 TensorFlow 없이 네이티브 C++ 코어에서 작동합니다. 라이브러리는 네 개의 계층으로 구조화되어 있습니다:

1. 텐서 및 BLAS: 행 우선 (row-major) float32 데이터를 저장하는 Tensor 클래스입니다. 행렬 곱셈을 macOS 의 AMX 보조 프로세서를 사용하는 Apple Accelerate 나 Linux 의 OpenBLAS 로 디스패치합니다. Cholesky 분해는 LAPACK spotrf 를 활용합니다.
2. 커널: 추상적인 Kernel 기본 클래스는 연산자 (+, *, Scale) 를 통한 조합을 지원합니다. 구체적인 구현으로는 RBF, Matérn, Periodic, Linear 커널이 포함됩니다.
3. GPU 백엔드:
   • Metal (Apple Silicon): 타일화된 커널 행렬 구성, 융합된 행렬 없는 커널 - 벡터 곱, 타일화된 스레드그룹 메모리 GEMM 을 통한 블록화된 Cholesky, 그리고 삼각 해법을 구현합니다.
   • CUDA (NVIDIA): 행렬 없는 연산을 위한 cuBLAS, cuSOLVER, 그리고 사용자 정의 커널을 사용하여 Metal 기능을 반영합니다.
   • 자동 선택: Backend::Auto 휴리스틱은 문제 크기와 차원성에 따라 CPU 또는 GPU 를 선택합니다.
4. 추론 모델:
   • GPExact: 세 가지 솔버를 지원합니다: 정확한 Cholesky ($O(N^3)$), 확률적 Lanczos 로그-행렬식을 통한 켤레 기울기법 (CG, 행렬 없는, $O(N^2k)$), 그리고 FFT 를 통한 구조화된 커널 보간법 (SKI, $O(N \log N)$).
   • GPSparse: 가장 먼 지점 유도 초기화와 Adam 최적화를 통한 Titsias 의 변분 자유 에너지 (VFE) 를 구현합니다. 반복 피팅 시 재초기화를 건너뛰는 "warm-refit" 기능을 포함합니다.

이 라이브러리는 외부 종속성이 없는 단일 정적 라이브러리로 컴파일되며 pip install lightgp를 통해 설치할 수 있습니다.

주요 기여

1. 네이티브 C++ 구현: Python 이나 프레임워크 종속성 없이 배포 가능한 네 가지 GP 추론 경로 (정확한, CG, 희소 VFE, SKI) 와 조합 가능한 커널에 대한 완전한 구현.
2. 하드웨어 최적화 커널: Apple Silicon 을 위한 손으로 튜닝된 Metal 컴퓨트 쉐이더와 NVIDIA GPU 를 위한 CUDA 커널. 특히 PyTorch 의 Metal 백엔드가 네이티브로 지원하지 않는 기능을 가능하게 하는 융합된 행렬 없는 커널 - 벡터 곱을 도입하여 CG 에 대해 $O(N)$ 메모리 추론을 가능하게 합니다.
3. FFT 가속 SKI: Apple Accelerate vDSP 와 cuFFT 를 사용한 SKI 구현으로 대규모에서 $O(N \log N)$ 행렬 - 벡터 곱셈을 달성합니다.
4. 크로스 플랫폼 벤치마크: Apple Silicon 에서 CPU 기반 AMX 보조 프로세서 (Accelerate 를 통해) 가 밀집 Cholesky 분해에 대해 Metal 보다 현저히 우수하다는 체계적인 평가로, 이는 백엔드 선택 휴리스틱에 반영됩니다.
5. 종속성 제로 Python 바인딩: pybind11 을 통한 NumPy 상호 운용성과 함께 완전한 API 노출로, 딥러닝 프레임워크의 무거운 풋프린트 없이 설치가 가능합니다.

결과

벤치마크는 Apple M4 (Metal 3) 와 Intel i7-12700 및 NVIDIA RTX 3060 (CUDA 12.0) 이 탑재된 Linux 데스크탑에서 수행되었습니다.

• GPyTorch 대비 성능:
  • Apple M4 (CPU): LightGP 는 모든 테스트된 규모에서 정확한 GP 추론의 경우 GPyTorch 보다 2.6–8.7 배 빠르며, 희소 GP 의 경우 약 1.5 배 빠릅니다. 이 속도 향상은 Python 과 PyTorch 디스패치 오버헤드를 제거한 데 기인합니다.
  • Linux (CUDA): LightGP 는 $N=2,048$까지 정확한 GP 의 경우 GPyTorch 보다 2.3–6.7 배 빠릅니다. $N=4,096$에서는 GPyTorch 가 융합된 autograd 파이프라인이 오버헤드를 분산시켜 격차를 좁힙니다. 희소 GP 의 경우 LightGP 는 $N=10,000$에서 더 빠르지만 (1.7 배), 영구 디바이스 텐서와 컴파일된 그래디언트 계산으로 인해 GPyTorch 가 $N=50,000$에서 승리합니다.
• 행렬 없는 연산: Apple M4 에서 융합된 Metal 행렬 없는 커널 - 벡터 곱은 $N=20,000$에서 명시적 CPU 경로 대비 32 배의 속도 향상을 달성하며, 16GB 메모리에서 $N=30,000$을 초과하는 유일한 실현 가능한 옵션입니다. Linux 에서 CUDA 는 $N=10,000$에서 최대 259 배의 속도 향상을 달성합니다.
• 메모리 확장성: 행렬 없는 CG, 희소 VFE, 그리고 SKI 는 $O(N)$ 메모리 사용량을 유지합니다. 반면, 정확한 Cholesky ($O(N^2)$) 는 Mac 에서 $N \approx 42,000$–45,000 에서 16GB RAM 을 소진하고 Linux 에서 $N \approx 35,000$에서 12GB VRAM 을 소진합니다.
• 정확도: 표준 데이터셋 (Motorcycle, Mauna Loa, kin40k) 에서 CPU 와 GPU 백엔드의 예측은 평균과 분산 모두에서 $10^{-5}$ 이내로 일치합니다.

중요성

LightGP 는 네이티브 C++ GP 추론이 외부 종속성 제로로 프레임워크 기반 라이브러리의 성능을 맞먹거나 능가할 수 있음을 보여줍니다. Python 과 딥러닝 프레임워크의 디스패치 오버헤드를 제거함으로써 LightGP 는 일관된 성능 향상을 제공하며, 특히 AMX 보조 프로세서를 직접 활용하는 Apple Silicon 에서 두드러집니다. 이 라이브러리는 $O(N)$ 메모리 행렬 없는 방법과 FFT 가속 SKI 를 통해 소비자 하드웨어에서 이전에는 비실용적이었던 규모 ($N > 100,000$) 의 GP 추론을 가능하게 합니다.

이 논문은 중요한 하드웨어 통찰을 강조합니다: Apple Silicon 에서 CPU (Accelerate/AMX 를 통해) 는 Cholesky 분해와 같은 밀집 선형 대수 작업에 대해 GPU (Metal) 보다 종종 우수하며, LightGP 의 휴리스틱은 이를 자동으로 관리합니다. GPyTorch 는 성숙한 autograd 와 영구 텐서 인프라로 인해 CUDA 에서 매우 큰 규모 ($N \ge 50,000$) 에서 경쟁력 있거나 우월하지만, LightGP 는 전체 딥러닝 스택의 부담 없이 경량, 고성능, 배포 가능한 GP 추론을 위한 틈새 시장을 채웁니다.