深度隐式层 | 可微优化

翻译自 https://implicit-layers-tutorial.org/differentiable_optimization/

在本教程的前几章中，我们考虑了诸如深度平衡模型和神经常微分方程这样的隐式层，它们都对层的性质施加了某种特定类型的结构。另一种常见的高度相关的结构类型是，将层定义为优化问题的解。这种处理需要事先做一些说明，因为在机器学习语境中，"优化"几乎总是指拟合数据参数的学习过程（例如，优化深度网络的参数，包括那些包含 DEQ 或 Neural ODE 层的网络）。但这并不是我们在本节中所指的内容。相反，我们谈论的是那些求解小规模（否则网络的前向传播将花费太长时间）优化问题作为层本身函数的层。

为了更形式化地描述，我们可以按如下方式定义泛化优化问题

\[\begin{aligned} \min_z \;\; & f(z,x) \\ \text{subject to} \;\; & z \in \mathcal{C}(x) \end{aligned} \]

其中 \(z \in \mathbb{R}^n\) 是优化变量；\(x \in \mathbb{R}^m\) 表示定义优化问题的附加问题数据；\(f : \mathbb{R}^n \times \mathbb{R}^m \rightarrow \mathbb{R}\) 是目标函数，\(C(x) \subseteq \mathbb{R}^n\) 表示约束集。换句话说，输入 \(x\) 同时定义了优化目标和约束集，然后我们找到最小化目标并满足约束的 \(z\) 变量。

现在我们可以将优化层的输出定义为给定输入 \(x\) 所定义的问题的最优解（即达到约束优化问题最优值的 \(z^\star\)）。换句话说，我们可以将层定义为

\[z^\star(x) = \argmin_{z \in \mathcal{C}(x)} f(z,x) \]

其中 \(\argmin\) 表示达到最小值的变量 \(z\) 的值。

可微优化层作为"传统"深度网络的组成部分以及其自身都有许多应用。双层优化领域；我们在 OptNet 及后续工作中的成果；Deep Declarative Networks（深度声明式网络）；控制中的可微优化。下面，我们还将展示一些简单的结果，表明深度学习中使用的许多常见（以及不那么常见）的非线性，事实上可以解释为优化问题的解；这样做通常会揭示深度网络中每一层精确"目标"的新视角。我们还将展示一些简单的示例，说明如何使用这些层从数据中学习"硬约束"，这是许多现代网络经常面临的挑战。

凸优化问题

为了使我们的任务比所有通用优化问题（通常可能极其困难）更简单，我们将专注于凸优化问题的特定子集。凸问题具有一些理想的特性，即它们（通常）允许使用高效（即多项式时间）的求解方法来找到问题的全局最优解。与本教程更相关的是，存在现成的求解器可用于许多类别的凸优化问题，并且存在领域特定的建模语言，允许人们以自然形式轻松指定这些问题，然后可以解析这些问题，以便通过现成求解器求解。

出于本教程的目的，我们将考虑一种相当通用的凸优化问题形式，它将涵盖我们感兴趣的大多数情况（注意，本节将以某种非正式的方式呈现这些内容，不讨论与这些问题相关的凸分析的微妙问题）。具体来说，我们考虑以下形式的凸优化问题

\[\begin{aligned} \min_z \;\; & f(z) \\ \text{subject to} \;\; & g(z) \leq 0 \\ & h(z) = 0, \end{aligned} \]

其中 \(z \in \mathbb{R}^n\) 是优化变量，\(f : \mathbb{R}^n \rightarrow \mathbb{R}\) 是目标函数，\(g : \mathbb{R}^n \rightarrow \mathbb{R}^m\) 是（向量值的）不等式约束集，\(h : \mathbb{R}^n \rightarrow \mathbb{R}^p\) 是（向量值的）等式约束集。为了使这个问题成为凸优化问题，需要满足：

1. 函数 \(f(z)\) 和每个分量 \(g_i(z)\) 必须是凸函数。简而言之，如果对任意 \(x,y \in \text{dom}(f)\) 和 \(\theta \in [0,1]\) 都满足如下条件，则 \(f\) 是凸的。我们假定读者对凸优化有基本了解，并建议读者参阅 [Boyd and Vandenberghe, 2004] 以获取更全面的概述。

