LinearOperator

LinearOperator是一个PyTorch包，用于抽象化结构矩阵（或算子）所需的线性代数例程。

此包处于测试阶段。 目前，大部分功能仅支持正半定和三角矩阵。包开发待办事项

• 支持正半定算子
• 支持三角算子
• 接口指定结构（例如，对称、三角、正半定等）
• 为对称算子添加代数例程
• 为通用平方算子添加代数例程
• 为通用矩形算子添加代数例程
• 添加稀疏算子

要开始，运行以下之一

pip install linear_operator
# or
conda install linear_operator -c gpytorch

或见下文以获取更详细的说明。

为什么使用LinearOperator

在描述线性算子是什么以及为什么它们是一个有用的抽象之前，最简单的方法是看一个例子。假设你想计算矩阵求解

$$\boldsymbol A^{-1} \boldsymbol b.$$

如果你对矩阵 $\boldsymbol A$ 一无所知，代码中完成此操作的最简单（也是最好的）方法是

# A = torch.randn(1000, 1000)
# b = torch.randn(1000)
torch.linalg.solve(A, b)  # computes A^{-1} b

虽然这很简单，但 solve 例程是 $\mathcal O(N^3)$，当 $N$ 增大时速度会非常慢。

然而，让我们设想我们已经知道矩阵 $\boldsymbol A$ 等于一个低秩矩阵加上一个对角矩阵（即 $\boldsymbol A = \boldsymbol C \boldsymbol C^\top + \boldsymbol D$，其中 $\boldsymbol C$ 是一个瘦矩阵，$\boldsymbol D$ 是一个对角矩阵）。现在有一个非常高效的 $\boldsymbol O(N)$ 算法来计算 $\boldsymbol A^{-1}$（即Woodbury 公式）。一般来说，如果我们知道 $\boldsymbol A$ 有某种结构，我们希望使用高效的线性代数例程——而不是通用例程——来利用这种结构。

不使用 LinearOperator

实现利用 $\boldsymbol A$ 的低秩加对角结构的高效求解看起来可能像这样

def low_rank_plus_diagonal_solve(C, d, b):
    # A = C C^T + diag(d)
    # A^{-1} b = D^{-1} b - D^{-1} C (I + C^T D^{-1} C)^{-1} C^T D^{-1} b
    #   where D = diag(d)

    D_inv_b = b / d
    D_inv_C = C / d.unsqueeze(-1)
    eye = torch.eye(C.size(-2))
    return (
        D_inv_b - D_inv_C @ torch.cholesky_solve(
            C.mT @ D_inv_b,
            torch.linalg.cholesky(eye + C.mT @ D_inv_C, upper=False),
            upper=False
        )
    )


# C = torch.randn(1000, 20)
# d = torch.randn(1000)
# b = torch.randn(1000)
low_rank_plus_diagonal_solve(C, d, b)  # computes A^{-1} b in O(N) time, instead of O(N^3)

虽然这是一个高效的代码，但它有多个方面的不足

1. 它比 torch.linalg.solve(A, b) 要复杂得多。
2. 没有代表 $\boldsymbol A$ 的对象。要对 $\boldsymbol A$ 执行任何数学运算，我们必须传递矩阵 C 和向量 d。

使用 LinearOperator

LinearOperator 包提供了两全其美的解决方案

from linear_operator.operators import DiagLinearOperator, LowRankRootLinearOperator
# C = torch.randn(1000, 20)
# d = torch.randn(1000)
# b = torch.randn(1000)
A = LowRankRootLinearOperator(C) + DiagLinearOperator(d)  # represents C C^T + diag(d)

它提供了一个接口，让我们可以将 $\boldsymbol A$ 作为通用张量来处理，使用标准的 PyTorch API

torch.linalg.solve(A, b)  # computes A^{-1} b efficiently!

在底层，LinearOperator 对象跟踪 $\boldsymbol A$ 的代数结构（低秩加对角），并确定最有效的例程来使用（Woodbury 公式）。这样，我们可以获得高效的 $\mathcal O(N)$ 求解，同时抽象出所有细节。

关键的是，$\boldsymbol A$ 从不会被明确实例化为矩阵，这使得我们可以扩展到非常大的算子而不会耗尽内存

