Meta-Learning with Implicit Gradients

摘要:

智能系统的一个核心功能是能够通过借鉴先前的经验来快速学习新任务的能力。最近，基于梯度（或优化）的元学习已成为一种有效的少样本学习方法。在此形式中，仅使用当前任务中的少量数据，即可在outer loop中学习元参数，而在inner-loop中学习特定于任务的模型。扩展这些方法的关键挑战是需要通过inner-loop学习过程计算微分，这可能会带来相当大的计算和内存负担。借助隐式微分，我们开发了隐式MAML算法，该算法仅取决于inner level优化的解，而不取决于inner loop优化器采用的路径。这有效地将元梯度计算与inner loop优化器的选择解耦。因此，我们的方法与inner loop优化器的选择无关，并且可以优雅地处理许多梯度步骤而不会梯度消失或内存限制。从理论上讲，我们证明隐式MAML可以使用不超过计算单个内循环梯度所需的内存占用量来计算准确的元梯度，而不会增加总的计算成本。从实验上，我们证明了隐式MAML的这些好处可转化为在少样本的图像识别benchmarks上的经验收益。

其他来源：Notes on iMAML: Meta-Learning with Implicit Gradients

对这篇笔记也进行了翻译，链接如下: iMAML笔记（翻译）

Problem Formulation and Notations

我们首先在少样本监督学习的背景下提出元学习问题，然后概括该概念以帮助本文的其余论述。

Review of Few-Shot Supervised Learning and MAML

在这种设置下，我们从$P(T)$中提取了一系列元训练任务$\{T_i\} ^M_{i = 1}$。每个任务$T_i$与一个数据集$D_i$相关联，我们可以从中采样两个不相交的集合：$D^{tr}_i$和$D^{test}_i$。这些数据集都由$K$个输入输出对组成。令$x\in X$和$y\in Y$分别表示输入和输出。数据集采用$D^{tr}_i = \{(x^k_i，y^k_i)\}^K_{k = 1}$的形式，对于$D^{test}_i$同样。我们对学习 $h_\phi(x): X \to Y $形式的模型感兴趣，由$\phi\in\Phi\equiv\mathbb {R}^{d}$参数化。任务的性能由损失函数指定，例如交叉熵或平方误差损失。我们将损失函数以$L(\phi，D)$的形式写成参数向量和数据集的函数。任务$T_i$的目标是使用$D^{tr}_ i$学习特定于任务的参数$\phi_i$，以便我们可以使任务的总体或测试损失$L(\phi_i，D^{test}_ i)$最小化。

在一般的双层元学习设置中，我们考虑使用一组元参数$\theta\in\Theta\equiv\mathbb{R}^{d}$和来自任务任务的训练数据集用于计算任务特定参数的算法空间，形式化如$\phi_i= Alg(\theta,D^{tr}_ i)$针对任务$T_i$。元学习的目标是学习适应后产生良好任务特定参数的元参数，如下所示：

$$\overbrace{\boldsymbol{\theta}_{\mathrm{ML}}^{*}:=\underset{\boldsymbol{\theta} \in \Theta}{\operatorname{argmin}} F(\boldsymbol{\theta})}^{\text{outer-level}}, \text { where } F(\boldsymbol{\theta})=\frac{1}{M} \sum_{i=1}^{M} \mathcal{L}\left(\overbrace{Alg\left(\boldsymbol{\theta}, \mathcal{D}_{i}^{\mathrm{tr}}\right)}^{\text{inner-level}}, \mathcal{D}_{i}^{\mathrm{test}}\right)$$

由于我们通常将$Alg(\theta,D^{tr}_i)$解释为显式或隐式解决潜在优化问题，因此我们将其视为双层优化问题。在元测试（部署）时，当提供与新任务$T_j\sim P(T)$相对应的数据集$D^{tr}_j$时，我们可以通过对学到的元参数使用适应过程来获得良好的泛化性能（即低test error），即得到$\phi_j= Alg(\theta^\ast_{ML},D^{tr}_j)$。

在MAML中，$Alg(\theta,D)$对应于以$\theta$初始化的一个或多个梯度下降步骤。例如，如果使用一个梯度下降步骤，我们得到：