\[f(\theta x + (1-\theta)y) \leq \theta f(x) + (1-\theta) f(y), \]

2. 函数 \(h(z)\) 必须是仿射的，即形式为 \(h(z) = A z + b\)，其中 \(A \in \mathbb{R}^{p \times n}\)，\(b \in \mathbb{R}^p\)。

凸优化最具实际吸引力的方面之一是，我们可以轻松描述保证底层优化问题全局最优性的条件。例如，在具有可微目标 \(f\) 的无约束优化问题的情况下，保证最优性的一个简单条件是零梯度条件，即如果我们找到一个点 \(z^\star\) 使得

\[\nabla f(z^\star) = 0, \]

那么 \(z^\star\) 就是优化问题的全局解。对于像上面这样的不等式和等式约束问题，存在这个条件的推广形式，称为 Karush-Kuhn-Tucker（KKT）条件。这些条件表明，当且仅当我们可以找到 \((z^\star,\lambda^\star, \nu^\star)\) 使得以下五个条件成立

\[\begin{aligned} g(z^\star) & \leq 0 \\ h(z^\star) & = 0 \\ \lambda^\star & \geq 0 \\ \lambda^\star \circ g(z^\star) & = 0 \\ \nabla f(z^\star) + \sum_{i=1}^m \lambda_i^\star \nabla g_i(z^\star) + \sum_{i=1}^p \nu_i^\star \nabla h_i(z^\star) & = 0 \end{aligned} \]

其中 \(\circ\) 表示 Hadamard 积（逐元素乘法）。这里 \(\lambda\) 和 \(\nu\) 被称为对偶变量（而 \(z\) 被称为原始变量），KKT 条件描述了最优原始-对偶对的必要且充分条件。

实践中的凸优化问题求解

如上所述，凸优化的关键方面之一是，尽管开发能够稳健求解不同类别凸问题的通用求解器可能相当复杂，但存在易于使用的库来快速指定问题，然后将这些问题传递给经过良好调优的求解器。我们将在这里使用的特定库是 cvxpy 库 [Diamond and Boyd,2016, Agarwal et al., 2018]。该库提供了在 Python 内指定和求解优化问题的简洁接口。例如，下面是一个使用 cvxpy 求解带 \(\ell_1\) 正则化的非负最小二乘问题的快速演示。

import cvxpy as cp
import numpy as npm,n = 50,10
A = np.random.randn(m,n)
b = np.random.randn(m)
x = cp.Variable(n)
lam = 0.1f = cp.sum_squares(A@x - b) + lam*cp.norm1(x)
cons = [x >= 0]
cp.Problem(cp.Minimize(f), cons).solve(verbose=False, eps_abs=1e-8, eps_rel=1e-8)
print(np.round(x.value,5))

[-0.       0.0294   0.       0.06528 -0.       0.23629 -0.      -0.-0.      -0.     ]

注意，使用 cvxpy 时，我们不需要像上面那样以 \(g\) 或 \(h\) 函数的形式精确指定优化形式，而通常只需以自然的方式将约束指定为一个列表。在幕后，cvxpy 将这些问题转换为所谓的锥规划（cone programs），可以根据所涉及的不同锥类型（如正锥、二阶锥、半正定锥或指数锥）由多种求解器求解。然而，这些转换的细节对于下面的讨论并不重要。此外，我们不仅可以得到原始解，还可以得到每个约束对应的对偶解（即 \(\lambda\) 变量）。

print(np.round(cons[0].dual_value, 5))

[ 3.18426  0.      10.78774  0.       3.6743   0.       0.93088 21.022455.42407 14.82586]

可微凸优化

现在我们能够求解凸优化问题（尽管只是通过调用现有求解器），我们可以讨论如何使这些问题可微。稍微融合前两节的符号，我们将考虑一个凸优化问题，其中每个函数 \(f,g,h\) 同时是优化变量 \(z\) 和层输入 \(x\) 的函数（即 \(x\) 实际上决定了函数的形式）。具体来说，我们考虑如下定义的层

