Meta-Learning with Implicit Gradient

NIPS 2019

8/18/2020

https://papers.nips.cc/paper/8306-meta-learning-with-implicit-gradients

Motivation

There are some limitations in MAML. The meta-learning process requires higher-order derivatives, imposes a non-trivial computational and memory burden, and can suffer from vanishing gradients. These limitations make it harder to scale optimization-based meta learning methods to tasks involving medium or large datasets, or those that require many inner-loop optimization steps.

Contributions

The development of the implicit MAML (iMAML) algorithm, an approach for optimization-based meta-learning with deep neural networks that removes the need for differentiating through the optimization path.

The algorithm aims to learn a set of parameters such that an optimization algorithm that is initialized at and regularized to this parameter vector leads to good generalization for a variety of learning tasks.

Few-shot supervised learning and MAML (Bi-level optimization)

The goal of meta-learning is to learn meta-parameters that produce good task specific parameters after adaptation.

Notations:

• $\theta^*_{ML}$ optimal meta-learned parameters

• $M$ the number of tasks in meta-train, $i$ is the index of task $i$

• $\mathcal{D}^{tr}_i$ support set, $\mathcal{D}^{test}_i$ query set in task $i$

• $\mathcal{L}(\phi, \mathcal{D})$ loss function with parameter vector and dataset

• $\phi_i = \mathcal{A}lg(\theta, \mathcal{D}^{tr}_{i})=\theta - \alpha\nabla_{\theta}\mathcal{L}(\theta,\mathcal{D}^{tr}_i)$ : one (or multiple) steps of gradient descent initialized at $\theta$. [inner-level of MAML]

Proximal regularization in the inner level

Two challenges in MAML in cases like ill-conditioned optimization landscapes and medium-shot learning, we may want to take many gradient steps:

• we need to store and differentiate through the long optimization path of $\mathcal{A}lg$, imposing a considerable computation and memory burden

• The dependence of the model-parameters {$\phi_i$} on the meta-parameters ( $\theta$ ) shrinks and vanishes as the number of gradient steps in $\mathcal{A}lg$ grows, making meta-learning difficult.

A more explicitly regularized algorithm is considered:

New Bi-level optimization problem after proximal regularization

Simplify the notation:

$\mathcal{L}_i(\phi):=\mathcal{L}(\phi,\mathcal{D}_i^{test}), \hat{\mathcal{L}}_i(\phi):=\mathcal{L}(\phi,\mathcal{D}_i^{tr}), \mathcal{A}lg(\theta):=\mathcal{A}lg(\theta,\mathcal{D}^{tr}_i)$

bi-level meta-learning problem (more general)

Total and Partial Derivatives:

( $\boldsymbol{d}$ denotes the total derivative, $\nabla$ denotes the partial derivative)

$$\boldsymbol{d}_{\theta}\mathcal{L}_i(\mathcal{A}lg_i(\theta)) = \frac{\boldsymbol{d}\mathcal{A}lg_i(\theta)}{d\theta}\nabla_{\phi}\mathcal{L}_i(\phi)|_{\phi=\mathcal{A}lg_i(\theta)} = \frac{\boldsymbol{d}\mathcal{A}lg_i(\theta)}{d\theta}\nabla_{\phi}\mathcal{L}_i(\mathcal{A}lg_i(\theta))$$

Implicit MAML

Goal: to solve the bi-level meta-learning problem in Eq (4) using an iterative gradient based algorithm of the form:

$$\theta \leftarrow \theta -\eta d_{\theta}F(\theta)$$

Specifically:

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

Note: many available iterative gradient based algorithm and other optimization methods could be used here:

• quasi-Newton or Newton methods

• Adam

• gradient descent with momentum

• etc.,

• $\nabla_{\phi} \mathcal{L}_{i}\left(\mathcal{A}lg _{i}^{\star}(\boldsymbol{\theta})\right)$ can be easily obtained in practice via automatic differentiation

• $\frac{d \mathcal{A}lg _{i}^{\star}(\boldsymbol{\theta})}{d \boldsymbol{\theta}}$ presents the primary challenge. $\mathcal{A}lg_i^{\star}(\theta)$ is implicitly defined as an optimization problem in Equ.4.

Meta-Gradient Computation

In theory

Theoretically we can calculate the meta-gradient computation $\frac{d \mathcal{A}lg _{i}^{\star}(\boldsymbol{\theta})}{d \boldsymbol{\theta}}$ exactly using the following lemma.

lemma 1: (Implicit Jacobian) Consider $\mathcal{A}lg_i^{\star}(\theta)$ as defined in Eq.4 for task $\mathcal{T}_i$. Let $\phi_i=\mathcal{A}lg_i^{\star}(\theta)$ be the results of $\mathcal{A}lg_i^{\star}(\theta)$. If $(I+\frac{1}{\lambda}\nabla^2_{\phi}\hat{\mathcal{L}_i}(\phi_i) )$ is invertible, then the derivative Jacobian is

$\frac{d \mathcal{A}lg _{i}^{\star}(\boldsymbol{\theta})}{d \boldsymbol{\theta}} = (I+\frac{1}{\lambda}\nabla^2_{\phi}\hat{\mathcal{L}_i}(\phi_i) )^{-1}$ (6)

Proof: We drop $i$ subscripts in the proof for convenience.

