参数化教程

作者: Mario Lezcano

正则化深度学习模型是一项极具挑战性的任务。传统的惩罚方法在应用于深度模型时往往表现不佳，因为优化函数的复杂性较高。这在处理病态模型时尤为明显，例如在长序列上训练的RNN和GAN。近年来，已经提出了许多技术来正则化这些模型并提高其收敛性。在循环模型上，有建议通过控制循环核的奇异值来使RNN保持良好的条件。例如，可以通过使循环核正交化来实现这一点。另一种正则化循环模型的方法是“权重归一化”。该方法建议将参数的学习与其范数的学习解耦。为此，参数被其Frobenius范数除，并学习一个单独的参数来编码其范数。类似的归一化方法也被提出用于GAN，称为“谱归一化”。该方法通过将参数除以其谱范数而不是Frobenius范数来控制网络的Lipschitz常数。

所有这些方法都有一个共同的模式：它们在使用参数之前，都会以适当的方式对其进行转换。在第一种情况下，它们通过使用一个将矩阵映射为正交矩阵的函数，使其变为正交矩阵。在权重和谱归一化的情况下，它们将原始参数除以其范数。

更一般地说，所有这些示例都使用一个函数来为参数添加额外的结构。换句话说，它们使用一个函数来约束参数。

在本教程中，您将学习如何实现并使用这种模式来对模型施加约束。实现这一点就像编写自己的 nn.Module 一样简单。

要求：torch>=1.9.0

手动实现参数化

假设我们希望有一个权重对称的方形线性层，即权重 X 满足 X = Xᵀ。实现这一目标的一种方法是将矩阵的上三角部分复制到其下三角部分。

importtorch
importtorch.nnasnn
importtorch.nn.utils.parametrizeasparametrize

defsymmetric(X):
    return X.triu() + X.triu(1).transpose(-1, -2)

X = torch.rand(3, 3)
A = symmetric(X)
assert torch.allclose(A, A.T)  # A is symmetric
print(A)                       # Quick visual check

tensor([[0.8823, 0.9150, 0.3829],
        [0.9150, 0.3904, 0.6009],
        [0.3829, 0.6009, 0.9408]])

我们可以利用这个思路来实现一个具有对称权重的线性层

classLinearSymmetric(nn.Module):
    def__init__(self, n_features):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(n_features, n_features))

    defforward(self, x):
        A = symmetric(self.weight)
        return x @ A

该层可以作为常规的线性层使用

layer = LinearSymmetric(3)
out = layer(torch.rand(8, 3))

该实现虽然正确且自包含，但存在若干问题：

它重新实现了这一层。我们不得不将线性层实现为 x @ A。这对于线性层来说并不是什么大问题，但如果要重新实现一个卷积神经网络（CNN）或 Transformer，情况就会变得复杂。
它没有将层和参数化分离开来。如果参数化过程更加复杂，我们不得不在每个想要使用它的层中重写其代码。
每次使用该层时，它都会重新计算参数化。如果我们在前向传播过程中多次使用该层（想象一下 RNN 的循环核），那么每次调用该层时都会计算相同的 A。

参数化简介

参数化可以解决所有这些问题以及其他问题。

让我们从使用 torch.nn.utils.parametrize 重新实现上述代码开始。我们唯一需要做的就是将参数化编写为一个常规的 nn.Module。

classSymmetric(nn.Module):
    defforward(self, X):
        return X.triu() + X.triu(1).transpose(-1, -2)

这就是我们需要做的全部工作。一旦我们实现了这一点，就可以通过以下方式将任何普通层转换为对称层：

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Symmetric())

ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Symmetric()
    )
  )
)

现在，线性层的矩阵是对称的

A = layer.weight
assert torch.allclose(A, A.T)  # A is symmetric
print(A)                       # Quick visual check

tensor([[ 0.2430,  0.5155,  0.3337],
        [ 0.5155,  0.3333,  0.1033],
        [ 0.3337,  0.1033, -0.5715]], grad_fn=<AddBackward0>)

我们可以对其他任何层进行相同的操作。例如，我们可以创建一个使用斜对称核的CNN。我们使用类似的参数化方法，将上三角部分复制到下三角部分，并反转符号。

classSkew(nn.Module):
    defforward(self, X):
        A = X.triu(1)
        return A - A.transpose(-1, -2)


