Meta-Learning with Implicit Gradients

Aravind Rajeswaran; Chelsea Finn; Sham Kakade; Sergey Levine

Source-linked AI summary

Meta-Learning with Implicit Gradients

Aravind Rajeswaran, Chelsea Finn, Sham Kakade, Sergey Levine

arXiv:1909.04630v1Sep 10, 2019cs.LGcs.AImath.OCstat.ML

TL;DR

Optimization-based meta-learning struggles with the cost and instability of differentiating through long inner-loop paths. The paper introduces implicit MAML, which uses implicit differentiation at the inner solution, and reports improved efficiency and competitive or better few-shot recognition results.

Problem
Differentiating through inner optimization imposes computational and memory burdens, can require difficult derivatives, and becomes problematic for long or non-differentiable optimization paths.
Method
Implicit MAML computes an analytical meta-gradient from the inner optimization solution, decoupling meta-gradient computation from the inner optimizer.
Results
iMAML is competitive with full MAML on Omniglot and performs better than MAML and FOMAML on Mini-ImageNet.
Takeaways & Limitations
The approach provides compute and memory efficiency while allowing a variety of inner optimization methods, including higher-order methods.
Takeaways & Limitations
The theoretical guarantees assume smoothness, Lipschitz continuity, and strong convexity of the inner problem, while practical solutions may be approximate.

Abstract

from arXiv · show

A core capability of intelligent systems is the ability to quickly learn new tasks by drawing on prior experience. Gradient (or optimization) based meta-learning has recently emerged as an effective approach for few-shot learning. In this formulation, meta-parameters are learned in the outer loop, while task-specific models are learned in the inner-loop, by using only a small amount of data from the current task. A key challenge in scaling these approaches is the need to differentiate through the inner loop learning process, which can impose considerable computational and memory burdens. By drawing upon implicit differentiation, we develop the implicit MAML algorithm, which depends only on the solution to the inner level optimization and not the path taken by the inner loop optimizer. This effectively decouples the meta-gradient computation from the choice of inner loop optimizer. As a result, our approach is agnostic to the choice of inner loop optimizer and can gracefully handle many gradient steps without vanishing gradients or memory constraints. Theoretically, we prove that implicit MAML can compute accurate meta-gradients with a memory footprint that is, up to small constant factors, no more than that which is required to compute a single inner loop gradient and at no overall increase in the total computational cost. Experimentally, we show that these benefits of implicit MAML translate into empirical gains on few-shot image recognition benchmarks.

1 Introduction

Optimization-based meta-learning enables rapid adaptation, but differentiating through long inner-loop optimization paths creates computational, memory, and gradient problems. Implicit MAML removes this path dependence by computing meta-gradients from the inner solution, improving scalability and optimizer flexibility.

Meta-learning uses prior experience across tasks to learn models that adapt to new tasks from very little data.
Differentiating through inner optimization requires higher-order derivatives, substantial computation and memory, and can produce vanishing gradients.
Implicit MAML computes the meta-gradient from the inner optimization solution rather than its optimization path.
Implicit MAML decouples meta-gradient computation from the inner optimizer, permitting higher-order and potentially non-differentiable optimization methods.
The method avoids storing or differentiating through the inner path, making it memory-efficient for many inner optimization steps.
˜O(log(1/ϵ)) gradient evaluations and ˜O(1) memory suffice for an ϵ–approximate meta-gradient.

2 Problem Formulation and Notations

The paper formulates few-shot meta-learning as a bi-level problem: inner adaptation produces task-specific parameters, while outer optimization learns meta-parameters. Proximal regularization preserves dependence on the initialization and supports implicit differentiation, whose total and partial derivatives must be distinguished.