\[\begin{aligned} z^\star(x) = \argmin_{z} \;\; & f(z,x) \\ \text{subject to} \;\; & g(z,x) \leq 0 \\ & h(z,x) = 0 \end{aligned} \]

其中函数的定义使得对任何输入 \(x\)，优化问题都是凸的。更一般地，由于优化问题实际上返回原始-对偶最优解 \((z^\star, \lambda^\star, \nu^\star)\)，我们可以将整个解写作 \(x\) 的函数

\[(z^\star, \lambda^\star, \nu^\star)(x). \]

与不动点迭代的情况一样，对这些问题进行微分的关键在于将凸优化求解器视为简单地寻找 KKT 条件的根（由于 KKT 条件也包含不等式约束，因此需要一些额外的细微处理）。具体来说，让我们将 KKT 条件的三个等式约束记为

\[G(z, \lambda, \nu) = \left [ \begin{array}{c} \nabla_z f(z,x) + \partial_z g(z,x)^T \lambda + \partial_z h(z,x)^T \nu \\ \lambda \circ g(z,x) \\ h(z) \end{array} \right ] \]

其中我们略微简化了之前的符号，用 \(\partial_z g(z,x)\) 表示 \(g\) 对其第一个参数 \(z\) 的 Jacobian。因此，我们可以将凸优化求解器视为一种寻根方法，它试图找到 \((z^\star, \lambda^\star, \nu^\star)\) 使得

\[G(z^\star, \lambda^\star, \nu^\star) = 0, \;\; g(z^\star,x) \leq 0, \;\; \lambda^\star \geq 0. \]

通过凸优化问题进行求导的关键方面在于，局部上，我们可以通过仅考虑 KKT 条件的等式部分 \(G(z^\star, \lambda^\star, \nu^\star) = 0\) 来计算变量的导数。这是因为互补条件 \(\lambda^\star \circ g(z^\star,x)\) 意味着对所有 \(i=1,\ldots,m\)，\(\lambda^\star_i = 0\) 或 \(g_i(z^\star,x) = 0\)；虽然可能出现 \(\lambda^\star_i = 0\) 和 \(g_i(z^\star,x) = 0\) 同时成立的情况，但这些点对应于约束"偶然"被拉紧的情况，即即使忽略该约束，优化问题也会满足相同的约束，而求解器通常会找到仅 \(\lambda^\star_i = 0\) 或 \(g_i(z^\star,x) = 0\) 中一个成立的解。假设情况如此（这个论证可以用优化问题的次可微性来形式化，但我们不在这里深入讨论细节），那么对于 \(\lambda^\star\) 或 \(z^\star\) 的任意"小"扰动，保证 \(\lambda_i \cdot g_i(z,x) = 0\) 的唯一方法是维持条件 \(\lambda_i \geq 0\) 和 \(g_i(z,x) \leq 0\)（即，如果 \(\lambda_i\) 被扰动为正值，那么由于 \(g_i(z,x) < 0\)，乘积将不为零）。

上述讨论的最终结论是，尽管包含不等式和等式约束，在解点 \((z^\star, \lambda^\star, \nu^\star)\) 处，我们可以将求解器视为简单地寻找非线性方程组的根

\[G(z^\star, \lambda^\star, \nu^\star) = 0. \]

现在我们可以使用前几节中使用的相同隐式微分方法来推导一个线性方程组，用于求解通过优化问题的导数（我们在这里快速过一下这个方法，因为它与之前章节使用的过程相同）。具体来说，使用简写符号 \((z^\star, \lambda^\star, \nu^\star)(x)\) 表示作为 \(x\) 函数的原始-对偶最优解，我们有

\[\begin{aligned} & \partial_{x} G(z^\star(x), \lambda^\star(x), \nu^\star(x), x) = 0 \\ \Longrightarrow \;\; & \partial_{z, \lambda, \nu} G(z^\star, \lambda^\star, \nu^\star, x) \partial_x (z^\star, \lambda^\star, \nu^\star)(x) + \partial_{x} G(z^\star, \lambda^\star, \nu^\star, x) = 0 \\ \Longrightarrow \;\; & \partial_x (z^\star, \lambda^\star, \nu^\star)(x) = -\left (\partial_{z, \lambda, \nu} G(z^\star, \lambda^\star, \nu^\star, x) \right)^{-1} \partial_{x} G(z^\star, \lambda^\star, \nu^\star, x). \end{aligned} \]

