Spaces:
Sleeping
Sleeping
File size: 3,564 Bytes
ce7bf5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# Copyright Generate Biomedicines, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Layers for linear algebra.
进行线性代数计算
This module contains additional pytorch layers for linear algebra operations,
such as a more parallelization-friendly implementation of eigvenalue estimation.
"""
import torch
def eig_power_iteration(A, num_iterations=50, eps=1e-5):
"""Estimate largest magnitude eigenvalue and associated eigenvector.
This uses a simple power iteration algorithm to estimate leading
eigenvalues, which can often be considerably faster than torch's built-in
eigenvalue routines. All steps are differentiable and small constants are
added to any division to preserve the stability of the gradients. For more
information on power iteration, see
https://en.wikipedia.org/wiki/Power_iteration.
Args:
A (tensor): Batch of square matrices with shape
`(..., num_dims, num_dims)`.
num_iterations (int, optional): Number of iterations for power
iteration. Default: 50.
eps (float, optional): Small number to prevent division by zero.
Default: 1E-5.
Returns:
lam (tensor): Batch of estimated highest-magnitude eigenvalues with
shape `(...)`.
v (tensor): Associated eigvector with shape `(..., num_dims)`.
"""
_safe = lambda x: x + eps
dims = list(A.size())[:-1]
v = torch.randn(dims, device=A.device).unsqueeze(-1)
for i in range(num_iterations):
v_prev = v
Av = torch.matmul(A, v)
v = Av / _safe(Av.norm(p=2, dim=-2, keepdim=True))
# Compute eigenvalue
v_prev = v_prev.transpose(-1, -2)
lam = torch.matmul(v_prev, Av) / _safe(torch.abs(torch.matmul(v_prev, v)))
# Reshape
v = v.squeeze(-1)
lam = lam.view(list(lam.size())[:-2])
return lam, v
def eig_leading(A, num_iterations=50):
"""Estimate largest positive eigenvalue and associated eigenvector.
This estimates the *most positive* eigenvalue of each matrix in a batch of
matrices by using two consecutive power iterations with spectral shifting.
Args:
A (tensor): Batch of square matrices with shape
`(..., num_dims, num_dims)`.
num_iterations (int, optional): Number of iterations for power
iteration. Default: 50.
Returns:
lam (tensor): Estimated most positive eigenvalue with shape `(...)`.
v (tensor): Associated eigenvectors with shape `(..., num_dims)`.
"""
batch_dims = list(A.size())[:-2]
# First pass gets largest magnitude
lam_1, vec_1 = eig_power_iteration(A, num_iterations)
# Second pass guaranteed to grab most positive eigenvalue
lam_1_abs = torch.abs(lam_1)
lam_I = lam_1_abs.reshape(batch_dims + [1, 1]) * torch.eye(4, device=A.device).view(
[1 for _ in batch_dims] + [4, 4]
)
A_shift = A + lam_I
lam_2, vec = eig_power_iteration(A_shift, num_iterations)
# Shift back to original specta
lam = lam_2 - lam_1_abs
return lam, vec
|