Meta-Learning with Adjoint Methods

TL;DR

This paper introduces Adjoint MAML, a novel method that uses adjoint ODE techniques to efficiently compute meta-learning gradients, reducing computational costs and improving accuracy for long training trajectories.

Contribution

The paper proposes Adjoint MAML, which applies adjoint methods to meta-learning, enabling efficient gradient computation without expanding computational graphs or approximations.

Findings

Efficient gradient computation for long trajectories

Reduced computational cost compared to traditional MAML

Effective on synthetic and real-world tasks

Abstract

Model Agnostic Meta Learning (MAML) is widely used to find a good initialization for a family of tasks. Despite its success, a critical challenge in MAML is to calculate the gradient w.r.t. the initialization of a long training trajectory for the sampled tasks, because the computation graph can rapidly explode and the computational cost is very expensive. To address this problem, we propose Adjoint MAML (A-MAML). We view gradient descent in the inner optimization as the evolution of an Ordinary Differential Equation (ODE). To efficiently compute the gradient of the validation loss w.r.t. the initialization, we use the adjoint method to construct a companion, backward ODE. To obtain the gradient w.r.t. the initialization, we only need to run the standard ODE solver twice -- one is forward in time that evolves a long trajectory of gradient flow for the sampled task; the other is backward and solves the adjoint ODE. We need not create or expand any intermediate computational graphs, adopt aggressive approximations, or impose proximal regularizers in the training loss. Our approach is cheap, accurate, and adaptable to different trajectory lengths. We demonstrate the advantage of our approach in both synthetic and real-world meta-learning tasks.