# C = torch.randn(10000000, 20)
# d = torch.randn(10000000)
# b = torch.randn(10000000)
A = LowRankRootLinearOperator(C) + DiagLinearOperator(d)  # represents a 10M x 10M matrix!
torch.linalg.solve(A, b)  # computes A^{-1} b efficiently!

什么是线性算子？

线性算子是矩阵的一种推广。它是一个线性函数，由其对向量的应用来定义。最常见的线性算子是（可能是有结构的）矩阵，其中对这些向量的应用是（可能是高效的）矩阵-向量乘法例程。

在代码中，LinearOperator 是一个类，它

1. 指定了定义线性算子所需的张量（
2. 指定了一个 _matmul 函数（如何将线性算子应用于向量），
3. 指定了一个 _size 函数（如果将线性算子表示为矩阵或矩阵批，它有多大），以及
4. 指定了一个 _transpose_nonbatch 函数（线性算子的共轭）。
5. （可选）定义其他函数（例如 logdet，eigh 等），以加速存在有效结构利用例程的计算。

例如

class DiagLinearOperator(linear_operator.LinearOperator):
    r"""
    A LinearOperator representing a diagonal matrix.
    """
    def __init__(self, diag):
        # diag: the vector that defines the diagonal of the matrix
        self.diag = diag

    def _matmul(self, v):
        return self.diag.unsqueeze(-1) * v

    def _size(self):
        return torch.Size([*self.diag.shape, self.diag.size(-1)])

    def _transpose_nonbatch(self):
        return self  # Diagonal matrices are symmetric

    # this function is optional, but it will accelerate computation
    def logdet(self):
        return self.diag.log().sum(dim=-1)
# ...