其中符号如 \(\partial_{(z,\lambda,\nu)} G\) 表示整个函数 \(G\) 对变量 \((z,\lambda,\nu)\) 拼接的 Jacobian。这个特定的 Jacobian 也值得关注，因为它是在反向传播中被求逆的项。通过将 \(G\) 函数的每个元素对三个输入中的每一个求导，我们可以用块矩阵形式推导出这一项的解析表达式，即

\[\partial_{z, \lambda, \nu} G(z^\star, \lambda^\star, \nu^\star, x) = \left [ \begin{array}{ccc} \nabla_z^2 f(z^\star,x) + \sum_{i=1}^m \lambda_i^\star \nabla_z^2 g_i(z^\star,x) & \partial_z g(z^\star, x) & \partial_z h(z^\star, x) \\ \partial_z g(z^\star, x)^T \mathrm{diag}(\lambda^\star) & \mathrm{diag}(g(z^\star, x)) & 0 \\ \partial_z h(z^\star, x)^T & 0 & 0 \end{array} \right ] \]

这样，给定优化问题的原始-对偶解（通过任何方法计算），我们就可以微分该解，即计算导数 \(\partial_x (z^\star, \lambda^\star, \nu^\star)(x)\)。注意，如果需要，我们也可以同样容易地微分最优对偶解；通常对偶变量在特定量方面具有直观的解释，我们可能也希望层以可微的方式输出这些对偶解。然而，出于本文的目的，我们将主要关注我们只需要对原始变量的输出求导的情况。

实现可微凸优化

与一开始求解凸优化问题一样，使用凸优化层的另一个好处是，同样有可用的库可以自动将优化问题转换为可微层，适用于 PyTorch 和 Tensorflow 库。即 cvxpylayer 库 [Agarwal el al., 2019] 提供了一个工具，可以轻松地从 cvxpy 中指定的凸优化问题生成层。在实践中，你可能只想直接使用这样的工具。但是，考虑到本教程的目标是描述可微优化的一些内部机制，下面我们将实现一个基本上是 cvxpylayer 库"精简版"的版本，只需大约 60 行 PyTorch 代码。[我们还应该强调，实际的 cvxpylayers 库以相当不同且更通用的方式工作，通过在 cvxpy 生成的"标准"锥形式中微分问题。这使得该库能在更广泛的情况下工作，例如半定规划、不可微的目标和约束，或其他直接用我们上面列出的形式建模优化问题有些繁琐的问题。然而，出于说明目的，使用我们之前在本章中描述的形式来工作是方便的。]

虽然我们在前一节中推导的 Jacobian 表达式可能看起来比前几章中的复杂得多，但在实践中我们可以再次使用自动微分来计算大部分项（这个过程在 PyTorch 等库中可能不是特别高效，因为我们将使用它们来计算整个 Jacobian/Hessian 矩阵，但对于小规模示例来说已经足够了）。再次从高层次来看，与前面章节中镜像的过程类似，该过程分为三个步骤。

1. 给定 \(x\)，建立并求解优化问题（使用任何方法，例如现成求解器），得到最优原始解和对偶解 \(z^\star, \lambda^\star, \nu^\star\)（在任何梯度带之外进行）。在我们的示例中，我们仍然依赖 cvxpy 库来提供解。
2. 通过按如下形式赋值，重新接合梯度带。

\[(z, \lambda, \nu) = (z^\star, \lambda^\star, \nu^\star) - G(z^\star, \lambda^\star, \nu^\star, x) \]

3. 添加一个后向钩子，执行乘法 \(\partial_{z, \lambda, \nu} G(z^\star, \lambda^\star, \nu^\star, x)^{-T} g\) 到反向传播，这里我们通过形成 Jacobian 并求解线性系统来精确求解。

以下是完整的实现。

