目录通过优化求导

第 11 章 通过优化求导Chapter 11 Differentiating through optimization

本章研究如何通过优化问题求导,更一般地说,如何通过非线性方程组求导。In this chapter, we study how to differentiate through optimization problems, and more generally through nonlinear systems of equations.

11.1 隐函数11.1 Implicit functions

隐函数无法显式分解为初等函数,因而第 8 章所讨论的自动微分不能直接施用于其上。本章介绍对这类函数求导的技术,以及如何将其纳入自动微分框架。Implicit functions are functions that do not enjoy an explicit decomposition into elementary functions, for which automatic differentiation, as studied in Chapter 8, can therefore not be directly applied. We describe in this chapter techniques to differentiate through such functions and how to integrate them into an autodiff framework.

形式上,我们将隐函数记为 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}),其中 w⋆ ⁣:Λ→W{\bm{w}}^{\star}\colon\Lambda\to\mathcal{W}。由此产生的一个问题是:如何计算雅可比矩阵 ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}})?作为初步应用,可以考虑系统的灵敏度分析。例如,w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 可对应某物理系统的平衡状态,此时 ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}) 便告诉我们该系统对某些参数 λ∈Λ{\bm{\lambda}}\in\Lambda 的灵敏度。Formally, we will denote an implicit function by w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}), where w⋆ ⁣:Λ→W{\bm{w}}^{\star}\colon\Lambda\to\mathcal{W}. One question is then how to compute the Jacobian ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}). As a first application one can consider sensitivity analysis of a system. For example, w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) could correspond to the equilibrium state of a physical system and in this case, ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}) would tell us about the sensitivity of the system to some parameters λ∈Λ{\bm{\lambda}}\in\Lambda.

11.1.1 优化问题Optimization problems

另一例子是隐式定义为某一优化问题的解(假定唯一)的函数Another example is a function implicitly defined as the solution (assumed unique) of an optimization problem

w⋆(λ):=arg max⁡w∈Wf(w,λ),{\bm{w}}^{\star}({\bm{\lambda}})\coloneqq\argmax_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}),

其中 f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R},W\mathcal{W} 表示一个约束集。注意,这里为方便起见使用了 arg max⁡\argmax,但改用 arg min⁡\argmin 时同样成立。where f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} and W\mathcal{W} denotes a constraint set. Note that we use an arg max⁡\argmax for convenience, but the same applies when using an arg min⁡\argmin.

11.1.2 非线性方程Nonlinear equations

更一般地,w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 可定义为某函数 F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W} 的根,即:w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 被隐式定义为满足下述(可能非线性的)方程组的取值More generally, w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) can be defined as the root of some function F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}, i.e., w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is implicitly defined as the value satisfying the (potentially nonlinear) system of equations

F(w,λ)=0F({\bm{w}},{\bm{\lambda}})={\mathbf{0}}

对一切 λ∈Λ{\bm{\lambda}}\in\Lambda。for all λ∈Λ{\bm{\lambda}}\in\Lambda.

11.1.3 在双层优化中的应用Application to bilevel optimization

除灵敏度分析外,双层优化是另一个应用实例。我们常常希望最小化的函数,是由某个固定函数与一个优化问题的解复合而成的。形式上,令 f,g ⁣:W×Λ→Rf,g\colon\mathcal{W}\times\Lambda\to\mathbb{R}。考察如下定义的复合 h(λ)h({\bm{\lambda}}):Besides sensitivity analysis, another example of application is bilevel optimization. Many times, we want to minimize a function defined as the composition of a fixed function and the solution of an optimization problem. Formally, let f,g ⁣:W×Λ→Rf,g\colon\mathcal{W}\times\Lambda\to\mathbb{R}. We consider the composition h(λ)h({\bm{\lambda}}) defined as

h(λ):=g(w⋆(λ),λ),wherew⋆(λ):=arg max⁡w∈Wf(w,λ).h({\bm{\lambda}})\coloneqq g({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}),\quad\text{where}\quad{\bm{w}}^{\star}({\bm{\lambda}})\coloneqq\argmax_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}). (11.1)

例如,这包括超参数优化:其中 ff 为内层的对数似然目标,gg 为外层的验证损失,w∈W{\bm{w}}\in\mathcal{W} 为模型参数,λ∈Λ{\bm{\lambda}}\in\Lambda 则为模型超参数(如正则化强度),如图 11.1 所示。要最小化 h(λ)h({\bm{\lambda}}),通常采用关于 λ{\bm{\lambda}} 的梯度下降方案,这就需要对 ∇h(λ)\nabla h({\bm{\lambda}}) 进行计算。设 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 在 λ{\bm{\lambda}} 处可微,则由链式法则可得雅可比矩阵This includes for instance hyperparameter optimization, where ff is an inner log-likelihood objective, gg is an outer validation loss, w∈W{\bm{w}}\in\mathcal{W} are model parameters and λ∈Λ{\bm{\lambda}}\in\Lambda are model hyperparameters, such as regularization strength, as illustrated in Fig. 11.1. To minimize h(λ)h({\bm{\lambda}}), one generally resorts to a gradient descent scheme w.r.t. λ{\bm{\lambda}}, which requires computing ∇h(λ)\nabla h({\bm{\lambda}}). Assuming that w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is differentiable at λ{\bm{\lambda}}, by the chain rule, we obtain the Jacobian

∂h(λ)=∂1g(w⋆(λ),λ)∂w⋆(λ)+∂2g(w⋆(λ),λ).\partial h({\bm{\lambda}})=\partial_{1}g({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})\partial{\bm{w}}^{\star}({\bm{\lambda}})+\partial_{2}g({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

借助 ∂h(λ)⊤=∇h(λ)\partial h({\bm{\lambda}})^{\top}=\nabla h({\bm{\lambda}})(见注 2.2.9),可得到梯度Using ∂h(λ)⊤=∇h(λ)\partial h({\bm{\lambda}})^{\top}=\nabla h({\bm{\lambda}}) (see Remark 2.2.9), we obtain the gradient

∇h(λ)=∂w⋆(λ)⊤∇1g(w⋆(λ),λ)+∇2g(w⋆(λ),λ).\nabla h({\bm{\lambda}})=\partial{\bm{w}}^{\star}({\bm{\lambda}})^{\top}\nabla_{1}g({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})+\nabla_{2}g({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

唯一棘手的一项是 ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}),因为它涉及对 argmax 求导。事实上,大多数情况下 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 并无显式公式,也无法分解为初等函数。The only problematic term is ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}), as it requires argmax differentiation. Indeed, most of the time, there is no explicit formula for w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) and it does not decompose into elementary functions.

图 11.1:非线性回归中的超参数优化可表述为双层优化问题。每条曲线对应以不同超参数 λ\lambda 拟合若干训练数据(蓝色圆点)所得的估计量。形式上,记 ff 为训练目标,则估计量为 w⋆(λ):=arg min⁡wf(w;λ)w^{\star}(\lambda)\coloneqq\argmin_{w}f(w;\lambda)。目标是找到能最好地拟合若干验证数据(此处为青色菱形)的最优超参数,即最小化 h(λ):=g(w⋆(λ),λ)h(\lambda)\coloneqq g(w^{\star}(\lambda),\lambda),其中 gg 为验证目标。λ1\lambda_{1} 过小会导致对训练目标过拟合,在验证目标上表现不佳;反之,λ2\lambda_{2} 较大则会导致对训练目标和验证目标都欠拟合。最优参数 λ⋆\lambda^{\star} 使验证目标最小,可通过对 λ\lambda 进行梯度下降迭代求得。这需要 h(λ)=g(w⋆(λ),λ)h(\lambda)=g(w^{\star}(\lambda),\lambda) 关于 λ\lambda 的梯度。Figure 11.1: Hyperparameter optimization in nonlinear regression can be cast as a bi-level optimization problem. Each line corresponds to the estimator obtained by fitting some training data (in blue circles) using a different hyperparameter λ\lambda. Formally, denoting ff the training objective, the estimators are w⋆(λ):=arg min⁡wf(w;λ)w^{\star}(\lambda)\coloneqq\argmin_{w}f(w;\lambda). The goal is to find the best hyperparameter that fits some validation data (here in cyan diamonds), that is, minimizing h(λ):=g(w⋆(λ),λ)h(\lambda)\coloneqq g(w^{\star}(\lambda),\lambda), where gg is the validation objective. A too small λ1\lambda_{1} leads to overfitting the training objective and performs badly on validation objective. Conversely, a larger λ2\lambda_{2} underfits both training and validation objectives. The optimal parameter λ⋆\lambda^{\star} minimizes the validation objective and may be obtained by iterating gradient descent w.r.t. λ\lambda. This requires gradients of h(λ)=g(w⋆(λ),λ)h(\lambda)=g(w^{\star}(\lambda),\lambda) w.r.t. λ\lambda.

11.2 包络定理11.2 Envelope theorems

在 g=fg=f 这一特殊情形下,式 11.1 所定义的复合 hh 简化为In the special case g=fg=f, the composition hh defined in Eq. 11.1 is simply given by

h(λ)=f(w⋆(λ),λ)=max⁡w∈Wf(w,λ).h({\bm{\lambda}})=f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})=\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}).

也就是说,此时不再需要对 argmax 求导,而只需对 max 求导;下文将会看到,后者要容易得多。函数 hh 常被称为值函数 [66]。“包络”这一名称的由来如图 11.2 所示。需要强调,包络定理并非只有一条,而是有多条,具体取决于对 ff 所作的假设。That is, we no longer need argmax differentiation, but only max differentiation, which, as we shall now see is much easier. The function hh is often called a value function [66]. The reason for the name “envelope” is illustrated in Fig. 11.2. We emphasize that there is not one, but several envelope theorems, depending on the assumptions on ff.

图 11.2:h(λ)=max⁡w∈Wf(w,λ)h({\bm{\lambda}})=\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}) 的图像是所有 w∈W{\bm{w}}\in\mathcal{W} 对应的函数 λ↦f(w,λ){\bm{\lambda}}\mapsto f({\bm{w}},{\bm{\lambda}}) 的图像的上包络。Figure 11.2: The graph of h(λ)=max⁡w∈Wf(w,λ)h({\bm{\lambda}})=\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}) is the upper-envelope of the graphs of the functions λ↦f(w,λ){\bm{\lambda}}\mapsto f({\bm{w}},{\bm{\lambda}}) for all w∈W{\bm{w}}\in\mathcal{W}.

11.2.1 Bertsekas 定理Bertsekas’ theorem

以下定理取自 [24, 命题 B.25]。The following theorem is due to [24, Proposition B.25].

定理 11.2.1(Bertsekas 定理)。Theorem 11.2.1 (Bertsekas’ theorem).
令Let h(λ):=max⁡w∈Wf(w,λ)h({\bm{\lambda}})\coloneqq\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}) 以及and w⋆(λ):=arg max⁡w∈Wf(w,λ),{\bm{w}}^{\star}({\bm{\lambda}})\coloneqq\argmax_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}), 其中 f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} 关于 w{\bm{w}} 是凹的,且 W\mathcal{W} 是一个紧凸集。若最大值 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 唯一,则函数 hh 可微,其梯度为where f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} is concave in w{\bm{w}} and W\mathcal{W} is a compact convex set. If the maximum w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is unique, then the function hh is differentiable with gradient ∇h(λ)=∇2f(w⋆(λ),λ).\nabla h({\bm{\lambda}})=\nabla_{2}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}). 若进一步 ff 关于 λ{\bm{\lambda}} 是凸的,则 h(λ)h({\bm{\lambda}}) 是凸的。If furthermore ff is convex in λ{\bm{\lambda}}, then h(λ)h({\bm{\lambda}}) is convex.

当 W=RP\mathcal{W}=\mathbb{R}^{P} 时,可以给出该定理的一个变体。A variant of the theorem can be stated when W=RP\mathcal{W}=\mathbb{R}^{P}.

定理 11.2.2(无界情形)。Theorem 11.2.2 (Unbounded case).
假设 W=RP\mathcal{W}=\mathbb{R}^{P}。若 ff 关于 w{\bm{w}} 是强凹的,则函数 hh 可微,其梯度为Suppose W=RP\mathcal{W}=\mathbb{R}^{P}. If ff is strongly concave in w{\bm{w}}, then the function hh is differentiable with gradient ∇h(λ)=∇2f(w⋆(λ),λ).\nabla h({\bm{\lambda}})=\nabla_{2}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}). 若进一步 ff 关于 λ{\bm{\lambda}} 是凸的,则 h(λ)h({\bm{\lambda}}) 也是凸的。If furthermore ff is convex in λ{\bm{\lambda}}, then h(λ)h({\bm{\lambda}}) is convex.

直观地说,定理 11.2.1 与定理 11.2.2 意味着我们可以把 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 当作 λ{\bm{\lambda}} 的一个常数来处理。换言之,我们不需要穿过 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 进行求导,即使它依赖于 λ{\bm{\lambda}}。Informally, Theorem 11.2.1 and Theorem 11.2.2 mean that we can treat w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) as if it were a constant of λ{\bm{\lambda}}. In other words, we do not need to differentiate through w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}), even though it depends on λ{\bm{\lambda}}.

注 11.2.3(历史注记)。Remark 11.2.3 (Historical note).
在机器学习文献中,Bertsekas 定理或其无界变体常被误引作 Danskin 定理,因为 Bertsekas 本人正是这样称呼它的 [24, Proposition B.25]。然而,两者的假设并不相同,因此更应将其视为 Danskin 原始定理的变体或推论。详见第 11.2.3 节。Bertsekas’ theorem or its unbounded variant are often incorrectly cited as Danskin’s theorem in the machine learning literature, because Bertsekas himself called it that way [24, Proposition B.25]. However, because the assumptions differ, they should rather be thought as variants or corollaries of Danskin’s original theorem. See Section 11.2.3 for more details.