D = DiagLinearOperator(torch.tensor([1., 2., 3.])
# Represents the matrix
#   [[1., 0., 0.],
#    [0., 2., 0.],
#    [0., 0., 3.]]
torch.matmul(D, torch.tensor([4., 5., 6.])
# Returns [4., 10., 18.]

虽然 _matmul，_size 和 _transpose_nonbatch 可能看起来像是一组有限的函数，但实际上，大多数 torch 和 torch.linalg 命名空间中的函数都可以仅使用这三个基本函数高效实现。

此外，因为 _matmul 是一个线性函数，所以很容易以各种方式组合线性算子。例如：添加两个线性算子（SumLinearOperator）只需要添加它们的 _matmul 函数的输出。这使得可以定义非常复杂的组合结构，同时仍然产生高效的线性代数例程。

最后，线性算子对象可以相互组合，产生新的 LinearOperator 对象，并在每次计算后自动跟踪代数结构。因此，用户无需推理应该使用什么高效的线性代数例程（只要用户定义的输入元素编码了已知输入结构）。

请参阅使用线性算子对象部分以获取更多详细信息。

用例

LinearOperator 包有几个用例。在这里，我们强调两个主题

结构化矩阵的模块化代码

例如，假设您有一个涉及从高维多元高斯分布中采样的生成模型。这种采样操作将需要存储和处理一个大的协方差矩阵，因此为了加快速度，您可能想尝试不同的协方差矩阵结构近似。使用LinearOperator包可以轻松做到这一点。

from gpytorch.distributions import MultivariateNormal

# variance = torch.randn(10000)
cov = DiagLinearOperator(variance)
# or
# cov = LowRankRootLinearOperator(...) + DiagLinearOperator(...)
# or
# cov = KroneckerProductLinearOperator(...)
# or
# cov = ToeplitzLinearOperator(...)
# or
# ...

mvn = MultivariateNormal(torch.zeros(cov.size(-1), cov) # 10000-dimensional MVN
mvn.rsample()  # returns a 10000-dimensional vector

高效复杂数学运算例程

LinearOperator中的许多高效线性代数例程是基于矩阵-向量乘法的迭代算法。由于矩阵-向量乘法遵循许多良好的组合属性，因此可以获得用于极其复杂组合线性算子的有效例程

from linear_operator.operators import KroneckerProductLinearOperator, RootLinearOperator, ToeplitzLinearOperator

# mat1 = 200 x 200 PSD matrix
# mat2 = 100 x 100 PSD matrix
# vec3 = 20000 vector

A = KroneckerProductLinearOperator(mat1, mat2) + RootLinearOperator(ToeplitzLinearOperator(vec3))
# represents a 20000 x 20000 matrix

torch.linalg.solve(A, torch.randn(20000))  # Sub O(N^3) routine!

使用LinearOperator对象

LinearOperator对象与torch.Tensor对象具有（大部分）相同的API。在底层，这些对象使用__torch_function__将所有有效的线性代数操作分派到torch和torch.linalg命名空间。这包括

• torch.add
• torch.cat
• torch.clone
• torch.diagonal
• torch.dim
• torch.div
• torch.expand
• torch.logdet
• torch.matmul
• torch.numel
• torch.permute
• torch.prod
• torch.squeeze
• torch.sub
• torch.sum
• torch.transpose
• torch.unsqueeze
• torch.linalg.cholesky
• torch.linalg.eigh
• torch.linalg.eigvalsh
• torch.linalg.solve
• torch.linalg.svd

这些函数中的每一个都会返回一个torch.Tensor，或者一个新创建的LinearOperator对象，具体取决于函数。例如

# A = RootLinearOperator(...)
# B = ToeplitzLinearOperator(...)
# d = vec

C = torch.matmul(A, B)  # A new LienearOperator representing the product of A and B
torch.linalg.solve(C, d)  # A torch.Tensor

有关更多示例，请参阅示例文件夹。

批处理支持和广播

LinearOperator对象自然支持批处理模式。例如，为了表示3个100 x 100的对角矩阵批次

# d = torch.randn(3, 100)
D = DiagLinearOperator(d)  # Reprents an operator of size 3 x 100 x 100

这些对象完全支持广播操作

D @ torch.randn(100, 2)  # Returns a tensor of size 3 x 100 x 2

D2 = DiagLinearOperator(torch.randn([2, 1, 100]))  # Represents an operator of size 2 x 1 x 100 x 100
D2 + D  # Represents an operator of size 2 x 3 x 100 x 100

索引

LinearOperator对象可以像torch张量一样进行索引。这包括

• 整数索引（获取行、列或批次）
• 切片索引（获取行、列或批次的子集）
• LongTensor索引（通过索引获取单个条目集合）
• 省略号（支持具有任意批次维度的索引操作）

D = DiagLinearOperator(torch.randn(2, 3, 100))  # Represents an operator of size 2 x 3 x 100 x 100
D[-1]  # Returns a 3 x 100 x 100 operator
D[..., :10, -5:]  # Returns a 2 x 3 x 10 x 5 operator
D[..., torch.LongTensor([0, 1, 2, 3]), torch.LongTensor([0, 1, 2, 3])]  # Returns a 2 x 3 x 4 tensor

组合和装饰

线性算子可以通过各种方式组合在一起。这包括

• 加法（LinearOpA + LinearOpB）
• 矩阵乘法（LinearOpA @ LinearOpB）
• 连接（torch.cat([LinearOpA, LinearOpB], dim=-2)）
• 克罗内克积（torch.kron(LinearOpA, LinearOpB)）

此外，还有许多方法可以“装饰”线性算子对象。这包括

• 逐元素乘以常数（torch.mul(2., LinearOpA)）
• 对批次求和（torch.sum(LinearOpA, dim=-3)）
• 批次内逐元素乘法（torch.prod(LinearOpA, dim=-3)）

有关支持的组合和装饰操作的完整列表，请参阅文档composition_decoration_operators.html。

安装

LinearOperator需要Python >= 3.8。

标准安装（最新稳定版本）

我们建议通过pip或Anaconda进行安装

pip install linear_operator
# or
conda install linear_operator -c gpytorch

安装需要以下软件包

• PyTorch >= 1.11
• Scipy

您可以通过遵循PyTorch安装说明来定制您的PyTorch安装（例如，CUDA版本，仅CPU选项）。

从main分支安装（最新不稳定版本）

要安装当前位于main分支上的内容（可能存在错误和不稳定）

pip install --upgrade git+https://github.com/cornellius-gp/linear_operator.git

开发安装

如果您正在提交拉取请求，则最好进行手动安装

git clone https://github.com/cornellius-gp/linear_operator.git
cd linear_operator
pip install -e ".[dev,docs,test]"

贡献

有关提交问题和拉取请求的信息，请参阅贡献指南CONTRIBUTING.md。

许可证

LinearOperator遵循MIT许可证。