import torch
import torch.nn as nn
import torch.autograd as autograd
from itertools import accumulateclass OptLayer(nn.Module):def __init__(self, variables, parameters, objective, inequalities, equalities, **cvxpy_opts):super().__init__()self.variables = variablesself.parameters = parametersself.objective = objectiveself.inequalities = inequalitiesself.equalities = equalitiesself.cvxpy_opts = cvxpy_opts# 创建包含目标、不等式、等式的 cvxpy 问题self.cp_inequalities = [ineq(*variables, *parameters) <= 0 for ineq in inequalities]self.cp_equalities = [eq(*variables, *parameters) == 0 for eq in equalities]self.problem = cp.Problem(cp.Minimize(objective(*variables, *parameters)),self.cp_inequalities + self.cp_equalities)def forward(self, *batch_params):out, J = [], []# 通过迭代处理小批量for batch in range(batch_params[0].shape[0]):# 求解优化问题并提取解 + 对偶变量params = [p[batch] for p in batch_params]with torch.no_grad():for i,p in enumerate(self.parameters):p.value = params[i].double().numpy()self.problem.solve(**self.cvxpy_opts)z = [torch.tensor(v.value).type_as(params[0]) for v in self.variables]lam = [torch.tensor(c.dual_value).type_as(params[0]) for c in self.cp_inequalities]nu = [torch.tensor(c.dual_value).type_as(params[0]) for c in self.cp_equalities]# 用于"展平"和"恢复"(z,lam,nu)的便捷函数def vec(z, lam, nu):return torch.cat([a.view(-1) for b in [z,lam,nu] for a in b])def mat(x):sz = [0] + list(accumulate([a.numel() for b in [z,lam,nu] for a in b]))val = [x[a:b] for a,b in zip(sz, sz[1:])]return ([val[i].view_as(z[i]) for i in range(len(z))],[val[i+len(z)].view_as(lam[i]) for i in range(len(lam))],[val[i+len(z)+len(lam)].view_as(nu[i]) for i in range(len(nu))])# 计算 KKT 残差def kkt(z, lam, nu, *params):g = [ineq(*z, *params) for ineq in self.inequalities]dnu = [eq(*z, *params) for eq in self.equalities]L = (self.objective(*z, *params) +sum((u*v).sum() for u,v in zip(lam,g)) + sum((u*v).sum() for u,v in zip(nu,dnu)))dz = autograd.grad(L, z, create_graph=True)dlam = [lam[i]*g[i] for i in range(len(lam))]return dz, dlam, dnu# 计算残差并重新接合自动微分带y = vec(z, lam, nu)y = y - vec(*kkt([z_.clone().detach().requires_grad_() for z_ in z], lam, nu, *params))# 计算 Jacobian 和后向钩子J.append(autograd.functional.jacobian(lambda x: vec(*kkt(*mat(x), *params)), y))y.register_hook(lambda grad,b=batch : torch.solve(grad[:,None], J[b].transpose(0,1))[0][:,0])out.append(mat(y)[0])out = [torch.stack(o, dim=0) for o in zip(*out)]return out[0] if len(out) == 1 else tuple(out)

让我们讨论一下这段代码的工作原理。OptLayer 层接受五个参数：variables，包含一组被优化的 cvxpy 变量，对应于上面的 \(z\) 变量（注意，我们实际上允许在 cvxpy 中定义一组 \(z\) 变量，因为这通常是编写优化问题更方便直观的方式，上面的"完整" \(z\) 变量将是所有这些变量的拼接）；parameters，类似地对应于参数化优化问题的所有输入 \(x\) 的集合；objective，指定目标函数的函数；最后是 inequalities 和 equalities，是指定不等式和等式约束的函数列表（这些列表中的每个函数可以是向量值的，但就像变量一样，我们允许多个函数来指定不等式和等式，因为 cvxpy 形式通常在这种形式下更自然）。注意，以这种方式指定目标和不等式与 cvxpy 略有不同，在 cvxpy 中人们直接在 cp.Problem 中指定所有约束；这是因为这里的函数需要同时服务于双重目的：既以 cvxpy 可以解释的方式指定问题，又作为 PyTorch 可以微分的函数。这意味着在某些情况下，我们需要根据输入是 cp.Variable 还是 torch.tensor 以不同的方式定义函数（我们稍后会看到这样的例子）。