现在,我们借助定理 11.2.2 说明如何穿过最小值 h(λ)=min⁡w∈Wf(w,λ)h({\bm{\lambda}})=\min_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}) 进行求导,其中假定 f(w,λ)f({\bm{w}},{\bm{\lambda}}) 关于 w{\bm{w}} 是强凸的。We now illustrate Theorem 11.2.2, to differentiate through a minimum, h(λ)=min⁡w∈Wf(w,λ)h({\bm{\lambda}})=\min_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}), assuming f(w,λ)f({\bm{w}},{\bm{\lambda}}) is strongly convex in w{\bm{w}}.

例 11.2.4(对凸二次函数的最小值求导)。Example 11.2.4 (Differentiating the minimum of a convex quadratic).
令 h(λ):=min⁡w∈Rf(w,λ)h(\lambda)\coloneqq\min_{w\in\mathbb{R}}f(w,\lambda),其中 f(w,λ):=λ2w2+bw+cf(w,\lambda)\coloneqq\frac{\lambda}{2}w^{2}+bw+c,λ>0\lambda>0。ff 关于 ww 的导数为 ∇1f(w,λ)=λw+b\nabla_{1}f(w,\lambda)=\lambda w+b;ff 关于 ww 的二阶导数为 ∇12f(w,λ)=λ\nabla_{1}^{2}f(w,\lambda)=\lambda。因此,ff 关于 ww 是强凸的。设 w⋆(λ)w^{\star}(\lambda) 为 ff 在 ww 上的最小值。令 ∇1f(w,λ)\nabla_{1}f(w,\lambda) 等于零,可得 w⋆(λ)=−bλw^{\star}(\lambda)=-\frac{b}{\lambda}。ff 关于 λ\lambda 的导数为 ∇2f(w,λ)=12w2\nabla_{2}f(w,\lambda)=\frac{1}{2}w^{2}。由定理 11.2.2 可得 h′(λ)=∇2f(w⋆(λ),λ)=12(w⋆(λ))2=12b2λ2h^{\prime}(\lambda)=\nabla_{2}f(w^{\star}(\lambda),\lambda)=\frac{1}{2}(w^{\star}(\lambda))^{2}=\frac{1}{2}\frac{b^{2}}{\lambda^{2}}。下面验证该导数确实无误:将 w⋆(λ)w^{\star}(\lambda) 代回 f(w,λ)f(w,\lambda),即得解析公式 h(λ)=−12b2λ+ch(\lambda)=-\frac{1}{2}\frac{b^{2}}{\lambda}+c。利用 (1λ)′=−1λ2(\frac{1}{\lambda})^{\prime}=-\frac{1}{\lambda^{2}},我们确实得到了 h′(λ)h^{\prime}(\lambda) 的同一表达式。Let us define h(λ):=min⁡w∈Rf(w,λ)h(\lambda)\coloneqq\min_{w\in\mathbb{R}}f(w,\lambda), where f(w,λ):=λ2w2+bw+cf(w,\lambda)\coloneqq\frac{\lambda}{2}w^{2}+bw+c and λ>0\lambda>0. The derivative of ff w.r.t. ww is ∇1f(w,λ)=λw+b\nabla_{1}f(w,\lambda)=\lambda w+b. The second derivative of ff w.r.t. ww is ∇12f(w,λ)=λ\nabla_{1}^{2}f(w,\lambda)=\lambda. Therefore ff is strongly convex in ww. Let w⋆(λ)w^{\star}(\lambda) be the minimum of ff in ww. Setting ∇1f(w,λ)\nabla_{1}f(w,\lambda) to zero, we get w⋆(λ)=−bλw^{\star}(\lambda)=-\frac{b}{\lambda}. The derivative of ff w.r.t. λ\lambda is ∇2f(w,λ)=12w2\nabla_{2}f(w,\lambda)=\frac{1}{2}w^{2}. From Theorem 11.2.2, we have h′(λ)=∇2f(w⋆(λ),λ)=12(w⋆(λ))2=12b2λ2h^{\prime}(\lambda)=\nabla_{2}f(w^{\star}(\lambda),\lambda)=\frac{1}{2}(w^{\star}(\lambda))^{2}=\frac{1}{2}\frac{b^{2}}{\lambda^{2}}. Let us check that the derivative is indeed correct. Plugging w⋆(λ)w^{\star}(\lambda) back into f(w,λ)f(w,\lambda), we get the analytical formula h(λ)=−12b2λ+ch(\lambda)=-\frac{1}{2}\frac{b^{2}}{\lambda}+c. Using (1λ)′=−1λ2(\frac{1}{\lambda})^{\prime}=-\frac{1}{\lambda^{2}}, we indeed obtain the same expression for h′(λ)h^{\prime}(\lambda).

如下所示,对于关于 λ{\bm{\lambda}} 为仿射的函数,定理 11.2.1 有一个简单的解释。Theorem 11.2.1 has a simple interpretation for functions that are affine in λ{\bm{\lambda}} as shown below.

例 11.2.5(凸共轭)。Example 11.2.5 (Convex conjugate).
设 f(w,λ):=⟨w,λ⟩−Ω(w)f({\bm{w}},{\bm{\lambda}})\coloneqq\langle{\bm{w}},{\bm{\lambda}}\rangle-\Omega({\bm{w}}),其中 Ω\Omega 为凸。于是有 h(λ)=max⁡w∈W ⟨w,λ⟩−Ω(w)=:Ω∗(λ)h({\bm{\lambda}})=\max_{{\bm{w}}\in\mathcal{W}} \ \langle{\bm{w}},{\bm{\lambda}}\rangle-\Omega({\bm{w}})\eqqcolon\Omega^{*}({\bm{\lambda}}),其中 Ω∗\Omega^{*} 表示 Ω\Omega 的凸共轭。由于 ff 是凸的,且假设 W\mathcal{W} 为凸紧集,定理 11.2.1 的条件得以满足。由 ∇2f(w,λ)=w\nabla_{2}f({\bm{w}},{\bm{\lambda}})={\bm{w}} 可得 ∇h(λ)=∇Ω∗(λ)=w⋆(λ)\nabla h({\bm{\lambda}})=\nabla\Omega^{*}({\bm{\lambda}})={\bm{w}}^{\star}({\bm{\lambda}})。换言之,在这一特殊情形下,最大值的梯度等于 argmax。其原因在于 f(w,λ)f({\bm{w}},{\bm{\lambda}}) 关于 λ{\bm{\lambda}} 是仿射的。Let f(w,λ):=⟨w,λ⟩−Ω(w)f({\bm{w}},{\bm{\lambda}})\coloneqq\langle{\bm{w}},{\bm{\lambda}}\rangle-\Omega({\bm{w}}) with Ω\Omega convex. We then have h(λ)=max⁡w∈W ⟨w,λ⟩−Ω(w)=:Ω∗(λ)h({\bm{\lambda}})=\max_{{\bm{w}}\in\mathcal{W}} \ \langle{\bm{w}},{\bm{\lambda}}\rangle-\Omega({\bm{w}})\eqqcolon\Omega^{*}({\bm{\lambda}}), where Ω∗\Omega^{*} denotes the convex conjugate of Ω\Omega. Since ff is convex and assuming W\mathcal{W} is convex compact, the conditions of Theorem 11.2.1 are met. Since ∇2f(w,λ)=w\nabla_{2}f({\bm{w}},{\bm{\lambda}})={\bm{w}}, we obtain ∇h(λ)=∇Ω∗(λ)=w⋆(λ)\nabla h({\bm{\lambda}})=\nabla\Omega^{*}({\bm{\lambda}})={\bm{w}}^{\star}({\bm{\lambda}}). In other words, in this special case, the gradient of the max is equal to the argmax. This is due to the fact that f(w,λ)f({\bm{w}},{\bm{\lambda}}) is affine in λ{\bm{\lambda}}.

另一个应用是鞍点优化。Another application is saddle point optimization.

例 11.2.6(鞍点问题)。Example 11.2.6 (Saddle point problem).
考虑鞍点问题 min⁡λ∈Λmax⁡w∈Wf(w,λ)\min_{{\bm{\lambda}}\in\Lambda}\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}})。若关于 λ{\bm{\lambda}} 求最小较为困难,而关于 w{\bm{w}} 求最大较为容易,我们便可把问题改写为 min⁡λ∈Λh(λ)\min_{{\bm{\lambda}}\in\Lambda}h({\bm{\lambda}}),其中 h(λ):=max⁡w∈Wf(w,λ)h({\bm{\lambda}})\coloneqq\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}),并借助 ∇h(λ)\nabla h({\bm{\lambda}}) 关于 λ{\bm{\lambda}} 执行(投影)梯度下降。Consider the saddle point problem min⁡λ∈Λmax⁡w∈Wf(w,λ)\min_{{\bm{\lambda}}\in\Lambda}\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}). If it is difficult to minimize w.r.t. λ{\bm{\lambda}} but easy to maximize w.r.t. w{\bm{w}}, we can rewrite the problem as min⁡λ∈Λh(λ)\min_{{\bm{\lambda}}\in\Lambda}h({\bm{\lambda}}), where h(λ):=max⁡w∈Wf(w,λ)h({\bm{\lambda}})\coloneqq\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}), and use ∇h(λ)\nabla h({\bm{\lambda}}) to perform (projected) gradient descent w.r.t. λ{\bm{\lambda}}.

11.2.2 Rockafellar 定理Rockafellar’s theorem

在与 ff 相关的不同假设下,可以证明一个相关定理 [169, 定理 10.31]。A related theorem [169, Theorem 10.31] can be proved under different assumptions on ff.

定理 11.2.7(Rockafellar 包络定理)。Theorem 11.2.7 (Rockafellar’s envelope theorem).
设 f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R},并设 W\mathcal{W} 为紧凸集。令Let f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} and W\mathcal{W} be a compact convex set. Let h(λ):=max⁡w∈Wf(w,λ)h({\bm{\lambda}})\coloneqq\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}) 以及and w⋆(λ):=arg max⁡w∈Wf(w,λ).{\bm{w}}^{\star}({\bm{\lambda}})\coloneqq\argmax_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}). 若对任意 w∈W{\bm{w}}\in\mathcal{W},ff 都在 λ{\bm{\lambda}} 上连续可微,∇1f\nabla_{1}f 连续,且最大值 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 唯一,则函数 hh 可微,其梯度为If ff is continuously differentiable in λ{\bm{\lambda}} for all w∈W{\bm{w}}\in\mathcal{W}, ∇1f\nabla_{1}f is continuous and the maximum w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is unique, then the function hh is differentiable with gradient ∇h(λ)=∇2f(w⋆(λ),λ).\nabla h({\bm{\lambda}})=\nabla_{2}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

与 Bertsekas 定理相比,Rockafellar 定理不要求 ff 关于 w{\bm{w}} 为凹,但对 ff 的可微性提出了更强的假设。Compared to Bertsekas’ theorem, Rockafellar’s theorem does not require ff to be concave in w{\bm{w}}, but it requires stronger assumptions on the differentiability of ff.

11.2.3 Danskin 定理Danskin’s theorem

最后,我们陈述 Danskin 的原始定理 [51]。Finally, we state Danskin’s original theorem [51].

定理 11.2.8(Danskin 定理)。Theorem 11.2.8 (Danskin’s theorem).
设 W\mathcal{W} 为一个紧拓扑空间。设 f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} 连续,且关于 λ∈Λ{\bm{\lambda}}\in\Lambda 连续可微。将值函数与最大值点集分别定义为Let W\mathcal{W} be a compact topological space. Let f ⁣:W×Λ→Rf\colon\mathcal{W}\times\Lambda\to\mathbb{R} be continuous and continuously differentiable with respect to λ∈Λ{\bm{\lambda}}\in\Lambda. Let the value function and the set of maximizers be defined as h(λ):=max⁡w∈Wf(w,λ)andW⋆(λ):=arg max⁡w∈Wf(w,λ).h({\bm{\lambda}})\coloneqq\max_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}})\quad\text{and}\quad\mathcal{W}^{\star}({\bm{\lambda}})\coloneqq\argmax_{{\bm{w}}\in\mathcal{W}}f({\bm{w}},{\bm{\lambda}}). 则 hh 在 λ∈Λ{\bm{\lambda}}\in\Lambda 处沿方向 d∈Λ{\bm{d}}\in\Lambda 的方向导数(记作 h′(λ;d)h^{\prime}({\bm{\lambda}};{\bm{d}}))存在,且由下式给出:Then, the directional derivative of hh at λ∈Λ{\bm{\lambda}}\in\Lambda in direction d∈Λ{\bm{d}}\in\Lambda, denoted h′(λ;d)h^{\prime}({\bm{\lambda}};{\bm{d}}), exists and is given by h′(λ;d)=max⁡w∈W⋆(λ)⟨∇2f(w,λ),d⟩.h^{\prime}({\bm{\lambda}};{\bm{d}})=\max_{{\bm{w}}\in\mathcal{W}^{\star}({\bm{\lambda}})}\langle\nabla_{2}f({\bm{w}},{\bm{\lambda}}),{\bm{d}}\rangle.

我们在表 11.1 中总结各包络定理及其假设。We sum up envelope theorems and their assumptions in Table 11.1.

表 11.1:不同包络定理的比较。Table 11.1: Comparison of different envelope theorems.
定理Theorem w⋆{\bm{w}}^{\star} 是否唯一?Unique w⋆{\bm{w}}^{\star}? 关于 w{\bm{w}} 是否为凹?Concave in w{\bm{w}}? 关于 w{\bm{w}} 是否光滑?Smooth in w{\bm{w}}? 结果Result
DanskinDanskin 否No 否No 否No 方向导数Dir. Deriv.
BertsekasBertsekas 是Yes 是Yes 否No 梯度Gradient
RockafellarRockafellar 是Yes 否No 是Yes 梯度Gradient