$$\phi_{i} \equiv \mathcal{A} l g\left(\boldsymbol{\theta}, \mathcal{D}_{i}^{\mathrm{tr}}\right)=\boldsymbol{\theta}-\alpha \nabla_{\boldsymbol{\theta}} \mathcal{L}\left(\boldsymbol{\theta}, \mathcal{D}_{i}^{\mathrm{tr}}\right) . \quad \text { (inner-level of MAML) }$$

通常，$\alpha$是标量超参数，但也可以是可学习的向量。因此，对于MAML，元学习参数（$\theta^\ast_{ML}$）具有学习归纳偏差，该偏差特别适合于使用$K$个样本对$P(T)$中的任务进行微调。为了使用基于梯度的方法解决外层问题，我们需要一种贯穿$Alg$计算微分的方法。在MAML中，这对应于通过梯度下降的动力学进行反向传播。

Proximal Regularization in the Inner Level

为了在内层水平上获得足够的学习，同时又避免过拟合，$Alg$需要纳入某种形式的正则化。由于MAML使用少量的梯度步骤，这对应于早停，可以解释为正则化和贝叶斯先验的一种形式。在病态条件数的优化情形和中等样本量学习的情况下，我们可能需要采用多个梯度下降步，这给MAML带来了两个挑战。首先，我们需要通过漫长的$Alg$优化路径进行存储和计算微分，这会带来相当大的计算和内存负担。其次，随着$Alg$中梯度步数的增加，模型参数$\{\phi_i\}$对元参数（$\theta $）的依赖关系会缩小和消失，从而使得元学习变得困难。为了克服这些限制，我们考虑使用更明确的正则化算法：

$$\mathcal{Alg}^{\star}\left(\boldsymbol{\theta}, \mathcal{D}_{i}^{\mathrm{tr}}\right)=\underset{\boldsymbol{\phi}^{\prime} \in \Phi}{\operatorname{argmin}} \mathcal{L}\left(\boldsymbol{\phi}^{\prime}, \mathcal{D}_{i}^{\mathrm{tr}}\right)+\frac{\lambda}{2}\left\|\boldsymbol{\phi}^{\prime}-\boldsymbol{\theta}\right\|^{2}$$

等式中的正则项促使$\phi_i$保持接近$\theta$，从而始终保持很强的依赖性。正则化强度（$\lambda$）的作用类似于MAML中的学习率（$\alpha$），控制先验（$\theta$）相对于数据（$D^{tr}_T$）的强度。像$\alpha$一样，正则强度$\lambda$也可以被学习。此外，$\alpha$和$\lambda$都可以是标量，向量或完整矩阵。为简单起见，我们将$\lambda$作为标量超参数。在等式中，我们用$\star$表示优化问题已完全解决。在实践中，我们使用迭代算法（由$Alg$表示）进行有限迭代，其返回近似最小化解。我们在分析中明确考虑了近似解与精确解之间的差异。

The Bi-Level Optimization Problem

为方便起见，有时我们会使用下标而不是参数来表达对任务$T_i$的依赖性，例如我们写：

$$\mathcal{L}_{i}(\phi):=\mathcal{L}\left(\phi, \mathcal{D}_{i}^{\text {test }}\right), \quad \hat{\mathcal{L}}_{i}(\phi):=\mathcal{L}\left(\phi, \mathcal{D}_{i}^{\text {tr }}\right), \quad \mathcal{A} l g_{i}(\boldsymbol{\theta}):=\mathcal{A} \lg \left(\boldsymbol{\theta}, \mathcal{D}_{i}^{\text {tr }}\right)$$

使用这种表示法，可以将二层元学习问题更一般地写为：

$$\begin{array}{l}{\boldsymbol{\theta}_{\mathrm{ML}}^{*}:=\underset{\boldsymbol{\theta} \in \Theta}{\operatorname{argmin}} F(\boldsymbol{\theta}), \text { where } F(\boldsymbol{\theta})=\frac{1}{M} \sum_{i=1}^{M} \mathcal{L}_{i}\left(\mathcal{A} l g_{i}^{\star}(\boldsymbol{\theta})\right), \text { and }} \\ {\mathcal{A} l g_{i}^{\star}(\boldsymbol{\theta}):=\underset{\boldsymbol{\phi}^{\prime} \in \Phi}{\operatorname{argmin}} G_{i}\left(\boldsymbol{\phi}^{\prime}, \boldsymbol{\theta}\right), \text { where } G_{i}\left(\boldsymbol{\phi}^{\prime}, \boldsymbol{\theta}\right)=\hat{\mathcal{L}}_{i}\left(\boldsymbol{\phi}^{\prime}\right)+\frac{\lambda}{2}\left\|\boldsymbol{\phi}^{\prime}-\boldsymbol{\theta}\right\|^{2}}\end{array} \tag{4}$$