在实现方面，该类使用上述过程操作。我们首先在梯度循环之外使用 cvxpy 建立并求解问题（即使没有 torch.no_grad() 块也没关系，因为 cvxpy 变量不会被微分，但我们仍将其包含在内以供说明）。然后我们创建 kkt 函数，它以列表形式输出 KKT 等式条件，并计算该函数的 Jacobian 并存储用于反向传播（mat 和 vec 例程是将 z、lambda 和 nu 变量列表展平为单个向量以便计算 Jacobian，以及将向量恢复为三个列表的便捷函数）。最后，我们为每个变量注册钩子，以便在反向传播中乘以转置 Jacobian 的逆矩阵。

让我们看看如何使用这个层来创建可微二次规划（quadratic program，QP）。可微 QP 在 OptNet 框架 [Amos and Kolter, 2017] 中已被考虑，下面的代码重新创建了这个层的一个（虽然相当低效的）版本。具体来说，二次规划是以下形式的优化层

\[\begin{aligned} \min_z \;\; & \frac{1}{2}z^T P z + q^T z \\ \text{subject to} \;\; & Gx \leq h \\ & Ax = b \end{aligned} \]

其中优化变量 \(z \in \mathbb{R}^n\)，问题参数 \(P \in \mathbb{R}^{n \times n} \succeq 0\) 和 \(q \in \mathbb{R}^n\) 指定目标，\(G \in \mathbb{R}^{m \times n}\) 和 \(h \in \mathbb{R}^m\) 指定不等式约束，\(A \in \mathbb{R}^{p \times n}\) 和 \(b \in \mathbb{R}^p\) 指定等式约束。这些问题参数作为层的输入（为确保 \(P\) 始终正定，我们实际上根据输入 \(P_{\mathrm{sqrt}}\) 参数化该层，并令 \(P = P_{\mathrm{sqrt}}^T P_{\mathrm{sqrt}}\)）。下面是我们如何使用上述优化问题实现这个可微层。

n,m,p = 10,4,5
z = cp.Variable(n)
Psqrt = cp.Parameter((n,n))
q = cp.Parameter(n)
G = cp.Parameter((m,n))
h = cp.Parameter(m)
A = cp.Parameter((p,n))
b = cp.Parameter()def f_(z,Psqrt,q,G,h,A,b):return 0.5*cp.sum_squares(Psqrt @ z) + q@z if isinstance(z, cp.Variable) else 0.5*torch.sum((Psqrt @ z)**2) + q@z
def g_(z,Psqrt,q,G,h,A,b):return G@z - h
def h_(z,Psqrt,q,G,h,A,b):return A@z - blayer = OptLayer(variables = [z], parameters = [Psqrt, q, G, h, A, b],objective = f_, inequalities = [g_], equalities=[h_],solver=cp.OSQP, eps_abs = 1e-8, eps_rel = 1e-8)

我们可以使用以下代码验证所有参数的梯度是否正确。

torch_params = [torch.randn(2,*p.shape, dtype=torch.double).requires_grad_() for p in layer.parameters]
autograd.gradcheck(lambda *x: layer(*x).sum(), tuple(torch_params), eps=1e-4, atol=1e-3, check_undefined_grad=False)

True

优化层的应用

现在我们已经定义了一个通用的凸优化层，让我们考虑这种方法的一些应用。这些示例大多来自 [Amos, 2019]，在一篇博客文章中还有更多的说明。然而，我们将在此重写这些示例以适应我们的框架。

重新创建常见的神经网络非线性

凸优化层的一个有趣且直接的应用（需要说明的是，这更多是解释性的而非实际有用的）是作为解释许多现有非线性的一种方式。例如，很容易看出 ReLU 非线性对应于一个点到正象限的投影，这是一个凸问题的解。可能不太为人所知的是，许多其他常见的非线性也是凸优化问题的解。

Sigmoid 例如，sigmoid 运算等于最小化输入的线性函数加上二元熵惩罚

\[\sigma(x) = \frac{1}{1+\exp(-x)} = \argmin_{z} \; -z\cdot x + z\log z + (1-z)\log(1-z). \]

我们可以使用上述框架轻松实现这个层。