11.3 隐函数定理11.3 Implicit function theorem

11.3.1 一元函数Univariate functions

隐函数定理(IFT)给出了如下条件:在该条件下,形如 F(w,λ)=0F(w,\lambda)=0 的隐式关系可以在局部改写为函数 w=w⋆(λ)w=w^{\star}(\lambda),并给出一种计算其关于 λ\lambda 的导数的方法。The implicit function theorem (IFT) provides conditions under which an implicit relationship of the form F(w,λ)=0F(w,\lambda)=0 can be rewritten as a function w=w⋆(λ)w=w^{\star}(\lambda) locally, and provides a way to compute its derivative w.r.t. λ\lambda.

定理 11.3.1(隐函数定理,一元情形)。Theorem 11.3.1 (Implicit function theorem, univariate case).
设 F ⁣:R×R→RF\colon\mathbb{R}\times\mathbb{R}\to\mathbb{R}。假设 F(w,λ)F(w,\lambda) 是 (w0,λ0)(w_{0},\lambda_{0}) 的某邻域 U\mathcal{U} 内的连续可微函数,且满足 F(w0,λ0)=0F(w_{0},\lambda_{0})=0 与 ∂1F(w0,λ0)≠0\partial_{1}F(w_{0},\lambda_{0})\neq 0。则存在 (w0,λ0)(w_{0},\lambda_{0}) 的一个邻域 V⊆U\mathcal{V}\subseteq\mathcal{U},其中存在函数 w⋆(λ)w^{\star}(\lambda),使得Let
F ⁣:R×R→RF\colon\mathbb{R}\times\mathbb{R}\to\mathbb{R}. Assume F(w,λ)F(w,\lambda) is a continuously differentiable function in a neighborhood U\mathcal{U} of (w0,λ0)(w_{0},\lambda_{0}) such that F(w0,λ0)=0F(w_{0},\lambda_{0})=0 and ∂1F(w0,λ0)≠0\partial_{1}F(w_{0},\lambda_{0})\neq 0. Then there exists a neighborhood V⊆U\mathcal{V}\subseteq\mathcal{U} of (w0,λ0)(w_{0},\lambda_{0}) in which there is a function w⋆(λ)w^{\star}(\lambda) such that
• w⋆(λ0)=w0w^{\star}(\lambda_{0})=w_{0}, • 对邻域 V\mathcal{V} 内的任意 λ\lambda,均有 F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0,F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0 for all λ\lambda in the neighborhood V\mathcal{V}, • ∂w⋆(λ)=−∂2F(w⋆(λ),λ)∂1F(w⋆(λ),λ)\partial w^{\star}(\lambda)=-\frac{\partial_{2}F(w^{\star}(\lambda),\lambda)}{\partial_{1}F(w^{\star}(\lambda),\lambda)}.

我们将证明推迟到多元情形,先从一个经典的应用该定理的例子开始。We postpone the proof to the multivariate case and begin with a classical example of application of the theorem.

图 11.3:圆的方程 F(x,y):=x2+y2−1=0F(x,y)\coloneqq x^{2}+y^{2}-1=0 并不是从 y∈[−1,1]y\in[-1,1] 到 x∈[−1,1]x\in[-1,1] 的函数,因为总有两个可能的 xx 取值,即 x=1−y2x=\sqrt{1-y^{2}} 或 x=−1−y2x=-\sqrt{1-y^{2}}。然而,在某点 (x0,y0)(x_{0},y_{0}) 的局部邻域内,例如满足 x0>0x_{0}>0 和 y0>0y_{0}>0(右上象限)时,函数 x=x⋆(y)=1−y2x=x^{\star}(y)=\sqrt{1-y^{2}} 是良定义的。隐函数定理给出了这样的函数在局部存在的条件,并提供了一种计算其导数的方法。Figure 11.3: The circle equation F(x,y):=x2+y2−1=0F(x,y)\coloneqq x^{2}+y^{2}-1=0 is not a function from y∈[−1,1]y\in[-1,1] to x∈[−1,1]x\in[-1,1], as there are always two possible xx values, x=1−y2x=\sqrt{1-y^{2}} or x=−1−y2x=-\sqrt{1-y^{2}}. However, locally around some point (x0,y0)(x_{0},y_{0}), e.g., such that x0>0x_{0}>0 and y0>0y_{0}>0 (upper-right quadrant), the function x=x⋆(y)=1−y2x=x^{\star}(y)=\sqrt{1-y^{2}} is well-defined. The implicit function theorem gives conditions for such a function to exist locally and provides a way to compute its derivative.
例 11.3.2(单位圆的方程)。Example 11.3.2 (Equation of the unit circle).
为清晰起见,我们使用 w≡xw\equiv x 与 λ≡y\lambda\equiv y。设 F(x,y):=x2+y2−1F(x,y)\coloneqq x^{2}+y^{2}-1。一般而言,我们无法把单位圆方程 F(x,y)=0F(x,y)=0 改写为从 yy 到 xx 的函数,因为对每个 y∈[−1,1]y\in[-1,1],总有两个可能的 xx 值,即 x=1−y2x=\sqrt{1-y^{2}} 或 x=−1−y2x=-\sqrt{1-y^{2}}。但在某一点 (x0,y0)(x_{0},y_{0}) 的局部附近,例如在满足 x0>0x_{0}>0 与 y0>0y_{0}>0 时(右上象限),函数 x=x⋆(y)=1−y2x=x^{\star}(y)=\sqrt{1-y^{2}} 有良定义。借助 ∂1F(x,y)=2x\partial_{1}F(x,y)=2x 与 ∂2F(x,y)=2y\partial_{2}F(x,y)=2y,我们在该邻域内(此例中即右上象限)得到 ∂x⋆(y)=−∂2F(x⋆(y),y)∂1F(x⋆(y),y)=−2y2x⋆(y)=−y1−y2\partial x^{\star}(y)=-\frac{\partial_{2}F(x^{\star}(y),y)}{\partial_{1}F(x^{\star}(y),y)}=-\frac{2y}{2x^{\star}(y)}=-\frac{y}{\sqrt{1-y^{2}}}。这与对 1−y2\sqrt{1-y^{2}} 施用链式法则所得的导数表达式完全相同,且在 y∈[0,1)y\in[0,1) 上有良定义。We use w≡xw\equiv x and λ≡y\lambda\equiv y for clarity. Let F(x,y):=x2+y2−1F(x,y)\coloneqq x^{2}+y^{2}-1. In general, we cannot rewrite the unit circle equation F(x,y)=0F(x,y)=0 as a function from yy to xx, because for every y∈[−1,1]y\in[-1,1], there are always two possible xx values, namely, x=1−y2x=\sqrt{1-y^{2}} or x=−1−y2x=-\sqrt{1-y^{2}}. However, locally around some point (x0,y0)(x_{0},y_{0}), e.g., such that x0>0x_{0}>0 and y0>0y_{0}>0 (upper-right quadrant), the function x=x⋆(y)=1−y2x=x^{\star}(y)=\sqrt{1-y^{2}} is well-defined. Using ∂1F(x,y)=2x\partial_{1}F(x,y)=2x and ∂2F(x,y)=2y\partial_{2}F(x,y)=2y, we get ∂x⋆(y)=−∂2F(x⋆(y),y)∂1F(x⋆(y),y)=−2y2x⋆(y)=−y1−y2\partial x^{\star}(y)=-\frac{\partial_{2}F(x^{\star}(y),y)}{\partial_{1}F(x^{\star}(y),y)}=-\frac{2y}{2x^{\star}(y)}=-\frac{y}{\sqrt{1-y^{2}}} in that neighborhood (the upper right quadrant in this case). This is indeed the same derivative expression as if we used the chain rule on 1−y2\sqrt{1-y^{2}} and is well-defined on y∈[0,1)y\in[0,1).

在上述简单例子中,我们可以容易地在给定邻域内推导出联系 yy 与 xx 的显式函数,但这并非总是可行。隐函数定理给出了保证这样的函数存在的条件以及对其求导的方法,但并不给出构造该函数的方法。事实上,寻找使 F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0 成立的 w⋆(λ)w^{\star}(\lambda),通常涉及求根算法、优化算法、非线性系统求解器等。In the above simple example, we can easily derive an explicit function relating yy to xx in a given neighborhood, but this is not always the case. The IFT gives us conditions guaranteeing that such function exists and a way to differentiate it, but not a way to construct such a function. In fact, finding w⋆(λ)w^{\star}(\lambda) such that F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0 typically involves a root finding algorithm, an optimization algorithm, a nonlinear system solver, etc.

例 11.3.3(多项式)。Example 11.3.3 (Polynomial).
设 F(w,λ)=w5+w3+w−λF(w,\lambda)=w^{5}+w^{3}+w-\lambda。根据 Abel-Ruffini 定理 [196],五次多项式(次数为 55 的多项式)的根不能用根式表示,只能借助数值求根。此外,奇次多项式必有实根;并且 ∂1F(w,λ)=5w4+3w2+1\partial_{1}F(w,\lambda)=5w^{4}+3w^{2}+1 严格为正。因此由介值定理,满足 F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0 的根 w⋆(λ)w^{\star}(\lambda) 必定唯一。这个唯一根例如可用二分法求得。利用 IFT 可得其导数为 ∂w⋆(λ)=(5w⋆(λ)4+3w⋆(λ)2+1)−1\partial w^{\star}(\lambda)=(5w^{\star}(\lambda)^{4}+3w^{\star}(\lambda)^{2}+1)^{-1}。Let F(w,λ)=w5+w3+w−λF(w,\lambda)=w^{5}+w^{3}+w-\lambda. According to the Abel-Ruffini theorem [196], quintics (polynomials of degree 55) do not enjoy roots in terms of radicals and one must resort to numerical root finding. In addition, odd-degree polynomials have real roots. Moreover, ∂1F(w,λ)=5w4+3w2+1\partial_{1}F(w,\lambda)=5w^{4}+3w^{2}+1 is strictly positive. Therefore, by the intermediate value theorem, there must be only one root w⋆(λ)w^{\star}(\lambda) such that F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0. This unique root can for example be found by bisection. Using the IFT, its derivative is found to be ∂w⋆(λ)=(5w⋆(λ)4+3w⋆(λ)2+1)−1\partial w^{\star}(\lambda)=(5w^{\star}(\lambda)^{4}+3w^{\star}(\lambda)^{2}+1)^{-1}.

若 IFT 的假设在某点的邻域内成立,则隐函数在该点可微;但反之不成立:IFT 假设不成立并不必然意味着隐函数不可微,下面我们对此加以说明。While an implicit function is differentiable at a point if the assumptions of the IFT hold in a neighborhood of that point, the reciprocal is not true: failure of the IFT assumptions does not necessarily mean that the implicit function is not differentiable, as we now illustrate.

例 11.3.4(隐函数定理的条件并非可微性的必要条件)。Example 11.3.4 (IFT conditions are not necessary for differentiability).
考虑 F(w,λ)=(w−λ)2F(w,\lambda)=(w-\lambda)^{2}。若将 w⋆(λ)=λw^{\star}(\lambda)=\lambda 定义为恒等函数,则显然有 F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0。它对所有 λ\lambda 显然可微,但 IFT 的假设却不成立,因为我们有 ∂1F(w,λ)=2(w−λ)\partial_{1}F(w,\lambda)=2(w-\lambda),从而 ∂1F(0,0)=0\partial_{1}F(0,0)=0。Consider F(w,λ)=(w−λ)2F(w,\lambda)=(w-\lambda)^{2}. We clearly have that F(w⋆(λ),λ)=0F(w^{\star}(\lambda),\lambda)=0 if we define w⋆(λ)=λw^{\star}(\lambda)=\lambda, the identity function. It is clearly differentiable for all λ\lambda, yet the assumptions of the IFT fail, since we have ∂1F(w,λ)=2(w−λ)\partial_{1}F(w,\lambda)=2(w-\lambda) and therefore ∂1F(0,0)=0\partial_{1}F(0,0)=0.

11.3.2 多元函数Multivariate functions

现在我们给出一般多元情形下的隐函数定理。非正式地说,若 F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}},则由链式法则,我们有We now present the IFT in the general multivariate setting. Informally, if F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}}, then by the chain rule, we have

∂1F(w⋆(λ),λ)∂w⋆(λ)+∂2F(w⋆(λ),λ)=0,\partial_{1}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})\partial{\bm{w}}^{\star}({\bm{\lambda}})+\partial_{2}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}},

这意味着雅可比矩阵 ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}})(假设存在)满足meaning that the Jacobian ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}), assuming that it exists, satisfies

−∂1F(w⋆(λ),λ)∂w⋆(λ)=∂2F(w⋆(λ),λ).-\partial_{1}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})\partial{\bm{w}}^{\star}({\bm{\lambda}})=\partial_{2}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

隐函数定理给出了 ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}) 存在的条件。The IFT gives us conditions for the existence of ∂w⋆(λ)\partial{\bm{w}}^{\star}({\bm{\lambda}}).

