本章研究如何通过优化问题求导,更一般地说,如何通过非线性方程组求导。In this chapter, we study how to differentiate through optimization problems, and more generally through nonlinear systems of equations.
隐函数无法显式分解为初等函数,因而第 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.
形式上,我们将隐函数记为 ,其中 。由此产生的一个问题是:如何计算雅可比矩阵 ?作为初步应用,可以考虑系统的灵敏度分析。例如, 可对应某物理系统的平衡状态,此时 便告诉我们该系统对某些参数 的灵敏度。Formally, we will denote an implicit function by , where . One question is then how to compute the Jacobian . As a first application one can consider sensitivity analysis of a system. For example, could correspond to the equilibrium state of a physical system and in this case, would tell us about the sensitivity of the system to some parameters .
另一例子是隐式定义为某一优化问题的解(假定唯一)的函数Another example is a function implicitly defined as the solution (assumed unique) of an optimization problem
其中 , 表示一个约束集。注意,这里为方便起见使用了 ,但改用 时同样成立。where and denotes a constraint set. Note that we use an for convenience, but the same applies when using an .
更一般地, 可定义为某函数 的根,即: 被隐式定义为满足下述(可能非线性的)方程组的取值More generally, can be defined as the root of some function , i.e., is implicitly defined as the value satisfying the (potentially nonlinear) system of equations
对一切 。for all .
除灵敏度分析外,双层优化是另一个应用实例。我们常常希望最小化的函数,是由某个固定函数与一个优化问题的解复合而成的。形式上,令 。考察如下定义的复合 :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 . We consider the composition defined as
| (11.1) |
例如,这包括超参数优化:其中 为内层的对数似然目标, 为外层的验证损失, 为模型参数, 则为模型超参数(如正则化强度),如图 11.1 所示。要最小化 ,通常采用关于 的梯度下降方案,这就需要对 进行计算。设 在 处可微,则由链式法则可得雅可比矩阵This includes for instance hyperparameter optimization, where is an inner log-likelihood objective, is an outer validation loss, are model parameters and are model hyperparameters, such as regularization strength, as illustrated in Fig. 11.1. To minimize , one generally resorts to a gradient descent scheme w.r.t. , which requires computing . Assuming that is differentiable at , by the chain rule, we obtain the Jacobian
借助 (见注 2.2.9),可得到梯度Using (see Remark 2.2.9), we obtain the gradient
唯一棘手的一项是 ,因为它涉及对 argmax 求导。事实上,大多数情况下 并无显式公式,也无法分解为初等函数。The only problematic term is , as it requires argmax differentiation. Indeed, most of the time, there is no explicit formula for and it does not decompose into elementary functions.
在 这一特殊情形下,式 11.1 所定义的复合 简化为In the special case , the composition defined in Eq. 11.1 is simply given by
也就是说,此时不再需要对 argmax 求导,而只需对 max 求导;下文将会看到,后者要容易得多。函数 常被称为值函数 [66]。“包络”这一名称的由来如图 11.2 所示。需要强调,包络定理并非只有一条,而是有多条,具体取决于对 所作的假设。That is, we no longer need argmax differentiation, but only max differentiation, which, as we shall now see is much easier. The function 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 .
以下定理取自 [24, 命题 B.25]。The following theorem is due to [24, Proposition B.25].
当 时,可以给出该定理的一个变体。A variant of the theorem can be stated when .
直观地说,定理 11.2.1 与定理 11.2.2 意味着我们可以把 当作 的一个常数来处理。换言之,我们不需要穿过 进行求导,即使它依赖于 。Informally, Theorem 11.2.1 and Theorem 11.2.2 mean that we can treat as if it were a constant of . In other words, we do not need to differentiate through , even though it depends on .
现在,我们借助定理 11.2.2 说明如何穿过最小值 进行求导,其中假定 关于 是强凸的。We now illustrate Theorem 11.2.2, to differentiate through a minimum, , assuming is strongly convex in .
如下所示,对于关于 为仿射的函数,定理 11.2.1 有一个简单的解释。Theorem 11.2.1 has a simple interpretation for functions that are affine in as shown below.
另一个应用是鞍点优化。Another application is saddle point optimization.
在与 相关的不同假设下,可以证明一个相关定理 [169, 定理 10.31]。A related theorem [169, Theorem 10.31] can be proved under different assumptions on .
与 Bertsekas 定理相比,Rockafellar 定理不要求 关于 为凹,但对 的可微性提出了更强的假设。Compared to Bertsekas’ theorem, Rockafellar’s theorem does not require to be concave in , but it requires stronger assumptions on the differentiability of .
最后,我们陈述 Danskin 的原始定理 [51]。Finally, we state Danskin’s original theorem [51].
我们在表 11.1 中总结各包络定理及其假设。We sum up envelope theorems and their assumptions in Table 11.1.
| 定理Theorem | 是否唯一?Unique ? | 关于 是否为凹?Concave in ? | 关于 是否光滑?Smooth in ? | 结果Result |
|---|---|---|---|---|
| DanskinDanskin | 否No | 否No | 否No | 方向导数Dir. Deriv. |
| BertsekasBertsekas | 是Yes | 是Yes | 否No | 梯度Gradient |
| RockafellarRockafellar | 是Yes | 否No | 是Yes | 梯度Gradient |
隐函数定理(IFT)给出了如下条件:在该条件下,形如 的隐式关系可以在局部改写为函数 ,并给出一种计算其关于 的导数的方法。The implicit function theorem (IFT) provides conditions under which an implicit relationship of the form can be rewritten as a function locally, and provides a way to compute its derivative w.r.t. .
我们将证明推迟到多元情形,先从一个经典的应用该定理的例子开始。We postpone the proof to the multivariate case and begin with a classical example of application of the theorem.
在上述简单例子中,我们可以容易地在给定邻域内推导出联系 与 的显式函数,但这并非总是可行。隐函数定理给出了保证这样的函数存在的条件以及对其求导的方法,但并不给出构造该函数的方法。事实上,寻找使 成立的 ,通常涉及求根算法、优化算法、非线性系统求解器等。In the above simple example, we can easily derive an explicit function relating to 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 such that typically involves a root finding algorithm, an optimization algorithm, a nonlinear system solver, etc.
若 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.
现在我们给出一般多元情形下的隐函数定理。非正式地说,若 ,则由链式法则,我们有We now present the IFT in the general multivariate setting. Informally, if , then by the chain rule, we have
这意味着雅可比矩阵 (假设存在)满足meaning that the Jacobian , assuming that it exists, satisfies
隐函数定理给出了 存在的条件。The IFT gives us conditions for the existence of .
我们从一个简单的无约束优化算法开始。We begin with a simple unconstrained optimization algorithm.
接下来,我们通过允许优化问题带有约束来推广前一例子。Next, we generalize the previous example, by allowing constraints in the optimization problem.
要将隐函数 集成到自动微分框架中,我们需要能够计算其 JVP 或 VJP。这正是下一个命题的目的。To integrate an implicit function in an autodiff framework, we need to be able to compute its JVP or VJP. This is the purpose of the next proposition.
注意,在上述线性系统中,我们可以把 和 作为线性映射来访问,它们正是 的 JVP。它们的伴随 和 对应于 的 VJP。因此,为求解这些系统,我们可以使用 9.4 节中详述的无矩阵求解器。例如,当 为对称半正定时,我们可以使用共轭梯度法 [93];当 不是对称正定时,则可以使用 GMRES [172] 或 BiCGSTAB [205]。Note that in the above linear systems, we can access and as linear maps, the JVPs of . Their adjoints, and , correspond to the VJPs of . To solve these systems, we can therefore use matrix-free solvers as detailed in Section 9.4. For example, when is symmetric positive semi-definite, we can use the conjugate gradient method [93]. When is not symmetric positive definite, we can use GMRES [172] or BiCGSTAB [205].
我们借助定理 11.5.1 中给出的逆函数定理来证明该定理。定义We prove the theorem using the inverse function theorem presented in Theorem 11.5.1. Define
其从 映到 。 的雅可比矩阵为which goes from onto . The Jacobian of is
因此在 处有 ,因为我们假定了 可逆。由逆函数定理,函数 在 的邻域 内可逆。特别地,它在 中可逆。于是隐式方程在 邻域内的解为 。由逆函数定理, 是连续可微的逆,因此 亦然。导数 由逆的微分得到,即So at , we have since we assumed invertible. By the inverse function theorem, the function is then invertible in a neighborhood of . In particular, it is invertible in . The solution of the implicit equation in a neighborhood of is then . By the inverse function theorem, is continuously differentiable inverse and so is . The derivative is obtained from the differential of the inverse as
并由逆函数定理,我们有and by the inverse function theorem, we have
因此,使用分块矩阵求逆公式,我们得到Using the block matrix inversion formula, we therefore obtain
我们便得到了所断言的表达式。尽管我们以雅可比矩阵和矩阵的形式表述了该证明,但这一结果自然对相应的线性算子、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.
本节介绍伴随状态法(又称伴随法、伴随方法、伴随灵敏度法)。该方法可用于计算显式函数与隐函数的复合的梯度,其中隐函数由等式约束(例如非线性方程)定义。该方法可追溯至 [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].
假设变量 (它对应最优控制中的状态)由(可能非线性的)方程 在给定某些参数 的条件下隐式地定义,其中 。假定对所有的 , 都唯一确定,于是便定义了一个从 到 的隐函数 ,使得 。给定目标函数 ,伴随状态法的目标便是计算以下函数的梯度Suppose a variable (which corresponds to a state in optimal control) is implicitly defined given some parameters through the (potentially nonlinear) equation , where . Assuming is uniquely determined for all , this defines an implicit function from to such that . Given an objective function , the goal of the adjoint state method is then to compute the gradient of
然而,这并不简单,因为 是一个隐函数。例如,这可用于将等式约束问题However, this is not trivial as is an implicit function. For instance, this can be used to convert the equality-constrained problem
转化为无约束问题into the unconstrained problem
能够访问 使我们得以通过梯度下降求解该问题。Access to allows us to solve this problem by gradient descent.
如下面的证明所示, 对应于一个拉格朗日乘子。该线性系统可以用无矩阵求解器求解。As shown in the proof below, corresponds to a Lagrange multiplier. The linear system can be solved using matrix-free solvers.
由于对任意的 , 都由 唯一确定,我们可以将 等价地改写为平凡的最小化或最大化,Because is uniquely determined for any by , we can alternatively rewrite as the trivial minimization or maximization,
因此,伴随状态法可视为一种用于计算 的包络定理,适用于 同时出现在目标函数与等式约束中的情形。Therefore, the adjoint state method can be seen as an envelope theorem for computing , for the case when is involved in both the objective function and in the equality constraint.
经典地,伴随状态法是借助拉格朗日乘子法推导出来的。让我们引入与 和 相关联的拉格朗日函数,Classically, the adjoint state method is derived using the method of Lagrange multipliers. Let us introduce the Lagrangian associated with and ,
其中 是与等式约束 相关联的拉格朗日乘子。在最优控制的文献中, 常被称为伴随变量或伴随状态。拉格朗日函数的梯度为where is the Lagrange multiplier associated with the equality constraint . In the optimal control literature, is often called the adjoint variable or adjoint state. The gradients of the Lagrangian are
其中 是伴随算子。令 为零,即得到约束 。令 为零,则得到所谓的伴随状态方程where are the adjoint operators. Setting to zero gives the constraint . Setting to zero gives the so-called adjoint state equation
在 处关于 求解该线性系统,便得到伴随变量 。于是我们得到Solving this linear system w.r.t. at gives the adjoint variable . We then get
这就完成了证明。which concludes the proof.
借助隐函数定理(11.3 节),我们可以给出一个更直接的证明。运用链式法则,得到A more direct proof is possible thanks to the implicit function theorem (Section 11.3). Using the chain rule, we get
其中 是 的 VJP,即一个从 到 的线性映射。where is the VJP of , a linear map from to .
在计算上,主要困难在于把 作用于向量 。对隐函数 应用隐函数定理(11.3 节)与命题 11.3.8,可得线性系统 ,其中 是从 到 的线性映射。解出 后,得到 ,其中 是从 到 的线性映射。将以上各部分合在一起,便得到Computationally, the main difficulty is to apply to the vector . Using the implicit function theorem (Section 11.3) on the implicit function , and Proposition 11.3.8, we get the linear system , where is a linear map from to . After solving for , we get , where is a linear map from to . Putting everything together, we get
本节从伴随状态法的角度重新审视反向模式自动微分。为叙述清晰起见,我们只考虑输入为 、网络权重为 的前馈网络,其定义为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 and network weights , defined by
| (11.2) |
这里我们关注的是关于参数 的梯度,故记作 。借助伴随状态法,可以重新得到反向模式自动微分,并在此过程中证明其正确性。虽然为简便起见我们只讨论前馈网络,但这一论述可推广至计算图。Here we focus on gradients with respect to the parameters , hence the notation . 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.
尽管我们把中间计算的集合 定义为一串运算,它也可以定义为非线性方程 的唯一解,其中While we defined the set of intermediate computations as a sequence of operations, they can also be defined as the unique solution of the nonlinear equation , where
由此定义了一个隐函数 ,即该非线性系统的解,它由式 11.2 中定义的变量 给出。于是前馈网络的输出为 。This defines an implicit function , the solution of this nonlinear system, which is given by the variables defined in Eq. 11.2. The output of the feedforward network is then .
在机器学习中,最后一层 通常会被送入损失函数 ,从而定义In machine learning, the final layer is typically fed into a loss , to define
注意,另一种做法是把 写成Note that an alternative is to write as
更一般地,若只想计算 沿某个方向 的 VJP,我们可以定义标量值函数More generally, if we just want to compute the VJP of in some direction , we can define the scalar-valued function
于是so that
我们令 表示 (即损失 的梯度情形)或 (即 沿方向 的 VJP 情形)。根据伴随状态法可知,该目标的梯度可由下式得到Let us define as (gradient of the loss case) or (VJP of in the direction case). Using the adjoint state method, we know that the gradient of this objective is obtained as
其中 为线性方程组的解for the solution of the linear system
约束函数 在 处的 JVP 具体化为矩阵后,呈分块下三角矩阵的形式The JVP of the constraint function at , materialized as a matrix, takes the form of a block lower-triangular matrix
其中 。关键在于,JVP 的三角结构源于如下事实:每个中间激活值只依赖于此前各中间激活值。因此,与雅可比矩阵各行对应的约束,无法在其对角线之外引入非零值。VJP 则呈分块上三角矩阵的形式where . 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
于是,像 这样的上三角系统,便可借助回代高效求解。从最后一个伴随状态 出发,我们可由在 处算出的伴随状态,逐一计算各伴随状态 。也就是说,对于 ,有Solving an upper triangular system like can then be done efficiently by backsubstitution. Starting from the last adjoint state , we can compute each adjoint state from that computed at . Namely, for , we have
而关于参数的 VJP 则可由下式得到The VJPs with respect to the parameters are then obtained by
重新得到反向模式自动微分。recovering reverse-mode autodiff.
对于具有独立参数 的网络,反向传播的拉格朗日视角是众所周知的;例如参见 [116] 或 [166]。对于具有共享参数 的网络,通过时间反向传播 [209] 的拉格朗日视角例如由 [70] 讨论。我们的阐述使用伴随状态法,该方法本身既可以用拉格朗日乘子法(第 11.4.3 节)证明,也可以用隐函数定理(第 11.4.4 节)证明,并结合回代来求解上三角线性系统。过去的工作常常在 上最小化,但我们不要求这样做,因为梯度不一定用于优化。我们的阐述还支持计算任意向量值函数 的 VJP,而现有工作推导的是标量值损失函数的梯度。The Lagrangian perspective of backpropagation for networks with separate parameters is well-known; see for instance [116] or [166]. The Lagrangian perspective of backpropagation through time [209] for networks with shared parameter 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 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 , while existing works derive the gradient of a scalar-valued loss function.
在某些情形下(例如参见 12.4.4 节),计算逆函数 的雅可比矩阵会很有用。下面的逆函数定理使我们能够把 的雅可比矩阵与 的雅可比矩阵联系起来。In some cases (see for instance Section 12.4.4), it is useful to compute the Jacobian of an inverse function . The inverse function theorem below allows us to relate the Jacobian of with the Jacobian of .
逆函数定理可用于证明隐函数定理,参见定理 11.3.5 的证明。反过来,回顾一下:要用隐函数定理,需选取一个根目标 。若令 、,其中 ,则满足 的根 恰为 。并且 、。以该 应用隐函数定理,便确实重新得到了逆函数定理。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 . If we set and , with , then we have that the root satisfying is exactly . Moreover, and . By applying the implicit function theorem with this , we indeed recover the inverse function theorem.
我们首先在 于 的某邻域内良定义且连续可微的假设下,给出该公式的一个证明。此时,对 邻域内的任意 ,有We first give a proof of the formula assuming that is well-defined and continuously differentiable in a neighborhood of . In that case, we have for any in a neighborhood of ,
将两边关于 求导,可得Differentiating both sides w.r.t. , we get
其中 为 上的恒等函数。特别地,取 ,便重新得到定理陈述中所给的公式。where is the identity function in . In particular, for we recover the formula presented in the statement.
接下来只需证明:JVP 的可逆性保证了函数在 的某邻域内可逆,且其逆连续可微。为此,记 ,使得按定义 良定义。 具有连续可微的逆,当且仅当 具有连续可微的逆。于是不失一般性,考虑 、、。Now, it remains to show that invertibility of the JVP ensures that the function is invertible in a neighborhood of and that the inverse is continuously differentiable. For that, denote such that is well-defined by definition. is invertible with continuously differentiable inverse, if and only if is invertible with continuously differentiable inverse. So without loss of generality, we consider , , .
由于 连续可微,存在邻域 ,在其上 。由中值定理, 在该邻域内为压缩映射,压缩因子为 。对任意满足 的 ,序列 始终落在 内,并收敛(因 的压缩性,它是柯西序列)到唯一的不动点 ,且满足 。由此可知,逆映射在 的邻域 上存在,并映入 。As is continuously differentiable, there exists a neighborhood on which we have . In this neighborhood, the function is contractive by the mean value theorem with contraction factor . For any such that , the sequence remains in and converges (since it is a Cauchy sequence by the contraction of ) to a unique fixed point satisfying . This shows the existence of the inverse in the neighborhood of onto .
我们现在来处理 的可微性(从而连续性)。对 邻域内的任意 ,其逆为 ;由假设, 在 处的 JVP 满足 。因此, 定义了一个收敛级数 ,且容易验证 ,即 可逆,并且 。为计算逆的 JVP,我们随后把 视为候选 JVP,并考察We tackle now the differentiability (hence the continuity) of . For any in the neighborhood of with inverse , the JVP of at satisfies by assumption . Hence, defines a convergent series and one verifies easily that , that is is invertible and . To compute the JVP of the inverse, we consider then as the candidate JVP and examine
再记 ,使得 。由于 在 中是 -压缩的,故有 。于是 。由此得到Denote then such that . As is -contractive in , we have . So . We then get
由于 ,我们有 ,从而 。因此, 可微,其 JVP 为 。这表明 连续,因而作为连续函数的组合, 也连续。As , we have and so . Hence, is differentiable with JVP . This shows that is continuous and so is continuous as a composition of continuous functions.
隐函数是无法分解为基本操作、因而无法直接应用自动微分的函数。优化问题与非线性方程即为例子。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.