z = cp.Variable(1)
x = cp.Parameter(1)
def f_(z,x):if isinstance(z, cp.Variable):return -cp.multiply(z,x) - cp.entr(z) - cp.entr(1-z)else:return -z*x + z*torch.log(z) + (1-z)*torch.log(1-z)layer = OptLayer([z], [x], f_, [], [])

import matplotlib.pyplot as plt
%matplotlib inlinex = torch.linspace(-5,5,100).requires_grad_()
y = layer(x[:,None])[:,0]
plt.figure(dpi=150)
plt.plot(x.detach().numpy(), y.detach().numpy())

正如你所期望的，各项的梯度也能正常工作（尽管仅使用上述实现不支持二阶导数）。

y.sum().backward()
plt.figure(dpi=150)
plt.plot(x.detach().numpy(), x.grad.detach().numpy())

Softmax softmax 函数

\[\mathrm{softmax}(x) = \frac{\exp(x)}{\sum_{i=1}^n \exp(x_i)} \]

也由一个优化问题的解给出，即一个类似熵正则化的线性目标

\[\begin{aligned} \mathrm{softmax}(x) = \argmin_{z : \mathbf{1}^T z = 1} \;\; & -z^T x + \sum_{i=1}^n z_i \log z_i. \end{aligned} \]

这同样很容易写成一个优化层。

z = cp.Variable(10)
x = cp.Parameter(10)f_ = lambda z,x: -z@x - cp.sum(cp.entr(z)) if isinstance(z, cp.Variable) else -z@x + z@torch.log(z)
h_ = lambda z,x: cp.sum(z) - 1 if isinstance(z, cp.Variable) else z.sum() - 1layer = OptLayer([z], [x], f_, [], [h_])

让我们测试该层及其梯度与 softmax 解析形式的对比。

import torch.nn.functional as F
x = torch.randn(20,10).requires_grad_()
y = layer(x)
y0 = nn.Softmax(dim=1)(x)
print("Mean absolute error of forward pass:", (y - y0).abs().mean().item())

Mean absolute error of forward pass: 3.955818829126656e-05

r = torch.randn(*y0.shape)
(y*r).sum().backward()
print("Mean absolute error of backward pass:", (x.grad - autograd.grad((y0*r).sum(), x)[0]).abs().mean().item())

Mean absolute error of backward pass: 4.4939119106857106e-05

Sparsemax 虽然为重新创建具有简单解析形式的常见非线性提供了一些有趣的见解，但其实际价值有限，因为解析形式比使用通用凸优化工具快得多。然而，凸优化方法的一个优点是，它可以轻松扩展到更一般的非线性形式，这在实际中可能有若干优势。例如，[Martins and Astudillo, 2016] 提出了 sparsemax 激活函数，它由以下优化问题的解给出

\[\begin{aligned} \mathrm{sparsemax}(x) = \argmin_{z} \;\; & \|z - x\|^2 \\ \text{subject to} \;\; & \mathbf{1}^T z = 1 \\ & z \geq 0. \end{aligned} \]

按照原始论文中的方法实现这个函数，需要为该特定优化问题构建（直接但仍有相当参与量的）自定义求解器，并推导 Jacobian。然而，使用我们的优化层，我们可以用少量代码测试这样的层。

z = cp.Variable(10)
x = cp.Parameter(10)f_ = lambda z,x : cp.sum_squares(z - x) if isinstance(z, cp.Variable) else torch.sum((x-z)**2)
g_ = lambda z,x : -z
h_ = lambda z,x: cp.sum(z) - 1 if isinstance(z, cp.Variable) else z.sum() - 1layer = OptLayer([z], [x], f_, [g_], [h_])

x = torch.randn(3,10).requires_grad_()
y = layer(x)
print(np.round(y.detach().numpy(),4))

[[-0.     -0.     -0.     -0.     -0.     -0.     -0.      1.     -0.-0.    ][ 0.7871  0.      0.      0.      0.2129  0.      0.      0.      0.0.    ][ 0.      0.      0.      0.      0.      0.      0.      0.9695  0.03050.    ]]

当然，如果你打算在实际网络中频繁使用某个层，那么用自定义的快速求解器重写它是合理的（最可能使用原生 GPU 实现），但对于可微层的快速实验，像上面的代码（或更实际地，cvxpylayers 库）可能非常有用。事实上，当确实想编写更专门的、更快的版本时，相同的技术和编程方法也能很好地工作（在适当的特化下）。