定理 11.3.5(隐函数定理,多元情形)。Theorem 11.3.5 (Implicit function theorem, multivariate case).
我们定义 F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}。设 F(w,λ)F({\bm{w}},{\bm{\lambda}}) 是 (w0,λ0)({\bm{w}}_{0},{\bm{\lambda}}_{0}) 的某邻域内的连续可微函数,满足 F(w0,λ0)=0F({\bm{w}}_{0},{\bm{\lambda}}_{0})={\mathbf{0}},且 ∂1F(w0,λ0)\partial_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}) 可逆,即其行列式非零。则存在 λ0{\bm{\lambda}}_{0} 的一个邻域,其中存在函数 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}),使得Let us define F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}. Assume F(w,λ)F({\bm{w}},{\bm{\lambda}}) is a continuously differentiable function in a neighborhood of (w0,λ0)({\bm{w}}_{0},{\bm{\lambda}}_{0}) such that F(w0,λ0)=0F({\bm{w}}_{0},{\bm{\lambda}}_{0})={\mathbf{0}} and ∂1F(w0,λ0)\partial_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}) is invertible, i.e., its determinant is nonzero. Then there exists a neighborhood of λ0{\bm{\lambda}}_{0} in which there is a function w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) such that • w⋆(λ0)=w0{\bm{w}}^{\star}({\bm{\lambda}}_{0})={\bm{w}}_{0}, • 对邻域内的所有 λ{\bm{\lambda}},均有 F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}},F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}} for all λ{\bm{\lambda}} in the neighborhood, • −∂1F(w⋆(λ),λ)∂w⋆(λ)=∂2F(w⋆(λ),λ)-\partial_{1}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})\partial{\bm{w}}^{\star}({\bm{\lambda}})=\partial_{2}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})   ⟺  ∂w⋆(λ)=−∂1F(w⋆(λ),λ)−1∂2F(w⋆(λ),λ).\iff\partial{\bm{w}}^{\star}({\bm{\lambda}})=-\partial_{1}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})^{-1}\partial_{2}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

我们从一个简单的无约束优化算法开始。We begin with a simple unconstrained optimization algorithm.

例 11.3.6(无约束优化)。Example 11.3.6 (Unconstrained optimization).
假设我们希望穿过 w⋆(λ)=arg min⁡w∈RPf(w,λ){\bm{w}}^{\star}({\bm{\lambda}})=\argmin_{{\bm{w}}\in\mathbb{R}^{P}}f({\bm{w}},{\bm{\lambda}}) 求导,其中 ff 关于 w{\bm{w}} 严格凸,从而保证解唯一。由驻值条件,若定义 F(w,λ):=∇1f(w,λ)F({\bm{w}},{\bm{\lambda}})\coloneqq\nabla_{1}f({\bm{w}},{\bm{\lambda}}),则 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 可唯一地刻画为 FF 在第一个自变量上的根,即 F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}}。式中 ∂1F(w,λ)=∇12f(w,λ)\partial_{1}F({\bm{w}},{\bm{\lambda}})=\nabla^{2}_{1}f({\bm{w}},{\bm{\lambda}}) 是 ff 关于 w{\bm{w}} 的 Hessian 矩阵,∂2F(w,λ)=∂2∇1f(w,λ)\partial_{2}F({\bm{w}},{\bm{\lambda}})=\partial_{2}\nabla_{1}f({\bm{w}},{\bm{\lambda}}) 则是 ff 关于 w{\bm{w}} 与 λ{\bm{\lambda}} 的交叉导数。因此,只要 Hessian 矩阵在 (w⋆(λ),λ)({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}) 处有良定义且可逆,我们便可利用 IFT 穿过 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 求导,得到 ∂w⋆(λ)=−(∇12f(w⋆(λ),λ))−1∂2∇1f(w⋆(λ),λ)\partial{\bm{w}}^{\star}({\bm{\lambda}})=-(\nabla^{2}_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}))^{-1}\partial_{2}\nabla_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})。Assume we want to differentiate through w⋆(λ)=arg min⁡w∈RPf(w,λ){\bm{w}}^{\star}({\bm{\lambda}})=\argmin_{{\bm{w}}\in\mathbb{R}^{P}}f({\bm{w}},{\bm{\lambda}}), where ff is strictly convex in w{\bm{w}}, which ensures that the solution is unique. From the stationary conditions, if we define F(w,λ):=∇1f(w,λ)F({\bm{w}},{\bm{\lambda}})\coloneqq\nabla_{1}f({\bm{w}},{\bm{\lambda}}), then w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is uniquely characterized as the root of FF in the first argument, i.e., F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}}. We have ∂1F(w,λ)=∇12f(w,λ)\partial_{1}F({\bm{w}},{\bm{\lambda}})=\nabla^{2}_{1}f({\bm{w}},{\bm{\lambda}}), the Hessian of ff in w{\bm{w}}, and ∂2F(w,λ)=∂2∇1f(w,λ)\partial_{2}F({\bm{w}},{\bm{\lambda}})=\partial_{2}\nabla_{1}f({\bm{w}},{\bm{\lambda}}), the cross derivatives of ff in w{\bm{w}} and λ{\bm{\lambda}}. Therefore, assuming that the Hessian is well-defined and invertible at (w⋆(λ),λ)({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}), we can use the IFT to differentiate through w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) and obtain ∂w⋆(λ)=−(∇12f(w⋆(λ),λ))−1∂2∇1f(w⋆(λ),λ)\partial{\bm{w}}^{\star}({\bm{\lambda}})=-(\nabla^{2}_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}))^{-1}\partial_{2}\nabla_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}).

接下来,我们通过允许优化问题带有约束来推广前一例子。Next, we generalize the previous example, by allowing constraints in the optimization problem.

例 11.3.7(约束优化)。Example 11.3.7 (Constrained optimization).
现在假设我们希望穿过 w⋆(λ)=arg min⁡w∈Cf(w,λ){\bm{w}}^{\star}({\bm{\lambda}})=\argmin_{{\bm{w}}\in\mathcal{C}}f({\bm{w}},{\bm{\lambda}}) 求导,其中 ff 关于 w{\bm{w}} 严格凸,且 C⊆W\mathcal{C}\subseteq\mathcal{W} 为凸集。此时解由不动点方程 w⋆(λ)=PC(w⋆(λ)−η∇1f(w⋆(λ),λ)){\bm{w}}^{\star}({\bm{\lambda}})=P_{\mathcal{C}}({\bm{w}}^{\star}({\bm{\lambda}})-\eta\nabla_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})) 刻画:对任意 η>0\eta>0,P(y):=arg min⁡x∈C∥x−y∥22P({\bm{y}})\coloneqq\argmin_{{\bm{x}}\in\mathcal{C}}\|{\bm{x}}-{\bm{y}}\|_{2}^{2} 是 y{\bm{y}} 到 C\mathcal{C} 的欧几里得投影。因此,w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 是 F(w,λ)=w−PC(w−η∇1f(w,λ))F({\bm{w}},{\bm{\lambda}})={\bm{w}}-P_{\mathcal{C}}({\bm{w}}-\eta\nabla_{1}f({\bm{w}},{\bm{\lambda}})) 的根(见第 16 章)。在定理条件成立的前提下,我们可以利用 IFT 穿过 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 求导。注意,∂1F(w,λ)\partial_{1}F({\bm{w}},{\bm{\lambda}}) 需要给出雅可比矩阵 ∂PC(y)\partial P_{\mathcal{C}}({\bm{y}}) 的表达式。所幸对许多集合 C\mathcal{C} 而言,PC(y)P_{\mathcal{C}}({\bm{y}}) 及其雅可比矩阵都易于计算 [27]。Now, assume we want to differentiate through w⋆(λ)=arg min⁡w∈Cf(w,λ){\bm{w}}^{\star}({\bm{\lambda}})=\argmin_{{\bm{w}}\in\mathcal{C}}f({\bm{w}},{\bm{\lambda}}), where ff is strictly convex in w{\bm{w}} and C⊆W\mathcal{C}\subseteq\mathcal{W} is a convex set. A solution is characterized by the fixed point equation w⋆(λ)=PC(w⋆(λ)−η∇1f(w⋆(λ),λ)){\bm{w}}^{\star}({\bm{\lambda}})=P_{\mathcal{C}}({\bm{w}}^{\star}({\bm{\lambda}})-\eta\nabla_{1}f({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})), for any η>0\eta>0, where P(y):=arg min⁡x∈C∥x−y∥22P({\bm{y}})\coloneqq\argmin_{{\bm{x}}\in\mathcal{C}}\|{\bm{x}}-{\bm{y}}\|_{2}^{2} is the Euclidean projection of y{\bm{y}} onto C\mathcal{C}. Therefore, w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) is the root of F(w,λ)=w−PC(w−η∇1f(w,λ))F({\bm{w}},{\bm{\lambda}})={\bm{w}}-P_{\mathcal{C}}({\bm{w}}-\eta\nabla_{1}f({\bm{w}},{\bm{\lambda}})) (see Chapter 16). We can differentiate through w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) using the IFT, assuming that the conditions of the theorem apply. Note that ∂1F(w,λ)\partial_{1}F({\bm{w}},{\bm{\lambda}}) requires the expression of the Jacobian ∂PC(y)\partial P_{\mathcal{C}}({\bm{y}}). Fortunately, PC(y)P_{\mathcal{C}}({\bm{y}}) and its Jacobian are easy to compute for many sets C\mathcal{C} [27].

11.3.3 隐函数的 JVP 与 VJPJVP and VJP of implicit functions

要将隐函数 w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) 集成到自动微分框架中,我们需要能够计算其 JVP 或 VJP。这正是下一个命题的目的。To integrate an implicit function w⋆(λ){\bm{w}}^{\star}({\bm{\lambda}}) in an autodiff framework, we need to be able to compute its JVP or VJP. This is the purpose of the next proposition.

命题 11.3.8(隐函数的 JVP 与 VJP)。Proposition 11.3.8 (JVP and VJP of implicit functions).
设 w⋆ ⁣:Λ→W{\bm{w}}^{\star}\colon\Lambda\to\mathcal{W} 是由 F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}} 的解隐式定义的函数,其中 F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W} 为某个函数。定义Let w⋆ ⁣:Λ→W{\bm{w}}^{\star}\colon\Lambda\to\mathcal{W} be a function implicitly defined as the solution of F(w⋆(λ),λ)=0F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}})={\mathbf{0}}, for some function F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}. Define A\displaystyle A :=−∂1F(w⋆(λ),λ)\displaystyle\coloneqq-\partial_{1}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}) B\displaystyle B :=∂2F(w⋆(λ),λ).\displaystyle\coloneqq\partial_{2}F({\bm{w}}^{\star}({\bm{\lambda}}),{\bm{\lambda}}). 假设 IFT 的假设成立。则输入方向 v∈Λ{\bm{v}}\in\Lambda 上的 JVP t:=∂w⋆(λ)v{\bm{t}}\coloneqq\partial{\bm{w}}^{\star}({\bm{\lambda}}){\bm{v}} 可通过求解如下线性方程组得到Assume the assumptions of the IFT hold. The JVP t:=∂w⋆(λ)v{\bm{t}}\coloneqq\partial{\bm{w}}^{\star}({\bm{\lambda}}){\bm{v}} in the input direction v∈Λ{\bm{v}}\in\Lambda is obtained by solving the linear system At=Bv.A{\bm{t}}=B{\bm{v}}. 而输出方向 u∈W{\bm{u}}\in\mathcal{W} 上的 VJP ∂w⋆(λ)∗u\partial{\bm{w}}^{\star}({\bm{\lambda}})^{*}{\bm{u}} 则可通过求解如下线性方程组得到The VJP ∂w⋆(λ)∗u\partial{\bm{w}}^{\star}({\bm{\lambda}})^{*}{\bm{u}} in the output direction u∈W{\bm{u}}\in\mathcal{W} is obtained by solving the linear system A∗r=u.A^{*}{\bm{r}}={\bm{u}}. 利用解 r{\bm{r}},我们得到Using the solution r{\bm{r}}, we get ∂w⋆(λ)∗u=∂w⋆(λ)∗A∗r=B∗r.\partial{\bm{w}}^{\star}({\bm{\lambda}})^{*}{\bm{u}}=\partial{\bm{w}}^{\star}({\bm{\lambda}})^{*}A^{*}{\bm{r}}=B^{*}{\bm{r}}.

注意,在上述线性系统中,我们可以把 AA 和 BB 作为线性映射来访问,它们正是 FF 的 JVP。它们的伴随 A∗A^{*} 和 B∗B^{*} 对应于 FF 的 VJP。因此,为求解这些系统,我们可以使用 9.4 节中详述的无矩阵求解器。例如,当 AA 为对称半正定时,我们可以使用共轭梯度法 [93];当 AA 不是对称正定时,则可以使用 GMRES [172] 或 BiCGSTAB [205]。Note that in the above linear systems, we can access AA and BB as linear maps, the JVPs of FF. Their adjoints, A∗A^{*} and B∗B^{*}, correspond to the VJPs of FF. To solve these systems, we can therefore use matrix-free solvers as detailed in Section 9.4. For example, when AA is symmetric positive semi-definite, we can use the conjugate gradient method [93]. When AA is not symmetric positive definite, we can use GMRES [172] or BiCGSTAB [205].

11.3.4 隐函数定理的证明Proof of the implicit function theorem

我们借助定理 11.5.1 中给出的逆函数定理来证明该定理。定义We prove the theorem using the inverse function theorem presented in Theorem 11.5.1. Define

f(λ,w)=(λ,F(w,λ))f({\bm{\lambda}},{\bm{w}})=({\bm{\lambda}},F({\bm{w}},{\bm{\lambda}}))

其从 RQ×RP\mathbb{R}^{Q}\times\mathbb{R}^{P} 映到 RQ×RP\mathbb{R}^{Q}\times\mathbb{R}^{P}。ff 的雅可比矩阵为which goes from RQ×RP\mathbb{R}^{Q}\times\mathbb{R}^{P} onto RQ×RP\mathbb{R}^{Q}\times\mathbb{R}^{P}. The Jacobian of ff is

∂f(λ,w)=(I⁡0∂2F(w,λ)∂1F(w,λ)).\bm{\partial}f({\bm{\lambda}},{\bm{w}})=\begin{pmatrix}\operatorname{I}&0\\ \bm{\partial}_{2}F({\bm{w}},{\bm{\lambda}})&\bm{\partial}_{1}F({\bm{w}},{\bm{\lambda}})\end{pmatrix}.