2.1 Review of Few-Shot Supervised Learning and MAML: Few-shot supervised learning uses task datasets containing input-output pairs to learn models h_φ(x): X → Y and minimize task loss.
2.1 Review of Few-Shot Supervised Learning and MAML: Meta-learning learns meta-parameters θ that an adaptation algorithm applies to each task's training data to produce task-specific parameters φ_i.
2.1 Review of Few-Shot Supervised Learning and MAML: The bi-level formulation treats the adaptation algorithm as solving an underlying optimization problem, with deployment applying it to a new task's training dataset.
2.1 Review of Few-Shot Supervised Learning and MAML: MAML uses one or more gradient-descent steps initialized at θ, requiring differentiation through the resulting optimization dynamics.
2.2 Proximal Regularization in the Inner Level: Many inner-loop steps require storing and differentiating through a long path, while dependence on θ can shrink and vanish as optimization proceeds.
2.2 Proximal Regularization in the Inner Level: Proximal regularization keeps φ_i close to θ, with λ controlling the relative strength of the prior and task data.
2.2 Proximal Regularization in the Inner Level: The regularized inner problem is ideally solved exactly, but practical iterative algorithms return approximate minimizers whose discrepancy is analyzed explicitly.
2.4 Total and Partial Derivatives: Total derivatives propagate through the adaptation algorithm, whereas partial gradients evaluate the loss gradient without differentiating through that algorithm.

3 The Implicit MAML Algorithm

Implicit MAML computes meta-gradients from the final inner-level solution rather than the optimization path, decoupling meta-gradient computation from the inner optimizer. Its practical algorithm uses approximate inner solutions, iterative Hessian-vector-product solvers, and theoretical guarantees for accuracy, computation, and memory.

Implicit meta-gradient: Implicit differentiation computes the meta-gradient from the final inner-level solution, avoiding differentiation through the optimization path.The resulting Jacobian depends on the final result rather than the path used to obtain it.
Implicit meta-gradient: The approach can use iterative inner optimization methods and is not restricted to a particular optimizer, provided they find an approximate inner-level solution.The paper discusses approximate solutions and optimization methods including gradient-based, quasi-Newton, Newton, Adam, and momentum methods.
Practical algorithm: The practical algorithm obtains task parameters to accuracy δ, computes the partial outer gradient, and uses an iterative solver with reverse-mode differentiation to approximate each task meta-gradient.Conjugate gradient is used with Hessian-vector products rather than explicitly forming or storing the Hessian.
Theory: The method’s approximation analysis assumes smooth outer losses, strongly convex regularized inner objectives, and controlled Hessian regularity, with κ measuring inner-problem conditioning.Theorem 1 provides an accuracy guarantee under these assumptions, while the analysis does not require the task loss to be convex.
Theory: For an ϵ-accurate meta-gradient, iMAML uses logarithmic dependence on 1/ϵ in gradient evaluations and memory proportional to a single inner gradient up to constant factors.The stated memory bound is at most 2 · Mem(∇ˆLi), and the method can match the path-based approach’s minimax computational complexity while using substantially less memory.

4 Experimental Results and Discussion

The experiments evaluate iMAML's meta-gradient accuracy, compute and memory use, and few-shot performance against MAML, FOMAML, and Reptile. Across synthetic and image-recognition tasks, iMAML provides accurate gradients with lower memory requirements and competitive or better performance.

Meta-gradient accuracy: Both iMAML and MAML asymptotically match the exact meta-gradient, while iMAML is more accurate with finite iterations.With two conjugate-gradient iterations, iMAML has a small terminal error that vanishes with five iterations.
Computation and memory: 100 inner gradient steps give iMAML significantly lower computation and memory costs than MAML.One conjugate-gradient step has computational cost comparable to one MAML inner gradient step because both require a Hessian-vector product.
Computation and memory: On 20-way-5-shot Omniglot, iMAML memory is independent of inner-loop gradient steps and conjugate-gradient iterations, whereas MAML memory grows linearly and reaches 12 GB after approximately 16 steps.iMAML has a constant conjugate-gradient overhead in computation, while FOMAML has lower cost but does not compute an accurate meta-gradient.
Few-shot benchmarks: The evaluation compares iMAML with prior methods on Omniglot and Mini-ImageNet using identical convolutional architectures and few-shot protocols.The study includes different N-way, K-shot settings and compares GD and Hessian-free iMAML variants.
Few-shot benchmarks: On Omniglot, GD iMAML is competitive with MAML and better than FOMAML and Reptile, while Hessian-free iMAML performs substantially better than the other methods.The reported comparisons are especially favorable on harder 20-way tasks.
Few-shot benchmarks: On Mini-ImageNet, iMAML performs better than MAML and FOMAML.The experiments use λ = 0.5, 10 inner-loop gradient steps, and five conjugate-gradient steps; the authors did not perform an extensive hyperparameter sweep.