学习凸多面体 (convex polytopes)

让我们考虑凸优化问题的最后一个（同时也是具有很好可视化效果的）示例：从数据中学习多面体。具体来说，考虑优化层

\[\begin{aligned} \min_{z} \;\; &\|z-x\|^2 \\ \text{subject to} \;\; & G z \leq h \end{aligned} \]

即该层计算多面体 \(G z \leq h\) 内距离 \(x\) 最近的点。然而，在这里 \(G\) 和 \(h\) 可以被视为层的未知参数：对于给定的数据，我们并不知道一个实际的多面体（事实上，可能不存在"真实的"底层多面体），而是假设我们给定成对的 \((x^{(i)},y^{(i)})\)，其中 \(x^{(i)}\) 是某个原始点，\(y^{(i)}\) 是其到某个未知集合的投影。由于该层是可微的，我们可以使用梯度下降来拟合 \(G\) 和 \(h\) 参数，以在底层数据上通过多面体逼近这个集合。为方便起见，我们可以轻松地将其包装到另一个处理细节的 PyTorch 模块中。

class PolytopeProjection(nn.Module):def __init__(self, n, m):super().__init__()self.G = nn.Parameter(torch.randn(m,n))self.h = nn.Parameter(torch.ones(m))obj = lambda z,x,G,h : cp.sum_squares(z - x) if isinstance(z, cp.Variable) else torch.sum((x-z)**2)ineq = lambda z,x,G,h : G@z - hself.layer = OptLayer([cp.Variable(n)], [cp.Parameter(n), cp.Parameter((m,n)), cp.Parameter(m)],obj, [ineq], [])def forward(self, x):return self.layer(x, self.G.expand(x.shape[0], *self.G.shape), self.h.expand(x.shape[0], *self.h.shape))

现在让我们创建一些数据，其中 \(x^{(i)}\) 是一个二维点，\(y^{(i)}\) 是其到单位 \(\ell_2\) 球的投影。当然，我们无法将单位球精确表示为多面体，但使用梯度下降和我们的层，我们可以尽可能地在给定数据上用该层逼近球面。

p = 100
X = torch.randn(p, 2)
Y = X / X.norm(dim=1).clamp(min=1)[:,None]
plt.figure(dpi=150)
plt.plot(Y[:,0], Y[:,1],'bx')
plt.axis("equal")

torch.manual_seed(0)
layer = PolytopeProjection(2, 20)

让我们先看看初始化时（随机初始化但保证非空）的多面体。

from scipy.spatial import HalfspaceIntersection
from matplotlib.patches import Polygondef plot_polytope(G, h, X):hs = HalfspaceIntersection(np.hstack((G.detach(),  -h[:,None].detach())), np.array([0.,0.]))pts = hs.intersections - hs.interior_pointpts = pts[np.argsort(np.arctan2(pts[:,1],pts[:,0])),:] + hs.interior_pointplt.fill(pts[:,0], pts[:,1], color="orange", alpha=0.5)plt.plot(X[:,0], X[:,1], 'bx')plt.figure(dpi=150)
plot_polytope(layer.G, layer.h, Y)
plt.axis('equal');

现在让我们使用 Adam 优化多面体，以最小化投影到多面体的点与投影到单位球的点之间的距离。

import torch.optim as optim
opt = optim.Adam(layer.parameters(), lr=1e-1)
for i in range(40):if i == 20:opt.param_groups[0]["lr"] = 1e-2loss = nn.MSELoss()(layer(X), Y)opt.zero_grad()loss.backward()opt.step()

plt.figure(dpi=150)
plot_polytope(layer.G, layer.h, Y)
plt.axis('equal');

回顾一下，根据我们的损失函数，我们并不是试图将点"包含"在多面体内（如凸包那样），而是最小化投影到多面体的点与投影到球的点之间的平方误差，这将导致一些点位于表面内部，一些位于外部。我们可以轻松地使用基于梯度的技术来执行这样的操作，并将它们集成到深度架构中，这是一个强大的特性，在广泛的领域中都非常有用。