因此在 w0,λ0{\bm{w}}_{0},{\bm{\lambda}}_{0} 处有 det⁡(∂f(λ0,w0))=det⁡(I⁡)det⁡(∂1F(w0,λ0))≠0\det(\bm{\partial}f({\bm{\lambda}}_{0},{\bm{w}}_{0}))=\det(\operatorname{I})\det(\bm{\partial}_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}))\neq 0,因为我们假定了 ∂1F(w0,λ0)\partial_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}) 可逆。由逆函数定理,函数 ff 在 f(λ0,w0)=(λ0,0)f({\bm{\lambda}}_{0},{\bm{w}}_{0})=({\bm{\lambda}}_{0},{\mathbf{0}}) 的邻域 NN 内可逆。特别地,它在 N∩{(λ,0),λ∈RQ}N\cap\{({\bm{\lambda}},{\mathbf{0}}),{\bm{\lambda}}\in\mathbb{R}^{Q}\} 中可逆。于是隐式方程在 λ0{\bm{\lambda}}_{0} 邻域内的解为 (λ,w∗(λ))=f−1(λ,0)({\bm{\lambda}},{\bm{w}}^{*}({\bm{\lambda}}))=f^{-1}({\bm{\lambda}},{\mathbf{0}})。由逆函数定理,f−1f^{-1} 是连续可微的逆,因此 w∗(λ){\bm{w}}^{*}({\bm{\lambda}}) 亦然。导数 ∂w∗(λ)\partial{\bm{w}}^{*}({\bm{\lambda}}) 由逆的微分得到,即So at w0,λ0{\bm{w}}_{0},{\bm{\lambda}}_{0}, we have det⁡(∂f(λ0,w0))=det⁡(I⁡)det⁡(∂1F(w0,λ0))≠0\det(\bm{\partial}f({\bm{\lambda}}_{0},{\bm{w}}_{0}))=\det(\operatorname{I})\det(\bm{\partial}_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}))\neq 0 since we assumed ∂1F(w0,λ0)\partial_{1}F({\bm{w}}_{0},{\bm{\lambda}}_{0}) invertible. By the inverse function theorem, the function ff is then invertible in a neighborhood NN of f(λ0,w0)=(λ0,0)f({\bm{\lambda}}_{0},{\bm{w}}_{0})=({\bm{\lambda}}_{0},{\mathbf{0}}). In particular, it is invertible in N∩{(λ,0),λ∈RQ}N\cap\{({\bm{\lambda}},{\mathbf{0}}),{\bm{\lambda}}\in\mathbb{R}^{Q}\}. The solution of the implicit equation in a neighborhood of λ0{\bm{\lambda}}_{0} is then (λ,w∗(λ))=f−1(λ,0)({\bm{\lambda}},{\bm{w}}^{*}({\bm{\lambda}}))=f^{-1}({\bm{\lambda}},{\mathbf{0}}). By the inverse function theorem, f−1f^{-1} is continuously differentiable inverse and so is w∗(λ){\bm{w}}^{*}({\bm{\lambda}}). The derivative ∂w∗(λ)\partial{\bm{w}}^{*}({\bm{\lambda}}) is obtained from the differential of the inverse as

(∼∼∂w∗(λ)∼)=∂f−1(λ,0),\begin{pmatrix}\sim&\sim\\ \partial{\bm{w}}^{*}({\bm{\lambda}})&\sim\end{pmatrix}=\partial f^{-1}({\bm{\lambda}},{\mathbf{0}}),

并由逆函数定理,我们有and by the inverse function theorem, we have

∂f−1(λ,0)=(∂f(λ,w∗(λ)))−1.\partial f^{-1}({\bm{\lambda}},{\mathbf{0}})=(\partial f({\bm{\lambda}},{\bm{w}}^{*}({\bm{\lambda}})))^{-1}.

因此,使用分块矩阵求逆公式,我们得到Using the block matrix inversion formula, we therefore obtain

(ABCD)−1=(∼∼−(D−CA−1B)−1CA−1∼),\begin{pmatrix}{\bm{A}}&{\bm{B}}\\ {\bm{C}}&{\bm{D}}\end{pmatrix}^{-1}=\begin{pmatrix}\sim&\sim\\ -({\bm{D}}-{\bm{C}}{\bm{A}}^{-1}{\bm{B}})^{-1}{\bm{C}}{\bm{A}}^{-1}&\sim\end{pmatrix},

我们便得到了所断言的表达式。尽管我们以雅可比矩阵和矩阵的形式表述了该证明,但这一结果自然对相应的线性算子、JVP、VJP 及其逆同样成立。we get the claimed expression. Though we expressed the proof in terms of Jacobians and matrices, the result naturally holds for the corresponding linear operators, JVPs, VJPs, and their inverses.

11.4 伴随状态法11.4 Adjoint state method

11.4.1 对非线性方程求导Differentiating nonlinear equations

本节介绍伴随状态法(又称伴随法、伴随方法、伴随灵敏度法)。该方法可用于计算显式函数与隐函数的复合的梯度,其中隐函数由等式约束(例如非线性方程)定义。该方法可追溯至 [41]。We describe in this section the adjoint state method (a.k.a. adjoint method, method of adjoints, adjoint sensitivity method). The method can be used to compute the gradient of the composition of an explicit function and an implicit function, defined through an equality constraint (e.g., a nonlinear equation). The method dates back to [41].

假设变量 s∈S{\bm{s}}\in\mathcal{S}(它对应最优控制中的状态)由(可能非线性的)方程 c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}} 在给定某些参数 w∈W{\bm{w}}\in\mathcal{W} 的条件下隐式地定义,其中 c ⁣:S×W→Sc\colon\mathcal{S}\times\mathcal{W}\to\mathcal{S}。假定对所有的 w∈W{\bm{w}}\in\mathcal{W},s{\bm{s}} 都唯一确定,于是便定义了一个从 W\mathcal{W} 到 S\mathcal{S} 的隐函数 s⋆(w){\bm{s}}^{\star}({\bm{w}}),使得 c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}}。给定目标函数 L ⁣:S×W→RL\colon\mathcal{S}\times\mathcal{W}\to\mathbb{R},伴随状态法的目标便是计算以下函数的梯度Suppose a variable s∈S{\bm{s}}\in\mathcal{S} (which corresponds to a state in optimal control) is implicitly defined given some parameters w∈W{\bm{w}}\in\mathcal{W} through the (potentially nonlinear) equation c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}, where c ⁣:S×W→Sc\colon\mathcal{S}\times\mathcal{W}\to\mathcal{S}. Assuming s{\bm{s}} is uniquely determined for all w∈W{\bm{w}}\in\mathcal{W}, this defines an implicit function s⋆(w){\bm{s}}^{\star}({\bm{w}}) from W\mathcal{W} to S\mathcal{S} such that c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}}. Given an objective function L ⁣:S×W→RL\colon\mathcal{S}\times\mathcal{W}\to\mathbb{R}, the goal of the adjoint state method is then to compute the gradient of

L(w):=L(s⋆(w),w).L({\bm{w}})\coloneqq L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}}).

然而,这并不简单,因为 s⋆(w){\bm{s}}^{\star}({\bm{w}}) 是一个隐函数。例如,这可用于将等式约束问题However, this is not trivial as s⋆(w){\bm{s}}^{\star}({\bm{w}}) is an implicit function. For instance, this can be used to convert the equality-constrained problem

min⁡w∈WL(s,w)s.t.c(s,w)=0.\displaystyle\min_{{\bm{w}}\in\mathcal{W}}L({\bm{s}},{\bm{w}})\quad\text{s.t.}\quad c({\bm{s}},{\bm{w}})={\mathbf{0}}.

转化为无约束问题into the unconstrained problem

min⁡w∈WL(s⋆(w),w).\displaystyle\min_{{\bm{w}}\in\mathcal{W}}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}}).

能够访问 ∇L(w)\nabla L({\bm{w}}) 使我们得以通过梯度下降求解该问题。Access to ∇L(w)\nabla L({\bm{w}}) allows us to solve this problem by gradient descent.

命题 11.4.1(伴随状态法)。Proposition 11.4.1 (Adjoint state method).
设 c:S×W→Sc:\mathcal{S}\times\mathcal{W}\rightarrow\mathcal{S} 为一映射,定义形如 c(s,w)c({\bm{s}},{\bm{w}}) 的约束。假设对每个 w∈W{\bm{w}}\in\mathcal{W},都存在唯一的 s⋆(w){\bm{s}}^{\star}({\bm{w}}) 满足 c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}},且 s⋆(w){\bm{s}}^{\star}({\bm{w}}) 可微。则Let c:S×W→Sc:\mathcal{S}\times\mathcal{W}\rightarrow\mathcal{S} be a mapping defining constraints of the form c(s,w)c({\bm{s}},{\bm{w}}). Assume that for each w∈W{\bm{w}}\in\mathcal{W}, there exists a unique s⋆(w){\bm{s}}^{\star}({\bm{w}}) satisfying c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}} and that s⋆(w){\bm{s}}^{\star}({\bm{w}}) is differentiable. The gradient of L(w):=L(s⋆(w),w),L({\bm{w}})\coloneqq L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}}), 对某个可微函数 L:S×W→RL:\mathcal{S}\times\mathcal{W}\rightarrow\mathbb{R},其梯度由下式给出for some differentiable function L:S×W→RL:\mathcal{S}\times\mathcal{W}\rightarrow\mathbb{R}, is given by ∇L(w)=∇2L(s⋆(w),w)+∂2c(s⋆(w),w)∗r⋆(w),\nabla L({\bm{w}})=\nabla_{2}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})+\partial_{2}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}^{\star}({\bm{w}}), 其中 r⋆(w){\bm{r}}^{\star}({\bm{w}}) 是如下线性方程组的解where r⋆(w){\bm{r}}^{\star}({\bm{w}}) is the solution of the linear system ∂1c(s⋆(w),w)∗r=−∇1L(s⋆(w),w).\partial_{1}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}=-\nabla_{1}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}}).

如下面的证明所示,r⋆(w){\bm{r}}^{\star}({\bm{w}}) 对应于一个拉格朗日乘子。该线性系统可以用无矩阵求解器求解。As shown in the proof below, r⋆(w){\bm{r}}^{\star}({\bm{w}}) corresponds to a Lagrange multiplier. The linear system can be solved using matrix-free solvers.

11.4.2 与包络定理的关系Relation with envelope theorems

由于对任意的 w∈W{\bm{w}}\in\mathcal{W},s{\bm{s}} 都由 c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}} 唯一确定,我们可以将 L(w)L({\bm{w}}) 等价地改写为平凡的最小化或最大化,Because s{\bm{s}} is uniquely determined for any w∈W{\bm{w}}\in\mathcal{W} by c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}, we can alternatively rewrite L(w)L({\bm{w}}) as the trivial minimization or maximization,

L(w)\displaystyle L({\bm{w}}) =min⁡s∈SL(s,w)s.t.c(s,w)=0\displaystyle=\min_{{\bm{s}}\in\mathcal{S}}L({\bm{s}},{\bm{w}})\quad\text{s.t.}\quad c({\bm{s}},{\bm{w}})={\mathbf{0}}
=max⁡s∈SL(s,w)s.t.c(s,w)=0.\displaystyle=\max_{{\bm{s}}\in\mathcal{S}}L({\bm{s}},{\bm{w}})\quad\text{s.t.}\quad c({\bm{s}},{\bm{w}})={\mathbf{0}}.

因此,伴随状态法可视为一种用于计算 ∇L(w)\nabla L({\bm{w}}) 的包络定理,适用于 w{\bm{w}} 同时出现在目标函数与等式约束中的情形。Therefore, the adjoint state method can be seen as an envelope theorem for computing ∇L(w)\nabla L({\bm{w}}), for the case when w{\bm{w}} is involved in both the objective function and in the equality constraint.

11.4.3 使用拉格朗日乘子法的证明Proof using the method of Lagrange multipliers

经典地,伴随状态法是借助拉格朗日乘子法推导出来的。让我们引入与 LL 和 cc 相关联的拉格朗日函数,Classically, the adjoint state method is derived using the method of Lagrange multipliers. Let us introduce the Lagrangian associated with LL and cc,

L(s,w,r):=L(s,w)+⟨r,c(s,w)⟩,\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}})\coloneqq L({\bm{s}},{\bm{w}})+\langle{\bm{r}},c({\bm{s}},{\bm{w}})\rangle,

其中 r∈S{\bm{r}}\in\mathcal{S} 是与等式约束 c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}} 相关联的拉格朗日乘子。在最优控制的文献中,r{\bm{r}} 常被称为伴随变量或伴随状态。拉格朗日函数的梯度为where r∈S{\bm{r}}\in\mathcal{S} is the Lagrange multiplier associated with the equality constraint c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}. In the optimal control literature, r{\bm{r}} is often called the adjoint variable or adjoint state. The gradients of the Lagrangian are

∇sL(s,w,r)\displaystyle\nabla_{\bm{s}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) =∇1L(s,w)+∂1c(s,w)∗r\displaystyle=\nabla_{1}L({\bm{s}},{\bm{w}})+\partial_{1}c({\bm{s}},{\bm{w}})^{*}{\bm{r}}
∇wL(s,w,r)\displaystyle\nabla_{\bm{w}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) =∇2L(s,w)+∂2c(s,w)∗r\displaystyle=\nabla_{2}L({\bm{s}},{\bm{w}})+\partial_{2}c({\bm{s}},{\bm{w}})^{*}{\bm{r}}
∇rL(s,w,r)\displaystyle\nabla_{\bm{r}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) =c(s,w),\displaystyle=c({\bm{s}},{\bm{w}}),

