[논문 리뷰] Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization

Geoff Pleiss, Martin Jankowiak외 3명|arXiv (Cornell University)|2020. 06. 19.

Gaussian Processes and Bayesian Inference참고 문헌 79인용 수 11

한 줄 요약

이 논문은 고차원 가우시안 프로세스와 베이지안 최적화에서 효율적인 샘플링과 화이트닝을 가능하게 하기 위해, 행렬-벡터 곱셈(MVM)을 통한 행렬 제곱근과 그 역행렬을 계산하는 빠르고 행렬을 직접 생성하지 않는 알고리즘을 제안한다. 유리 근사와 전처리가 가미된 다중 이동(mini-shift) MINRES 솔버를 조합함으로써, 100번 이내의 MVM으로 4~5자리 소수 정확도를 달성한다. 이는 GPU 가속을 통해 50,000×50,000 행렬까지 스케일링 가능하며, 확장 가능한 변분 추론과 깁스 샘플링을 가능하게 한다.

ABSTRACT

Matrix square roots and their inverses arise frequently in machine learning, e.g., when sampling from high-dimensional Gaussians $\mathcal{N}(\mathbf 0, \mathbf K)$ or whitening a vector $\mathbf b$ against covariance matrix $\mathbf K$. While existing methods typically require $O(N^3)$ computation, we introduce a highly-efficient quadratic-time algorithm for computing $\mathbf K^{1/2} \mathbf b$, $\mathbf K^{-1/2} \mathbf b$, and their derivatives through matrix-vector multiplication (MVMs). Our method combines Krylov subspace methods with a rational approximation and typically achieves $4$ decimal places of accuracy with fewer than $100$ MVMs. Moreover, the backward pass requires little additional computation. We demonstrate our method's applicability on matrices as large as $50,\!000 imes 50,\!000$ - well beyond traditional methods - with little approximation error. Applying this increased scalability to variational Gaussian processes, Bayesian optimization, and Gibbs sampling results in more powerful models with higher accuracy.

연구 동기 및 목표

• 고차원 가우시안 프로세스와 베이지안 최적화에서 K^{±1/2}b 계산의 계산 병목 현상을 해결하기 위해.
• 최대 10,000개의 유도점(inducing points)을 가진 대규모 가우시안 프로세스 모델의 확장 가능한 추론을 가능하게 하기 위해.
• 행렬 제곱근을 활용해 고차원 문제(예: 25,600차원)에서 효율적인 깁스 샘플링을 지원하기 위해.
• 학습 및 최적화 파ip라인에서 자동 미분을 위한 백워드 패스를 개발하기 위해.
• 명시적인 콜레프스 분해를 피하고 GPU 가속을 활용한 MVM을 활용함으로써 메모리 및 계산 비용을 절감하기 위해.

제안 방법

• Hale 등(2020)의 기반으로, 행렬 제곱근의 유리 근사를 이동된 행렬 역행렬의 합으로 사용한다.
• 다중 이동 MINRES(msMINRES) 알고리즘을 수정하여, t_q I + K의 역행렬을 하나의 반복에서 다중 이동 시스템 (t_q I + K)^{-1}b를 동시에 해결한다. 이 과정에서 MVM을 공유한다.
• 다양한 이동 파rameter를 가짐에도 불구하고, 단일 전처리자를 적용하여 msMINRES의 수렴 속도를 가속화한다.
• 전체 행렬 저장을 피하기 위해, 분할된 MVM을 사용하여 O(N) 메모리 사용을 유지한다.
• 연역 감도 방법을 활용하여 ∂(K^{±1/2}b)/∂K 및 ∂(K^{±1/2}b)/∂b에 대한 확장 가능한 백워드 패스를 유도한다.
• 정적 적분 기반 유리 근사와 반복적 킬로프 부분공간 방법을 조합하여 정확도와 효율성의 균형을 이루도록 한다.

실험 결과

연구 질문

• RQ1 명시적인 콜레프스 분해 없이 고차원 환경에서 행렬 제곱근을 효율적으로 계산할 수 있는가?
• RQ2 다중 이동 킬로프 솔버는 다수의 이동에 걸쳐 단일 전처리자를 통해 가속화될 수 있는가?
• RQ3 100번 이내의 행렬-벡터 곱셈으로 4자리 이상의 정확도(예: 4+ 소수 자릿수)를 달성할 수 있는가?
• RQ4 50,000×50,000 크기의 행렬에 대해 근사 오차가 낮은 상태로 알고리즘이 스케일링 가능한가?
• RQ5 변분 가우시안 프로세스와 베이지안 최적화를 포함한 엔드 투 엔드 학습 파이프라인에 통합 가능하며, 백프로파게이션까지 지원하는가?

주요 결과

• 이 방법은 100번 이내의 행렬-벡터 곱셈으로 행렬 제곱근 계산에서 4~5자리 소수 정확도를 달성한다.
• 알고리즘이 50,000×50,000 행렬까지 근사 오차가 거의 없는 상태로 스케일링 가능하며, 기존 콜레프스 기반 방법보다 훨씬 뛰어난 성능을 보인다.
• K^{±1/2}b에 대한 백워드 패스는 추가 계산 비용이 최소한이어서, 기울기 기반 기계학습 파이프라인에서 효율적으로 활용할 수 있다.
• O(M²) MVM 기반 자연 경사 하강법 업데이트를 통해 최대 10,000개의 유도점이 포함된 변분 가우시안 프로세스 추론이 가능해진다.
• 제안된 K^{-1/2}b 루틴을 활용해 25,600차원의 이미지 복원 문제에서 깁스 샘플링이 성공적으로 수행되었다.
• 이론적 분석을 통해 수렴 속도는 K의 조건수에 의존하며, 반복 횟수에 따라 오차 경계가 지수적으로 감소하는 것으로 나타났다.