5 Related Work

The related work situates iMAML within optimization-based meta-learning and implicit-gradient methods. Its distinction is finite-time analysis for approximately solved inner problems rather than only exact inner-level solutions.

Meta-learning approaches: Meta-learning methods include metric-learning, black-box, optimization-based, and hybrid approaches.Optimization-based methods formulate learning procedures such as gradient descent within a bi-level meta-optimization problem.
Optimization-based methods: iMAML builds on optimization-based meta-learning, particularly MAML, which meta-learns an initial parameter set.Prior work also studies implicit gradients and methods that facilitate backpropagation through optimization paths.
Implicit differentiation: Prior implicit-differentiation methods assume an exact inner-level solution and therefore provide only asymptotic guarantees.The paper contrasts this with practical iterative optimizers that return approximate solutions after finite iterations.
Implicit differentiation: iMAML provides finite-time guarantees for approximately solved inner problems.This analysis addresses the finite-iteration setting associated with practical methods such as gradient descent.

6 Conclusion

The paper concludes that implicit differentiation removes dependence on the inner optimization path and separates outer meta-gradient computation from the inner optimizer. It reports efficiency gains and identifies broader optimizer classes and richer regularization as future directions.

Conclusion: iMAML removes the need to differentiate through the inner optimization path and decouples meta-gradient computation from the inner optimization algorithm.The paper presents this as the basis for compute and memory efficiency gains.
Future directions: iMAML can in principle support inner-loop procedures including Q-learning, GANs, energy-based models, actor-critic reinforcement learning, and trajectory optimization.These examples extend beyond the gradient-based optimization methods studied experimentally.
Future directions: Learning vector- or matrix-valued λ is proposed as an extension for co-adapting and co-regularizing different model parameters.The paper's experiments used simple ℓ2 regularization toward the parameter initialization.

A Relationship between iMAML and Prior Algorithms

This section relates iMAML to MAML, FOMAML, and Reptile, emphasizing differences in meta-gradient computation and the role of explicit regularization. It also reviews the Hessian-vector products used for iMAML's complexity guarantees.

Regularization: Unlike conventional early-stopped formulations, iMAML uses an explicitly regularized inner-level problem.The regularizer is discussed as distinct from the implicit regularization produced by taking only a few gradient steps.
MAML and iMAML: MAML and iMAML target the same idealized update, but MAML backpropagates through the optimization path while iMAML uses an implicit Jacobian.iMAML is therefore more memory-efficient and can accommodate higher-order optimizers and non-differentiable components.
First-order MAML: FOMAML ignores the dependence of task parameters on meta-parameters, whereas iMAML retains this dependence through conjugate-gradient computation.With zero conjugate-gradient steps, iMAML reduces to FOMAML and interpolates toward full MAML as more steps are used.
Reptile: Reptile and FOMAML are identical in this explicitly regularized formulation up to the learning-rate choice.The equivalence follows from the proximal-point relation between task parameters and task-loss gradients.
Complexity assumptions: For smooth strongly convex objectives, accelerated gradient descent supplies a finite-iteration approximate inner solver with complexity governed by the condition number.The section introduces Lipschitz, smoothness, strong-convexity, and first-order complexity assumptions for the analysis.
Hessian-vector products: Reverse-mode automatic differentiation computes Hessian-vector products through gradients of scalar gradient-vector products.The stated assumptions bound Hessian-vector-product computation by a constant factor of gradient computation and memory by at most twice the gradient memory.

D Additional Discussion About Compute and Memory Complexity

The complexity analysis distinguishes exact-solve error from approx-solve error and compares how MAML, truncated backpropagation, and iMAML compute meta-gradients. iMAML provides exact-solve guarantees without requiring an exact inner solver and uses memory comparable to a single gradient computation.