其中 ∂ic(s,w)∗\partial_{i}c({\bm{s}},{\bm{w}})^{*} 是伴随算子。令 ∇rL(s,w,r)\nabla_{\bm{r}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) 为零,即得到约束 c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}。令 ∇sL(s,w,r)\nabla_{\bm{s}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) 为零,则得到所谓的伴随状态方程where ∂ic(s,w)∗\partial_{i}c({\bm{s}},{\bm{w}})^{*} are the adjoint operators. Setting ∇rL(s,w,r)\nabla_{\bm{r}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) to zero gives the constraint c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}. Setting ∇sL(s,w,r)\nabla_{\bm{s}}\mathcal{L}({\bm{s}},{\bm{w}},{\bm{r}}) to zero gives the so-called adjoint state equation

∂1c(s,w)∗r=−∇1L(s,w).\partial_{1}c({\bm{s}},{\bm{w}})^{*}{\bm{r}}=-\nabla_{1}L({\bm{s}},{\bm{w}}).

在 s=s⋆(w){\bm{s}}={\bm{s}}^{\star}({\bm{w}}) 处关于 r{\bm{r}} 求解该线性系统,便得到伴随变量 r⋆(w){\bm{r}}^{\star}({\bm{w}})。于是我们得到Solving this linear system w.r.t. r{\bm{r}} at s=s⋆(w){\bm{s}}={\bm{s}}^{\star}({\bm{w}}) gives the adjoint variable r⋆(w){\bm{r}}^{\star}({\bm{w}}). We then get

∇L(w)\displaystyle\nabla L({\bm{w}}) =∇2L(s⋆(w),w,r⋆(w))\displaystyle=\nabla_{2}\mathcal{L}({\bm{s}}^{\star}({\bm{w}}),{\bm{w}},{\bm{r}}^{\star}({\bm{w}}))
=∇2L(s⋆(w),w)+∂2c(s⋆(w),w)∗r⋆(w),\displaystyle=\nabla_{2}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})+\partial_{2}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}^{\star}({\bm{w}}),

这就完成了证明。which concludes the proof.

11.4.4 使用隐函数定理的证明Proof using the implicit function theorem

借助隐函数定理(11.3 节),我们可以给出一个更直接的证明。运用链式法则,得到A more direct proof is possible thanks to the implicit function theorem (Section 11.3). Using the chain rule, we get

∇L(w)=∇2L(s⋆(w),w)+∂s⋆(w)∗∇1L(s⋆(w),w),\nabla L({\bm{w}})=\nabla_{2}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})+\partial{\bm{s}}^{\star}({\bm{w}})^{*}\nabla_{1}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}}),

其中 ∂s⋆(w)∗\partial{\bm{s}}^{\star}({\bm{w}})^{*} 是 s⋆{\bm{s}}^{\star} 的 VJP,即一个从 S\mathcal{S} 到 W\mathcal{W} 的线性映射。where ∂s⋆(w)∗\partial{\bm{s}}^{\star}({\bm{w}})^{*} is the VJP of s⋆{\bm{s}}^{\star}, a linear map from S\mathcal{S} to W\mathcal{W}.

在计算上,主要困难在于把 ∂s⋆(w)∗\partial{\bm{s}}^{\star}({\bm{w}})^{*} 作用于向量 u=∇1L(s⋆(w),w)∈S{\bm{u}}=\nabla_{1}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})\in\mathcal{S}。对隐函数 c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}} 应用隐函数定理(11.3 节)与命题 11.3.8,可得线性系统 A∗r=uA^{*}{\bm{r}}={\bm{u}},其中 A∗:=∂1c(s⋆(w),w)∗A^{*}\coloneqq\partial_{1}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*} 是从 S\mathcal{S} 到 S\mathcal{S} 的线性映射。解出 r{\bm{r}} 后,得到 ∂s⋆(w)∗u=B∗r\partial{\bm{s}}^{\star}({\bm{w}})^{*}{\bm{u}}=B^{*}{\bm{r}},其中 B∗:=∂2c(s⋆(w),w)∗B^{*}\coloneqq\partial_{2}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*} 是从 S\mathcal{S} 到 W\mathcal{W} 的线性映射。将以上各部分合在一起,便得到Computationally, the main difficulty is to apply ∂s⋆(w)∗\partial{\bm{s}}^{\star}({\bm{w}})^{*} to the vector u=∇1L(s⋆(w),w)∈S{\bm{u}}=\nabla_{1}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})\in\mathcal{S}. Using the implicit function theorem (Section 11.3) on the implicit function c(s⋆(w),w)=0c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})={\mathbf{0}}, and Proposition 11.3.8, we get the linear system A∗r=uA^{*}{\bm{r}}={\bm{u}}, where A∗:=∂1c(s⋆(w),w)∗A^{*}\coloneqq\partial_{1}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*} is a linear map from S\mathcal{S} to S\mathcal{S}. After solving for r{\bm{r}}, we get ∂s⋆(w)∗u=B∗r\partial{\bm{s}}^{\star}({\bm{w}})^{*}{\bm{u}}=B^{*}{\bm{r}}, where B∗:=∂2c(s⋆(w),w)∗B^{*}\coloneqq\partial_{2}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*} is a linear map from S\mathcal{S} to W\mathcal{W}. Putting everything together, we get

∇L(w)=∇2L(s⋆(w),w)+∂2c(s⋆(w),w)∗r.\nabla L({\bm{w}})=\nabla_{2}L({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})+\partial_{2}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}.

11.4.5 反向模式作为带回代的伴随方法Reverse mode as adjoint method with backsubstitution

本节从伴随状态法的角度重新审视反向模式自动微分。为叙述清晰起见,我们只考虑输入为 x∈X{\bm{x}}\in\mathcal{X}、网络权重为 w=(w1,…,wk)∈W1×…×WK{\bm{w}}=({\bm{w}}_{1},\ldots,{\bm{w}}_{k})\in\mathcal{W}_{1}\times\ldots\times\mathcal{W}_{K} 的前馈网络,其定义为In this section, we revisit reverse-mode autodiff from the perspective of the adjoint state method. For clarity, we focus our exposition on feedforward networks with input x∈X{\bm{x}}\in\mathcal{X} and network weights w=(w1,…,wk)∈W1×…×WK{\bm{w}}=({\bm{w}}_{1},\ldots,{\bm{w}}_{k})\in\mathcal{W}_{1}\times\ldots\times\mathcal{W}_{K}, defined by

s0\displaystyle{\bm{s}}_{0} :=x∈X\displaystyle\coloneqq{\bm{x}}\in\mathcal{X}
s1\displaystyle{\bm{s}}_{1} :=f1(s0,w1)∈S1\displaystyle\coloneqq f_{1}({\bm{s}}_{0},{\bm{w}}_{1})\in\mathcal{S}_{1}
⋮\displaystyle\hskip 6.0pt\vdots
sK\displaystyle{\bm{s}}_{K} :=fK(sK−1,wK)∈SK\displaystyle\coloneqq f_{K}({\bm{s}}_{K-1},{\bm{w}}_{K})\in\mathcal{S}_{K}
f(w)\displaystyle f({\bm{w}}) :=sK.\displaystyle\coloneqq{\bm{s}}_{K}. (11.2)

这里我们关注的是关于参数 w{\bm{w}} 的梯度,故记作 f(w)f({\bm{w}})。借助伴随状态法,可以重新得到反向模式自动微分,并在此过程中证明其正确性。虽然为简便起见我们只讨论前馈网络,但这一论述可推广至计算图。Here we focus on gradients with respect to the parameters w{\bm{w}}, hence the notation f(w)f({\bm{w}}). We can use the adjoint state method to recover reverse-mode autodiff, and prove its correctness in the process. While we focus for simplicity on feedforward networks, our exposition can be generalized to computation graphs.

将前馈网络视为非线性方程的解Feedforward networks as the solution of a nonlinear equation

尽管我们把中间计算的集合 s=(s1,…,sK)∈S1×…×SK{\bm{s}}=({\bm{s}}_{1},\ldots,{\bm{s}}_{K})\in\mathcal{S}_{1}\times\ldots\times\mathcal{S}_{K} 定义为一串运算,它也可以定义为非线性方程 c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}} 的唯一解,其中While we defined the set of intermediate computations s=(s1,…,sK)∈S1×…×SK{\bm{s}}=({\bm{s}}_{1},\ldots,{\bm{s}}_{K})\in\mathcal{S}_{1}\times\ldots\times\mathcal{S}_{K} as a sequence of operations, they can also be defined as the unique solution of the nonlinear equation c(s,w)=0c({\bm{s}},{\bm{w}})={\mathbf{0}}, where

c(s,w):=(s1−f1(x,w1)s2−f2(s1,w2)⋮sK−fK(sK−1,wK)).c({\bm{s}},{\bm{w}})\coloneqq\begin{pmatrix}{\bm{s}}_{1}-f_{1}({\bm{x}},{\bm{w}}_{1})\\ {\bm{s}}_{2}-f_{2}({\bm{s}}_{1},{\bm{w}}_{2})\\ \vdots\\ {\bm{s}}_{K}-f_{K}({\bm{s}}_{K-1},{\bm{w}}_{K})\end{pmatrix}.

由此定义了一个隐函数 s⋆(w)=(s1⋆(w),…,sK⋆(w)){\bm{s}}^{\star}({\bm{w}})=({\bm{s}}_{1}^{\star}({\bm{w}}),\ldots,{\bm{s}}_{K}^{\star}({\bm{w}})),即该非线性系统的解,它由式 11.2 中定义的变量 s1,…,sK{\bm{s}}_{1},\ldots,{\bm{s}}_{K} 给出。于是前馈网络的输出为 f(w)=sK⋆(w)f({\bm{w}})={\bm{s}}_{K}^{\star}({\bm{w}})。This defines an implicit function s⋆(w)=(s1⋆(w),…,sK⋆(w)){\bm{s}}^{\star}({\bm{w}})=({\bm{s}}_{1}^{\star}({\bm{w}}),\ldots,{\bm{s}}_{K}^{\star}({\bm{w}})), the solution of this nonlinear system, which is given by the variables s1,…,sK{\bm{s}}_{1},\ldots,{\bm{s}}_{K} defined in Eq. 11.2. The output of the feedforward network is then f(w)=sK⋆(w)f({\bm{w}})={\bm{s}}_{K}^{\star}({\bm{w}}).

在机器学习中,最后一层 sK⋆(w){\bm{s}}^{\star}_{K}({\bm{w}}) 通常会被送入损失函数 ℓ\ell,从而定义In machine learning, the final layer sK⋆(w){\bm{s}}^{\star}_{K}({\bm{w}}) is typically fed into a loss ℓ\ell, to define

L(w):=ℓ(sK⋆(w);y).L({\bm{w}})\coloneqq\ell({\bm{s}}_{K}^{\star}({\bm{w}});{\bm{y}}).

注意,另一种做法是把 L(w)L({\bm{w}}) 写成Note that an alternative is to write L(w)L({\bm{w}}) as

L(w)=min⁡s∈Sℓ(s;y)s.t.c(s,w)=0.L({\bm{w}})=\min_{{\bm{s}}\in\mathcal{S}}\ell({\bm{s}};{\bm{y}})\quad\text{s.t.}\quad c({\bm{s}},{\bm{w}})={\mathbf{0}}.

更一般地,若只想计算 sK⋆(w){\bm{s}}_{K}^{\star}({\bm{w}}) 沿某个方向 uK∈SK{\bm{u}}_{K}\in\mathcal{S}_{K} 的 VJP,我们可以定义标量值函数More generally, if we just want to compute the VJP of sK⋆(w){\bm{s}}_{K}^{\star}({\bm{w}}) in some direction uK∈SK{\bm{u}}_{K}\in\mathcal{S}_{K}, we can define the scalar-valued function

L(w):=ℓ(sK⋆(w);uK):=⟨sK⋆(w),uK⟩L({\bm{w}})\coloneqq\ell({\bm{s}}_{K}^{\star}({\bm{w}});{\bm{u}}_{K})\coloneqq\langle{\bm{s}}_{K}^{\star}({\bm{w}}),{\bm{u}}_{K}\rangle

于是so that

∂f(w)∗uK=∇L(w).\partial f({\bm{w}})^{*}{\bm{u}}_{K}=\nabla L({\bm{w}}).

我们令 u∈S1×⋯×SK−1×SK{\bm{u}}\in\mathcal{S}_{1}\times\dots\times\mathcal{S}_{K-1}\times\mathcal{S}_{K} 表示 u:=(0,…,0,∇1ℓ(f(w);y)){\bm{u}}\coloneqq({\mathbf{0}},\ldots,{\mathbf{0}},\nabla_{1}\ell(f({\bm{w}});{\bm{y}}))(即损失 ℓ\ell 的梯度情形)或 u:=(0,…,0,uK){\bm{u}}\coloneqq({\mathbf{0}},\ldots,{\mathbf{0}},{\bm{u}}_{K})(即 ff 沿方向 uK{\bm{u}}_{K} 的 VJP 情形)。根据伴随状态法可知,该目标的梯度可由下式得到Let us define u∈S1×⋯×SK−1×SK{\bm{u}}\in\mathcal{S}_{1}\times\dots\times\mathcal{S}_{K-1}\times\mathcal{S}_{K} as u:=(0,…,0,∇1ℓ(f(w);y)){\bm{u}}\coloneqq({\mathbf{0}},\ldots,{\mathbf{0}},\nabla_{1}\ell(f({\bm{w}});{\bm{y}})) (gradient of the loss ℓ\ell case) or u:=(0,…,0,uK){\bm{u}}\coloneqq({\mathbf{0}},\ldots,{\mathbf{0}},{\bm{u}}_{K}) (VJP of ff in the direction uK{\bm{u}}_{K} case). Using the adjoint state method, we know that the gradient of this objective is obtained as