cnn = nn.Conv2d(in_channels=5, out_channels=8, kernel_size=3)
parametrize.register_parametrization(cnn, "weight", Skew())
# Print a few kernels
print(cnn.weight[0, 1])
print(cnn.weight[2, 2])

tensor([[ 0.0000,  0.0457, -0.0311],
        [-0.0457,  0.0000, -0.0889],
        [ 0.0311,  0.0889,  0.0000]], grad_fn=<SelectBackward0>)
tensor([[ 0.0000, -0.1314,  0.0626],
        [ 0.1314,  0.0000,  0.1280],
        [-0.0626, -0.1280,  0.0000]], grad_fn=<SelectBackward0>)

检查参数化模块

当模块被参数化时，我们发现模块在三个方面发生了变化：

model.weight 现在是一个属性
它有一个新的 module.parametrizations 属性
未参数化的权重已移动到 module.parametrizations.weight.original

在将 weight 参数化之后，layer.weight 被转换为一个 Python 属性。每当我们请求 layer.weight 时，该属性都会计算 parametrization(weight)，就像我们在上面的 LinearSymmetric 实现中所做的那样。

注册的参数化存储在模块的 parametrizations 属性下。

layer = nn.Linear(3, 3)
print(f"Unparametrized:\n{layer}")
parametrize.register_parametrization(layer, "weight", Symmetric())
print(f"\nParametrized:\n{layer}")

Unparametrized:
Linear(in_features=3, out_features=3, bias=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Symmetric()
    )
  )
)

这个 parametrizations 属性是一个 nn.ModuleDict，可以像这样访问它

print(layer.parametrizations)
print(layer.parametrizations.weight)

ModuleDict(
  (weight): ParametrizationList(
    (0): Symmetric()
  )
)
ParametrizationList(
  (0): Symmetric()
)

nn.ModuleDict 的每个元素都是一个 ParametrizationList，它的行为类似于 nn.Sequential。这个列表允许我们在一个权重上串联多个参数化。由于这是一个列表，我们可以通过索引访问其中的参数化。这就是我们的 Symmetric 参数化所在的位置。

print(layer.parametrizations.weight[0])

Symmetric()

我们注意到的另一件事是，如果我们打印参数，会发现参数 weight 已经被移动了。

print(dict(layer.named_parameters()))

{'bias': Parameter containing:
tensor([-0.0730, -0.2283,  0.3217], requires_grad=True), 'parametrizations.weight.original': Parameter containing:
tensor([[-0.4328,  0.3425,  0.4643],
        [ 0.0937, -0.1005, -0.5348],
        [-0.2103,  0.1470,  0.2722]], requires_grad=True)}

现在它位于 layer.parametrizations.weight.original 下

print(layer.parametrizations.weight.original)

Parameter containing:
tensor([[-0.4328,  0.3425,  0.4643],
        [ 0.0937, -0.1005, -0.5348],
        [-0.2103,  0.1470,  0.2722]], requires_grad=True)

除了这三处小差异外，参数化的实现与我们的手动实现完全一致。

symmetric = Symmetric()
weight_orig = layer.parametrizations.weight.original
print(torch.dist(layer.weight, symmetric(weight_orig)))

tensor(0., grad_fn=<DistBackward0>)

参数化是一等公民

由于 layer.parametrizations 是一个 nn.ModuleList，这意味着参数化被正确地注册为原始模块的子模块。因此，在模块中注册参数的相同规则也适用于注册参数化。例如，如果参数化有参数，当调用 model = model.cuda() 时，这些参数将从 CPU 移动到 CUDA。

缓存参数化配置的值

参数化通过上下文管理器 parametrize.cached() 自带内置的缓存系统。

classNoisyParametrization(nn.Module):
    defforward(self, X):
        print("Computing the Parametrization")
        return X

layer = nn.Linear(4, 4)
parametrize.register_parametrization(layer, "weight", NoisyParametrization())
print("Here, layer.weight is recomputed every time we call it")
foo = layer.weight + layer.weight.T
bar = layer.weight.sum()
with parametrize.cached():
    print("Here, it is computed just the first time layer.weight is called")
    foo = layer.weight + layer.weight.T
    bar = layer.weight.sum()

