10 行 Python 代码实现令陶哲轩惊叹的数学公式

最近，陶神的一篇小论文引起了广泛关注，我们也来跟风吃个瓜。

为了代码实现，有必要先来解读一下数学公式。陶神的研究领域需要比较深的数学背景知识，一般人是很难看懂的。而这篇里的公式涉及的数学知识大家在本科时都已经学过了，所以只要稍作展开就能看懂哦。

论文里的证明比较简练，为了让更多的吃瓜群众看明白，本文按照陶神论文里的思路将证明作了详细展开。如果不想啃公式，那就直接跳到后面看代码吧。

理论证明

我们先看一下公式，$A$ 是一个 $Hermitian$ 矩阵（简单点，可以只考虑元素是实数的情况，就是实对称矩阵），$\lambda_i(A)$ 表示它的第 $i$ 个特征值（实数），$v_{i}$ 是它第 $i$ 个特征向量，$v_{i,j}$ 就是 $v_{i}$ 的第 $j$ 个分量，以及由 $A$ 删掉第 $j$ 行和第 $j$ 列得到的子矩阵 $M_j$（$n\!-\!\small1$ 阶主子式），那么有如下公式成立，

$\begin{aligned} \left|v_{i, j}\right|^{2} &\prod_{k=1 ; k \neq i}^{n}\left(\lambda_{i}(A)-\lambda_{k}(A)\right)\\ = &\prod_{k=1}^{n-1}\left(\lambda_{i}(A)-\lambda_{k}\left(M_{j}\right)\right). \end{aligned}$

具体一点，举个例子看看。例如我们将 $Hermitian$ 矩阵 $A$ 分解为，

$\displaystyle A = \begin{pmatrix} a & X^* \\ X & M \end{pmatrix},$

其中，$a$ 是实数，$X$ 是 ${n\!-\!\small 1}$ 维向量，以及 $M$ 是 ${n\!-\!\small{1}\: \! \times \normalsize{n}\!-\!\small{1}}$ 的 $Hermitian$ 矩阵，则有，

$\displaystyle |v_{i,1}|^2 = \frac{\prod_{k=1}^{n-1} (\lambda_i(A) - \lambda_k(M))}{\prod_{k=1; k \neq i}^n (\lambda_i(A) - \lambda_k(A))},$

当然这里假设分母不为零。

简单地说，就是由原矩阵的特征值以及它所有 $n\!-\!\small 1$ 阶主子式的特征值，可以计算出原矩阵的所有特征向量分量的模平方。

下面我们根据论文 https://arxiv.org/pdf/1908.03795.pdf 来理一下陶神的证明思路。

引理$\;$ 假设 $Hermitian$ 矩阵 $A$ 的一个特征值为 $0$，不失一般性，令 $\lambda_n(A)=0$，那么对于任意 $n$ 行 $n\!-\!\small1$ 列的矩阵 $B$，下式均成立，

$\begin{aligned} & \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([B,v_{n}]\right)\right|^{2} \\ &=\operatorname{det}\left(B^{*} A B\right), \end{aligned} \tag{1}$

其中 $v_n$ 表示 $A$ 的第 $n$ 个特征向量，$[B , v_{n}]$ 表示由一个 $n$ 行 $n\!-\!\small1$ 列的矩阵 $B$ 和一个有 $n$ 个元素的列向量 $v_n$ 组成的 $n$ 行 $n$ 列的新矩阵。

证明: 对 $Hermitian$ 矩阵 $A$ 作特征分解（对角化 $A$），得 $A=V D V^{*}$，其中 $D$ 为对角矩阵

$D \equiv \operatorname{diag}\left(\lambda_{1}(A), \ldots, \lambda_{n-1}(A), 0\right).$

将分解代入等式右端，并令 $U = V^{*}B$，得