∇L(w)=∂2c(s(w),w)∗r⋆(w),\nabla L({\bm{w}})=\partial_{2}c({\bm{s}}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}^{\star}({\bm{w}}),

其中 r⋆(w){\bm{r}}^{\star}({\bm{w}}) 为线性方程组的解for r⋆(w){\bm{r}}^{\star}({\bm{w}}) the solution of the linear system

∂1c(s(w),w)∗r=−u.\partial_{1}c({\bm{s}}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}=-{\bm{u}}.

用回代法求解该线性方程组Solving the linear system using backsubtitution

约束函数 cc 在 s⋆(w){\bm{s}}^{\star}({\bm{w}}) 处的 JVP 具体化为矩阵后,呈分块下三角矩阵的形式The JVP of the constraint function cc at s⋆(w){\bm{s}}^{\star}({\bm{w}}), materialized as a matrix, takes the form of a block lower-triangular matrix

∂1c(s⋆(w),w)=(I⁡0……0−A1I⁡⋱⋮0−A2I⁡⋱⋮⋮⋱⋱⋱00…0−AKI⁡),\bm{\partial}_{1}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})=\begin{pmatrix}\operatorname{I}&0&\dots&\dots&0\\ -A_{1}&\operatorname{I}&\ddots&&\vdots\\ 0&-A_{2}&\operatorname{I}&\ddots&\vdots\\ \vdots&\ddots&\ddots&\ddots&0\\ 0&\dots&0&-A_{K}&\operatorname{I}\end{pmatrix},

其中 Ak:=∂1fk(sk−1,wk)A_{k}\coloneqq\bm{\partial}_{1}f_{k}({\bm{s}}_{k-1},{\bm{w}}_{k})。关键在于,JVP 的三角结构源于如下事实:每个中间激活值只依赖于此前各中间激活值。因此,与雅可比矩阵各行对应的约束,无法在其对角线之外引入非零值。VJP 则呈分块上三角矩阵的形式where Ak:=∂1fk(sk−1,wk)A_{k}\coloneqq\bm{\partial}_{1}f_{k}({\bm{s}}_{k-1},{\bm{w}}_{k}). Crucially the triangular structure of the JVP stems from the fact that each intermediate activation only depends on the past intermediate activations. Therefore, the constraints, corresponding to the lines of the Jacobian, cannot introduce non-zero values beyond its diagonal. The VJP takes the form of a block upper-triangular matrix

∂1c(s⋆(w),w)∗=(I⁡−A1∗0…00I⁡−A2∗⋱⋮⋮⋱I⁡⋱0⋮⋱⋱−AK∗0……0I⁡).\bm{\partial}_{1}c({\bm{s}}^{\star}({\bm{w}}),{\bm{w}})^{*}=\begin{pmatrix}\operatorname{I}&-A_{1}^{*}&0&\dots&0\\ 0&\operatorname{I}&-A_{2}^{*}&\ddots&\vdots\\ \vdots&\ddots&\operatorname{I}&\ddots&0\\ \vdots&&\ddots&\ddots&-A_{K}^{*}\\ 0&\dots&\dots&0&\operatorname{I}\end{pmatrix}.

于是,像 ∂1c(s(w),w)∗r=u\partial_{1}c({\bm{s}}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}={\bm{u}} 这样的上三角系统,便可借助回代高效求解。从最后一个伴随状态 rK=u{\bm{r}}_{K}={\bm{u}} 出发,我们可由在 k+1k+1 处算出的伴随状态,逐一计算各伴随状态 rk{\bm{r}}_{k}。也就是说,对于 k∈(K−1,…,1)k\in(K-1,\dots,1),有Solving an upper triangular system like ∂1c(s(w),w)∗r=u\partial_{1}c({\bm{s}}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}={\bm{u}} can then be done efficiently by backsubstitution. Starting from the last adjoint state rK=u{\bm{r}}_{K}={\bm{u}}, we can compute each adjoint state rk{\bm{r}}_{k} from that computed at k+1k+1. Namely, for k∈(K−1,…,1)k\in(K-1,\dots,1), we have

rk−Ak+1∗rk+1=0  ⟺  rk=∂fk+1(sk,wk+1)∗rk+1.\displaystyle{\bm{r}}_{k}-A_{k+1}^{*}{\bm{r}}_{k+1}={\mathbf{0}}\iff{\bm{r}}_{k}=\partial f_{k+1}({\bm{s}}_{k},{\bm{w}}_{k+1})^{*}{\bm{r}}_{k+1}.

而关于参数的 VJP 则可由下式得到The VJPs with respect to the parameters are then obtained by

∂2c(s(w),w)∗r=(∂2f1(x,w1)∗r1∂2f2(s1(w),w2)∗r2⋮∂2fK(sK−1(w),wK)∗rK),\partial_{2}c({\bm{s}}({\bm{w}}),{\bm{w}})^{*}{\bm{r}}=\begin{pmatrix}\partial_{2}f_{1}({\bm{x}},{\bm{w}}_{1})^{*}{\bm{r}}_{1}\\ \partial_{2}f_{2}({\bm{s}}_{1}({\bm{w}}),{\bm{w}}_{2})^{*}{\bm{r}}_{2}\\ \vdots\\ \partial_{2}f_{K}({\bm{s}}_{K-1}({\bm{w}}),{\bm{w}}_{K})^{*}{\bm{r}}_{K}\end{pmatrix},

重新得到反向模式自动微分。recovering reverse-mode autodiff.

对于具有独立参数 w=(w1,…,wK){\bm{w}}=({\bm{w}}_{1},\dots,{\bm{w}}_{K}) 的网络,反向传播的拉格朗日视角是众所周知的;例如参见 [116] 或 [166]。对于具有共享参数 w{\bm{w}} 的网络,通过时间反向传播 [209] 的拉格朗日视角例如由 [70] 讨论。我们的阐述使用伴随状态法,该方法本身既可以用拉格朗日乘子法(第 11.4.3 节)证明,也可以用隐函数定理(第 11.4.4 节)证明,并结合回代来求解上三角线性系统。过去的工作常常在 w{\bm{w}} 上最小化,但我们不要求这样做,因为梯度不一定用于优化。我们的阐述还支持计算任意向量值函数 ff 的 VJP,而现有工作推导的是标量值损失函数的梯度。The Lagrangian perspective of backpropagation for networks with separate parameters w=(w1,…,wK){\bm{w}}=({\bm{w}}_{1},\dots,{\bm{w}}_{K}) is well-known; see for instance [116] or [166]. The Lagrangian perspective of backpropagation through time [209] for networks with shared parameter w{\bm{w}} is discussed for instance by [70]. Our exposition uses the adjoint state method, which can itself be proved either using the method of Lagrange multipliers (Section 11.4.3) or by the implicit function theorem (Section 11.4.4), combined with backsubtitution for solving the upper-triangular linear system. Past works often minimize over w{\bm{w}} but we do not require this, as gradients are not necessarily used for optimization. Our exposition also supports computing the VJP of any vector-valued function ff, while existing works derive the gradient of a scalar-valued loss function.

11.5 反函数定理11.5 Inverse function theorem

11.5.1 对逆函数求导Differentiating inverse functions

在某些情形下(例如参见 12.4.4 节),计算逆函数 f−1f^{-1} 的雅可比矩阵会很有用。下面的逆函数定理使我们能够把 f−1f^{-1} 的雅可比矩阵与 ff 的雅可比矩阵联系起来。In some cases (see for instance Section 12.4.4), it is useful to compute the Jacobian of an inverse function f−1f^{-1}. The inverse function theorem below allows us to relate the Jacobian of f−1f^{-1} with the Jacobian of ff.

定理 11.5.1(逆函数定理)。Theorem 11.5.1 (Inverse function theorem).
假设 f ⁣:W→Wf\colon\mathcal{W}\to\mathcal{W} 连续可微,且在 w0{\bm{w}}_{0} 处的雅可比矩阵 ∂f(w0)\partial f({\bm{w}}_{0}) 可逆。则 ff 是 w0{\bm{w}}_{0} 的某个邻域到 f(w0)f({\bm{w}}_{0}) 的某个邻域的双射。此外,逆映射 f−1f^{-1} 在 ω0=f(w0){\bm{\omega}}_{0}=f({\bm{w}}_{0}) 附近连续可微,且逆映射的雅可比矩阵 ∂f−1(ω)\partial f^{-1}({\bm{\omega}}) 为Assume f ⁣:W→Wf\colon\mathcal{W}\to\mathcal{W} is continuously differentiable with invertible Jacobian ∂f(w0)\partial f({\bm{w}}_{0}) at w0{\bm{w}}_{0}. Then ff is bijective from a neighborhood of w0{\bm{w}}_{0} to a neighborhood of f(w0)f({\bm{w}}_{0}). Moreover, the inverse f−1f^{-1} is continuously differentiable near ω0=f(w0){\bm{\omega}}_{0}=f({\bm{w}}_{0}) and the Jacobian of the inverse ∂f−1(ω)\partial f^{-1}({\bm{\omega}}) is ∂f(w)∂f−1(ω)=I⇔∂f−1(ω)=(∂f(w))−1,\partial f({\bm{w}})\partial f^{-1}({\bm{\omega}})=I\Leftrightarrow\partial f^{-1}({\bm{\omega}})=(\partial f({\bm{w}}))^{-1}, 其中 w=f−1(ω){\bm{w}}=f^{-1}({\bm{\omega}})。with w=f−1(ω){\bm{w}}=f^{-1}({\bm{\omega}}).

11.5.2 与隐函数定理的联系Link with the implicit function theorem

逆函数定理可用于证明隐函数定理,参见定理 11.3.5 的证明。反过来,回顾一下:要用隐函数定理,需选取一个根目标 F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}。若令 W=Λ=RQ\mathcal{W}=\Lambda=\mathbb{R}^{Q}、F(w,ω)=f(w)−ωF({\bm{w}},{\bm{\omega}})=f({\bm{w}})-{\bm{\omega}},其中 f ⁣:RQ→RQf\colon\mathbb{R}^{Q}\to\mathbb{R}^{Q},则满足 F(w⋆(ω),ω)=0F({\bm{w}}^{\star}({\bm{\omega}}),{\bm{\omega}})={\mathbf{0}} 的根 w⋆(ω){\bm{w}}^{\star}({\bm{\omega}}) 恰为 w⋆(ω)=f−1(ω){\bm{w}}^{\star}({\bm{\omega}})=f^{-1}({\bm{\omega}})。并且 ∂1F(w,ω)=∂f(w)\partial_{1}F({\bm{w}},{\bm{\omega}})=\partial f({\bm{w}})、∂2F(w,ω)=−I\partial_{2}F({\bm{w}},{\bm{\omega}})=-I。以该 FF 应用隐函数定理,便确实重新得到了逆函数定理。The inverse function theorem can be used to prove the implicit function theorem; see proof of Theorem 11.3.5. Conversely, recall that, in order to use the implicit function theorem, we need to choose a root objective F ⁣:W×Λ→WF\colon\mathcal{W}\times\Lambda\to\mathcal{W}. If we set W=Λ=RQ\mathcal{W}=\Lambda=\mathbb{R}^{Q} and F(w,ω)=f(w)−ωF({\bm{w}},{\bm{\omega}})=f({\bm{w}})-{\bm{\omega}}, with f ⁣:RQ→RQf\colon\mathbb{R}^{Q}\to\mathbb{R}^{Q}, then we have that the root w⋆(ω){\bm{w}}^{\star}({\bm{\omega}}) satisfying F(w⋆(ω),ω)=0F({\bm{w}}^{\star}({\bm{\omega}}),{\bm{\omega}})={\mathbf{0}} is exactly w⋆(ω)=f−1(ω){\bm{w}}^{\star}({\bm{\omega}})=f^{-1}({\bm{\omega}}). Moreover, ∂1F(w,ω)=∂f(w)\partial_{1}F({\bm{w}},{\bm{\omega}})=\partial f({\bm{w}}) and ∂2F(w,ω)=−I\partial_{2}F({\bm{w}},{\bm{\omega}})=-I. By applying the implicit function theorem with this FF, we indeed recover the inverse function theorem.

11.5.3 逆函数定理的证明Proof of inverse function theorem

我们首先在 f−1f^{-1} 于 f(w0)f({\bm{w}}_{0}) 的某邻域内良定义且连续可微的假设下,给出该公式的一个证明。此时,对 f(w0)f({\bm{w}}_{0}) 邻域内的任意 ω{\bm{\omega}},有We first give a proof of the formula assuming that f−1f^{-1} is well-defined and continuously differentiable in a neighborhood of f(w0)f({\bm{w}}_{0}). In that case, we have for any ω{\bm{\omega}} in a neighborhood of f(w0)f({\bm{w}}_{0}),

f∘f−1(ω)=ω.f\circ f^{-1}({\bm{\omega}})={\bm{\omega}}.

将两边关于 ω{\bm{\omega}} 求导,可得Differentiating both sides w.r.t. ω{\bm{\omega}}, we get

∂f(f−1(ω))∂f−1(ω)=I⁡,\partial f(f^{-1}({\bm{\omega}}))\partial f^{-1}({\bm{\omega}})=\operatorname{I},

其中 I⁡\operatorname{I} 为 RQ\mathbb{R}^{Q} 上的恒等函数。特别地,取 w=f−1(ω){\bm{w}}=f^{-1}({\bm{\omega}}),便重新得到定理陈述中所给的公式。where I⁡\operatorname{I} is the identity function in RQ\mathbb{R}^{Q}. In particular, for w=f−1(ω){\bm{w}}=f^{-1}({\bm{\omega}}) we recover the formula presented in the statement.

