Python中大型对称正半定矩阵的有效矩阵平方根

Question

我有一个大的（150,000 x 150,000）对称正半定样本协方差矩阵，我希望在Python中有效地计算其矩阵平方根。

鉴于矩阵是对称的 psd，有什么方法可以加快平方根计算速度吗？scipy.linalg.sqrtm 对于我的目的来说很慢。

Answer 1

根据您的应用程序，如果足以找到Bs @ Bs.T ~ B，您可以使用 Cholesky 分解。如果没有，您可以根据特征值分解得到平方根。

import numpy as np;
import scipy.linalg
A = np.random.randn(1500, 1500)

%%time
Bs = scipy.linalg.sqrtm(B)

Wall time: 4.4 s- 我们的基线

%%time
Bs = scipy.linalg.cholesky(B)

Wall time: 52 ms乔斯基速度更快

D, V = scipy.linalg.eigh(B)
Bs = (V * np.sqrt(D)) @ V.T

Wall time: 1.62 s快两倍多（它探索对称性）

使用火炬

Pytorch 支持一些线性代数函数，并且它们可以跨多个 CPU 进行矢量化

import torch.linalg
B_cpu = torch.tensor(B, device='cpu')

使用 8 个的平方根（12 个逻辑 / 6 个物理 CPU）

%%time
D, V = torch.linalg.eigh(B_cpu)
Bs = (V * torch.sqrt(D)) @ V.T

Wall time: 400 ms

或 Cholesky 分解

Bs = torch.linalg.cholesky(B_cpu)

Wall time: 27 ms