$\begin{aligned} &\quad\;\;\!\!\!\!\;\operatorname{det}\left(B^{*} A B\right) \\[1.5mm] &= \operatorname{det}\left(B^{*} V D V^{*} B\right) \\[1.5mm] &= {\operatorname{det}\left(U^{*} D U\right)} \\[1.5mm] &= {\operatorname{det}\left(U^{*}_{\large n\small{-1} \times {\large n\small{-1}}} D_{\large n\small{-1} \times {\large n\small{-1}}} U_{\large n\small{-1} \times {\large n\small{-1}}}\right)} \\[1.5mm] &= \color{#a3a}{\operatorname{det}\left(D_{\large n\small{-1} \times {\large n\small{-1}}} \right) \left|\operatorname{det}\left(U_{\large n\small{-1} \times {\large n\small{-1}}} \right)\right|^2}, \end{aligned}$

其中 $D_{\large n\small{-1} \times {\large n\small{-1}}}$ 表示矩阵 $D$ 的前 $\large n\small{-1}$ 行和前 $\large n\small{-1}$ 列构成的子矩阵，$U^{*}_{\large n\small{-1} \times {\large n\small{-1}}}$ 和 $U_{\large n\small{-1} \times {\large n\small{-1}}}$ 亦同理。

由 $U = V^{*}B$ 得 $B = VU$，将之代入等式左端得，

$\begin{aligned} &\quad\;\;\!\!\!\!\; \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([B, v_{n}]\right)\right|^{2} \\ &= \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([VU, v_{n}]\right)\right|^{2} \\ &= \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([VU, VV^{*}v_{n}]\right)\right|^{2} \\ &= \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left(V[U, V^{*}v_{n}]\right)\right|^{2} \\ &= \prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([U, V^{*}v_{n}]\right)\right|^{2} \\ &= \color{#a3a}{\prod_{i=1}^{n-1} \lambda_{i}(A)}\color{#a33}{\left|\operatorname{det}\left([U, e_{n}]\right)\right|^{2}} \\ &= \color{#a3a}{\operatorname{det}\left(D_{\large n\small{-1} \times {\large n\small{-1}}} \right)}\color{#a33}{\left|\operatorname{det}\left([U, e_{n}]\right)\right|^{2}} \\[2.15mm] &= \color{#a3a}{\operatorname{det}\left(D_{\large n\small{-1} \times {\large n\small{-1}}} \right)}\color{#a33}{ \left|\operatorname{det}\left(U_{\large n\small{-1} \times {\large n\small{-1}}} \right)\right|^2}. \end{aligned}$

可见，左右两端相等，证毕 $\boxdot$。

其实上面证明过程中很多中间步骤可以省略掉，但为了尽可能让更多人看明白，这里尽量保留了详细步骤。

定理$\;$ 矩阵 $A$ 的特征向量各元素的平方与 $A$ 的特征值以及各个子矩阵 $\{M_j\}$ 的特征值之间有如下关系，

$\begin{aligned} \left|v_{i, j}\right|^{2} & \prod_{k=1 ; k \neq i}^{n}\left(\lambda_{i}(A)-\lambda_{k}(A)\right) \\ = & \prod_{k=1}^{n-1}\left(\lambda_{i}(A)-\lambda_{k}\left(M_{j}\right)\right). \end{aligned} \tag{2}$

不失一般性，我们设 $j=1$ 和 $i=n$，即第 $n$ 个特征向量的第 $1$ 个元素。为了使用引理，让矩阵第 $n$ 个特征值 $\lambda_{n}(A) = 0$，因此让 $A$ 减去 $\lambda_n (A)I_n$，得 $A = A - \lambda_n (A)I_n$，为了简化符号，变换后的新 $A$ 仍然记作 $A$，此时 新 $A$ 的特征值与原来 $A$ 的特征值之间的关系刚好是 $\lambda_{new} = \lambda_{old} - \lambda_{n}$。因此现在有 $\lambda_n(A) = 0$；当然 $A$ 的其他特征值以及 $M_j$ 的特征值也都作了相同偏差。关于 $M_j$ 的特征值，我们可以来验证一下。将变换后的新 $M_j$ 记为 $\hat{M}_j$, 即有

$\hat{M}_j = M_j - \lambda_n I_{\large n\small{-1} \times {\large n\small{-1}}},$

它的特征多项式为，

$\begin{aligned} &\quad\; \! \left| \hat{M}_j - \lambda_{new} I_{\large n\small{-1} \times {\large n\small{-1}}}\right| \\[1.98mm] &= \left| M_j - \lambda_n I_{\large n\small{-1} \times {\large n\small{-1}}} - \lambda_{new} I_{\large n\small{-1} \times {\large n\small{-1}}}\right| \\[1.98mm] &= \left| M_j - (\lambda_n + \lambda_{new}) I_{\large n\small{-1} \times {\large n\small{-1}}}\right|. \end{aligned}$

可见 $\lambda_{new} = \lambda_{old} - \lambda_{n}$。同样的，将变换后的新 $M_j$ 仍然记为 $M_j$。

回到前面，此时矩阵 $A$ 满足引理的条件，且式 $(2)$ 可简化为，

$\left|v_{n, 1}\right|^{2} \prod_{k=1}^{n-1} \lambda_{k}(A)=\prod_{k=1}^{n-1} \lambda_{k}\left(M_{1}\right). \tag{3}$

式 $(3)$ 的右端就是 $\operatorname{det}\left(M_{1}\right)$。

接下来使用引理来证明上面式 $(3)$ 成立。我们先构造一个矩阵，

$B=\left(\begin{array}{c}{0} \\ {I_{n\small{-1}}}\end{array}\right) = \begin{pmatrix} 0 & \cdots & 0 \\ 1 & \cdots & 0 \\ \vdots & \ddots & \vdots \\ 0 & \cdots & 1 \end{pmatrix}_{\large n \times {\large n\small{-1}}}.$

将 $B$ 代入式 $(1)$

$\prod_{i=1}^{n-1} \lambda_{i}(A)\left|\operatorname{det}\left([B,v_{n}]\right)\right|^{2}=\operatorname{det}\left(B^{*} A B\right)$

得左端为，

$\prod_{i=1}^{n-1} \lambda_{i}(A)\left|v_{n, 1}\right|^{2}.$

易知右端为 $\operatorname{det}\left(M_{1}\right)$，所以式 $(3)$ 的左端 $\left|v_{n, 1}\right|^{2} \prod_{k=1}^{n-1} \lambda_{k}(A)$ 也等于 $\operatorname{det}\left(M_{1}\right)$，因此式 $(3)$ 得证 $\boxdot$。

思维导图

这里联想一下上面定理的证明思路，画了一个简单的导图。

最后构造了一个类似柯西–比内公式（Cauchy–Binet formula）的引理，而柯西–比内公式将行列式的可乘性推广到非方块矩阵，即非方块矩阵乘积的行列式转化为方块矩阵行列式的乘积之和。

另外，论文里还提供了使用伴随矩阵的另一种证明方法，

$\operatorname{adj}\left(\lambda I_{n}-A\right)=\operatorname{det}\left(\lambda I_{n}-A\right)\left(\lambda I_{n}-A\right)^{-1}.$

有兴趣的话可以前往 arxiv 下载论文学习。

经过网络传播，各路网友纷纷加入吃瓜，后来大家知道这个公式其实早在 1968 年的线性代数书籍上就出现了，它并不是什么新公式。那么这次重新提出是不是完全就没有意义呢？

可以这么认为，公式早就诞生了，但是可能并没有很好地应用开。这次几个物理学家从他们从事的中微子研究中重新发现了它，让这个公式出现在了具体的应用领域，也具有了相应的物理意义。虽然目前应用不是很广泛，但对于该公式来说，意义并不小，因为后面必将会受到更多的关注。

至于该公式在机器学习或更广的应用数学中如何应用，值得大家思考。

下面我们来用 Python 实现一下，看看它的性能如何。

代码实现

再回顾下公式，

$\begin{aligned} \left|v_{i, j}\right|^{2} \prod_{k=1 ; k \neq i}^{n}\left(\lambda_{i}(A)-\lambda_{k}(A)\right) = \prod_{k=1}^{n-1}\left(\lambda_{i}(A)-\lambda_{k}\left(M_{j}\right)\right), \end{aligned}$

或者

$\displaystyle |v_{i,j}|^2 = \frac{\prod_{k=1}^{n-1} (\lambda_i(A) - \lambda_k(M_j))}{\prod_{k=1; k \neq i}^n (\lambda_i(A) - \lambda_k(A))}.$

因此，要计算矩阵 $A$ 的第 $i$ 个特征向量 $v_{i}$ 的第 $j$ 个元素 $v_{i,j}$ 的模，需要求出矩阵 $A$ 和 子矩阵 $M_j$ 的特征值。

本文实现了 NumPy 和 PyTorch 两个版本的代码。我们只处理实对称矩阵的情况，$Hermitian$ 矩阵只需要稍微修改一下代码即可。由于直接根据公式来实现，所以写法比较直接，代码也很容易懂。

NumPy 版本

import numpy as np
from numpy import linalg as LA

np.__version__

'1.17.2'

# 取 A 的子矩阵 M_j，多个版本，可以比较性能
def sub_matrix_np0(A, n, j):
    row, row[j] = np.ones(n, dtype=bool), False
    return A[row][:,row]

def sub_matrix_np1(A, n, j):
    row, row[j:] = np.arange(n-1), np.arange(n-1)[j:]+1
    return A[row][:,row]

def sub_matrix_np2(A, n, j):
    row = np.arange(n)
    row = np.concatenate([row[:j], row[j+1:]])
    return A[row][:,row]

# 此处用到 NumPy 只计算特征值，不计算特征向量
def eigv_ij_tao_np(A, i, j):
    eigvals = LA.eigvals(A)
    M_j = sub_matrix_np0(A, eigvals.shape[0], j)
    eigvals_M_j = LA.eigvals(M_j)
    eigvals_M_j = eigvals_M_j - eigvals[i]
    eigvals, eigvals[i] = eigvals - eigvals[i], 1.0
    return np.prod(eigvals_M_j) / np.prod(eigvals)

# 我们使用如下实对称矩阵作测试
A = np.array([[3.20233843, 2.06764157, 2.88108853, 3.04719701, 1.35948144,
        2.56403503, 1.45954379, 2.13631733, 1.53986513, 3.25995899],
       [2.06764157, 2.29127686, 2.1014315 , 2.70700189, 1.39188518,
        2.29620216, 1.40820018, 1.53013976, 1.63189934, 2.82578892],
       [2.88108853, 2.1014315 , 3.50400477, 3.47397131, 1.76088485,
        3.23409601, 1.32428533, 2.77130533, 2.23036096, 3.42124947],
       [3.04719701, 2.70700189, 3.47397131, 4.73687997, 2.20972905,
        3.77378375, 2.34901207, 2.48563082, 2.19559837, 4.13322167],
       [1.35948144, 1.39188518, 1.76088485, 2.20972905, 1.59570009,
        2.14854992, 1.1252818 , 1.94769495, 1.34582313, 2.1726746 ],
       [2.56403503, 2.29620216, 3.23409601, 3.77378375, 2.14854992,
        3.90925026, 2.11631829, 2.99858233, 2.44964915, 3.61299442],
       [1.45954379, 1.40820018, 1.32428533, 2.34901207, 1.1252818 ,
        2.11631829, 2.14092743, 1.13700682, 1.22321095, 2.05145548],
       [2.13631733, 1.53013976, 2.77130533, 2.48563082, 1.94769495,
        2.99858233, 1.13700682, 3.41056761, 2.08224044, 3.02608265],
       [1.53986513, 1.63189934, 2.23036096, 2.19559837, 1.34582313,
        2.44964915, 1.22321095, 2.08224044, 2.12302856, 2.41838461],
       [3.25995899, 2.82578892, 3.42124947, 4.13322167, 2.1726746 ,
        3.61299442, 2.05145548, 3.02608265, 2.41838461, 4.55396155]])

n = A.shape[0]

# 测试，计算第 n 个特征向量分量的模平方
[eigv_ij_tao_np(A, i=n-1, j=j) for j in range(n)]

[0.010078898355312161,
 0.008526097435704025,
 0.03552221239921944,
 0.021411824331722774,
 0.17690013381018135,
 0.5189551148981736,
 0.06059635621137234,
 0.00246512246354466,
 0.10612109694238948,
 0.059423143152378614]

%%timeit
[eigv_ij_tao_np(A, i=n-1, j=j) for j in range(n)]

766 µs ± 3.02 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# 设置 numpy 输出的精度，用于展示精度对比
np.set_printoptions(precision=18)

# 使用 numpy.linalg 直接计算结果，比较用
n = A.shape[0]
(LA.eig(A)[1]**2)[:,n-1].reshape(-1,1)

array([[0.010078898355311885],
       [0.008526097435703766],
       [0.035522212399219155],
       [0.021411824331722753],
       [0.17690013381018374 ],
       [0.5189551148981745  ],
       [0.0605963562113716  ],
       [0.002465122463544841],
       [0.10612109694238844 ],
       [0.059423143152379225]])

%%timeit
(LA.eig(A)[1]**2)[:,n-1].reshape(-1,1)

39.8 µs ± 227 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

看到没有，NumPy 直接计算特征向量的效率明显比上面 Tao 算法的实现要高效得多。

PyTorch 版本

import torch

torch.__version__

'1.1.0'

# 设置 torch 输出的精度，用于展示精度对比
torch.set_printoptions(precision=18)

# 取 A 的子矩阵 M_j，多个版本，可以比较性能
def sub_matrix_torch0(A, n, j):
    row, row[j] = torch.ones(n, dtype=torch.bool), False
    return A[row][:,row]

def sub_matrix_torch1(A, n, j):
    row, row[j:] = torch.arange(n-1), torch.arange(n-1)[j:]+1
    return A[row][:,row]

def sub_matrix_torch2(A, n, j):
    row = torch.arange(n)
    row = torch.cat([row[:j], row[j+1:]])
    return A[row][:,row]

def eigv_ij_tao_torch(A, i, j):
    eigvals, _ = torch.symeig(A)
    M = sub_matrix_torch2(A, eigvals.shape[0], j)
    eigvals_M_j, _ = torch.symeig(M)
    eigvals_M_j = eigvals_M_j - eigvals[i]
    eigvals, eigvals[i] = eigvals - eigvals[i], 1.0
    return torch.prod(eigvals_M_j) / torch.prod(eigvals)

A_torch = torch.tensor(A)

n = A_torch.shape[0]

# 由于特征值次序不同，所以这里 i 跟上面 numpy 版本取值不同
[eigv_ij_tao_torch(A_torch, i=n-8, j=j) for j in range(n)]

[tensor(0.010078898355311965, dtype=torch.float64),
 tensor(0.008526097435700507, dtype=torch.float64),
 tensor(0.035522212399216449, dtype=torch.float64),
 tensor(0.021411824331723325, dtype=torch.float64),
 tensor(0.176900133810184573, dtype=torch.float64),
 tensor(0.518955114898170700, dtype=torch.float64),
 tensor(0.060596356211371390, dtype=torch.float64),
 tensor(0.002465122463545969, dtype=torch.float64),
 tensor(0.106121096942389470, dtype=torch.float64),
 tensor(0.059423143152376595, dtype=torch.float64)]

可以跟上面 NumPy 版本的结果对比一下看是否一致。

%%timeit
[eigv_ij_tao_torch(A_torch, i=n-8, j=j) for j in range(n)]

828 µs ± 2.27 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

从这个小矩阵来看，PyTorch 实现的效率还不如 NumPy 版本。
在实现之际，发现网络上已经有一些版本，本文的实现跟这个版本比较接近，但比它的更加简短且性能更好一些。

性能对比

上面使用一个小矩阵测试代码，下面用比较大的矩阵测试对比下 NumPy 和 PyTorch 实现的两个版本的性能。

A = np.random.rand(200, 200)
A = A.dot(A.T)

NumPy 版本 Tao 算法:

3.08 s ± 19.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

PyTorch 版本 Tao 算法:

539 ms ± 1.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

NumPy 直接求解特征向量:

10.8 ms ± 33.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

PyTorch 直接求解特征向量:

10.5 ms ± 53.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

从陶等人的公式本身可想而知，计算效率并不高。不过这个公式并不是为了计算来的，而是揭示了特征值与特征向量分量的模之间的关系。

另外，本人其实更喜欢 Julia 语言，但考虑到目前了解 Python 的人更多。上面代码要改写成 Julia 版本也方便。

最后留个作业，用该公式计算 $Hermitian$ 矩阵的特征向量的时间复杂度为多少呢？

下载本文代码请前往仓库。
文章首发于下面的公众号。