接下来只需证明:JVP 的可逆性保证了函数在 f(w0)f({\bm{w}}_{0}) 的某邻域内可逆,且其逆连续可微。为此,记 l=∂f(w0){\bm{l}}=\partial f({\bm{w}}_{0}),使得按定义 l−1{\bm{l}}^{-1} 良定义。ff 具有连续可微的逆,当且仅当 l−1(f(w))−f(w0){\bm{l}}^{-1}(f({\bm{w}}))-f({\bm{w}}_{0}) 具有连续可微的逆。于是不失一般性,考虑 ∂f(w0)=I⁡\partial f({\bm{w}}_{0})=\operatorname{I}、f(w0)=0f({\bm{w}}_{0})={\mathbf{0}}、w0=0{\bm{w}}_{0}={\mathbf{0}}。Now, it remains to show that invertibility of the JVP ensures that the function is invertible in a neighborhood of f(w0)f({\bm{w}}_{0}) and that the inverse is continuously differentiable. For that, denote l=∂f(w0){\bm{l}}=\partial f({\bm{w}}_{0}) such that l−1{\bm{l}}^{-1} is well-defined by definition. ff is invertible with continuously differentiable inverse, if and only if l−1(f(w))−f(w0){\bm{l}}^{-1}(f({\bm{w}}))-f({\bm{w}}_{0}) is invertible with continuously differentiable inverse. So without loss of generality, we consider ∂f(w0)=I⁡\partial f({\bm{w}}_{0})=\operatorname{I}, f(w0)=0f({\bm{w}}_{0})={\mathbf{0}}, w0=0{\bm{w}}_{0}={\mathbf{0}}.

由于 ff 连续可微,存在邻域 N:={w:∥w−w0∥2≤δ}\mathcal{N}\coloneqq\{{\bm{w}}:\|{\bm{w}}-{\bm{w}}_{0}\|_{2}\leq\delta\},在其上 ∥∂f(w)−I⁡∥2≤1/2\|\partial f({\bm{w}})-\operatorname{I}\|_{2}\leq 1/2。由中值定理,g(w)=f(w)−wg({\bm{w}})=f({\bm{w}})-{\bm{w}} 在该邻域内为压缩映射,压缩因子为 1/21/2。对任意满足 ∥ω−f(w0)∥2≤δ/2\|{\bm{\omega}}-f({\bm{w}}_{0})\|_{2}\leq\delta/2 的 ω{\bm{\omega}},序列 wk+1=wk−f(wk)−ω′{\bm{w}}_{k+1}={\bm{w}}_{k}-f({\bm{w}}_{k})-{\bm{\omega}}^{\prime} 始终落在 N\mathcal{N} 内,并收敛(因 gg 的压缩性,它是柯西序列)到唯一的不动点 w∞{\bm{w}}_{\infty},且满足 w∞=w∞−f(w∞)−ω  ⟺  f(w∞)=ω{\bm{w}}_{\infty}={\bm{w}}_{\infty}-f({\bm{w}}_{\infty})-{\bm{\omega}}\iff f({\bm{w}}_{\infty})={\bm{\omega}}。由此可知,逆映射在 ω0=f(w0){\bm{\omega}}_{0}=f({\bm{w}}_{0}) 的邻域 M:={ω:∥ω−ω0∥2≤δ/2}\mathcal{M}\coloneqq\{{\bm{\omega}}:\|{\bm{\omega}}-{\bm{\omega}}_{0}\|_{2}\leq\delta/2\} 上存在,并映入 N\mathcal{N}。As ff is continuously differentiable, there exists a neighborhood N:={w:∥w−w0∥2≤δ}\mathcal{N}\coloneqq\{{\bm{w}}:\|{\bm{w}}-{\bm{w}}_{0}\|_{2}\leq\delta\} on which we have ∥∂f(w)−I⁡∥2≤1/2\|\partial f({\bm{w}})-\operatorname{I}\|_{2}\leq 1/2. In this neighborhood, the function g(w)=f(w)−wg({\bm{w}})=f({\bm{w}})-{\bm{w}} is contractive by the mean value theorem with contraction factor 1/21/2. For any ω{\bm{\omega}} such that ∥ω−f(w0)∥2≤δ/2\|{\bm{\omega}}-f({\bm{w}}_{0})\|_{2}\leq\delta/2, the sequence wk+1=wk−f(wk)−ω′{\bm{w}}_{k+1}={\bm{w}}_{k}-f({\bm{w}}_{k})-{\bm{\omega}}^{\prime} remains in N\mathcal{N} and converges (since it is a Cauchy sequence by the contraction of gg) to a unique fixed point w∞{\bm{w}}_{\infty} satisfying w∞=w∞−f(w∞)−ω  ⟺  f(w∞)=ω{\bm{w}}_{\infty}={\bm{w}}_{\infty}-f({\bm{w}}_{\infty})-{\bm{\omega}}\iff f({\bm{w}}_{\infty})={\bm{\omega}}. This shows the existence of the inverse in the neighborhood M:={ω:∥ω−ω0∥2≤δ/2}\mathcal{M}\coloneqq\{{\bm{\omega}}:\|{\bm{\omega}}-{\bm{\omega}}_{0}\|_{2}\leq\delta/2\} of ω0=f(w0){\bm{\omega}}_{0}=f({\bm{w}}_{0}) onto N\mathcal{N}.

我们现在来处理 f−1f^{-1} 的可微性(从而连续性)。对 ω0{\bm{\omega}}_{0} 邻域内的任意 ω{\bm{\omega}},其逆为 w:=f−1(ω)∈N{\bm{w}}\coloneqq f^{-1}({\bm{\omega}})\in N;由假设,ff 在 w{\bm{w}} 处的 JVP 满足 ∥∂f(w)−I⁡∥2≤1/2\|\partial f({\bm{w}})-\operatorname{I}\|_{2}\leq 1/2。因此,a=∂f(w)−I⁡{\bm{a}}=\partial f({\bm{w}})-\operatorname{I} 定义了一个收敛级数 b=∑k=0+∞ak{\bm{b}}=\sum_{k=0}^{+\infty}{\bm{a}}^{k},且容易验证 b=∂f(w)−1{\bm{b}}=\partial f({\bm{w}})^{-1},即 ∂f(w)\partial f({\bm{w}}) 可逆,并且 ∥∂f(w)−1∥≤2\|\partial f({\bm{w}})^{-1}\|\leq 2。为计算逆的 JVP,我们随后把 ∂f(w)−1\partial f({\bm{w}})^{-1} 视为候选 JVP,并考察We tackle now the differentiability (hence the continuity) of f−1f^{-1}. For any ω{\bm{\omega}} in the neighborhood of ω0{\bm{\omega}}_{0} with inverse w:=f−1(ω)∈N{\bm{w}}\coloneqq f^{-1}({\bm{\omega}})\in N, the JVP of ff at w{\bm{w}} satisfies by assumption ∥∂f(w)−I⁡∥2≤1/2\|\partial f({\bm{w}})-\operatorname{I}\|_{2}\leq 1/2. Hence, a=∂f(w)−I⁡{\bm{a}}=\partial f({\bm{w}})-\operatorname{I} defines a convergent series b=∑k=0+∞ak{\bm{b}}=\sum_{k=0}^{+\infty}{\bm{a}}^{k} and one verifies easily that b=∂f(w)−1{\bm{b}}=\partial f({\bm{w}})^{-1}, that is ∂f(w)\partial f({\bm{w}}) is invertible and ∥∂f(w)−1∥≤2\|\partial f({\bm{w}})^{-1}\|\leq 2. To compute the JVP of the inverse, we consider then ∂f(w)−1\partial f({\bm{w}})^{-1} as the candidate JVP and examine

∥f−1(ω+η)−f(ω)−(∂f(w))−1η∥2∥η∥2.\frac{\|f^{-1}({\bm{\omega}}+{\bm{\eta}})-f({\bm{\omega}})-(\partial f({\bm{w}}))^{-1}{\bm{\eta}}\|_{2}}{\|{\bm{\eta}}\|_{2}}.

再记 v{\bm{v}},使得 f−1(ω+η)=w+vf^{-1}({\bm{\omega}}+{\bm{\eta}})={\bm{w}}+{\bm{v}}。由于 g(w)=f(w)−wg({\bm{w}})=f({\bm{w}})-{\bm{w}} 在 N\mathcal{N} 中是 1/21/2-压缩的,故有 ∥v−η∥2=∥g(w+v)−g(w)∥2≤1/2∥v∥2\|{\bm{v}}-{\bm{\eta}}\|_{2}=\|g({\bm{w}}+{\bm{v}})-g({\bm{w}})\|_{2}\leq 1/2\|{\bm{v}}\|_{2}。于是 ∥v∥2≥∥η∥/2\|{\bm{v}}\|_{2}\geq\|{\bm{\eta}}\|/2。由此得到Denote then v{\bm{v}} such that f−1(ω+η)=w+vf^{-1}({\bm{\omega}}+{\bm{\eta}})={\bm{w}}+{\bm{v}}. As g(w)=f(w)−wg({\bm{w}})=f({\bm{w}})-{\bm{w}} is 1/21/2-contractive in N\mathcal{N}, we have ∥v−η∥2=∥g(w+v)−g(w)∥2≤1/2∥v∥2\|{\bm{v}}-{\bm{\eta}}\|_{2}=\|g({\bm{w}}+{\bm{v}})-g({\bm{w}})\|_{2}\leq 1/2\|{\bm{v}}\|_{2}. So ∥v∥2≥∥η∥/2\|{\bm{v}}\|_{2}\geq\|{\bm{\eta}}\|/2. We then get

∥f−1(ω+η)−f(ω)−(∂f(w))−1η∥2∥η∥2\displaystyle\frac{\|f^{-1}({\bm{\omega}}+{\bm{\eta}})-f({\bm{\omega}})-(\partial f({\bm{w}}))^{-1}{\bm{\eta}}\|_{2}}{\|{\bm{\eta}}\|_{2}}
=∥v−(∂f(w))−1(f(w+v)−f(w))∥2∥η∥2\displaystyle=\frac{\|{\bm{v}}-(\partial f({\bm{w}}))^{-1}(f({\bm{w}}+{\bm{v}})-f({\bm{w}}))\|_{2}}{\|{\bm{\eta}}\|_{2}}
≤4∥f(w+v)−f(w)−∂f(w)v∥2∥v∥2\displaystyle\leq 4\frac{\|f({\bm{w}}+{\bm{v}})-f({\bm{w}})-\partial f({\bm{w}}){\bm{v}}\|_{2}}{\|{\bm{v}}\|_{2}}

由于 ∥η∥2→0\|{\bm{\eta}}\|_{2}\rightarrow 0,我们有 ∥v∥2→0\|{\bm{v}}\|_{2}\rightarrow 0,从而 ∥f(w+v)−f(w)−∂f(w)v∥2/∥v∥2→0\|f({\bm{w}}+{\bm{v}})-f({\bm{w}})-\partial f({\bm{w}}){\bm{v}}\|_{2}/\|{\bm{v}}\|_{2}\rightarrow 0。因此,f−1f^{-1} 可微,其 JVP 为 ∂f−1(ω)=(∂f(w))−1=(∂f(f−1(ω)))−1\partial f^{-1}({\bm{\omega}})=(\partial f({\bm{w}}))^{-1}=(\partial f(f^{-1}({\bm{\omega}})))^{-1}。这表明 f−1f^{-1} 连续,因而作为连续函数的组合,∂f−1(ω)\partial f^{-1}({\bm{\omega}}) 也连续。As ∥η∥2→0\|{\bm{\eta}}\|_{2}\rightarrow 0, we have ∥v∥2→0\|{\bm{v}}\|_{2}\rightarrow 0 and so ∥f(w+v)−f(w)−∂f(w)v∥2/∥v∥2→0\|f({\bm{w}}+{\bm{v}})-f({\bm{w}})-\partial f({\bm{w}}){\bm{v}}\|_{2}/\|{\bm{v}}\|_{2}\rightarrow 0. Hence, f−1f^{-1} is differentiable with JVP ∂f−1(ω)=(∂f(w))−1=(∂f(f−1(ω)))−1\partial f^{-1}({\bm{\omega}})=(\partial f({\bm{w}}))^{-1}=(\partial f(f^{-1}({\bm{\omega}})))^{-1}. This shows that f−1f^{-1} is continuous and so ∂f−1(ω)\partial f^{-1}({\bm{\omega}}) is continuous as a composition of continuous functions.

11.6 小结11.6 Summary

  • •

    隐函数是无法分解为基本操作、因而无法直接应用自动微分的函数。优化问题与非线性方程即为例子。Implicit functions are functions that cannot be decomposed into elementary operations and for which autodiff can therefore not be directly applied. Examples are optimization problems and nonlinear equations.

  • •

    包络定理可用于对函数的最小值或最大值(而非其解)求导。Envelope theorems can be used for differentiating through the min or max value (not solution) of a function.

  • •

    更一般地,隐函数定理使我们能够对隐函数求导。它给出了导数存在的条件,以及求得导数的方法。More generally, the implicit function theorem allows us to differentiate through implicit functions. It gives conditions for the existence of derivatives and how to obtain them.

  • •

    伴随状态法可用于求显式函数与由等式约束指定的隐函数之复合的梯度,并可用于证明反向模式自动微分的正确性。The adjoint state method can be used to obtain the gradient of the composition of an explicit function and of an implicit function, specified by equality constraints. It can be used to prove the correctness of reverse-mode autodiff.

  • •

    逆函数定理可用于对函数的逆求导。The inverse function theorem can be used to differentiate function inverses.

  • •

    在某种意义上,隐函数定理可视为母定理,因为它可用于证明包络定理、伴随状态法与逆函数定理。In a sense, the implicit function theorem can be thought of as the mother theorem, as it can be used to prove envelope theorems, the adjoint state method and the inverse function theorem.