Error definitions: The analysis distinguishes exact-solve error from approx-solve error when evaluating computed meta-gradients.Exact-solve error concerns the gradient of F(θ), whereas approx-solve error concerns gradients associated with a δ-accurate inner solution.
Comparison with MAML: MAML has zero approx-solve error because it backpropagates through the entire inner optimization path, but its exact-solve error is not known.The exact-solve guarantee is unclear even as δ approaches zero without additional regularity conditions.
Comparison with truncated backpropagation: Truncated backpropagation can reduce memory while retaining an ϵ-accurate approximation of the full-backpropagation gradient, but its accuracy with accelerated methods is unclear.The cited guarantee applies to the full-backpropagation gradient and does not establish the same accuracy for accelerated inner optimization.
iMAML guarantees: iMAML guarantees small exact-solve error without requiring access to the exact inner solver.Its analysis explicitly separates this guarantee from the use of an approximate inner optimization algorithm.
Memory complexity: iMAML uses memory comparable, up to small constant factors, to computing a single gradient of the inner objective.The memory requirement does not scale with the number of inner-loop gradient steps.

E Proofs

The proofs establish implicit-gradient accuracy under regularity conditions, using the inner optimum's stationary equation and iterative procedures for computing both the inner solution and implicit gradient.

Implicit differentiation: The inner optimum is characterized by a stationary equation that can be differentiated to obtain the implicit Jacobian.The resulting equation is noted as one that often arises in proximal point methods.
Assumptions: The analysis assumes each task loss is Lipschitz and smooth, while the regularized inner objective is smooth and strongly convex.These conditions support uniqueness, stability, and the subsequent implicit-gradient bounds.
Gradient accuracy: Lemma 3 bounds implicit-gradient error when the inner solution and computed gradient satisfy their respective accuracy conditions.The proof decomposes the error and uses smoothness, strong convexity, and a Lipschitz Hessian.
Proof completion: The proof concludes by substituting the established bounds into the target error guarantee.The intermediate proof steps complete the bounds used by the theorem.
Iteration complexity: Theorem 2 specifies iteration requirements for Nesterov acceleration or conjugate gradients to compute the inner solution and implicit gradient accurately.The proof sets separate tolerances for inner-solution error and gradient-subproblem error.

F Experiment Details

The experiments were conducted on a single NVIDIA Titan Xp GPU.

Hardware: All training runs used a single NVIDIA Titan Xp GPU.The statement describes the hardware used for the experimental setup.

F.1 Synthetic Experiments

The synthetic experiments use linear regression with quadratic task objectives and compare exact and approximate meta-gradients from iMAML and MAML under a difficult conditioning setting.

Synthetic setup: The synthetic tasks use linear predictors hφ(x) = φT x, with raw inputs or feature representations such as Fourier features.Each task can be expressed with a quadratic objective.
Analytic construction: The inner-level minimizer has a closed-form expression involving the task matrix, regularization, and task vector.The exact meta-gradient is then written from this minimizer.

F.2 Omniglot and Mini-ImageNet experiments

The experiments follow prior standard training and evaluation protocols and compare gradient-descent and Hessian-free variants across Omniglot and Mini-ImageNet settings.

Training and evaluation followed the standard protocol used in prior works.
Omniglot Experiments: Omniglot's gradient-descent iMAML used 16 or 25 inner-loop steps, λ = 2.0, and five conjugate-gradient steps for meta-gradient computation.The 16-step setting covered 5-way tasks, while 20-way tasks used 25 steps; Adam averaged task meta-gradients in the outer loop.
Omniglot Experiments: The Omniglot Hessian-free MAML variant used a local quadratic objective approximation and conjugate gradient to approximately compute the Newton direction.Five conjugate-gradient steps were used to compute the search direction in the experiments.
Mini-ImageNet: On Mini-ImageNet, gradient-descent iMAML used 10 inner-loop steps with λ = 0.5 and five conjugate-gradient steps for the meta-gradient.The meta-gradients were computed with conjugate gradient after the inner-loop gradient-descent steps.
Mini-ImageNet: On Mini-ImageNet, the Hessian-free variant used five conjugate-gradient steps per search direction, line search, and three repetitions to solve the inner optimization.Five additional conjugate-gradient steps were used to compute the meta-gradient.