Total and Partial Derivatives

我们用$d$表示全导数，用$\nabla$表示偏导数。对于形式为$L_i(\phi_i)$的嵌套函数，其中$\phi_i= Alg_i(\theta)$，我们有链式法则

$$\boldsymbol{d}_{\boldsymbol{\theta}} \mathcal{L}_{i}\left(\mathcal{A} l g_{i}(\boldsymbol{\theta})\right)=\frac{d \mathcal{A} l g_{i}(\boldsymbol{\theta})}{d \boldsymbol{\theta}} \nabla_{\boldsymbol{\phi}} \mathcal{L}_{i}(\boldsymbol{\phi})|_{\boldsymbol{\phi}=\mathcal{A} l g_{i}(\boldsymbol{\theta})}=\frac{d \mathcal{A} l g_{i}(\boldsymbol{\theta})}{d \boldsymbol{\theta}} \nabla_{\boldsymbol{\phi}} \mathcal{L}_{i}\left(\mathcal{A} l g_{i}(\boldsymbol{\theta})\right)$$

注意$d_{\theta} \mathcal{L}_{i}(Alg_{i}(\theta))$和$\nabla_{\phi} \mathcal{L}_{i}(Alg_{i}(\theta))$之间的重要区别。前者传递导数穿过$Alg_i(\theta)$，而后者则不穿过。$\nabla_{\phi} \mathcal{L}_{i}(Alg_{i}(\theta))$只是梯度函数，即$\nabla_{\phi} \mathcal{L}_{i}(\phi)$，以$\phi=Alg_{i}(\theta)$求值。还要注意，$d_{\theta} \mathcal{L}_{i}(Alg_{i}(\theta))$和$\nabla_{\phi} \mathcal{L}_{i}(Alg_{i}(\theta))$是d维向量，而$\frac{d Alg_{i}(\theta)}{d \theta}$是（d×d）大小的Jacobian矩阵。在本文中，我们还将无差别地使用$d_\theta$和$\frac{d}{d\theta}$。

The Implicit MAML Algorithm

我们的目的是使用形如$\theta \gets \theta - \eta d_\theta F(\theta)$的基于迭代梯度的算法解决公式4中的双层元学习问题。尽管为简单起见，我们基于标准梯度下降法导出了我们的方法，但也可以使用任何其他优化方法，例如准牛顿法或牛顿法，Adam或带动量的梯度下降法，而无需进行任何修改。使用链式法则扩展梯度下降更新为

$$\boldsymbol{\theta} \leftarrow \boldsymbol{\theta}-\eta \frac{1}{M} \sum_{i=1}^{M} \frac{d \mathcal{Alg}_{i}^{\star}(\boldsymbol{\theta})}{d \boldsymbol{\theta}} \nabla_{\phi} \mathcal{L}_{i}\left(\mathcal{A} l g_{i}^{\star}(\boldsymbol{\theta})\right)$$

在此，$\nabla_{\phi} \mathcal{L}_{i}(\mathcal{Alg}_{i}^{\star}(\theta))$即$\nabla_{\phi} \mathcal{L}_{i}(\phi)|_{\phi=\mathcal{Alg}_{i}^{\star}(\theta)}$，操作上可以通过自动微分而容易地求出。对于此更新规则，我们必须计算$\frac{d Alg^\star_{i}(\theta)}{d \theta}$，其中$Alg^\star_{i}$被隐式定义为优化问题(公式4)，这带来了主要挑战。现在，我们提出一种有效的算法（在计算和内存中）以计算元梯度。

Meta-Gradient Computation

如果将$Alg^\star_{i}(\theta)$实现为迭代算法（例如梯度下降），则计算$\frac{d Alg^\star_{i}(\theta)}{d \theta}$的一种方法是通过迭代过程传播导数，无论是正向还是反向。但是，这样做的缺点是必须明确依赖于优化路径，而优化路径必须完全存储在内存中，而当所需的梯度下降步数较多时，该路径很快变得难以处理。此外，对于诸如牛顿法的二阶优化方法，需要难以获得的三阶导数。此外，当使用不可微分的操作（如行搜索）时，此方法变得不可能。但是，通过认识到$Alg^\star_{i}$是隐式定义的优化问题的解决方案，我们可以采用一个不同的策略，该策略不需要考虑优化路径，而只需考虑最终结果。这源自以下引理。

