> For the complete documentation index, see [llms.txt](https://lichangbin.gitbook.io/paper_notes/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://lichangbin.gitbook.io/paper_notes/meta-learning-with-implicit-gradient.md). # Meta-Learning with Implicit Gradient ## 8/18/2020 ## 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.](/files/-MF8z_YBtXgQpcIW6UBe) #### 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: ![](/files/-MF99ZAgcf7njQhze9s-) ## 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)](/files/-MF9HE5HjDSsIDZ2Uavh) #### 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. {% hint style="info" %} 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) {% endhint %} {% hint style="success" %} 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}$$ {% endhint %} 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** ![](/files/-MFCpAHjl-k0lzZ7MtfZ) ## Some Questions: {% tabs %} {% tab title="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. {% endtab %} {% tab title="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. {% endtab %} {% tab title="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. {% endtab %} {% tab title="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). {% endtab %} {% tab title="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)$$ {% endtab %} {% tab title="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. {% endtab %} {% endtabs %} **Reference:** * *