Computing the Parametrization
Here, layer.weight is recomputed every time we call it
Computing the Parametrization
Computing the Parametrization
Computing the Parametrization
Here, it is computed just the first time layer.weight is called
Computing the Parametrization

参数化串联

将两个参数化连接起来就像在同一个张量上注册它们一样简单。我们可以利用这一点，从更简单的参数化创建更复杂的参数化。例如，凯莱映射将斜对称矩阵映射到具有正行列式的正交矩阵。我们可以将 Skew 和实现凯莱映射的参数化连接起来，得到一个具有正交权重的层。

classCayleyMap(nn.Module):
    def__init__(self, n):
        super().__init__()
        self.register_buffer("Id", torch.eye(n))

    defforward(self, X):
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Skew())
parametrize.register_parametrization(layer, "weight", CayleyMap(3))
X = layer.weight
print(torch.dist(X.T @ X, torch.eye(3)))  # X is orthogonal

tensor(2.8527e-07, grad_fn=<DistBackward0>)

这也可以用于修剪参数化模块，或者重用参数化配置。例如，矩阵指数将对称矩阵映射到对称正定（SPD）矩阵，但它也将斜对称矩阵映射到正交矩阵。利用这两个事实，我们可以有效地重用之前的参数化配置。

classMatrixExponential(nn.Module):
    defforward(self, X):
        return torch.matrix_exp(X)

layer_orthogonal = nn.Linear(3, 3)
parametrize.register_parametrization(layer_orthogonal, "weight", Skew())
parametrize.register_parametrization(layer_orthogonal, "weight", MatrixExponential())
X = layer_orthogonal.weight
print(torch.dist(X.T @ X, torch.eye(3)))         # X is orthogonal

layer_spd = nn.Linear(3, 3)
parametrize.register_parametrization(layer_spd, "weight", Symmetric())
parametrize.register_parametrization(layer_spd, "weight", MatrixExponential())
X = layer_spd.weight
print(torch.dist(X, X.T))                        # X is symmetric
print((torch.linalg.eigvalsh(X) > 0.).all())  # X is positive definite

tensor(1.9066e-07, grad_fn=<DistBackward0>)
tensor(4.2147e-08, grad_fn=<DistBackward0>)
tensor(True)

初始化参数化

参数化自带初始化机制。如果我们实现一个签名如下的方法 right_inverse

defright_inverse(self, X: Tensor) -> Tensor

它将在赋值给参数化张量时使用。\ 让我们升级 Skew 类的实现以支持此功能。

classSkew(nn.Module):
    defforward(self, X):
        A = X.triu(1)
        return A - A.transpose(-1, -2)

    defright_inverse(self, A):
        # We assume that A is skew-symmetric
        # We take the upper-triangular elements, as these are those used in the forward
        return A.triu(1)

我们现在可以初始化一个由 Skew 参数化的层

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Skew())
X = torch.rand(3, 3)
X = X - X.T                             # X is now skew-symmetric
layer.weight = X                        # Initialize layer.weight to be X
print(torch.dist(layer.weight, X))      # layer.weight == X

tensor(0., grad_fn=<DistBackward0>)

当我们将参数化连接在一起时，这个 right_inverse 按预期工作。为了验证这一点，让我们升级 Cayley 参数化，使其也支持初始化。

classCayleyMap(nn.Module):
    def__init__(self, n):
        super().__init__()
        self.register_buffer("Id", torch.eye(n))

    defforward(self, X):
        # Assume X skew-symmetric
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)

    defright_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (A - I)(A + I)^{-1}
        return torch.linalg.solve(A + self.Id, self.Id - A)

layer_orthogonal = nn.Linear(3, 3)
parametrize.register_parametrization(layer_orthogonal, "weight", Skew())
parametrize.register_parametrization(layer_orthogonal, "weight", CayleyMap(3))
# Sample an orthogonal matrix with positive determinant
X = torch.empty(3, 3)
nn.init.orthogonal_(X)
if X.det() < 0.:
    X[0].neg_()
layer_orthogonal.weight = X
print(torch.dist(layer_orthogonal.weight, X))  # layer_orthogonal.weight == X

tensor(2.2141, grad_fn=<DistBackward0>)

此初始化步骤可以更简洁地写为