引理 1. （隐式雅克比）考虑公式4中对于任务$T_i$定义的$Alg^\star_{i}(\theta)$。另$\phi_i=\mathcal{Alg}_{i}^{\star}(\theta)$为$Alg^\star_{i}(\theta)$的解。如果$\left(\boldsymbol{I}+\frac{1}{\lambda} \nabla_{\phi}^{2} \hat{\mathcal{L}}_{i}\left(\boldsymbol{\phi}_{i}\right)\right)$是可逆的，则导数雅可比行列式为

$$\frac{d \mathcal{Alg}_{i}^{\star}(\boldsymbol{\theta})}{d \boldsymbol{\theta}}=\left(\boldsymbol{I}+\frac{1}{\lambda} \nabla_{\phi}^{2} \hat{\mathcal{L}}_{i}\left(\boldsymbol{\phi}_{i}\right)\right)^{-1} \tag{6}$$

请注意，导数（Jacobian）仅取决于算法的最终结果，而不取决于算法采用的路径。因此，原则上，任何算法方法都可用于计算$Alg^\star_{i}(\theta)$，从而将元梯度计算与内部级优化器的选择解耦。

具体算法：虽然引理1提供了一种理想的方法来计算$Alg^\star_{i}$的雅可比行列式，因此通过扩展元梯度，在实践中可能很难直接使用它。有两个问题特别相关。首先，元梯度要求计算$Alg^\star_{i}(\theta)$，这是内部优化问题的精确解。实际上，我们可能只能获得近似解。第二，对于大型深度神经网络而言，显式地构造和求逆等式6中的矩阵，用于计算雅可比行列式，可能是棘手的。为了解决这些困难，我们考虑接近理想方法的可行算法。

首先，我们考虑内部优化问题的近似解，可以通过迭代优化算法（例如梯度下降）获得。

定义1. （$\delta$-近似算法）另$Alg_i(\theta)$为$Alg^\star_{i}(\theta)$的$\delta$-近似，如：

$$\|Alg_{i}(\theta)-Alg_{i}^{\star}(\theta)\| \leq \delta$$

其次，我们将执行部分或近似矩阵求逆：

定义2. （$\delta’$-近似雅可比向量积）另$g_i$为一个向量满足

$$\|g_{i}-\left(I+\frac{1}{\lambda} \nabla_{\phi}^{2} \hat{\mathcal{L}}_{i}\left(\phi_{i}\right)\right)^{-1} \nabla_{\phi} \mathcal{L}_{i}\left(\phi_{i}\right)\| \leq \delta^{\prime}$$

这里$\phi_i=Alg_i(\theta)$且$Alg_i$基于定义1.

注意，定义2中的$g_i$是任务$T_i$的元梯度的近似值。观察到可以将$g_i$作为优化问题的近似解：

$$\min _{\boldsymbol{w}} \boldsymbol{w}^{\top}\left(\boldsymbol{I}+\frac{1}{\lambda} \nabla_{\boldsymbol{\phi}}^{2} \hat{\mathcal{L}}_{i}\left(\boldsymbol{\phi}_{i}\right)\right) \boldsymbol{w}-\boldsymbol{w}^{\top} \nabla_{\phi} \mathcal{L}_{i}\left(\boldsymbol{\phi}_{i}\right)$$

共轭梯度（CG）算法由于其出色的迭代复杂性和仅要求$\nabla_\phi^2\hat{\mathcal{L}}_i(\phi_i)$形式的Hessian矢量积的要求而特别适合于此问题。无需显式构造或存储Hessian矩阵即可容易地获得这样的hessian矢量积（如我们在附录C中所述）。这种基于CG的求逆已成功地应用在用于深度学习的Hessian-free或Newton-CG方法[36，44]和强化学习[52，47]中的信任区域方法。算法1提出了完整的实用算法。注意，这些近似值用于开发实用的算法，会在元梯度计算中引入误差。我们将在3.2节中分析这些误差的影响，并证明它们是可控制的。有关iMAML如何推广了基于先验梯度优化的元学习算法，请参阅附录A。