$\phi$ is the minimizer of $G(\phi',\theta)$ , namely:

$\phi = \mathcal{A}lg^{\star}(\boldsymbol{\theta}):=\underset{\boldsymbol{\phi}^{\prime} \in \Phi}{\operatorname{argmin}} G\left(\boldsymbol{\phi}^{\prime}, \boldsymbol{\theta}\right), \text { where } \phi = G\left(\boldsymbol{\phi}^{\prime}, \boldsymbol{\theta}\right)=\hat{\mathcal{L}}\left(\boldsymbol{\phi}^{\prime}\right)+\frac{\lambda}{2}\left\|\boldsymbol{\phi}^{\prime}-\boldsymbol{\theta}\right\|^{2}$

According to the stationary point conditions, we have:

$\nabla_{\phi'}G(\phi',\theta)|_{\phi'=\phi}=0 \\ \implies \nabla_{\phi'}(\hat{\mathcal{L}}(\phi')+\frac{\lambda}{2}\left\|\phi'-\theta \right\|^2 )|_{\phi'=\phi}=\nabla\hat{\mathcal{L}}(\phi)+\lambda(\phi-\theta) \\ \implies \phi=\theta-\frac{1}{\lambda}\nabla\hat{\mathcal{L}}(\theta)$

which is an implicit equation.

When the derivative exists:

$\frac{d\phi}{d\theta}=I-\frac{1}{\lambda}\nabla^2(\hat{\mathcal{L}}(\theta)) \frac{d\phi}{d\theta}\\ \implies (I+\frac{1}{\lambda}\nabla_2\hat{\mathcal{L}}(\theta)) \frac{d\phi}{d\theta}=I \\ \implies \frac{d\phi}{d\theta} = (I+\frac{1}{\lambda}\nabla_2\hat{\mathcal{L}}(\theta))^{-1}$

Implicit Jacobian

In practice

Two issues of theory solution in practice:

• The meta-gradients require computation of $\mathcal{A}lg_{i}^{\star}(\theta),$ which is the exact solution to the inner optimization problem. Only approximation could be obtained in practice.

• Explicitly forming and inverting the matrix Eq.6 for computing the Jacobian may be intractable for large deep learning network.

1), we consider an approximate solution to the inner optimization problem, that can be obtained with iterative optimization algorithms like gradient descent. Red

2), we will perform a partial or approximate matrix inversion. Green

Some Questions:

Q1

Use Figure 1 to explain the differences between MAML, first-order MAML, and implicit MAML. Appendix A might be helpful for this.

I will write a new answer for this question independently later including intuition and math details.

Q2

In Section 3.1, it talks about the high memory cost when the number of gradient steps is large for MAML. What is the memory cost with respect to the number of gradient steps?

Using iterative algorithm (Gradient Descent) for the optimization of inner loop has drawback: depending explicitly on the path of the optimization, which has to be fully stored in memory, quickly becoming intractable when the number of gradient steps needed is large.

According to Theorem 1: Algorithm 2 can be implemented using at most

• $\tilde{O}\left(\sqrt{\kappa} \log \left(\frac{poly(\kappa, D, B, L, \rho, \mu, \lambda)}{\epsilon}\right)\right)$ gradient computations of $\hat{\mathcal{L}}_i(\cdot)$

• $2\cdot Mem(\nabla\hat{\mathcal{L}}_i)$ memory.

Detailed proof will be given later.

Q3

Please provide a review of the conjugate gradient algorithm that is used in Section 3.1.

Another independent tutorial of conjugate gradient will be given.

Q4

Please explain why $g_i$ can be obtained as an approximate solution to Problem (7).

According to definition 2:

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

$\boldsymbol{g}_{i}$ is an approximation of the meta-gradient for task $i$.

In order to solve the following problem:

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

we have:

$\frac{df(\boldsymbol{w})}{d\boldsymbol{w}} = \left(I+\frac{1}{\lambda}\nabla^2_{\phi}\hat{\mathcal{L}_i}(\phi_i)\right)\boldsymbol{w}-\nabla_{\phi}\mathcal{L}_i(\phi_i)=0 \\ \implies \boldsymbol{w} = \left(I+\frac{1}{\lambda}\nabla^2_{\phi}\hat{\mathcal{L}_i}(\phi_i)\right)^{-1}\nabla_{\phi}\mathcal{L}_i(\phi_i)$

So $\boldsymbol{g}_i$ can be obtained as an approximate solution to optimization problem (7).

Q5

About Line 3 in Algorithm 2, what iterative optimization solver can be used and how many iterations are enough such that the error $\leq \delta$ .

Potential possible iterative optimization solver here could be:

• gradient descent

• Nesterov's accelerated gradient descent

If Nesterov's accelerated gradient descent algorithm is used to compute $\phi$, the number of iterations could be: (Theorem 2)

$2\sqrt{\kappa}\log\left(8\kappa D\left(\frac{B_1}{\epsilon} + \frac{\rho}{\mu}\right)\right)$

Q6

What is the relationships between $g_i$ in Definition 2 and the derivative Jacobian?

According to definition 2:

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

$\boldsymbol{g}_{i}$ is an approximation of the meta-gradient for task $i$.

$\left(\boldsymbol{I}+\frac{1}{\lambda} \nabla_{\boldsymbol{\phi}}^{2} \hat{\mathcal{L}}_{i}\left(\boldsymbol{\phi}_{i}\right)\right)^{-1}$ is the derivative Jacobian.

Reference:

• https://www.inference.vc/notes-on-imaml-meta-learning-without-differentiating-through/

• https://www.youtube.com/watch?v=u5BkO8XMS2I