layer_orthogonal.weight = nn.init.orthogonal_(layer_orthogonal.weight)

此方法的名称源于我们通常期望 forward(right_inverse(X)) == X。这是一种直接的表达方式，表明在使用值 X 初始化后，forward 方法应返回值 X。实际上，这一约束并未被严格执行。事实上，有时放松这一关系可能是有意义的。例如，考虑以下随机剪枝方法的实现：

classPruningParametrization(nn.Module):
    def__init__(self, X, p_drop=0.2):
        super().__init__()
        # sample zeros with probability p_drop
        mask = torch.full_like(X, 1.0 - p_drop)
        self.mask = torch.bernoulli(mask)

    defforward(self, X):
        return X * self.mask

    defright_inverse(self, A):
        return A

在这种情况下，对于每个矩阵 A，forward(right_inverse(A)) == A 并不成立。只有矩阵 A 在掩码相同的位置上为零时，这一等式才成立。即便如此，如果我们将一个张量分配给一个被修剪的参数，那么这个张量实际上会被修剪，这并不令人意外。

layer = nn.Linear(3, 4)
X = torch.rand_like(layer.weight)
print(f"Initialization matrix:\n{X}")
parametrize.register_parametrization(layer, "weight", PruningParametrization(layer.weight))
layer.weight = X
print(f"\nInitialized weight:\n{layer.weight}")

Initialization matrix:
tensor([[0.3513, 0.3546, 0.7670],
        [0.2533, 0.2636, 0.8081],
        [0.0643, 0.5611, 0.9417],
        [0.5857, 0.6360, 0.2088]])

Initialized weight:
tensor([[0.3513, 0.3546, 0.7670],
        [0.2533, 0.0000, 0.8081],
        [0.0643, 0.5611, 0.9417],
        [0.5857, 0.6360, 0.0000]], grad_fn=<MulBackward0>)

移除参数化

我们可以使用parametrize.remove_parametrizations()来移除模块中某个参数或缓冲区的所有参数化设置。

layer = nn.Linear(3, 3)
print("Before:")
print(layer)
print(layer.weight)
parametrize.register_parametrization(layer, "weight", Skew())
print("\nParametrized:")
print(layer)
print(layer.weight)
parametrize.remove_parametrizations(layer, "weight")
print("\nAfter. Weight has skew-symmetric values but it is unconstrained:")
print(layer)
print(layer.weight)

Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0669, -0.3112,  0.3017],
        [-0.5464, -0.2233, -0.1125],
        [-0.4906, -0.3671, -0.0942]], requires_grad=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Skew()
    )
  )
)
tensor([[ 0.0000, -0.3112,  0.3017],
        [ 0.3112,  0.0000, -0.1125],
        [-0.3017,  0.1125,  0.0000]], grad_fn=<SubBackward0>)

After. Weight has skew-symmetric values but it is unconstrained:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0000, -0.3112,  0.3017],
        [ 0.3112,  0.0000, -0.1125],
        [-0.3017,  0.1125,  0.0000]], requires_grad=True)

移除参数化时，我们可以通过设置标志 leave_parametrized=False 来选择保留原始参数（即 layer.parametriations.weight.original 中的参数），而不是其参数化版本。

layer = nn.Linear(3, 3)
print("Before:")
print(layer)
print(layer.weight)
parametrize.register_parametrization(layer, "weight", Skew())
print("\nParametrized:")
print(layer)
print(layer.weight)
parametrize.remove_parametrizations(layer, "weight", leave_parametrized=False)
print("\nAfter. Same as Before:")
print(layer)
print(layer.weight)

Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[-0.3447, -0.3777,  0.5038],
        [ 0.2042,  0.0153,  0.0781],
        [-0.4640, -0.1928,  0.5558]], requires_grad=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Skew()
    )
  )
)
tensor([[ 0.0000, -0.3777,  0.5038],
        [ 0.3777,  0.0000,  0.0781],
        [-0.5038, -0.0781,  0.0000]], grad_fn=<SubBackward0>)

After. Same as Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0000, -0.3777,  0.5038],
        [ 0.0000,  0.0000,  0.0781],
        [ 0.0000,  0.0000,  0.0000]], requires_grad=True)

下载 Python 源代码: parametrizations.py

下载 Jupyter 笔记本: parametrizations.ipynb