# Continuous-Time Meta-Learning with
Forward Mode Differentiation

Tristan Deleu David Kanaa Leo Feng Giancarlo Kerg
Yoshua Bengio 1,2 Guillaume Lajoie 2 Pierre-Luc Bacon 2
Mila – Université de Montréal Correspondence to: Tristan Deleu
<EMAIL_ADDRESS>
1CIFAR Senior Fellow, 2CIFAR AI Chair
Code is available at: https://github.com/tristandeleu/jax-comln

###### Abstract

Drawing inspiration from gradient-based meta-learning methods with infinitely
small gradient steps, we introduce Continuous-Time Meta-Learning (COMLN), a
meta-learning algorithm where adaptation follows the dynamics of a gradient
vector field. Specifically, representations of the inputs are meta-learned
such that a task-specific linear classifier is obtained as a solution of an
ordinary differential equation (ODE). Treating the learning process as an ODE
offers the notable advantage that the length of the trajectory is now
continuous, as opposed to a fixed and discrete number of gradient steps. As a
consequence, we can optimize the amount of adaptation necessary to solve a new
task using stochastic gradient descent, in addition to learning the initial
conditions as is standard practice in gradient-based meta-learning.
Importantly, in order to compute the exact meta-gradients required for the
outer-loop updates, we devise an efficient algorithm based on forward mode
differentiation, whose memory requirements do not scale with the length of the
learning trajectory, thus allowing longer adaptation in constant memory. We
provide analytical guarantees for the stability of COMLN, we show empirically
its efficiency in terms of runtime and memory usage, and we illustrate its
effectiveness on a range of few-shot image classification problems.

##  1 Introduction

Among the existing meta-learning algorithms, gradient-based methods as
popularized by Model-Agnostic Meta-Learning (MAML, Finn et al., 2017) have
received a lot of attention over the past few years. They formulate the
problem of learning a new task as an inner optimization problem, typically
based on a few steps of gradient descent. An outer _meta_ -optimization
problem is then responsible for updating the meta-parameters of this learning
process, such as the initialization of the gradient descent procedure. However
since the updates at the outer level typically require backpropagating through
the learning process, this class of methods has often been limited to only a
few gradient steps of adaptation, due to memory constraints. Although
solutions have been proposed to alleviate the memory requirements of these
algorithms, including checkpointing (Baranchuk, 2019), using implicit
differentiation (Rajeswaran et al., 2019), or reformulating the meta-learning
objective (Flennerhag et al., 2018), they are generally either more
computationally demanding, or only approximate the gradients of the meta-
learning objective (Nichol et al., 2018; Flennerhag et al., 2020).

In this work, we propose a continuous-time formulation of gradient-based meta-
learning, called _Continuous-Time Meta-Learning_ (COMLN), where the adaptation
is the solution of a differential equation (see Fig. 1). Moving to continuous
time allows us to devise a novel algorithm, based on forward mode
differentiation, to efficiently compute the exact gradients for meta-
optimization, no matter how long the adaptation to a new task might be. We
show that using forward mode differentiation leads to a stable algorithm,
unlike the counterpart of backpropagation in continuous time called the
adjoint method (frequently used in the Neural ODE literature; Chen et al.,
2018) which tends to be unstable in conjunction with gradient vector fields.
Moreover as the length of the adaptation trajectory is a continuous quantity,
as opposed to a discrete number of gradient steps fixed ahead of time, we can
treat the amount of adaptation in COMLN as a meta-parameter—on par with the
initialization—which we can meta-optimize using stochastic gradient descent.
We verify empirically that our method is both computationally and memory
efficient, and we show that COMLN outperforms other standard meta-learning
algorithms on few-shot image classification datasets.

##  2 Background

In this work, we consider the problem of few-shot classification, that is the
problem of learning a classification model with only a small number of
training examples. More precisely for a classification task $\tau$, we assume
that we have access to a (small) training dataset
$\mathcal{D}_{\tau}^{\mathrm{train}}=\\{({\bm{x}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
to fit a model on task $\tau$, and a distinct test dataset
$\mathcal{D}_{\tau}^{\mathrm{test}}$ to evaluate how well this adapted model
generalizes on that task. In the few-shot learning literature, it is standard
to consider the problem of $k$-shot $N$-way classification, meaning that the
model has to classify among $N$ possible classes, and there are only $k$
examples of each class in $\mathcal{D}_{\tau}^{\mathrm{train}}$, so that
overall the number of training examples is $M=kN$. We use the convention that
the target labels ${\bm{y}}_{m}\in\\{0,1\\}^{N}$ are one-hot vectors.

figures/comln

Figure 1: Illustration of the adaptation process in (a) a gradient-based meta-
learning algorithm, such as ANIL (Raghu et al., 2019), where the adapted
parameters ${\bm{W}}_{T}$ are given after $T$ steps of gradient descent, and
in (b) Continuous-Time Meta-Learning (COMLN), where the adapted parameters
${\bm{W}}(T)$ are the result of following the dynamics of a gradient vector
field up to time $T$.

###  2.1 Gradient-based meta-learning

Gradient-based meta-learning methods aim to learn an initialization such that
the model is able to adapt to a new task via gradient descent. Such methods
are often cast as a bi-level optimization process: adapting the task-specific
parameters ${\bm{\theta}}$ in the inner loop, and training the (task-agnostic)
meta-parameters $\Phi$ and initialization ${\bm{\theta}}_{0}$ in the outer
loop. The meta-learning objective is:

$\displaystyle\min_{{\bm{\theta}}_{0},\Phi}\ \ $
$\displaystyle\mathbb{E}_{\tau}\big{[}{\mathcal{L}}({\bm{\theta}}^{\tau}_{T},\Phi;\mathcal{D}_{\tau}^{\mathrm{test}})\big{]}$
(1) s.t.
$\displaystyle{\bm{\theta}}^{\tau}_{t+1}={\bm{\theta}}^{\tau}_{t}-\alpha\nabla_{{\bm{\theta}}}{\mathcal{L}}({\bm{\theta}}^{\tau}_{t},\Phi;\mathcal{D}_{\tau}^{\mathrm{train}})$
$\displaystyle{\bm{\theta}}_{0}^{\tau}$ $\displaystyle={\bm{\theta}}_{0}$
$\displaystyle\forall\tau\sim p(\tau),$ (2)

where $T$ is the number of inner loop updates. For example, in the case of
MAML (Finn et al., 2017), there is no additional meta-parameter other than the
initialization ($\Phi\equiv\emptyset$); in ANIL (Raghu et al., 2019),
${\bm{\theta}}$ are the parameters of the last layer, and $\Phi$ are the
parameters of the shared embedding network; in CAVIA (Zintgraf et al., 2019),
${\bm{\theta}}$ are referred to as context parameters.

During meta-training, the model is trained over many tasks $\tau$. The task-
specific parameters ${\bm{\theta}}$ are learned via gradient descent on
$\mathcal{D}_{\tau}^{\mathrm{train}}$. The meta-parameters are then updated by
evaluating the error of the trained model on the test dataset
$\mathcal{D}_{\tau}^{\mathrm{test}}$. At meta-testing time, the meta-trained
model is adapted on $\mathcal{D}_{\tau}^{\mathrm{train}}$—i.e. applying Eq. 2
with the learned meta-parameters ${\bm{\theta}}_{0}$ and $\Phi$.

###  2.2 Local sensitivity analysis of Ordinary Differential Equations

Consider the following (autonomous) Ordinary Differential Equation (ODE):

$\displaystyle\frac{d{\bm{z}}}{dt}$
$\displaystyle=g\big{(}{\bm{z}}(t);{\bm{\theta}}\big{)}$
$\displaystyle{\bm{z}}(0)$ $\displaystyle={\bm{z}}_{0}({\bm{\theta}}),$ (3)

where the dynamics $g$ and initial value ${\bm{z}}_{0}$ may depend on some
external parameters ${\bm{\theta}}$, and integration is carried out from $0$
to some time $T$. _Local sensitivity analysis_ is the study of how the
solution of this dynamical system responds to local changes in
${\bm{\theta}}$; this effectively corresponds to calculating the derivative
$d{\bm{z}}(t)/d{\bm{\theta}}$. We present here two methods to compute this
derivative, with a special focus on their memory efficiency.

#### Adjoint sensitivity method

Based on the adjoint state (Pontryagin, 2018), and taking its root in control
theory (Lions & Magenes, 2012), the _adjoint sensitivity method_ (Bryson & Ho,
1969; Chavent et al., 1974) provides an efficient approach for evaluating
derivatives of ${\mathcal{L}}\big{(}{\bm{z}}(T);{\bm{\theta}}\big{)}$, a
function of ${\bm{z}}(T)$ the solution of the ODE in Eq. 3. This method,
popularized lately by the literature on Neural ODEs (Chen et al., 2018),
requires the integration of the adjoint equation

$\displaystyle\frac{d{\bm{a}}}{dt}$ $\displaystyle=-{\bm{a}}(t)\frac{\partial
g\big{(}{\bm{z}}(t);{\bm{\theta}}\big{)}}{\partial{\bm{z}}(t)}$
$\displaystyle{\bm{a}}(T)$
$\displaystyle=\frac{d{\mathcal{L}}\big{(}{\bm{z}}(T);{\bm{\theta}}\big{)}}{d{\bm{z}}(T)},$
(4)

backward in time. The adjoint sensitivity method can be viewed as a
continuous-time counterpart to backpropagation, where the forward pass would
correspond to integrating Eq. 3 forward in time from $0$ to $T$, and the
backward pass to integrating Eq. 4 backward in time from $T$ to $0$.

One possible implementation, reminiscent of backpropagation through time
(BPTT), is to store the intermediate values of ${\bm{z}}(t)$ during the
forward pass, and reuse them to compute the adjoint state during the backward
pass. While several sophisticated checkpointing schemes have been proposed
(Serban & Hindmarsh, 2003; Gholami et al., 2019), with different
compute/memory trade-offs, the memory requirements of this approach typically
grow with $T$; this is similar to the memory limitations standard gradient-
based meta-learning methods suffer from as the number of gradient steps
increases. An alternative is to augment the adjoint state ${\bm{a}}(t)$ with
the original state ${\bm{z}}(t)$, and to solve this augmented dynamical system
backward in time (Chen et al., 2018). This has the notable advantage that the
memory requirements are now independent of $T$, since ${\bm{z}}(t)$ are no
longer stored during the forward pass, but they are recomputed on the fly
during the backward pass.

#### Forward sensitivity method

While the adjoint method is related to reverse-mode automatic differentiation
(backpropagation), the _forward sensitivity method_ (Feehery et al., 1997;
Leis & Kramer, 1988; Maly & Petzold, 1996; Caracotsios & Stewart, 1985), on
the other hand, can be viewed as the continuous-time counterpart to forward
(tangent-linear) mode differentiation (Griewank & Walther, 2008). This method
is based on the fact that the derivative ${\mathcal{S}}(t)\triangleq
d{\bm{z}}(t)/d{\bm{\theta}}$ is the solution of the so-called _forward
sensitivity equation_

$\displaystyle\frac{d{\mathcal{S}}}{dt}$ $\displaystyle=\frac{\partial
g\big{(}{\bm{z}}(t);{\bm{\theta}}\big{)}}{\partial{\bm{z}}(t)}{\mathcal{S}}(t)+\frac{\partial
g\big{(}{\bm{z}}(t);{\bm{\theta}}\big{)}}{\partial{\bm{\theta}}}$
$\displaystyle{\mathcal{S}}(0)$
$\displaystyle=\frac{\partial{\bm{z}}_{0}}{\partial{\bm{\theta}}}.$ (5)

This equation can be found throughout the literature in optimal control and
system identification (Betts, 2010; Biegler, 2010). Unlike the adjoint method,
which requires an explicit forward _and_ backward pass, the forward
sensitivity method only requires the integration forward in time of the
original ODE in Eq. 3, augmented by the sensitivity state ${\mathcal{S}}(t)$
with the dynamics above. The memory requirements of the forward sensitivity
method do not scale with $T$ either, but it now requires storing
${\mathcal{S}}(t)$, which may be very large; we will come back to this problem
in Section 3.2. We will simply note here that in discrete-time, this is the
same issue afflicting forward-mode training of RNNs with real-time recurrent
learning (RTRL; Williams & Zipser, 1989), or other meta-learning algorithms
(Sutton, 1992; Franceschi et al., 2017; Xu et al., 2018).

##  3 Continuous-time adaptation

In the limit of infinitely small steps, some optimization algorithms can be
viewed as the solution trajectory of a differential equation. This point of
view has often been taken to analyze their behavior (Platt & Barr, 1988;
Wilson et al., 2016; Su et al., 2014; Orvieto & Lucchi, 2019). In fact, some
optimization algorithms such as gradient descent with momentum have even been
introduced initially from the perspective of dynamical systems (Polyak, 1964).
As the simplest example, gradient descent with a constant step size
$\alpha\rightarrow 0^{+}$ (i.e. $\alpha$ tends to $0$ by positive values)
corresponds to following the dynamics of an autonomous ODE called a _gradient
vector field_

$\displaystyle{\bm{z}}_{t+1}$ $\displaystyle={\bm{z}}_{t}-\alpha\nabla
f({\bm{z}}_{t})$ $\displaystyle\underset{\alpha\rightarrow
0^{+}}{\longrightarrow}$ $\displaystyle\frac{d{\bm{z}}}{dt}$
$\displaystyle=-\nabla f\big{(}{\bm{z}}(t)\big{)},$ (6)

where the iterate ${\bm{z}}(t)$ is now a continuous function of time $t$. The
solution of this dynamical system is uniquely defined by the choice of the
initial condition ${\bm{z}}(0)={\bm{z}}_{0}$.

###  3.1 Continuous-time meta-learning

In gradient-based meta-learning, the task-specific adaptation with gradient
descent may also be replaced by a gradient vector field in the limit of
infinitely small steps. Inspired by prior work in meta-learning (Raghu et al.,
2019; Javed & White, 2019), we assume that an embedding network $f_{\Phi}$
with meta-parameters $\Phi$ is shared across tasks, and only the parameters
${\bm{W}}$ of a task-specific linear classifier are adapted, starting at some
initialization ${\bm{W}}_{0}$. Instead of being the result of a few steps of
gradient descent though, the final parameters ${\bm{W}}(T)$ now correspond to
integrating an ODE similar to Eq. 6 up to a certain horizon $T$, with the
initial conditions ${\bm{W}}(0)={\bm{W}}_{0}$. We call this new meta-learning
algorithm Continuous-Time Meta-Learning111COMLN is pronounced _chameleon_.
(COMLN).

Treating the learning algorithm as a continuous-time process has the notable
advantage that the adapted parameter ${\bm{W}}(T)$ is now differentiable wrt.
the time horizon $T$ (Wiggins, 2003, Chap. 7), in addition to being
differentiable wrt. the initial conditions ${\bm{W}}_{0}$—which plays a
central role in gradient-based meta-learning, as described in Section 2.1.
This allows us to view $T$ as a meta-parameter on par with $\Phi$ and
${\bm{W}}_{0}$, and to effectively optimize the amount of adaptation using
stochastic gradient descent (SGD). The meta-learning objective of COMLN can be
written as

$\displaystyle\min_{\Phi,{\bm{W}}_{0},T}\ \ $
$\displaystyle\mathbb{E}_{\tau}\big{[}{\mathcal{L}}\big{(}{\bm{W}}_{\tau}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}\big{]}$
(7) s.t.
$\displaystyle\frac{d{\bm{W}}_{\tau}}{dt}=-\nabla{\mathcal{L}}\big{(}{\bm{W}}_{\tau}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle{\bm{W}}_{\tau}(0)$ $\displaystyle={\bm{W}}_{0}$
$\displaystyle\forall\tau\sim p(\tau),$ (8)

where
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{(f_{\Phi}({\bm{x}}_{m}),{\bm{y}}_{m})\mid({\bm{x}}_{m},{\bm{y}}_{m})\in\mathcal{D}_{\tau}^{\mathrm{train}}\\}$
is the embedded training dataset, and
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})$ is defined similarly for
$\mathcal{D}_{\tau}^{\mathrm{test}}$. In practice, adaptation is implemented
using a numerical integration scheme based on an iterative discretization of
the problem, such as Runge-Kutta methods. Although a complete discussion of
numerical solvers is outside of the scope of this paper, we recommend
(Butcher, 2008) for a comprehensive overview of numerical methods for solving
ODEs.

###  3.2 The challenges of optimizing the meta-learning objective

${\bm{W}}(0)$${\bm{W}}(T)$${\bm{W}}^{\star}$ForwardBackward

Figure 2: Numerical instability of the adjoint method applied to the gradient
vector field of a quadratic loss function. The trajectory in green starting at
${\bm{W}}(0)$ corresponds to the integration of the dynamical system in Eq. 8
forward in time up to $T$, and the trajectory in red starting at ${\bm{W}}(T)$
corresponds to its integration backward in time. Note that $T$ was chosen so
that ${\bm{W}}(T)$ does not reach the equilibrium/minimum of the loss
${\bm{W}}^{\star}$.

In order to minimize the meta-learning objective of COMLN, it is common
practice to use (stochastic) gradient methods; that requires computing its
derivatives wrt. the meta-parameters, which we call _meta-gradients_. Our
primary goal is to devise an algorithm whose memory requirements do not scale
with the amount of adaptation $T$; this would contrast with standard gradient-
based meta-learning methods that backpropagate through a sequence of gradient
steps (similar to BPTT), where the intermediate parameters are stored during
adaptation (i.e. ${\bm{\theta}}^{\tau}_{t}$ for all $t$ in Eq. 2). Since this
objective involves the solution ${\bm{W}}(T)$ of an ODE, we can use either the
adjoint method, or the forward sensitivity method, in order to compute the
derivatives wrt. $\Phi$ and ${\bm{W}}_{0}$ (see Section 2.2).

Although the adjoint method has proven to be an effective strategy for
learning Neural ODEs, in practice computing the state ${\bm{W}}(t)$ backward
in time is numerically unstable when applied to a gradient vector field like
the one in Eq. 8, even for convex loss functions. Fig. 2 shows an example
where the trajectory of ${\bm{W}}(t)$ recomputed backward in time (in red)
diverges significantly from the original trajectory (in green) on a quadratic
loss function, even though the two should match exactly in theory since they
follow the same dynamics. Intuitively, recomputing ${\bm{W}}(t)$ backward in
time for a gradient vector field requires doing gradient _ascent_ on the loss
function, which is prone to compounding numerical errors; this is closely
related to the loss of entropy observed by Maclaurin et al. (2015). This
divergence makes the backward trajectory of ${\bm{W}}(t)$ unreliable to find
the adjoint state, ruling out the adjoint sensitivity method for computing the
meta-gradients in COMLN.

The forward sensitivity method addresses this shortcoming by avoiding the
backward pass altogether. However, it can also be particularly expensive here
in terms of memory requirements, since the sensitivity state
${\mathcal{S}}(t)$ in Section 2.2 now corresponds to Jacobian matrices, such
as $d{\bm{W}}(t)/d{\bm{W}}_{0}$. As the size $d$ of the feature vectors
returned by $f_{\Phi}$ may be very large, this $Nd\times Nd$ Jacobian matrix
would be almost impossible to store in practice; for example in our
experiments, it can be as large as $d=16,\\!000$ for a ResNet-12 backbone. In
Section 4.1, we will show how to apply forward sensitivity in a memory-
efficient way, by leveraging the structure of the loss function. This is
achieved by carefully decomposing the Jacobian matrices into smaller pieces
that follow specific dynamics. We show in Appendix D that unlike the adjoint
method, this process is stable.

###  3.3 Connection with Almost No Inner-Loop (ANIL)

Similarly to ANIL (Raghu et al., 2019), COMLN only adapts the parameters
${\bm{W}}$ of the last linear layer of the neural network. There is a deeper
connection between both algorithms though: while our description of the
adaptation in COMLN (Eq. 8) was independent of the choice of the ODE solver
used to find the solution ${\bm{W}}(T)$ in practice, if we choose an explicit
Euler scheme (Euler, 1913; roughly speaking, discretizing Eq. 6 from right to
left), then the adaptation of COMLN becomes functionally equivalent to ANIL.
However, this equivalence can greatly benefit from the memory-efficient
algorithm to compute the meta-gradients described in Section 4, based on the
forward sensitivity method. This means that using the methods devised here for
COMLN, we can effectively compute the meta-gradients of ANIL with a constant
memory cost wrt. the number of gradient steps of adaptation, instead of
relying on backpropagation (see also Section 4.2).

##  4 Memory-efficient meta-gradients

For some fixed task $\tau$ and
$({\bm{x}}_{m},{\bm{y}}_{m})\in\mathcal{D}_{\tau}^{\mathrm{train}}$, let
${\bm{\phi}}_{m}=f_{\Phi}({\bm{x}}_{m})\in{\mathbb{R}}^{d}$ be the embedding
of input ${\bm{x}}_{m}$ through the feature extractor $f_{\Phi}$. Since we are
confronted with a classification problem, the loss function of choice
${\mathcal{L}}({\bm{W}})$ is typically the cross-entropy loss. Böhning (1992)
showed that the gradient of the cross-entropy loss wrt. ${\bm{W}}$ can be
written as

$\displaystyle\nabla{\mathcal{L}}\big{(}{\bm{W}};f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=\frac{1}{M}\sum_{m=1}^{M}({\bm{p}}_{m}-{\bm{y}}_{m}){\bm{\phi}}_{m}^{\top},$
(9)

where ${\bm{p}}_{m}=\mathrm{softmax}({\bm{W}}{\bm{\phi}}_{m})$ is the vector
of probabilities returned by the neural network. The key observation here is
that the gradient can be decomposed as a sum of $M$ rank-one matrices, where
the feature vectors ${\bm{\phi}}_{m}$ are independent of ${\bm{W}}$. Therefore
we can fully characterize the gradient of the cross-entropy loss with $M$
vectors ${\bm{p}}_{m}-{\bm{y}}_{m}\in{\mathbb{R}}^{N}$, as opposed to the full
$N\times d$ matrix. This is particularly useful in the context of few-shot
classification, where the number of training examples $M$ is small, and
typically significantly smaller than the embedding size $d$.

###  4.1 Decomposition of the meta-gradients

We saw in Section 3.2 that the forward sensitivity method was the only stable
option to compute the meta-gradients of COMLN. However, naively applying the
forward sensitivity equation would involve quantities that typically scale
with $d^{2}$, which can be too expensive in practice. Using the structure of
Eq. 9, the Jacobian matrices appearing in the computation of the meta-
gradients for COMLN can be decomposed in such a way that only small quantities
will depend on time.

#### Meta-gradients wrt. ${\bm{W}}_{0}$

By the chain rule of derivatives, it is sufficient to compute the Jacobian
matrix $d{\bm{W}}(T)/d{\bm{W}}_{0}$ in order to obtain the meta-gradient wrt.
${\bm{W}}_{0}$. We show in Section B.2 that the sensitivity state
$d{\bm{W}}(t)/d{\bm{W}}_{0}$ can be decomposed as:

$\frac{d{\bm{W}}(t)}{d{\bm{W}}_{0}}={\bm{I}}-\sum_{i=1}^{M}\sum_{j=1}^{M}{\bm{B}}_{t}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top},$
(10)

where $\otimes$ is the Kronecker product, and each ${\bm{B}}_{t}[i,j]$ is an
$N\times N$ matrix, solution of the following system of ODEs222Here we used
the notation ${\bm{B}}_{t}[i,j]$ to make the dependence on $t$ explicit,
without overloading the notation. A more precise notation would be
${\bm{B}}[i,j](t)$.

$\displaystyle\frac{d{\bm{B}}_{t}[i,j]}{dt}$
$\displaystyle=\mathbbm{1}(i=j){\bm{A}}_{i}(t)-{\bm{A}}_{i}(t)\sum_{m=1}^{M}({\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{m}){\bm{B}}_{t}[m,j]$
$\displaystyle{\bm{B}}_{0}[i,j]$ $\displaystyle={\bm{0}},$ (11)

and ${\bm{A}}_{i}(t)$, defined in Section B.2, is also an $N\times N$ matrix
that only depends on ${\bm{W}}(t)$ and ${\bm{\phi}}_{i}$. The main consequence
of this decomposition is that we can simply integrate the augmented ODE in
$\big{\\{}{\bm{W}}(t),{\bm{B}}_{t}[i,j]\big{\\}}$ up to $T$ to obtain the
desired Jacobian matrix, along with the adapted parameters ${\bm{W}}(T)$.
Furthermore, in contrast to naively applying the forward sensitivity method
(see Section 3.2), the $M^{2}$ matrices ${\bm{B}}_{t}[i,j]$ are significantly
smaller than the full Jacobian matrix. In fact, we show in Appendix C that we
can compute vector-Jacobian products—required for the chain rule—using only
these smaller matrices, and without ever having to explicitly construct the
full $Nd\times Nd$ Jacobian matrix $d{\bm{W}}(t)/d{\bm{W}}_{0}$ with Eq. 10.

#### Meta-gradients wrt. $\Phi$

To backpropagate the error through the embedding network $f_{\Phi}$, we need
to first compute the gradients of the outer-loss wrt. the feature vectors
${\bm{\phi}}_{m}$. Again, by the chain rule, we can get these gradients with
the Jacobian matrices $d{\bm{W}}(T)/d{\bm{\phi}}_{m}$. Similar to Eq. 10, we
can show that these Jacobian matrices can be decomposed as:

$\frac{d{\bm{W}}(t)}{d{\bm{\phi}}_{m}}=-\bigg{[}{\bm{s}}_{m}(t)\otimes{\bm{I}}+\sum_{i=1}^{M}{\bm{B}}_{t}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}+\sum_{i=1}^{M}\sum_{j=1}^{M}{\bm{z}}_{t}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\bigg{]},$
(12)

where ${\bm{s}}_{m}(t)$ and ${\bm{z}}_{t}[i,j,m]$ are vectors of size $N$,
that follow some dynamics; the exact form of this system of ODEs, as well as
the proof of this decomposition, are given in Section B.3. Crucially, the only
quantities that depend on time are small objects independent of the embedding
size $d$. Following the same strategy as above, we can incorporate these
vectors in the augmented ODE, and integrate it to get the necessary Jacobians.
Once all the $d{\bm{W}}(t)/d{\bm{\phi}}_{m}$ are known, for all the training
datapoints, we can apply standard backpropagation through $f_{\Phi}$ to obtain
the meta-gradients wrt. $\Phi$.

#### Meta-gradient wrt. $T$

One of the major novelties of COMLN is the capacity to meta-learn the amount
of adaptation using stochastic gradient descent. To compute the meta-gradient
wrt. the time horizon $T$, we can directly borrow the results derived by Chen
et al. (2018) in the context of Neural ODEs, and apply it to our gradient
vector field in Eq. 8 responsible for adaptation:

$\frac{d{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{dT}=-\bigg{[}\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}\bigg{]}^{\top}\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}}{\partial{\bm{W}}(T)}.$
(13)

The proof is available in Section B.4. Interestingly, we find that this
involves the alignment between the vectors of partial derivatives of the
inner-loss and the outer-loss at ${\bm{W}}(T)$, which appeared in various
contexts in the meta-learning and the multi-task learning literature (Li et
al., 2018; Rothfuss et al., 2019; Yu et al., 2020; Von Oswald et al., 2021).

###  4.2 Memory efficiency

Although naively applying the forward sensitivity method would be memory
intensive, we have shown in Section 4.1 that the Jacobians can be carefully
decomposed into smaller pieces. It turns out that even the parameters
${\bm{W}}(t)$ can be expressed using the vectors ${\bm{s}}_{m}(t)$ from the
decomposition in Eq. 12; see Section B.1 for details. As a consequence, to
compute the adapted parameters ${\bm{W}}(T)$ as well as all the necessary
meta-gradients, it is sufficient to integrate a dynamical system in
$\big{\\{}{\bm{B}}_{t}[i,j],{\bm{s}}_{m}(t),{\bm{z}}_{t}[i,j,m]\big{\\}}$ (see
Algorithms 1 and 2 in App. A.1), involving exclusively quantities that are
independent of the embedding size $d$. Instead, the size of that system scales
with $M$ the total number of training examples, which is typically much
smaller than $d$ for few-shot classification.

Table 1: Memory required to compute meta-gradients for different algorithms. Exact: the method returns the exact meta-gradients. Full net.: the whole network is adapted, with a number of meta-parameters $|{\bm{\theta}}|$. The requirements for checkpointing are taken from (Shaban et al., 2019). Note that typically $M\ll d$ in few-shot learning. Model | Exact | Full net. | Memory
---|---|---|---
MAML (Finn et al., 2017) | ✓ | ✓ | ${\mathcal{O}}(|{\bm{\theta}}|\cdot T)$
ANIL (Raghu et al., 2019) | ✓ | ✗ | ${\mathcal{O}}(Nd\cdot T)$
Checkpointing (every $\sqrt{T}$ steps) | ✓ | ✓ | ${\mathcal{O}}(|{\bm{\theta}}|\cdot\sqrt{T})$
iMAML (Rajeswaran et al., 2019) | ✗ | ✓ | ${\mathcal{O}}(|{\bm{\theta}}|)$
Forward sensitivity (naive) | ✓ | ✗ | ${\mathcal{O}}(N^{2}d^{2}+MNd^{2})$
COMLN | ✓ | ✗ | ${\mathcal{O}}(M^{2}N^{2}+M^{3}N)$

Table 1 shows a comparison of the memory cost for different algorithms. It is
important to note that contrary to other standard gradient-based meta-learning
methods, the memory requirements of COMLN do not scale with the amount of
adaptation $T$ (i.e. the number of gradient steps in MAML & ANIL), while still
returning the exact meta-gradients—unlike iMAML (Rajeswaran et al., 2019),
which only returns an approximation of the meta-gradients. We verified
empirically this efficiency, both in terms of memory and computation costs, in
Section 5.2.

##  5 Experiments

For our embedding network $f_{\Phi}$, we consider two commonly used
architectures in meta-learning: Conv-4, a convolutional neural network with 4
convolutional blocks, and ResNet-12, a 12-layer residual network (He et al.,
2016). Note that following Lee et al. (2019), ResNet-12 does not include a
global pooling layer at the end of the network, leading to feature vectors
with embedding dimension $d=16,\\!000$. Additional details about these
architectures are given in Appendix E. To compute the adapted parameters and
the meta-gradients in COMLN, we integrate the dynamical system described in
Section 4.2 with a 4th order Runge-Kutta method with a Dormand Prince adaptive
step size (Runge, 1895; Dormand & Prince, 1980); we will come back to the
choice of this numerical solver in Section 5.2. Furthermore to ensure that
$T>0$, we parametrized it with an exponential activation.

###  5.1 Few-shot image classification

We evaluate COMLN on two standard few-shot image classification benchmarks:
the miniImageNet (Vinyals et al., 2016) and the tieredImageNet datasets (Ren
et al., 2018), both datasets being derived from ILSVRC-2012 (Russakovsky et
al., 2015). The process for creating tasks follows the standard procedure from
the few-shot classification literature (Santoro et al., 2016), with distinct
classes between the different splits. miniImagenet consists of 100 classes,
split into 64 training classes, 16 validation classes, and 20 test classes.
tieredImageNet consists of 608 classes grouped into 34 high-level categories
from ILSVRC-2012, split into 20 training, 6 validation, and 8 testing
categories—corresponding to 351/97/160 classes respectively; Ren et al. (2018)
argue that separating data according to high-level categories results in a
more difficult and more realistic regime.

Table 2: Few-shot classification on miniImageNet & tieredImageNet. The average
accuracy (%) on 1,000 held-out meta-test tasks is reported with $95\%$
confidence interval. ✓ denotes gradient-based meta-learning algorithms.
$\star$ denotes baseline results we executed using the official
implementations.

Model Backbone miniImageNet 5-way tieredImageNet 5-way 1-shot 5-shot 1-shot
5-shot MAML (Finn et al., 2017) ✓ Conv-4 $48.70\pm 1.84^{\phantom{\star}}$
$63.11\pm 0.92^{\phantom{\star}}$ $51.67\pm 1.81^{\phantom{\star}}$ $70.30\pm
1.75^{\phantom{\star}}$ ANIL (Raghu et al., 2019) ✓ Conv-4 $46.30\pm
0.40^{\phantom{\star}}$ $61.00\pm 0.60^{\phantom{\star}}$ $49.35\pm
0.26^{\phantom{\star}}$ $65.82\pm 0.12^{\phantom{\star}}$ Meta-SGD (Li et al.,
2017) ✓ Conv-4 $50.47\pm 1.87^{\phantom{\star}}$ $64.03\pm
0.94^{\phantom{\star}}$ $52.80\pm 0.44^{\phantom{\star}}$ $62.35\pm
0.26^{\phantom{\star}}$ CAVIA (Zintgraf et al., 2019) ✓ Conv-4 $51.82\pm
0.65^{\phantom{\star}}$ $65.85\pm 0.55^{\phantom{\star}}$ $52.41\pm
2.64^{\star}$ $67.55\pm 2.05^{\star}$ iMAML (Rajeswaran et al., 2019) ✓ Conv-4
$49.30\pm 1.88^{\phantom{\star}}$ $59.77\pm 0.73^{\star}$ $38.54\pm
1.37^{\star}$ $60.24\pm 0.76^{\star}$ MetaOptNet-RR (Lee et al., 2019) Conv-4
$\mathbf{53.23\pm 0.59^{\phantom{\star}}}$ $69.51\pm 0.48^{\phantom{\star}}$
$54.63\pm 0.67^{\phantom{\star}}$ $\mathbf{72.11\pm 0.59^{\phantom{\star}}}$
MetaOptNet-SVM (Lee et al., 2019) Conv-4 $52.87\pm 0.57^{\phantom{\star}}$
$68.76\pm 0.48^{\phantom{\star}}$ $\mathbf{54.71\pm 0.67^{\phantom{\star}}}$
$71.79\pm 0.59^{\phantom{\star}}$ COMLN (Ours) ✓ Conv-4 $53.01\pm
0.62^{\phantom{\star}}$ $\mathbf{70.54\pm 0.54^{\phantom{\star}}}$ $54.30\pm
0.69^{\phantom{\star}}$ $71.35\pm 0.57^{\phantom{\star}}$ MAML (Finn et al.,
2017) ✓ ResNet-12 $49.92\pm 0.65^{\phantom{\star}}$ $63.93\pm
0.59^{\phantom{\star}}$ $55.37\pm 0.74^{\phantom{\star}}$ $72.93\pm
0.60^{\phantom{\star}}$ ANIL (Raghu et al., 2019) ✓ ResNet-12 $49.65\pm
0.65^{\phantom{\star}}$ $59.51\pm 0.56^{\phantom{\star}}$ $54.77\pm
0.76^{\phantom{\star}}$ $69.28\pm 0.67^{\phantom{\star}}$ MetaOptNet-RR (Lee
et al., 2019) ResNet-12 $61.41\pm 0.61^{\phantom{\star}}$ $77.88\pm
0.46^{\phantom{\star}}$ $65.36\pm 0.71^{\phantom{\star}}$ $81.34\pm
0.52^{\phantom{\star}}$ MetaOptNet-SVM (Lee et al., 2019) ResNet-12
$\mathbf{62.64\pm 0.61^{\phantom{\star}}}$ $\mathbf{78.63\pm
0.46^{\phantom{\star}}}$ $\mathbf{65.99\pm 0.72^{\phantom{\star}}}$
$\mathbf{81.56\pm 0.53^{\phantom{\star}}}$ COMLN (Ours) ✓ ResNet-12 $59.26\pm
0.65^{\phantom{\star}}$ $77.26\pm 0.49^{\phantom{\star}}$ $62.93\pm
0.71^{\phantom{\star}}$ $81.13\pm 0.53^{\phantom{\star}}$

Table 2 shows the average accuracies of COMLN compared to various meta-
learning methods, be it gradient-based or not. For both backbones, COMLN
decisively outperforms all other gradient-based meta-learning methods.
Compared to methods that explicitly backpropagate through the learning
process, such as MAML or ANIL, the performance gain shown by COMLN could be
credited to the longer adaptation $T$ it learns, as opposed to a small number
of gradient steps—usually about $10$ steps; this was fully enabled by our
memory-efficient method to compute meta-gradients, which does not scale with
the length of adaptation anymore (see Section 4.2). We analyse the evolution
of $T$ during meta-training for these different settings in Section E.3. In
almost all settings, COMLN is even closing the gap with a strong non-gradient-
based method like MetaOptNet; the remainder may be explained in part by the
training choices made by Lee et al. (2019) (see Appendix E for details).

###  5.2 Empirical efficiency of COMLN

In Section 4.2, we showed that our algorithm to compute the meta-gradients,
based on forward differentiation, had a memory cost independent of the length
of adaptation $T$. We verify this empirically in Fig. 3, where we compare the
memory required by COMLN and other methods to compute the meta-gradients on a
single task, with a Conv-4 backbone (Fig. 6 in Section E.2 shows similar
results for ResNet-12). To ensure an aligned comparison between discrete and
continuous time, we use a conversion corresponding to a learning rate
$\alpha=0.01$ in Eq. 2; see Section E.2 for a justification. As expected, the
memory cost increases for both MAML and ANIL as the number of gradient steps
increases, while it remains constant for iMAML and COMLN. Interestingly, we
observe that the cost of COMLN is equivalent to the cost of running ANIL for a
small number of steps, showing that the additional cost of integrating the
augmented ODE in Section 4.2 to compute the meta-gradients is minimal.

Increasing the length of adaptation also has an impact on the time it takes to
compute the adapted parameters, and the meta-gradients. Fig. 3 (right) shows
how the runtime increases with the amount of adaptation for different
algorithms. We see that the efficiency of COMLN depends on the numerical
solver used. When we use a simple explicit-Euler scheme, the time taken to
compute the meta-gradients matches the one of ANIL; this behavior empirically
confirms our observation in Section 3.3. When we use an adaptive numerical
solver, such as Runge-Kutta (RK) with a Dormand Prince step size, this
computation can be significantly accelerated, thanks to the smaller number of
function evaluations. In practice, we show in Section E.1 that the choice of
the ODE solver has a very minimal impact on the accuracy.

[width=]figures/pgfplots/conv4_memory_runtime

Figure 3: Empirical efficiency of COMLN on a single 5-shot 5-way task, with a
Conv-4 backbone. (Left) Memory usage for computing the meta-gradients as a
function of the number of inner-gradient steps. The extrapolated dashed lines
correspond to the method reaching the memory capacity of a Tesla V100 GPU with
32Gb of memory. (Right) Average time taken (in ms) to compute the exact meta-
gradients. The extrapolated dashed lines correspond to the method taking over
2 seconds.

##  6 Related work

We are interested in meta-learning (Bengio et al., 1991; Schmidhuber, 1987;
Thrun & Pratt, 2012), and in particular we focus on gradient-based meta-
learning methods (Finn, 2018), where the learning rule is based on gradient
descent. While in MAML (Finn et al., 2017) the whole network was updated
during this process, follow-up works have shown that it is generally
sufficient to share most parts of the neural network, and to only adapt a few
layers (Raghu et al., 2019; Chen et al., 2020b; Tian et al., 2020). Even
though this hypothesis has been challenged recently (Arnold & Sha, 2021),
COMLN also updates only the last layer of a neural network, and therefore can
be viewed as a continuous-time extension of ANIL (Raghu et al., 2019); see
also Section 3.3. With its shared embedding network across tasks, COMLN is
also connected to metric-based meta-learning methods (Vinyals et al., 2016;
Snell et al., 2017; Sung et al., 2018; Bertinetto et al., 2018; Lee et al.,
2019).

Zhang et al. (2021) introduced a formulation where the adaptation of
prototypes follows a gradient vector field, but finally opted for modeling it
as a Neural ODE (Chen et al., 2018). Concurrent to our work, Li et al. (2021)
also propose a formulation of adaptation based on a gradient vector field, and
use the adjoint method to compute the meta-gradients, despite the challenges
we identified in Sec. 3.2; Li et al. (2021) also acknowledge these challenges,
and they limit their analysis to relatively small values of $T$ (in comparison
to the ones learned by COMLN), hence further from the equilibrium, to
circumvent this issue altogether. Zhou et al. (2021) also uses a gradient
vector field to motivate a novel method with a closed-form adaptation; COMLN
still explicitly updates the parameters following the gradient vector field,
since there is no closed-form solution of Eq. 8. As mentioned in Section 3,
treating optimization as a continuous-time process has been used to analyze
the convergence of different optimization algorithms, including the meta-
optimization of MAML (Xu et al., 2021), or to introduce new meta-optimizers
based on different integration schemes (Im et al., 2019). Guo et al. (2021)
also uses meta-learning to learn new integration schemes for ODEs. Although
this is a growing literature at the intersection of meta-learning and
dynamical systems, our work is the first algorithm that uses a gradient vector
field for adaptation in meta-learning (see also Li et al. (2021)).

Beyond the memory efficiency of our method, one of the main benefits of the
continuous-time perspective is that COMLN is capable of learning when to stop
the adaptation, as opposed to taking a number of gradient steps fixed ahead of
time. However unlike Chen et al. (2020a), where the number of gradient steps
are optimized (up to a maximal number) with variational methods, we
incorporate the amount of adaptation as a (continuous) meta-parameter that can
be learned using SGD. To compute the meta-gradients, which is known to be
challenging for long sequences in gradient-based meta-learning, we use
forward-mode differentiation as an alternative to backpropagation through the
learning process, similar to prior work in meta-learning (Franceschi et al.,
2017; Jiwoong Im et al., 2021) and hyperparameter optimization over long
horizons (Micaelli & Storkey, 2021). This yields the exact meta-gradients in
constant memory, without any assumption on the optimality of the inner
optimization problem, which is necessary when using the normal equations
(Bertinetto et al., 2018), or to apply implicit differentiation (Rajeswaran et
al., 2019).

##  7 Conclusion and Future Work

In this paper, we have introduced _Continuous-Time Meta-Learning_ (COMLN), a
novel algorithm that treats the adaptation in meta-learning as a continuous-
time process, by following the dynamics of a gradient vector field up to a
certain time horizon $T$. One of the major novelties of treating adaptation in
continuous time is that the amount of adaptation $T$ is now a continuous
quantity, that can be viewed as a meta-parameter and can be learned using SGD,
alongside the initial conditions and the parameters of the embedding network.
In order to learn these meta-parameters, we have also introduced a novel
practical algorithm based on forward mode automatic differentiation, capable
of efficiently computing the exact meta-gradients using an augmented dynamical
system. We have verified empirically that this algorithm was able to compute
the meta-gradients in constant memory, making it the first gradient-based
meta-learning approach capable of computing the exact meta-gradients with long
sequences of adaptation using gradient methods. In practice, we have shown
that COMLN significantly outperforms other standard gradient-based meta-
learning algorithms.

In addition to having a single meta-parameter $T$ that drives the adaptation
of all possible tasks, the fact that the time horizon can be learned with SGD
opens up new possibilities for gradient-based methods. For example, we could
imagine treating $T$ not as a shared meta-parameters, but as a task-specific
parameter. This would allow the learning process to be more _adaptive_ ,
possibly with different behaviors depending on the difficulty of the task.
This is left as a future direction of research.

## Acknowledgements

The authors are grateful to Samsung Electronics Co., Ldt., CIFAR, and IVADO
for their funding and Calcul Québec and Compute Canada for providing us with
the computing resources.

## Reproducibility statement

We provide in Section A.1 a full description in pseudo-code of the meta-
training procedure (Algorithm 1), along with the exact dynamics of the ODE
(Algorithm 2) and the projection operations (Algorithms 3 and 4) to avoid
explicitly building the Jacobian matrices to compute Jacobian-vector products
(see Section 4.1).

We also provide in Section A.2 a snippet of code in JAX (Bradbury et al.,
2018) to compute the adapted parameters ${\bm{W}}(T)$, as well as all the
necessary objects
$\big{\\{}{\bm{B}}_{t}[i,j],{\bm{s}}_{m}(t),{\bm{z}}_{t}[i,j,m]\big{\\}}$ to
compute all the meta-gradients (see Section 4.2). We also give in Code Snippet
5 the code to compute the meta-gradients wrt. the initialization
${\bm{W}}_{0}$ and the integration time $T$. Computing the meta-gradients wrt.
$\Phi$ involves non-minimal dependencies on Haiku (Hennigan et al., 2020), and
therefore is omitted here. The full code is available at
https://github.com/tristandeleu/jax-comln.

#### Data generation & hyperparameters

We used the miniImageNet and tieredImageNet datasets provided by Lee et al.
(2019) in order to create the 1-shot 5-way and 5-shot 5-way tasks for both
datasets. During evaluation, a fixed set of $1,\\!000$ tasks was sampled for
each setting; this means that both architectures for COMLN have been evaluated
using exactly the same data, to ensure direct comparison across backbones. A
full description of all the hyperparameters used in COMLN is given in Appendix
E.

#### Reproducibility of baseline results

To the best of our ability, we have tried to report baseline results from
existing work, to limit as much as possible the bias induced by running our
own baseline experiments. The references of those works are given in Table 3.
We still had to run CAVIA and iMAML on the remaining settings, since these
results have not been reported in the literature. For both methods, we used
the data generation described above.

  * •

_CAVIA_ : We used the official
implementation333https://github.com/lmzintgraf/cavia/. We used the
hyperparameters reported in (Zintgraf et al., 2019) for miniImageNet, and an
architecture with $64$ filters.

  * •

_iMAML_ : We used the official
implementation444https://github.com/aravindr93/imaml_dev. We used the
hyperparameters reported in (Rajeswaran et al., 2019) for miniImageNet 1-shot
5-way.

Table 3: References for the results provided in Table 2:  (Liu et al., 2019),
(Oh et al., 2021),  (Aimen et al., 2021),  (Arnold et al., 2021), and  are
reported in their respective references (under _Model_). Recall that $\star$
denotes baseline results we executed using the official implementations.

Model Backbone miniImageNet 5-way tieredImageNet 5-way 1-shot 5-shot 1-shot
5-shot MAML (Finn et al., 2017) ✓ Conv-4 $48.70\pm 1.84^{\phantom{\star}}$
$63.11\pm 0.92^{\phantom{\star}}$ $51.67\pm 1.81^{\phantom{\star}}$ $70.30\pm
1.75^{\phantom{\star}}$ ANIL (Raghu et al., 2019) ✓ Conv-4 $46.30\pm
0.40^{\phantom{\star}}$ $61.00\pm 0.60^{\phantom{\star}}$ $49.35\pm
0.26^{\phantom{\star}}$ $65.82\pm 0.12^{\phantom{\star}}$ Meta-SGD (Li et al.,
2017) ✓ Conv-4 $50.47\pm 1.87^{\phantom{\star}}$ $64.03\pm
0.94^{\phantom{\star}}$ $52.80\pm 0.44^{\phantom{\star}}$ $62.35\pm
0.26^{\phantom{\star}}$ CAVIA (Zintgraf et al., 2019) ✓ Conv-4 $51.82\pm
0.65^{\phantom{\star}}$ $65.85\pm 0.55^{\phantom{\star}}$ $52.41\pm
2.64^{\star}$ $67.55\pm 2.05^{\star}$ iMAML (Rajeswaran et al., 2019) ✓ Conv-4
$49.30\pm 1.88^{\phantom{\star}}$ $59.77\pm 0.73^{\star}$ $38.54\pm
1.37^{\star}$ $60.24\pm 0.76^{\star}$ MetaOptNet-RR (Lee et al., 2019) Conv-4
$\mathbf{53.23\pm 0.59^{\phantom{\star}}}$ $69.51\pm 0.48^{\phantom{\star}}$
$54.63\pm 0.67^{\phantom{\star}}$ $\mathbf{72.11\pm 0.59^{\phantom{\star}}}$
MetaOptNet-SVM (Lee et al., 2019) Conv-4 $52.87\pm 0.57^{\phantom{\star}}$
$68.76\pm 0.48^{\phantom{\star}}$ $\mathbf{54.71\pm 0.67^{\phantom{\star}}}$
$71.79\pm 0.59^{\phantom{\star}}$ COMLN (Ours) ✓ Conv-4 $53.01\pm
0.62^{\phantom{\star}}$ $\mathbf{70.54\pm 0.54^{\phantom{\star}}}$ $54.30\pm
0.69^{\phantom{\star}}$ $71.35\pm 0.57^{\phantom{\star}}$ MAML (Finn et al.,
2017) ✓ ResNet-12 $49.92\pm 0.65^{\phantom{\star}}$ $63.93\pm
0.59^{\phantom{\star}}$ $55.37\pm 0.74^{\phantom{\star}}$ $72.93\pm
0.60^{\phantom{\star}}$ ANIL (Raghu et al., 2019) ✓ ResNet-12 $49.65\pm
0.65^{\phantom{\star}}$ $59.51\pm 0.56^{\phantom{\star}}$ $54.77\pm
0.76^{\phantom{\star}}$ $69.28\pm 0.67^{\phantom{\star}}$ MetaOptNet-RR (Lee
et al., 2019) ResNet-12 $61.41\pm 0.61^{\phantom{\star}}$ $77.88\pm
0.46^{\phantom{\star}}$ $65.36\pm 0.71^{\phantom{\star}}$ $81.34\pm
0.52^{\phantom{\star}}$ MetaOptNet-SVM (Lee et al., 2019) ResNet-12
$\mathbf{62.64\pm 0.61^{\phantom{\star}}}$ $\mathbf{78.63\pm
0.46^{\phantom{\star}}}$ $\mathbf{65.99\pm 0.72^{\phantom{\star}}}$
$\mathbf{81.56\pm 0.53^{\phantom{\star}}}$ COMLN (Ours) ✓ ResNet-12 $59.26\pm
0.65^{\phantom{\star}}$ $77.26\pm 0.49^{\phantom{\star}}$ $62.93\pm
0.71^{\phantom{\star}}$ $81.13\pm 0.53^{\phantom{\star}}$

## References

  * Aimen et al. (2021) Aroof Aimen, Sahil Sidheekh, and Narayanan C Krishnan.  Task Attended Meta-Learning for Few-Shot Learning.  _arXiv preprint_ , 2021.
  * Arnold & Sha (2021) Sébastien MR Arnold and Fei Sha.  Embedding Adaptation is Still Needed for Few-Shot Learning.  _arXiv preprint_ , 2021.
  * Arnold et al. (2021) Sébastien MR Arnold, Guneet S Dhillon, Avinash Ravichandran, and Stefano Soatto.  Uniform Sampling over Episode Difficulty.  _arXiv preprint_ , 2021.
  * Baranchuk (2019) Dmitry Baranchuk.  Memory Efficient MAML, 2019.  URL https://github.com/dbaranchuk/memory-efficient-maml.
  * Bengio et al. (1991) Yoshua Bengio, Samy Bengio, Jocelyn Cloutier, and Jan Gecsei.  Learning a Synaptic Learning Rule.  _International Joint Conference on Neural Networks_ , 1991.
  * Bertinetto et al. (2018) Luca Bertinetto, Joao F Henriques, Philip HS Torr, and Andrea Vedaldi.  Meta-learning with Differentiable Closed-Form Solvers.  _arXiv preprint_ , 2018.
  * Betts (2010) John T Betts.  _Practical Methods for Optimal Control and Estimation Using Nonlinear Programming_.  SIAM, 2010.
  * Biegler (2010) Lorenz T. Biegler.  _Nonlinear Programming_.  Society for Industrial and Applied Mathematics, January 2010.
  * Böhning (1992) Dankmar Böhning.  Multinomial Logistic Regression Algorithm.  _Annals of the Institute of Statistical Mathematics_ , 1992.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Bryson & Ho (1969) A. E. Bryson and Y. C. Ho.  _Applied Optimal Control_.  Blaisdell, New York, 1969.
  * Butcher (2008) John Charles Butcher.  _Numerical Methods for Ordinary Differential Equations_.  Wiley, 2008.
  * Caracotsios & Stewart (1985) Makis Caracotsios and Warren E Stewart.  Sensitivity analysis of initial value problems with mixed odes and algebraic equations.  _Computers & Chemical Engineering_, 9(4):359–365, 1985.
  * Chavent et al. (1974) G Chavent, RE Goodson, and M Polis.  Identification of parameter distributed systems.  _Identification of function parameters in partial differential equations_ , pp. 31–48, 1974.
  * Chen et al. (2018) Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud.  Neural Ordinary Differential Equations.  _Advances in Neural Information Processing Systems_ , 2018.
  * Chen et al. (2020a) Xinshi Chen, Hanjun Dai, Yu Li, Xin Gao, and Le Song.  Learning To Stop While Learning To Predict.  In _International Conference on Machine Learning_ , 2020a.
  * Chen et al. (2020b) Yutian Chen, Abram L Friesen, Feryal Behbahani, Arnaud Doucet, David Budden, Matthew W Hoffman, and Nando de Freitas.  Modular Meta-Learning with Shrinkage.  _Neural Information Processing Systems_ , 2020b.
  * Dormand & Prince (1980) John R Dormand and Peter J Prince.  A family of embedded Runge-Kutta formulae.  _Journal of computational and applied mathematics_ , 1980.
  * Euler (1913) Leonhard Euler.  De integratione aequationum differentialium per approximationem.  _Opera Omnia_ , 1913.
  * Feehery et al. (1997) William F Feehery, John E Tolsma, and Paul I Barton.  Efficient sensitivity analysis of large-scale differential-algebraic systems.  _Applied Numerical Mathematics_ , 25(1):41–54, 1997.
  * Finn (2018) Chelsea Finn.  _Learning to Learn with Gradients_.  PhD thesis, UC Berkeley, 2018.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks.  _International Conference on Machine Learning (ICML)_ , 2017.
  * Flennerhag et al. (2018) Sebastian Flennerhag, Pablo G Moreno, Neil D Lawrence, and Andreas Damianou.  Transferring knowledge across learning processes.  _arXiv preprint_ , 2018.
  * Flennerhag et al. (2020) Sebastian Flennerhag, Andrei A Rusu, Razvan Pascanu, Francesco Visin, Hujun Yin, and Raia Hadsell.  Meta-Learning with Warped Gradient Descent.  _International Conference on Learning Representations_ , 2020.
  * Franceschi et al. (2017) Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil.  Forward and reverse gradient-based hyperparameter optimization.  In _International Conference on Machine Learning_ , 2017.
  * Ghiasi et al. (2018) Golnaz Ghiasi, Tsung-Yi Lin, and Quoc V Le.  Dropblock: A regularization method for convolutional networks.  In _Neural Information Processing Systems_ , 2018.
  * Gholami et al. (2019) Amir Gholami, Kurt Keutzer, and George Biros.  ANODE: unconditionally accurate memory-efficient gradients for neural odes.  _CoRR_ , abs/1902.10298, 2019.  URL http://arxiv.org/abs/1902.10298.
  * Griewank & Walther (2008) Andreas Griewank and Andrea Walther.  _Evaluating Derivatives_.  Society for Industrial and Applied Mathematics, January 2008.
  * Guo et al. (2021) Yue Guo, Felix Dietrich, Tom Bertalan, Danimir T Doncevic, Manuel Dahmen, Ioannis G Kevrekidis, and Qianxiao Li.  Personalized Algorithm Generation: A Case Study in Meta-Learning ODE Integrators.  _arXiv preprint_ , 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep Residual Learning for Image Recognition.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2016.
  * Hennigan et al. (2020) Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin.  Haiku: Sonnet for JAX, 2020.  URL http://github.com/deepmind/dm-haiku.
  * Im et al. (2019) Daniel Jiwoong Im, Yibo Jiang, and Nakul Verma.  Model-Agnostic Meta-Learning using Runge-Kutta Methods.  _arXiv preprint_ , 2019.
  * Javed & White (2019) Khurram Javed and Martha White.  Meta-Learning Representations for Continual Learning.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Jiwoong Im et al. (2021) Daniel Jiwoong Im, Cristina Savin, and Kyunghyun Cho.  Online hyperparameter optimization by Real-Time Recurrent Learning.  _arXiv preprint_ , 2021.
  * Lee et al. (2019) Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto.  Meta-learning with Differentiable Convex Optimization.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019.
  * Leis & Kramer (1988) Jorge R Leis and Mark A Kramer.  The simultaneous solution and sensitivity analysis of systems described by ordinary differential equations.  _ACM Transactions on Mathematical Software (TOMS)_ , 14(1):45–60, 1988.
  * Li et al. (2018) Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales.  Learning to Generalize: Meta-Learning for Domain Generalization.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2018.
  * Li et al. (2021) Shibo Li, Zheng Wang, Akil Narayan, Robert Kirby, and Shandian Zhe.  Meta-Learning with Adjoint Methods.  _arXiv preprint_ , 2021.
  * Li et al. (2017) Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li.  Meta-SGD: Learning to learn quickly for few-shot learning.  _arXiv preprint_ , 2017.
  * Lions & Magenes (2012) Jacques Louis Lions and Enrico Magenes.  _Non-homogeneous boundary value problems and applications: Vol. 1_ , volume 181.  Springer Science & Business Media, 2012.
  * Liu et al. (2019) Yanbin Liu, Juho Lee, Minseop Park, Saehoon Kim, Eunho Yang, Sung Ju Hwang, and Yi Yang.  Learning to Propagate Labels: Transductive Propagation Network for Few-Shot Learning.  _International Conference on Learning Representations_ , 2019.
  * Maclaurin et al. (2015) Dougal Maclaurin, David Duvenaud, and Ryan Adams.  Gradient-based hyperparameter optimization through reversible learning.  In _International conference on machine learning_ , pp. 2113–2122. PMLR, 2015.
  * Maly & Petzold (1996) Timothy Maly and Linda R Petzold.  Numerical methods and software for sensitivity analysis of differential-algebraic systems.  _Applied Numerical Mathematics_ , 20(1-2):57–79, 1996.
  * Micaelli & Storkey (2021) Paul Micaelli and Amos Storkey.  Gradient-based Hyperparameter Optimization Over Long Horizons.  In _Neural Information Processing Systems_ , 2021.
  * Nichol et al. (2018) Alex Nichol, Joshua Achiam, and John Schulman.  On First-Order Meta-Learning Algorithms.  _arXiv preprint_ , 2018.
  * Oh et al. (2021) Jaehoon Oh, Hyungjun Yoo, ChangHwan Kim, and Se-Young Yun.  BOIL: Towards Representation Change for Few-Shot Learning.  _International Conference on Learning Representations_ , 2021.
  * Orvieto & Lucchi (2019) Antonio Orvieto and Aurelien Lucchi.  Shadowing Properties of Optimization Algorithms.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Platt & Barr (1988) John Platt and Alan Barr.  Constrained Differential Optimization.  In _Neural Information Processing Systems_ , 1988.
  * Polyak (1964) Boris T Polyak.  Some methods of speeding up the convergence of iteration methods.  _USSR computational mathematics and mathematical physics_ , 1964\.
  * Pontryagin (2018) Lev Semenovich Pontryagin.  _Mathematical theory of optimal processes_.  Routledge, 2018.
  * Raghu et al. (2019) Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals.  Rapid learning or feature reuse? towards understanding the effectiveness of maml.  _arXiv preprint_ , 2019.
  * Rajeswaran et al. (2019) Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine.  Meta-Learning with Implicit Gradients.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Ren et al. (2018) Mengye Ren, Eleni Triantafillou, Sachin Ravi, Jake Snell, Kevin Swersky, Joshua B Tenenbaum, Hugo Larochelle, and Richard S Zemel.  Meta-learning for semi-supervised few-shot classification.  In _International Conference on Learning Representations_ , 2018.
  * Rothfuss et al. (2019) Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel.  ProMP: Proximal Meta-Policy Search.  _International Conference on Learning Representations_ , 2019.
  * Runge (1895) Carl Runge.  Über die numerische auflösung von differentialgleichungen.  _Mathematische Annalen_ , 1895.
  * Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  _International journal of computer vision_ , 115(3):211–252, 2015.
  * Rusu et al. (2018) Andrei A Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero, and Raia Hadsell.  Meta-learning with Latent Embedding Optimization.  _arXiv preprint_ , 2018.
  * Santoro et al. (2016) Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap.  One-shot learning with memory-augmented neural networks.  _arXiv preprint_ , 2016.
  * Schmidhuber (1987) Jürgen Schmidhuber.  _Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-… hook_.  PhD thesis, Technische Universität München, 1987.
  * Serban & Hindmarsh (2003) Radu Serban and Alan C Hindmarsh.  Cvodes: An ode solver with sensitivity analysis capabilities.  Technical report, Technical Report UCRL-JP-200039, Lawrence Livermore National Laboratory, 2003.
  * Shaban et al. (2019) Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots.  Truncated Back-propagation for Bilevel Optimization.  In _International Conference on Artificial Intelligence and Statistics_ , 2019.
  * Snell et al. (2017) Jake Snell, Kevin Swersky, and Richard Zemel.  Prototypical Networks for Few-shot Learning.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Su et al. (2014) Weijie Su, Stephen Boyd, and Emmanuel Candes.  A Differential Equation for Modeling Nesterov’s Accelerated Gradient Method: Theory and Insights.  _Advances in Neural Information Processing Systems_ , 2014.
  * Sung et al. (2018) Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang, Philip HS Torr, and Timothy M Hospedales.  Learning to compare: Relation network for few-shot learning.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2018.
  * Sutton (1992) Richard S. Sutton.  Adapting Bias by Gradient Descent: An Incremental Version of Delta-Bar-Delta.  In _Proceedings of the 10th National Conference on Artificial Intelligence_ , 1992.
  * Thrun & Pratt (2012) Sebastian Thrun and Lorien Pratt.  _Learning to learn_.  Springer Science & Business Media, 2012.
  * Tian et al. (2020) Yonglong Tian, Yue Wang, Dilip Krishnan, Joshua B Tenenbaum, and Phillip Isola.  Rethinking Few-Shot Image Classification: a Good Embedding Is All You Need?  2020\.
  * Vinyals et al. (2016) Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al.  Matching networks for one shot learning.  _Neural Information Processing Systems_ , 29:3630–3638, 2016.
  * Von Oswald et al. (2021) Johannes Von Oswald, Dominic Zhao, Seijin Kobayashi, Simon Schug, Massimo Caccia, Nicolas Zucchet, and João Sacramento.  Learning where to learn: Gradient sparsity in meta and continual learning.  _Advances in Neural Information Processing Systems_ , 2021.
  * Wiggins (2003) Stephen Wiggins.  _Introduction to Applied Nonlinear Dynamical Systems and Chaos_ , volume 2.  Springer, 2003.
  * Williams & Zipser (1989) Ronald J. Williams and David Zipser.  A Learning Algorithm for Continually Running Fully Recurrent Neural Networks.  _Neural Computation_ , 1989.
  * Wilson et al. (2016) Ashia C Wilson, Benjamin Recht, and Michael I Jordan.  A Lyapunov Analysis of Momentum Methods in Optimization.  _arXiv preprint_ , 2016.
  * Xu et al. (2021) Ruitu Xu, Lin Chen, and Amin Karbasi.  Meta Learning in the Continuous Time Limit .  In _Proceedings of The 24th International Conference on Artificial Intelligence and Statistics_. PMLR, 2021.
  * Xu et al. (2018) Zhongwen Xu, Hado van Hasselt, and David Silver.  Meta-Gradient Reinforcement Learning.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Yu et al. (2020) Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn.  Gradient Surgery for Multi-Task Learning.  _Neural Information Processing Systems_ , 2020.
  * Zhang et al. (2021) Baoquan Zhang, Xutao Li, Yunming Ye, Shanshan Feng, and Rui Ye.  MetaNODE: Prototype Optimization as a Neural ODE for Few-Shot Learning.  _arXiv preprint_ , 2021.
  * Zhou et al. (2021) Yufan Zhou, Zhenyi Wang, Jiayi Xian, Changyou Chen, and Jinhui Xu.  Meta-Learning with Neural Tangent Kernels.  _International Conference on Learning Representations_ , 2021.
  * Zintgraf et al. (2019) Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson.  Fast context adaptation via meta-learning.  In _International Conference on Machine Learning_ , pp. 7693–7702. PMLR, 2019.

Appendix

The Appendix is organized as follows: in Appendix A we give the pseudo-code
for meta-training and meta-testing COMLN, along with minimal code in JAX
(Bradbury et al., 2018). In Appendix B we prove the decomposition of the
Jacobians introduced in Section 4, and we give the exact dynamics of
${\bm{s}}_{m}(t)$ and ${\bm{z}}_{t}[i,j,m]$ in the decomposition of
$d{\bm{W}}(t)/d{\bm{\phi}}_{m}$, omitted from the main text for concision. In
Appendix C we show how the total derivatives may be computed from the Jacobian
matrices without ever having to form them explicitly, hence maintaining the
memory-efficiency described in Section 4.2. In Appendix D, we show that unlike
the adjoint method described in Section 2.2, the algorithm used in COMLN to
compute the meta-gradients based on the forward sensitivity method is stable
and guaranteed not to diverge. Finally in Appendix E we give additional
details regarding the experiments reported in Section 5, together with
additional analyses and results on the dataset introduced in Rusu et al.
(2018).

##  Appendix A Algorithmic details

###  A.1 Meta-training pseudo-code

We give in Algorithm 1 the pseudo-code for meta-training COMLN, based on a
distribution of tasks $p(\tau)$, with references to the relevant propositions
developed in Appendices B and C. Note that to simplify the presentation in
Algorithm 1, we introduce the notations $\nabla_{\tau}$ and
$\partial{\mathcal{L}}_{\tau}$ to denote the total derivative of the outer-
loss, and the partial derivative of the loss ${\mathcal{L}}$ respectively,
both computed at the adapted parameters ${\bm{W}}_{\tau}(T)$. For example:

$\displaystyle\nabla_{\tau}{\bm{W}}_{0}$
$\displaystyle=\frac{d{\mathcal{L}}\big{(}{\bm{W}}_{\tau}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{d{\bm{W}}_{0}}$
$\displaystyle\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{train}}$
$\displaystyle=\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}_{\tau}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}}{\partial{\bm{W}}(T)}$

Algorithm 1 COMLN – Meta-training

A task distribution $p(\tau)$

Initialize randomly $\Phi$ and ${\bm{W}}_{0}$. Initialize $T$ to a small value
$\varepsilon>0$.

loop

Sample a batch of tasks ${\mathcal{B}}\sim p(\tau)$

for all $\tau\in{\mathcal{B}}$ do

Embed the training set:
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$

${\bm{s}}(T),{\bm{B}}_{T},{\bm{z}}_{T}\leftarrow\textrm{ODESolve}([{\bm{0}},{\bm{0}},{\bm{0}}],\textrm{{Dynamics}},0,T)$
$\triangleright$ Algorithm 2

${\bm{W}}_{\tau}(T)\leftarrow{\bm{W}}_{0}-\sum_{m}{\bm{s}}_{m}(T){\bm{\phi}}_{m}^{\top}$
$\triangleright$ Proposition 1

Embed the test set: $f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})$

Compute the partial derivatives:
$\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{train}}$ &
$\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{test}}$

$\nabla_{\tau}{\bm{W}}_{0}\leftarrow\textrm{{Project}}_{{\bm{W}}_{0}}\big{(}\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{test}},{\bm{B}}_{T},{\bm{\phi}}\big{)}$$\triangleright$
Proposition 5

$\nabla_{\tau}{\bm{\phi}}_{m}\leftarrow\textrm{{Project}}_{{\bm{\phi}}_{m}}\big{(}\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{test}},{\bm{s}}(T),{\bm{B}}_{T},{\bm{z}}_{T},{\bm{\phi}}\big{)}+\partial_{{\bm{\phi}}_{m}}{\mathcal{L}}_{\tau}^{\mathrm{test}}$
$\triangleright$ Proposition 6

$\vphantom{\big{(}}\nabla_{\tau}\Phi\leftarrow$ Backpropagation through
$f_{\Phi}$, starting with $\nabla_{\tau}{\bm{\phi}}_{m}$ for all $m$

$\nabla_{\tau}T\leftarrow\big{(}\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{test}}\big{)}^{\top}\big{(}\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{train}}\big{)}$$\triangleright$
Proposition 4

end for

Update the meta-parameters:$\frac{\alpha}{|{\mathcal{B}}|}$

$\qquad{\bm{W}}_{0}\leftarrow{\bm{W}}_{0}-\frac{\alpha}{|{\mathcal{B}}|}\sum_{\tau\in{\mathcal{B}}}\nabla_{\tau}{\bm{W}}_{0}$

$\qquad\Phi\leftarrow\Phi-\frac{\alpha}{|{\mathcal{B}}|}\sum_{\tau\in{\mathcal{B}}}\nabla_{\tau}\Phi$

$\qquad T\leftarrow
T-\frac{\alpha}{|{\mathcal{B}}|}\sum_{\tau\in{\mathcal{B}}}\nabla_{\tau}T$

end loop

The dynamical system followed during adaptation is given in Propositions 1, 2
and 3 and is summarized in Algorithm 2. The solution of this dynamical system
is used not only to find the adapted parameters ${\bm{W}}_{\tau}(T)$, but also
to get the quantities necessary to compute the meta-gradients for the updates
of the meta-parameters ${\bm{W}}_{0}$ & $\Phi$. The integration of this
dynamical system, named ODESolve in Algorithm 1, is done in practice using a
numerical solver such as Runge-Kutta methods.

Algorithm 2 Dynamical system

function
Dynamics(${\bm{W}}_{0},f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}}),{\bm{s}}(t)$)

${\bm{W}}_{\tau}(t)\leftarrow{\bm{W}}_{0}-\sum_{m}{\bm{s}}_{m}(t){\bm{\phi}}_{m}^{\top}$
$\triangleright$ Proposition 1

${\bm{p}}_{m}(t)\leftarrow\mathrm{softmax}({\bm{W}}_{\tau}(t){\bm{\phi}}_{m})$

${\bm{A}}_{m}(t)\leftarrow\big{(}\operatorname{diag}({\bm{p}}_{m}(t))-{\bm{p}}_{m}(t){\bm{p}}_{m}(t)^{\top}\big{)}/M$

$\dfrac{d{\bm{s}}_{m}}{dt}\leftarrow\dfrac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}$
$\triangleright$ Proposition 1

$\displaystyle\dfrac{d{\bm{B}}_{t}[i,j]}{dt}\leftarrow\mathbbm{1}(i=j){\bm{A}}_{i}(t)-{\bm{A}}_{i}(t)\sum_{m=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{m}\big{)}{\bm{B}}_{t}[m,j]$
$\triangleright$ Propositions 2 and 3 ($\hookleftarrow$)

$\displaystyle\dfrac{d{\bm{z}}_{t}[i,j,m]}{dt}\leftarrow-{\bm{A}}_{i}(t)\bigg{[}\mathbbm{1}(i=j){\bm{s}}_{m}(t)+\mathbbm{1}(i=m){\bm{s}}_{j}(t)+\sum_{k=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{k}\big{)}{\bm{z}}_{t}[k,j,m]\bigg{]}$

return $d{\bm{s}}/dt,d{\bm{B}}_{t}/dt,d{\bm{z}}_{t}/dt$

end function

In Algorithms 3 and 4, we give the procedures responsible for the projection
of the partial derivatives
$\partial_{{\bm{W}}(T)}{\mathcal{L}}_{\tau}^{\mathrm{test}}$ onto the Jacobian
matrices $d{\bm{W}}(T)/d{\bm{W}}_{0}$ and $d{\bm{W}}(T)/d{\bm{\phi}}_{m}$
respectively, based on Propositions 5 and 6. Note that Algorithm 4 also
depends on the initial conditions ${\bm{W}}_{0}$, which is implicitly assumed
here for clarity of presentation. It is interesting to see that both functions
use the same matrix ${\bm{C}}$, and therefore this can be computed only once
and reused for both projections.

Algorithm 3 Projection onto $d{\bm{W}}(T)/d{\bm{W}}_{0}$

function Project${}_{{\bm{W}}_{0}}$(${\bm{V}}(T),{\bm{B}}_{T},{\bm{\phi}}$)

$\displaystyle{\bm{C}}_{j}\leftarrow\sum_{i=1}^{M}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,j]$

return ${\bm{V}}(T)-{\bm{C}}^{\top}{\bm{\phi}}$

end function

Algorithm 4 Projection onto $d{\bm{W}}(T)/d{\bm{\phi}}_{m}$

function
Project${}_{{\bm{\phi}}_{m}}$(${\bm{V}}(T),{\bm{s}}(T),{\bm{B}}_{T},{\bm{z}}_{T},{\bm{\phi}}$)

$\displaystyle{\bm{C}}_{j}\leftarrow\sum_{i=1}^{M}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,j]$

$\displaystyle{\bm{D}}_{m,j}\leftarrow\sum_{i=1}^{M}{\bm{z}}_{T}[i,j,m]^{\top}{\bm{V}}(T){\bm{\phi}}_{i}$

return
$-\big{[}{\bm{s}}_{m}(T)^{\top}{\bm{V}}(T)+{\bm{C}}_{m}{\bm{W}}_{0}+{\bm{D}}_{m}{\bm{\phi}}\big{]}$

end function

Finally in Algorithms 5 and 6, we show how to use COMLN at meta-test time on a
novel task, based on the learned meta-parameters ${\bm{W}}_{0}$, $\Phi$ and
$T$. This simply corresponds to integrating a dynamical system in
${\bm{W}}(t)$ (equivalently, in ${\bm{s}}_{m}(t)$, see Proposition 1). Note
that for adaptation during meta-testing, there is no need to compute either
${\bm{B}}_{t}[i,j]$ or ${\bm{z}}_{t}[i,j,m]$, since these are only necessary
to compute the meta-gradients during meta-training.

Algorithm 5 COMLN – Meta-test

A task $\tau$ with a dataset $\mathcal{D}_{\tau}^{\mathrm{train}}$

Meta-parameters $\Phi$, ${\bm{W}}_{0}$ & $T$

Embed the training dataset: $f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})$

${\bm{s}}(T)\leftarrow\mathrm{ODESolve}({\bm{0}},\mathrm{\textsc{Dynamics}},0,T)$

${\bm{W}}_{\tau}(T)\leftarrow{\bm{W}}_{0}-\sum_{m}{\bm{s}}_{m}(T){\bm{\phi}}_{m}^{\top}$

return ${\bm{W}}_{\tau}(T)\circ f_{\Phi}$

Algorithm 6 Dynamical system for adaptation

function
Dynamics(${\bm{W}}_{0},f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}}),{\bm{s}}(t)$)

${\bm{W}}_{\tau}(t)\leftarrow{\bm{W}}_{0}-\sum_{m}{\bm{s}}_{m}(t){\bm{\phi}}_{m}^{\top}$

${\bm{p}}_{m}(t)\leftarrow\mathrm{softmax}({\bm{W}}_{\tau}(t){\bm{\phi}}_{m})$

$d{\bm{s}}_{m}/dt\leftarrow\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}/M$

return $d{\bm{s}}/dt$

end function

###  A.2 Source code

We provide a snippet of code written in JAX (Bradbury et al., 2018) in order
to compute the adapted parameters, based on Algorithms 1 and 2. The dynamics
function not only computes the vectors ${\bm{s}}_{m}(t)$ necessary to compute
${\bm{W}}(t)$ (see Section B.1), but also ${\bm{B}}_{t}[i,j]$ and
${\bm{z}}_{t}[i,j,m]$ jointly in order to compute the meta-gradients. This
snippet shows that various necessary quantities can be precomputed ahead of
adaptation, such as the Gram matrix gram.

⬇

import jax.numpy as jnp

from jax import vmap, grad, nn, ops, tree_util

from jax.experimental import ode

from collections import namedtuple

State = namedtuple(’State’, [’s’, ’B’, ’z’])

M, N = train_inputs.shape[0], train_labels.shape[1]

gram = jnp.matmul(train_inputs, train_inputs.T)

logits_0 = jnp.matmul(train_inputs, W_0.T)

diag = jnp.diag_indices(M)

def dynamics(state, _):

preds = nn.softmax(logits_0 - jnp.matmul(gram, state.s), axis=1)

A = (vmap(jnp.diag)(preds) - vmap(jnp.outer)(preds, preds)) / M

# Update of s

ds = (preds - train_labels) / M

# Update of B

cross_prod = jnp.einsum(’ikn,im,mjnl->ijkl’, A, gram, state.B)

dB = ops.index_add(-cross_prod, diag, A)

# Update of z

cross_prod = jnp.einsum(’iln,ik,jmn->ijml’, A, gram, state.z)

A_s = jnp.einsum(’ikl,jl->ijk’, A, state.s)

dz = ops.index_add(cross_prod, diag, A_s)

dz = ops.index_add(dz, (diag[0], None, diag[1]), A_s)

return State(s=ds, B=dB, z=-dz)

state_0 = State(

s=jnp.zeros((M, N)),

B=jnp.zeros((M, M, N, N)),

z=jnp.zeros((M, M, M, N))

)

solution = ode.odeint(dynamics, state_0, jnp.array([0., T]))

state_T = tree_util.tree_map(lambda x: x[-1], solution)

W_T = W_0 - jnp.matmul(state_T.s.T, train_inputs)

Snippet of code in JAX to compute the adapted parameter W_T for a given task
specified by the embedded training set train_inputs & train_labels (note that
here train_inputs are the embedding vectors returned by $f_{\Phi}$) and an
initialization W_0.

We want to emphasize that all the information required to compute the meta-
gradients are available in state_T found after integration of the dynamics
function forward in time, and does not require any additional backward pass
(apart from backpropagation through the embedding network $f_{\Phi}$). In
particular, the meta-gradients wrt. the initial conditions ${\bm{W}}_{0}$ and
the wrt. the time horizon $T$ can be computed using the snippet of code
available in Code Snippet 5.

⬇

grads_test = grad(cross_entropy_loss)(W_T, test_inputs, test_labels)

grads_train = grad(cross_entropy_loss)(W_T, train_inputs, train_labels)

C = jnp.einsum(’in,kn,ijkl->jl’, train_inputs, grads_test, state_T.B)

grads_W_0 = grads_test - jnp.matmul(C.T, train_inputs)

grads_T = -jnp.vdot(grads_train, grads_test)

Snippet of code in JAX to compute the meta-gradients wrt. the meta-parameters
W_0 and T, based on state_T found above. See Section B.4 & Appendix C for
details.

The code to compute the meta-gradients wrt. the meta-parameters of the
embedding network $\Phi$ is not included here for clarity, as it involves non-
minimal dependencies on Haiku (Hennigan et al., 2020) in order to
backpropagate the error through the backbone network. The full code is
available at https://github.com/tristandeleu/jax-comln.

##  Appendix B Proofs of memory efficient meta-gradients

In this section, we will prove the results presented in Section 4 in a
slightly more general case, where the loss function ${\mathcal{L}}$ is the
cross-entropy loss regularized with a proximal term around the initialization
(Rajeswaran et al., 2019):

$\displaystyle{\mathcal{L}}\big{(}{\bm{W}};f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=-\frac{1}{M}\sum_{m=1}^{M}{\bm{y}}_{m}^{\top}\log{\bm{p}}_{m}+\frac{\lambda}{2}\|{\bm{W}}-{\bm{W}}_{0}\|_{F}^{2},$
(14)

where ${\bm{p}}_{m}=\mathrm{softmax}({\bm{W}}{\bm{\phi}}_{m})$,
$\|\cdot\|_{F}$ is the Frobenius norm, and $\lambda$ is a regularization
constant. Note that we can recover the setting presented in the main paper by
setting $\lambda=0$. The core idea of computing the meta-gradient efficiently
is based on the decomposition of the gradient (as well as the Hessian matrix)
of the cross-entropy loss (Böhning, 1992). We recall this decomposition as the
following lemma, extended to include the regularization term:

######  Lemma 1 (Böhning, 1992).

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
be the embedded training set through the embedding network $f_{\Phi}$, and
${\mathcal{L}}$ the regularized cross-entropy loss defined in Eq. 14. The
gradient $\nabla{\mathcal{L}}$ and the Hessian matrix
$\nabla^{2}{\mathcal{L}}$ of the regularized cross-entropy loss can be written
as

$\displaystyle\nabla{\mathcal{L}}\big{(}{\bm{W}};f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=\frac{1}{M}\sum_{m=1}^{M}\big{(}{\bm{p}}_{m}-{\bm{y}}_{m}\big{)}{\bm{\phi}}_{m}^{\top}+\lambda({\bm{W}}-{\bm{W}}_{0})$
$\displaystyle\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}};f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=\sum_{m=1}^{M}{\bm{A}}_{m}\otimes{\bm{\phi}}_{m}{\bm{\phi}}_{m}^{\top}+\lambda{\bm{I}},$

where
${\bm{A}}_{m}=\big{(}\operatorname{diag}({\bm{p}}_{m})-{\bm{p}}_{m}{\bm{p}}_{m}^{\top}\big{)}/M$
are $N\times N$ matrices, and $\otimes$ is the Kronecker product.

This lemma is particularly useful in the context of few-shot learning since it
reduces the characterization of the gradient and the Hessian from quantities
of size $N\times d$ and $Nd\times Nd$ respectively (where $d$ is the dimension
of the embedding vectors ${\bm{\phi}}$) to $M$ objects whose size is
independent of $d$—the embedding vectors ${\bm{\phi}}$ being independent of
${\bm{W}}$. Typically in few-shot learning, $M\ll d$.

In order to avoid higher-order tensors when defining the different Jacobians
and Hessians, we always consider them as matrices, with possibly an implicit
“flattening” operation. For example here even though ${\bm{W}}$ is a $N\times
d$ matrix, we treat the Hessian $\nabla^{2}{\mathcal{L}}$ as a $Nd\times Nd$
matrix, as opposed to a 4D tensor. When the context is required, we will make
this transformation from a higher-order tensor to a matrix explicit with an
encoding of indices. Moreover throughout this section, we will only consider
the computation of the meta-gradients for a single task $\tau$, and therefore
we will often drop the explicit dependence of the different objects on $\tau$
(e.g. we will write ${\bm{W}}(t)$ instead of ${\bm{W}}_{\tau}(t)$); the meta-
gradients are eventually averaged over a batch of tasks for the update of the
outer-loop, see Appendix A for details in the pseudo-code. Finally, since we
only consider adaptation of the last layer of the neural network, the
presentation here is always made in the context of an embedded training set
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
that went through the backbone $f_{\Phi}$.

###  B.1 Decomposition of $\bm{W}(t)$ for parameter adaptation

As a direct application of Lemma 1, we first decompose the parameters
${\bm{W}}(t)$ into smaller quantities ${\bm{s}}_{m}(t)$ that follow the
dynamics defined in Prop. 1. Although this decomposition is equivalent to
${\bm{W}}(t)$, in practice solving a smaller dynamical system in
${\bm{s}}_{m}(t)$ improves the efficiency of our method.

######  Proposition 1.

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
be the embedded training set through the embedding network $f_{\Phi}$, and
${\bm{W}}(t)$ be the solution of the following dynamical system

$\displaystyle\frac{d{\bm{W}}}{dt}$
$\displaystyle=-\nabla{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle{\bm{W}}(0)$ $\displaystyle={\bm{W}}_{0},$

where the loss function is the regularized cross-entropy loss defined in Eq.
14. The solution ${\bm{W}}(t)$ of this dynamical system can be written as

${\bm{W}}(t)={\bm{W}}_{0}-\sum_{m=1}^{M}{\bm{s}}_{m}(t){\bm{\phi}}_{m}^{\top},$

where for all $m$, ${\bm{s}}_{m}(t)$ is the solution of the following
dynamical system:

$\displaystyle\frac{d{\bm{s}}_{m}}{dt}$
$\displaystyle=\frac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}-\lambda{\bm{s}}_{m}(t)$
$\displaystyle{\bm{s}}_{m}(0)$ $\displaystyle={\bm{0}},$

and ${\bm{p}}_{m}(t)=\mathrm{softmax}({\bm{W}}(t){\bm{\phi}}_{m})$ is the
vector of predictions returned by the network at time $t$.

###### Proof.

Using Lemma 1, the function ${\bm{W}}(t)$ is the solution of the following
differential equation:

$\displaystyle\frac{d{\bm{W}}}{dt}$
$\displaystyle=-\nabla{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=-\frac{1}{M}\sum_{m=1}^{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}{\bm{\phi}}_{m}^{\top}-\lambda\big{(}{\bm{W}}(t)-{\bm{W}}_{0}\big{)}$
$\displaystyle{\bm{W}}(0)$ $\displaystyle={\bm{W}}_{0}.$

The proof relies on the unicity of the solution of a given autonomous
differential equation given a particular choice of initial conditions
(Wiggins, 2003, Prop. 7.4.2). In other words, if we can find another function
$\widetilde{{\bm{W}}}(t)$ that also satisfies the differential equation above
with the initial conditions $\widetilde{{\bm{W}}}(0)={\bm{W}}_{0}$, then it
means that for all $t$ we have $\widetilde{{\bm{W}}}(t)={\bm{W}}(t)$. Suppose
that this function $\widetilde{{\bm{W}}}(t)$ can be written as

$\widetilde{{\bm{W}}}(t)={\bm{W}}_{0}-\sum_{m=1}^{M}{\bm{s}}_{m}(t){\bm{\phi}}_{m}^{\top},$

where ${\bm{s}}_{m}$ satisfies the dynamical system defined in the statement
of Proposition 1. Then we have

$\displaystyle-\nabla{\mathcal{L}}\big{(}\widetilde{{\bm{W}}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle=-\frac{1}{M}\sum_{m=1}^{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}{\bm{\phi}}_{m}^{\top}-\lambda\big{(}\widetilde{{\bm{W}}}(t)-{\bm{W}}_{0}\big{)}$
$\displaystyle=-\sum_{m=1}^{M}\frac{d{\bm{s}}_{m}}{dt}{\bm{\phi}}_{m}^{\top}-\underbrace{\lambda\sum_{m=1}^{M}{\bm{s}}_{m}(t){\bm{\phi}}_{m}^{\top}-\lambda\big{(}\widetilde{{\bm{W}}}(t)-{\bm{W}}_{0}\big{)}}_{=\,{\bm{0}}}$
$\displaystyle=-\sum_{m=1}^{M}\frac{d{\bm{s}}_{m}}{dt}{\bm{\phi}}_{m}^{\top}=\frac{d\widetilde{{\bm{W}}}}{dt}$

We have shown that $\widetilde{{\bm{W}}}(t)$ follows the same dynamics as
${\bm{W}}(t)$. Moreover, using the initial conditions
${\bm{s}}_{m}(0)={\bm{0}}$, it is clear that
$\widetilde{{\bm{W}}}(0)={\bm{W}}_{0}$, which are the same initial conditions
as the equation satisfied by ${\bm{W}}(t)$. Therefore we have
$\widetilde{{\bm{W}}}(t)={\bm{W}}(t)$ for all $t$, showing the expected
decomposition of ${\bm{W}}(t)$. ∎

###  B.2 Decomposition of the Jacobian matrix $d\bm{W}(t)/d\bm{W}_{0}$

The core objective of gradient-based meta-learning methods is the capacity to
compute the meta-gradients wrt. the initial conditions ${\bm{W}}_{0}$. In
order to compute this meta-gradient using forward-mode automatic
differentiation, we want to first compute the Jacobian matrix
$d{\bm{W}}(t)/d{\bm{W}}_{0}$. We show that this Jacobian can be decomposed
into smaller quantities that follow the dynamics in Proposition 2.

######  Proposition 2.

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
be the embedded training set through the embedding network $f_{\Phi}$. The
Jacobian matrix $d{\bm{W}}(t)/d{\bm{W}}_{0}$ can be written as

$\frac{d{\bm{W}}(t)}{d{\bm{W}}_{0}}={\bm{I}}-\sum_{i=1}^{M}\sum_{j=1}^{M}{\bm{B}}_{t}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top},$

where $\otimes$ is the Kronecker product, and for all $i,j$,
${\bm{B}}_{t}[i,j]$ is a $N\times N$ matrix which the solution of the
following dynamical system

$\displaystyle\frac{d{\bm{B}}_{t}[i,j]}{dt}$
$\displaystyle=\mathbbm{1}(i=j){\bm{A}}_{i}(t)-\lambda{\bm{B}}_{t}[i,j]-{\bm{A}}_{i}(t)\sum_{m=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{m}\big{)}{\bm{B}}_{t}[m,j]$
$\displaystyle{\bm{B}}_{0}[i,j]$ $\displaystyle={\bm{0}},$

and
${\bm{A}}_{i}(t)=\big{(}\operatorname{diag}({\bm{p}}_{i}(t))-{\bm{p}}_{i}(t){\bm{p}}_{i}(t)^{\top}\big{)}/M$
are defined using the vectors of predictions at time $t$
${\bm{p}}_{i}(t)=\mathrm{softmax}\big{(}{\bm{W}}(t){\bm{\phi}}_{i}\big{)}$.

###### Proof.

We will use the forward sensitivity equation from Section 2.2 in order to
derive this new dynamical system over ${\bm{B}}_{t}[i,j]$. Recall that to
simplify the notations we can write the gradient vector field followed during
adaptation as

$\frac{d{\bm{W}}}{dt}=g\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}\triangleq-\nabla{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}.$

Introducing the matrix-valued function
${\mathcal{S}}(t)=d{\bm{W}}(t)/d{\bm{W}}_{0}$ as a sensitivity state, we can
use the forward sensitivity equations and see that ${\mathcal{S}}(t)$
satisfies the following equation

$\displaystyle\frac{d{\mathcal{S}}}{dt}$ $\displaystyle=\frac{\partial
g\big{(}{\bm{W}}(t)\big{)}}{\partial{\bm{W}}(t)}{\mathcal{S}}(t)+\frac{\partial
g\big{(}{\bm{W}}(t)\big{)}}{\partial{\bm{W}}_{0}}$
$\displaystyle{\mathcal{S}}(0)$ $\displaystyle={\bm{I}}$
$\displaystyle=-\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}{\mathcal{S}}(t)+\lambda{\bm{I}}.$

The rest of the proof is based on the unicity of the solution of a given
autonomous555While the dynamical system defined here for the sensitivity state
${\mathcal{S}}(t)$ alone is not exactly autonomous due to the dependence of
the Hessian matrix on ${\bm{W}}(t)$, we could augment ${\mathcal{S}}(t)$ with
${\bm{W}}(t)$ to obtain an autonomous system on the augmented state. We come
back to this distinction in Appendix D. The unicity argument still holds here
(the augmented solution would be unique). differential equation given a
particular choice of initial conditions (Wiggins, 2003, Prop. 7.4.2). In other
words, if we can find another function $\widetilde{{\mathcal{S}}}(t)$ that
also satisfies the above differential equation with the initial conditions
$\widetilde{{\mathcal{S}}}(0)={\bm{I}}$, then it means that for all $t$ we
have $\widetilde{{\mathcal{S}}}(t)={\mathcal{S}}(t)$. Suppose that this
function can be written as

$\widetilde{{\mathcal{S}}}(t)={\bm{I}}-\sum_{i=1}^{M}\sum_{j=1}^{M}{\bm{B}}_{t}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top},$

where ${\bm{B}}_{t}[i,j]$ satisfies the dynamical system defined in the
statement of Proposition 2. Then we have, using Lemma 1:

$\displaystyle-\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}\widetilde{{\mathcal{S}}}(t)+\lambda{\bm{I}}$
$\displaystyle\qquad=-\bigg{[}\sum_{i=1}^{M}{\bm{A}}_{i}(t)\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{i}^{\top}+\lambda{\bm{I}}\bigg{]}\bigg{[}{\bm{I}}-\sum_{m,j}{\bm{B}}_{t}[m,j]\otimes{\bm{\phi}}_{m}{\bm{\phi}}_{j}^{\top}\bigg{]}+\lambda{\bm{I}}$
$\displaystyle\qquad=-\sum_{i=1}^{M}{\bm{A}}_{i}(t)\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{i}^{\top}+\sum_{i,j,m}\big{(}{\bm{A}}_{i}(t){\bm{B}}_{t}[m,j]\big{)}\otimes\big{(}{\bm{\phi}}_{i}\underbrace{{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{m}}_{\in\,{\mathbb{R}}}{\bm{\phi}}_{j}^{\top}\big{)}+\lambda\sum_{i,j}{\bm{B}}_{t}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top}$
$\displaystyle\qquad=-\sum_{i,j}\bigg{[}\mathbbm{1}(i=j){\bm{A}}_{i}(t)-\lambda{\bm{B}}_{t}[i,j]-{\bm{A}}_{i}(t)\sum_{m=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{m}\big{)}{\bm{B}}_{t}[m,j]\bigg{]}\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top}$
$\displaystyle\qquad=-\sum_{i,j}\frac{d{\bm{B}}_{t}[i,j]}{dt}\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top}=\frac{d\widetilde{{\mathcal{S}}}}{dt}$

We have shown that $\widetilde{{\mathcal{S}}}(t)$ follows the same dynamics as
${\mathcal{S}}(t)$. Moreover, using the initial conditions
${\bm{B}}_{0}[i,j]={\bm{0}}$, it is clear that
$\widetilde{{\mathcal{S}}}(0)={\bm{I}}$, which are the same initial conditions
as the equation satisfied by ${\mathcal{S}}(t)$. Therefore, we have
$\widetilde{{\mathcal{S}}}(t)={\mathcal{S}}(t)$ for all $t$, showing the
expected decomposition of ${\mathcal{S}}(t)=d{\bm{W}}(t)/d{\bm{W}}_{0}$. ∎

###  B.3 Decomposition of the Jacobian matrix $d\bm{W}(t)/d\bm{\phi}_{m}$

Similar to Proposition 2, there exists a decomposition of the Jacobian matrix
$d{\bm{W}}(T)/d{\bm{\phi}}_{m}$. Recall that this Jacobian matrix appears in
the computation of the gradient of the outer-loss wrt. the embedding vectors
${\bm{\phi}}_{m}$, which is necessary in order to compute the meta-gradients
wrt. the meta-parameters of the embedding network $f_{\Phi}$ using
backpropagation from the last layer of $f_{\Phi}$.

######  Proposition 3.

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
be the embedded training set through the embedding network $f_{\Phi}$. For all
$m$, the Jacobian matrix $d{\bm{W}}(t)/d{\bm{\phi}}_{m}$ can be written as

$\frac{d{\bm{W}}(t)}{d{\bm{\phi}}_{m}}=-\bigg{[}{\bm{s}}_{m}(t)\otimes{\bm{I}}_{d}+\sum_{i=1}^{M}{\bm{B}}_{t}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}+\sum_{i=1}^{M}\sum_{j=1}^{M}{\bm{z}}_{t}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\bigg{]},$

where ${\bm{s}}_{m}(t)$ and ${\bm{B}}_{t}[i,j]$ are the solutions of the ODEs
defined in Propositions 1 and 2, and for all $i,j,m$, ${\bm{z}}_{t}[i,j,m]$ is
a vector of length $N$ solution of the following dynamical system:

$\frac{d{\bm{z}}_{t}[i,j,m]}{dt}=-{\bm{A}}_{i}(t)\bigg{[}\mathbbm{1}(i=j){\bm{s}}_{m}(t)+\mathbbm{1}(i=m){\bm{s}}_{j}(t)+\lambda{\bm{z}}_{t}[i,j,m]+\sum_{k=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{k}\big{)}{\bm{z}}_{t}[k,j,m]\bigg{]},$

with the initial conditions ${\bm{z}}_{0}[i,j,m]={\bm{0}}$.

###### Proof.

The outline of this proof follows the proof of Proposition 2: we first use the
forward sensitivity equations to get the dynamical system followed by the
Jacobian matrix $d{\bm{W}}(t)/d{\bm{\phi}}_{m}$, and then use a unicity
argument given that the new decomposition satisfies the same ODE. Recall that
we use the following notation to write the gradient vector field followed
during adaptation:

$\frac{d{\bm{W}}}{dt}=g\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}\triangleq-\nabla{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}.$

For a fixed $m$, we introduce the matrix valued function
${\mathcal{S}}(t)=d{\bm{W}}(t)/d{\bm{\phi}}_{m}$ as a sensitivity state. We
can use the forward sensitivity equations, and see that ${\mathcal{S}}(t)$
satisfies the following equation

$\displaystyle\frac{d{\mathcal{S}}}{dt}$ $\displaystyle=\frac{\partial
g\big{(}{\bm{W}}(t),{\bm{\phi}}_{m}\big{)}}{\partial{\bm{W}}(t)}{\mathcal{S}}(t)+\frac{\partial
g\big{(}{\bm{W}}(t),{\bm{\phi}}_{m}\big{)}}{\partial{\bm{\phi}}_{m}}$
$\displaystyle{\mathcal{S}}(0)$
$\displaystyle=\frac{d{\bm{W}}(0)}{d{\bm{\phi}}_{m}}={\bm{0}}$
$\displaystyle=-\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}{\mathcal{S}}(t)-{\bm{A}}_{m}(t){\bm{W}}_{0}\otimes{\bm{\phi}}_{m}$
$\displaystyle\qquad+{\bm{A}}_{m}(t)\sum_{i=1}^{M}{\bm{s}}_{i}(t){\bm{\phi}}_{i}^{\top}\otimes{\bm{\phi}}_{m}-\frac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}\otimes{\bm{I}}_{d}$

where we make the direct dependence of $g$ on ${\bm{\phi}}_{m}$ explicit, and
we used the decomposition of ${\bm{W}}(t)$ from Proposition 1. Suppose that we
define the function $\widetilde{{\mathcal{S}}}(t)$ as

$\widetilde{{\mathcal{S}}}(t)=-\bigg{[}{\bm{s}}_{m}(t)\otimes{\bm{I}}_{d}+\sum_{i=1}^{M}{\bm{B}}_{t}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}+\sum_{i,j}{\bm{z}}_{t}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\bigg{]},$

where ${\bm{z}}_{t}[i,j,m]$ satisfies the dynamical system defined in the
statement of Proposition 3. Then we have, using Lemma 1:

$\displaystyle-\\!\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}}(t);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}\widetilde{{\mathcal{S}}}(t)-{\bm{A}}_{m}(t)\bigg{[}{\bm{W}}_{0}-\sum_{i=1}^{M}{\bm{s}}_{i}(t){\bm{\phi}}_{i}^{\top}\bigg{]}\otimes{\bm{\phi}}_{m}-\frac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}\otimes{\bm{I}}_{d}$
$\displaystyle=\ $
$\displaystyle\bigg{[}\sum_{i=1}^{M}{\bm{A}}_{i}(t)\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{i}^{\top}+\lambda{\bm{I}}\bigg{]}\bigg{[}{\bm{s}}_{m}(t)\otimes{\bm{I}}_{d}+\sum_{i=1}^{M}{\bm{B}}_{t}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}+\sum_{i,j}{\bm{z}}_{t}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\bigg{]}$
$\displaystyle\qquad-{\bm{A}}_{m}(t){\bm{W}}_{0}\otimes{\bm{\phi}}_{m}+{\bm{A}}_{m}(t)\sum_{i=1}^{M}{\bm{s}}_{i}(t){\bm{\phi}}_{i}^{\top}\otimes{\bm{\phi}}_{m}-\frac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}\otimes{\bm{I}}_{d}$
$\displaystyle=\ $
$\displaystyle-\bigg{[}\underbrace{\vphantom{\sum_{i=1}^{M}}\frac{1}{M}\big{(}{\bm{p}}_{m}(t)-{\bm{y}}_{m}\big{)}-\lambda{\bm{s}}_{m}(t)}_{=\,d{\bm{s}}_{m}/dt}\bigg{]}\otimes{\bm{I}}_{d}$
$\displaystyle\qquad-\sum_{i=1}^{M}\bigg{[}\underbrace{\mathbbm{1}(i=m){\bm{A}}_{i}(t)-\lambda{\bm{B}}_{t}[i,m]-{\bm{A}}_{i}(t)\sum_{k=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{k}\big{)}{\bm{B}}_{t}[k,m]}_{=\,d{\bm{B}}_{t}[i,m]/dt}\bigg{]}{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}$
$\displaystyle\qquad+\sum_{i,j}\underbrace{{\bm{A}}_{i}(t)\bigg{[}\mathbbm{1}(i=j){\bm{s}}_{m}(t)+\mathbbm{1}(i=m){\bm{s}}_{j}(t)+\lambda{\bm{z}}_{t}[i,j,m]+\sum_{k=1}^{M}\big{(}{\bm{\phi}}_{i}^{\top}{\bm{\phi}}_{k}\big{)}{\bm{z}}_{t}[k,j,m]\bigg{]}}_{=\,-d{\bm{z}}_{t}[i,j,m]/dt}{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}$
$\displaystyle=\ $ $\displaystyle\frac{d\widetilde{{\mathcal{S}}}}{dt}$

We have shown that $\widetilde{{\mathcal{S}}}(t)$ follows the same dynamics as
${\mathcal{S}}(t)$. Moreover using the initial conditions
${\bm{s}}_{m}(0)={\bm{0}}$, ${\bm{B}}_{0}[i,j]={\bm{0}}$, and
${\bm{z}}_{0}[i,j,m]={\bm{0}}$, we have
$\widetilde{{\mathcal{S}}}(0)={\bm{0}}$, which are the same initial conditions
as the equation satisfied by ${\mathcal{S}}(t)$. Therefore we have
$\widetilde{{\mathcal{S}}}(t)={\mathcal{S}}(t)$ for all $t$, showing the
expected decomposition of the Jacobian matrix
${\mathcal{S}}(t)=d{\bm{W}}(t)/d{\bm{\phi}}_{m}$. ∎

###  B.4 Proof of the meta-gradient wrt. $T$

The novelty of COMLN over prior work on gradient-based meta-learning is the
ability to meta-learn the amount of adaptation $T$ using SGD. To compute the
meta-gradient wrt. the time horizon $T$, we can apply the result from Chen et
al. (2018).

######  Proposition 4.

Let $f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})$ and
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})$ be the embedding through the
network $f_{\Phi}$ of the training and test set respectively. The gradient of
the outer-loss wrt. the time horizon $T$ is given by:

$\frac{d{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})}{dT}=-\bigg{[}\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}\bigg{]}^{\top}\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}}{\partial{\bm{W}}(T)}.$

###### Proof.

We can directly apply the result from (Chen et al., 2018, App. B.2), which we
recall here for completeness. Given the ODE $d{\bm{z}}/dt=g({\bm{z}}(t))$, the
gradient of the loss $\widetilde{{\mathcal{L}}}\big{(}{\bm{z}}(T)\big{)}$ wrt.
$T$ is given by

$\displaystyle\frac{d\widetilde{{\mathcal{L}}}}{dT}$
$\displaystyle={\bm{a}}(T)^{\top}g({\bm{z}}(T))$
$\displaystyle\textrm{where}\quad{\bm{a}}(T)$
$\displaystyle=\frac{\partial\widetilde{{\mathcal{L}}}}{\partial{\bm{z}}(T)}$

is the _initial adjoint state_ (see Section 2.2). In our case we have

$\displaystyle g({\bm{z}}(T))$
$\displaystyle\equiv-\nabla{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})\big{)}$
$\displaystyle{\bm{a}}(T)$
$\displaystyle\equiv\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}.$

∎

##  Appendix C Projection onto the Jacobian matrices

Once we have computed the Jacobian $d{\bm{W}}(T)/d{\bm{\theta}}$ wrt. some
meta-parameter ${\bm{\theta}}$ (in our case, either the initial conditions
${\bm{W}}_{0}$ or the embedding vectors ${\bm{\phi}}_{m}$), we only have to
project the vector of partial derivatives onto it to obtain the meta-gradients
(vector-Jacobian product), using the chain rule:

$\frac{d{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{d{\bm{\theta}}}=\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}\frac{d{\bm{W}}(T)}{d{\bm{\theta}}}+\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{\theta}}}.$

While we have shown how to decompose the Jacobian matrices in such a way that
it only involves quantities that are independent of $d$ the dimension of the
embedding vectors ${\bm{\phi}}$, this final projection a priori requires us to
form the Jacobian explicitly. Even though this operation is done only once at
the end of the integration, this may be overly expensive since the Jacobian
matrices scale quadratically with $d$, which can be as high as $d=16,\\!000$
in our experiments.

Fortunately, we can perform this projection as a function of
${\bm{s}}_{m}(T)$, ${\bm{B}}_{T}[i,j]$, and ${\bm{z}}_{T}[i,j,m]$, without
having to explicitly form the full Jacobian matrix, by exchanging the order of
operations.

######  Proposition 5.

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
and $f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})$ be the embedded training
set and test set through the embedding network $f_{\Phi}$ respectively. Let
${\bm{B}}_{T}[i,j]$ be the solution at time $T$ of the differential equation
defined in Proposition 2, and ${\bm{V}}(T)\in{\mathbb{R}}^{N\times d}$ the
partial derivative of the outer-loss wrt. the adapted parameters
${\bm{W}}(T)$:

$\displaystyle{\bm{V}}(T)$
$\displaystyle=\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}$
$\displaystyle\textrm{so
that}\quad\frac{d{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{d{\bm{W}}_{0}}$
$\displaystyle=\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{T}\frac{d{\bm{W}}(T)}{d{\bm{W}}_{0}}.$

Let
${\bm{\phi}}=[{\bm{\phi}}_{1},\ldots,{\bm{\phi}}_{M}]^{\top}\in{\mathbb{R}}^{M\times
d}$ be the design matrix. The gradient of the outer-loss wrt. the initial
conditions ${\bm{W}}_{0}$ can be computed as

$\displaystyle\frac{d{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{d{\bm{W}}_{0}}$
$\displaystyle={\bm{V}}(T)-{\bm{C}}^{\top}{\bm{\phi}},$

where the rows of ${\bm{C}}\in{\mathbb{R}}^{M\times N}$ are defined by

${\bm{C}}_{j}=\sum_{i=1}^{M}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,j].$

###### Proof.

In this proof, we will make the encoding of the indices for the Jacobian
$d{\bm{W}}(T)/d{\bm{W}}_{0}$ more explicit, as mentioned in Appendix B. Recall
from Proposition 2 that this Jacobian matrix can be decomposed as

$\frac{d{\bm{W}}(T)}{d{\bm{W}}_{0}}={\bm{I}}-\sum_{i,j}{\bm{B}}_{T}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top}.$

Introducing the following notations

$\displaystyle{\bm{F}}$
$\displaystyle\triangleq\sum_{i,j}{\bm{B}}_{T}[i,j]\otimes{\bm{\phi}}_{i}{\bm{\phi}}_{j}^{\top}\in{\mathbb{R}}^{Nd\times
Nd}$ $\displaystyle{\bm{G}}$
$\displaystyle\triangleq\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{\top}{\bm{F}}\in{\mathbb{R}}^{Nd},$

for all $l\in\\{0,\ldots,N-1\\}$ and $y\in\\{0,\ldots,d-1\\}$:

$\displaystyle{\bm{G}}[dl+y]$
$\displaystyle=\sum_{k=0}^{N-1}\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{F}}[dk+x,dl+y]$
$\displaystyle=\sum_{i,j}\sum_{k=0}^{N-1}\bigg{[}\underbrace{\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{\phi}}_{i}[x]}_{=\,[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}}\bigg{]}{\bm{B}}_{T}[i,j,k,l]{\bm{\phi}}_{j}[y]$
$\displaystyle=\sum_{i,j}\bigg{[}\underbrace{\sum_{k=0}^{N-1}[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}{\bm{B}}_{T}[i,j,k,l]}_{=\,[{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,j]]_{l}}\bigg{]}{\bm{\phi}}_{j}[y]$
$\displaystyle=\sum_{j=1}^{M}\bigg{[}\underbrace{\sum_{i=1}^{M}\big{[}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,j]\big{]}_{l}}_{=\,{\bm{C}}_{j,l}}\bigg{]}{\bm{\phi}}_{j}[y]=\big{[}{\bm{C}}^{\top}{\bm{\phi}}\big{]}_{l,y}$

This shows that ${\bm{G}}$ is equal to ${\bm{C}}^{\top}{\bm{\phi}}$ up to
reshaping. Using the full form of the Jacobian, including ${\bm{I}}$,
concludes the proof. ∎

An interesting observation is that if the term ${\bm{C}}^{\top}{\bm{\phi}}$ is
ignored in the computation of the meta-gradient, we recover an equivalent of
the first-order approximation introduced in Finn et al. (2017). Similarly, we
can show that we can perform the projection onto the Jacobian matrix
$d{\bm{W}}(T)/d{\bm{\phi}}_{m}$ without having to form the explicit $Nd\times
d$ matrix.

######  Proposition 6.

Let
$f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{train}})=\\{({\bm{\phi}}_{m},{\bm{y}}_{m})\\}_{m=1}^{M}$
and $f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})$ be the embedded training
set and test set through the embedding network $f_{\Phi}$ respectively. Let
${\bm{s}}_{m}(T)$ be the solution at time $T$ of the differential equation
defined in Proposition 1, ${\bm{B}}_{T}[i,j]$ the solution of the one defined
in Proposition 2, and ${\bm{z}}_{T}[i,j,m]$ the solution of the one defined in
Proposition 3. Let ${\bm{V}}(T)\in{\mathbb{R}}^{N\times d}$ be the partial
derivative of the outer-loss wrt. the adapted parameters ${\bm{W}}(T)$:

${\bm{V}}(T)=\frac{\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}}{\partial{\bm{W}}(T)}.$

Let
${\bm{\phi}}=[{\bm{\phi}}_{1},\ldots,{\bm{\phi}}_{M}]^{\top}\in{\mathbb{R}}^{M\times
d}$ be the design matrix. The projection of ${\bm{V}}(T)$ onto the Jacobian
matrix $d{\bm{W}}(T)/d{\bm{\phi}}_{m}$ can be computed as

$\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{\top}\frac{d{\bm{W}}(T)}{d{\bm{\phi}}_{m}}=-\big{[}{\bm{s}}_{m}(T)^{\top}{\bm{V}}(T)+{\bm{C}}_{m}{\bm{W}}_{0}+{\bm{D}}_{m}{\bm{\phi}}\big{]},$

where ${\bm{C}}_{m}\in{\mathbb{R}}^{N}$ and ${\bm{D}}_{m}\in{\mathbb{R}}^{M}$
are defined by

$\displaystyle{\bm{C}}_{m}$
$\displaystyle=\sum_{i=1}^{M}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,m]$
$\displaystyle{\bm{D}}_{m,j}$
$\displaystyle=\sum_{i=1}^{M}{\bm{z}}_{T}[i,j,m]^{\top}{\bm{V}}(T){\bm{\phi}}_{i}$

Note that the definition of ${\bm{C}}$ in Proposition 6 matches exactly the
definition of ${\bm{C}}$ in Proposition 5, and therefore we only need to
compute this matrix once to perform the projections for both meta-gradients.

###### Proof.

Recall from Proposition 3 that the Jacobian $d{\bm{W}}(T)/d{\bm{\phi}}_{m}$
can be decomposed as

$\frac{d{\bm{W}}(T)}{d{\bm{\phi}}_{m}}=-\bigg{[}{\bm{s}}_{m}(T)\otimes{\bm{I}}_{d}+\sum_{i=1}^{M}{\bm{B}}_{T}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}+\sum_{i,j}{\bm{z}}_{T}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\bigg{]}.$

We will consider each of these 3 terms separately.

  * •

For the first term, let

$\displaystyle{\bm{F}}_{1}$
$\displaystyle\triangleq{\bm{s}}_{m}(T)\otimes{\bm{I}}_{d}\in{\mathbb{R}}^{Nd\times
d}$ $\displaystyle{\bm{G}}_{1}$
$\displaystyle\triangleq\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{\top}{\bm{F}}_{1}\in{\mathbb{R}}^{d},$

and for $y\in\\{0,\ldots,d-1\\}$:

$\displaystyle{\bm{G}}_{1}[y]$
$\displaystyle=\sum_{k=0}^{N-1}\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{F}}_{1}[dk+x,y]=\sum_{k=0}^{N-1}\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]\mathbbm{1}(x=y)\big{[}{\bm{s}}_{m}(T)\big{]}_{k}$
$\displaystyle=\sum_{k=0}^{N-1}\big{[}{\bm{s}}_{m}(T)\big{]}_{k}{\bm{V}}_{T}[k,y]=\big{[}{\bm{s}}_{m}(T)^{\top}{\bm{V}}(T)\big{]}_{y}$

Therefore, ${\bm{G}}_{1}$ the projection of ${\bm{V}}(T)$ onto the first term
of the Jacobian matrix is equal to ${\bm{s}}_{m}(T)^{\top}{\bm{V}}(T)$.

  * •

For the second term, let

$\displaystyle{\bm{F}}_{2}$
$\displaystyle\triangleq\sum_{i=1}^{M}{\bm{B}}_{T}[i,m]{\bm{W}}_{0}\otimes{\bm{\phi}}_{i}\in{\mathbb{R}}^{Nd\times
d}$ $\displaystyle{\bm{G}}_{2}$
$\displaystyle\triangleq\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{\top}{\bm{F}}_{2}\in{\mathbb{R}}^{d},$

and for $y\in\\{0,\ldots,d-1\\}$:

$\displaystyle{\bm{G}}_{2}[y]$
$\displaystyle=\sum_{k=0}^{N-1}\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{F}}_{2}[dk+x,y]$
$\displaystyle=\sum_{i=1}^{M}\sum_{l=0}^{N-1}\sum_{k=0}^{N-1}\bigg{[}\underbrace{\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{\phi}}_{i}[x]}_{=\,[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}}\bigg{]}{\bm{B}}_{T}[i,m,k,l]{\bm{W}}_{0}[l,y]$
$\displaystyle=\sum_{i=1}^{M}\sum_{l=0}^{N-1}\bigg{[}\underbrace{\sum_{k=0}^{N-1}[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}{\bm{B}}_{T}[i,m,k,l]}_{=\,\big{[}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,m]\big{]}_{l}}\bigg{]}{\bm{W}}_{0}[l,y]$
$\displaystyle=\sum_{l=0}^{N-1}\bigg{[}\underbrace{\sum_{i=1}^{M}\big{[}{\bm{\phi}}_{i}^{\top}{\bm{V}}(T)^{\top}{\bm{B}}_{T}[i,m]\big{]}_{l}}_{=\,{\bm{C}}_{m,l}}\bigg{]}{\bm{W}}_{0}[l,y]=\big{[}{\bm{C}}_{m}{\bm{W}}_{0}\big{]}_{y}$

${\bm{G}}_{2}$ the projection of ${\bm{V}}(T)$ onto the second term of the
Jacobian matrix is equal to ${\bm{C}}_{m}{\bm{W}}_{0}$.

  * •

Finally for the third term, let

$\displaystyle{\bm{F}}_{3}$
$\displaystyle\triangleq\sum_{i,j}{\bm{z}}_{T}[i,j,m]{\bm{\phi}}_{j}^{\top}\otimes{\bm{\phi}}_{i}\in{\mathbb{R}}^{Nd\times
d}$ $\displaystyle{\bm{G}}_{3}$
$\displaystyle\triangleq\mathrm{Vec}\big{(}{\bm{V}}(T)\big{)}^{\top}{\bm{F}}_{3}\in{\mathbb{R}}^{d}$

and for $y\in\\{0,\ldots,d-1\\}$:

$\displaystyle{\bm{G}}_{3}[y]$
$\displaystyle=\sum_{k=0}^{N-1}\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{F}}_{3}[dk+x,y]$
$\displaystyle=\sum_{i,j}\sum_{k=0}^{N-1}\bigg{[}\underbrace{\sum_{x=0}^{d-1}{\bm{V}}_{T}[k,x]{\bm{\phi}}_{i}[x]}_{=\,[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}}\bigg{]}{\bm{z}}_{T}[i,j,m,k]{\bm{\phi}}_{j}[y]$
$\displaystyle=\sum_{i,j}\bigg{[}\underbrace{\sum_{k=0}^{N-1}[{\bm{V}}(T){\bm{\phi}}_{i}]_{k}{\bm{z}}_{T}[i,j,m,k]}_{=\,{\bm{z}}_{T}[i,j,m]^{\top}{\bm{V}}(T){\bm{\phi}}_{i}\in{\mathbb{R}}}\bigg{]}{\bm{\phi}}_{j}[y]$
$\displaystyle=\sum_{j=1}^{M}\bigg{[}\underbrace{\sum_{i=1}^{M}{\bm{z}}_{T}[i,j,m]^{\top}{\bm{V}}(T){\bm{\phi}}_{i}}_{=\,{\bm{D}}_{m,j}}\bigg{]}{\bm{\phi}}_{j}[y]=\big{[}{\bm{D}}_{m}{\bm{\phi}}\big{]}_{y}$

${\bm{G}}_{3}$ the projection of ${\bm{V}}(T)$ onto the third term of the
Jacobian is equal to ${\bm{D}}_{m}{\bm{\phi}}$.

∎

Note that the projection of ${\bm{V}}(T)$ as defined in Proposition 6 is not
equal to the meta-gradient itself, as it is missing the term coming from the
partial derivative of the outer-loss wrt. the embedding vector
$\partial{\mathcal{L}}\big{(}{\bm{W}}(T);f_{\Phi}(\mathcal{D}_{\tau}^{\mathrm{test}})\big{)}/\partial{\bm{\phi}}_{m}$,
which is non-zero unlike the counterpart for ${\bm{W}}_{0}$. However, this
projection is the only expensive operation required to compute the total
gradient.

##  Appendix D Proof of stability

In this section, we would like to show that unlike the adjoint method (see
Sections 2.2 and 3.2), our method to compute the meta-gradients of COMLN is
guaranteed to be _stable_. In other words, even under some small perturbations
due to the ODE solver, the solution found by numerical integration is going to
stay close to the true solution of the dynamical system. To do so, we use the
concept of Lyapunov stability of the solution of an ODE, which we recall here
for completeness:

######  Definition 1 (Lyapunov stability; Wiggins, 2003).

The solution $\bar{{\bm{x}}}(t)$ of a dynamical system is said to be _Lyapunov
stable_ if, given $\varepsilon>0$, there exists $\delta$ such that for any
other solution ${\bm{x}}(t)$ such that
$\|\bar{{\bm{x}}}(0)-{\bm{x}}(0)\|<\delta$, we have
$\|\bar{{\bm{x}}}(t)-{\bm{x}}(t)\|<\varepsilon$ for all $t>0$.

We would like to understand the conditions needed for a function
$f:{\mathbb{R}}^{n}\rightarrow{\mathbb{R}}^{n}$, such that a trajectory
$\bar{{\bm{x}}}(t)$ satisfying the following autonomous differential equation
is (Lyapunov) stable:

$\frac{d{\bm{x}}}{dt}=f\big{(}{\bm{x}}(t)\big{)}$

If we assume for a moment that $f({\bm{x}})=-\nabla{\mathcal{L}}({\bm{x}})$
where ${\mathcal{L}}$ is a convex function, then for any two trajectories
${\bm{x}}_{1}(t)$ & ${\bm{x}}_{2}(t)$ of the dynamical system above, we can
see that the function $F$ defined by

$F(t)=\frac{1}{2}\|{\bm{x}}_{1}(t)-{\bm{x}}_{2}(t)\|^{2}$

satisfies $\dot{F}(t)\leq 0$, where we use the notation $\dot{F}(t)\triangleq
dF/dt$. Indeed, since $\nabla^{2}{\mathcal{L}}$ is positive semi-definite and
by defining ${\bm{h}}(t)={\bm{x}}_{2}(t)-{\bm{x}}_{1}(t)$, we get

$\displaystyle\dot{F}(t)$
$\displaystyle=\frac{dF}{dt}=\bigg{[}\frac{dF}{d{\bm{h}}}\bigg{]}^{\top}\frac{d{\bm{h}}}{dt}$
$\displaystyle=\big{(}{\bm{x}}_{2}(t)-{\bm{x}}_{1}(t)\big{)}^{\top}\big{(}\dot{{\bm{x}}}_{2}(t)-\dot{{\bm{x}}}_{1}(t)\big{)}$
$\displaystyle=\big{(}{\bm{x}}_{2}(t)-{\bm{x}}_{1}(t)\big{)}^{\top}\big{(}f({\bm{x}}_{2}(t))-f({\bm{x}}_{1}(t))\big{)}$
$\displaystyle={\bm{h}}(t)^{\top}\left[\int_{0}^{1}Df\big{(}{\bm{x}}_{1}(t)+s{\bm{h}}(t)\big{)}ds\right]{\bm{h}}(t)$
(15)
$\displaystyle=-\int_{0}^{1}\underbrace{{\bm{h}}(t)^{\top}\nabla^{2}{\mathcal{L}}\big{(}{\bm{x}}_{1}(t)+s{\bm{h}}(t)\big{)}{\bm{h}}(t)}_{\geq
0}ds$

where Eq. 15 follows from the fundamental theorem of calculus. Now since
$\dot{F}(t)\leq 0$, the distance between two trajectories decreases as time
$t$ increases, hence we have Lyapunov stability for all trajectories or
solutions of the above autonomous differential equation.

Now note that we didn’t actually need $f$ to be of the specific form
$f({\bm{x}})=-\nabla{\mathcal{L}}({\bm{x}})$, but having the Jacobian matrix
$Df$ negative semi-definite everywhere would have been sufficient. Hence we
get the following statement

######  Proposition 7.

If a continuously differentiable function
$f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}$ is such that the Jacobian matrix
$Df$ is negative semi-definite everywhere, then any solution
$\bar{{\bm{x}}}(t)$ of the autonomous differential equation

$\frac{d{\bm{x}}}{dt}=f\big{(}{\bm{x}}(t)\big{)}$

with initial condition $\bar{{\bm{x}}}(0)={\bm{x}}_{0}$ is (Lyapunov) stable.

We can extend the above result to the non-autonomous case. For this purpose,
let us define

$\frac{d{\bm{x}}}{dt}=f\big{(}t,{\bm{x}}(t)\big{)}$ (16)

which induces the parametrized function $f_{t}:{\bm{x}}\mapsto f(t,{\bm{x}})$.
With this notation we can state the following result.

######  Proposition 8.

If the Jacobian matrix $Df_{t}$ is negative semi-definite everywhere for all
$t\geq 0$, then any trajectory ${\bm{x}}(t)$, solution of Eq. 16 with initial
condition ${\bm{x}}(0)={\bm{x}}_{0}$, is (Lyapunov) stable.

###### Proof.

Let us start with two trajectories ${\bm{x}}_{1}(t)$ and ${\bm{x}}_{2}(t)$,
which are solutions to the above Eq. 16. By defining
${\bm{h}}(t)={\bm{x}}_{2}(t)-{\bm{x}}_{1}(t)$ and by applying the fundamental
theorem of calculus, we get

$\displaystyle
f\big{(}t,{\bm{x}}_{1}(t)\big{)}-f\big{(}t,{\bm{x}}_{2}(t)\big{)}$
$\displaystyle=f_{t}\big{(}{\bm{x}}_{1}(t)\big{)}-f_{t}\big{(}{\bm{x}}_{2}(t)\big{)}$
$\displaystyle=\left[\int_{0}^{1}Df_{t}\big{(}{\bm{x}}_{1}(t)+s{\bm{h}}(t)\big{)}ds\right]{\bm{h}}(t)$

Following the same idea as already previously outlined, let us define
$F(t)=\frac{1}{2}\|{\bm{x}}_{1}(t)-{\bm{x}}_{2}(t)\|^{2}$ and show that
$\dot{F}(t)\leq 0$:

$\displaystyle\dot{F}(t)$
$\displaystyle=({\bm{x}}_{2}(t)-{\bm{x}}_{1}(t))^{\top}\big{(}f(t,{\bm{x}}_{2}(t))-f(t,{\bm{x}}_{1}(t))\big{)}$
$\displaystyle={\bm{h}}(t)^{\top}\left[\int_{0}^{1}Df_{t}\big{(}{\bm{x}}_{1}(t)+s{\bm{h}}(t)\big{)}ds\right]{\bm{h}}(t)$
$\displaystyle=\int_{0}^{1}\underbrace{{\bm{h}}(t)^{T}Df_{t}\big{(}{\bm{x}}_{1}(t)+s{\bm{h}}(t)\big{)}{\bm{h}}(t)}_{\leq
0}ds\leq 0$

Hence the distance between two trajectories decreases as time $t$ increases,
and we have Lyapunov stability for all solutions of the non-autonomous
differential equation above. ∎

###  D.1 Stability of the forward sensitivity equations

In this subsection let us start by considering the general case of solving the
initial value problem to the following autonomous system of differential
equations:

$\centering\left\\{\begin{split}\frac{d{\bm{W}}}{dt}&=g({\bm{W}}(t),{\bm{\theta}})&{\bm{W}}(0)&={\bm{W}}_{0}\\\
\frac{d{\mathcal{S}}}{dt}&=\frac{\partial
g\big{(}{\bm{W}}(t),{\bm{\theta}}\big{)}}{\partial{\bm{W}}(t)}\cdot{\mathcal{S}}(t)+\frac{\partial
g\big{(}{\bm{W}}(t),{\bm{\theta}}\big{)}}{\partial{\bm{\theta}}}&{\mathcal{S}}(0)&=\frac{\partial{\bm{W}}(0)}{\partial{\bm{\theta}}}.\\\
\end{split}\right.\@add@centering$ (17)

######  Proposition 9.

There exists a solution of the autonomous system of differential equations in
Eq. 17, which is also unique.

###### Proof.

By Theorem 7.1.1 in Wiggins (2003), we know that there exists a solution of
the system of differential equations in Eq. 17. Then, by Theorem 7.4.1 in
Wiggins (2003), we can conclude that this solution is unique upon the choice
of initial conditions ${\bm{W}}(0)={\bm{W}}_{0}$ and
${\mathcal{S}}(0)=\partial{\bm{W}}(0)/\partial{\bm{\theta}}$.

Alternatively, one could also apply Theorems 7.1.1 & 7.4.1 in Wiggins (2003)
to conclude that the initial value problem
$d{\bm{W}}/dt=g({\bm{W}}(t),{\bm{\theta}})$ with initial condition
${\bm{W}}(0)={\bm{W}}_{0}$ has a unique solution. This existence and
uniqueness then gives rise to existence and uniqueness of a solution of the
entire system of differential equations via ${\mathcal{S}}(t)=d{\bm{W}}/dt$ by
applying Clairnaut’s theorem. ∎

In our case, we have
$g({\bm{W}}(t),\theta)=-\nabla{\mathcal{L}}\big{(}{\bm{W}}(T),{\bm{\theta}}\big{)}$,
where ${\mathcal{L}}$ is a convex function in ${\bm{W}}$. Hence the autonomous
system of differential equations can be rewritten as

$\centering\left\\{\begin{split}\frac{d{\bm{W}}}{dt}&=-\nabla{\mathcal{L}}\big{(}{\bm{W}}(T),{\bm{\theta}}\big{)}&{\bm{W}}(0)&={\bm{W}}_{0}\\\
\frac{d{\mathcal{S}}}{dt}&=-\nabla^{2}{\mathcal{L}}\big{(}{\bm{W}}(T),{\bm{\theta}}\big{)}{\mathcal{S}}(t)+\frac{\partial
g\big{(}{\bm{W}}(t),{\bm{\theta}}\big{)}}{\partial{\bm{\theta}}}&{\mathcal{S}}(0)&=\frac{\partial{\bm{W}}(0)}{\partial{\bm{\theta}}}.\\\
\end{split}\right.\@add@centering$ (18)

######  Proposition 10.

Let ${\bm{W}}(t)$ be the solution of

$\displaystyle\frac{d{\bm{W}}}{dt}$
$\displaystyle=-\nabla{\mathcal{L}}\big{(}{\bm{W}}(T),{\bm{\theta}}\big{)}$
$\displaystyle\qquad{\bm{W}}(0)$ $\displaystyle={\bm{W}}_{0}$

Then the trajectory ${\bm{W}}(t)$ is (Lyapunov) stable.

###### Proof.

Note that $Dg\big{(}{\bm{W}}(t)\big{)}=-\nabla^{2}{\mathcal{L}}$ here. Since
${\mathcal{L}}$ is convex, $\nabla^{2}{\mathcal{L}}$ is positive semi-
definite, and hence $Dg({\bm{W}}(t))$ is negative semi-definite. Hence the
result directly follows from Proposition 7. ∎

In addition we can now also conclude that the solution of Eq. 18 is also
(Lyapunov) stable.

######  Corollary 1.

The solution $\big{[}{\bm{W}}(t),{\mathcal{S}}(t)\big{]}$ of Eq. 18 is
(Lyapunov) stable.

This result is general to the application of the forward sensitivity equations
on a gradient vector field derived from a convex loss function. We can also
show that when the gradient $d{\bm{W}}^{\star}/d{\bm{\theta}}$ exists and is
finite (recall that ${\bm{W}}^{\star}$ is the minimizer of the loss function,
and the equilibrium of the gradient vector field), then the solution of the
system is also bounded, which guarantees us that our solution will not diverge
(unlike the adjoint method applied to a gradient vector field).

######  Proposition 11.

Assuming that $\big{\|}d{\bm{W}}^{\star}/d{\bm{\theta}}\big{\|}$ is finite,
the solution $\big{[}{\bm{W}}(t),{\mathcal{S}}(t)\big{]}$ of Eq. 18 is
bounded.

###### Proof.

Let us define

$\displaystyle V(t)$
$\displaystyle=\frac{1}{2}\|{\bm{W}}(t)\|^{2}+\frac{1}{2}\|{\mathcal{S}}(t)\|^{2}$

Since ${\bm{W}}(t)$ and ${\mathcal{S}}(t)$ are continuous functions in $t$,
verifying

$\displaystyle{\bm{W}}(t)$
$\displaystyle\underset{t\rightarrow\infty}{\longrightarrow}{\bm{W}}^{\star}$
and
$\displaystyle{\mathcal{S}}(t)\underset{t\rightarrow\infty}{\longrightarrow}\frac{d{\bm{W}}^{\star}}{d{\bm{\theta}}}$

Hence

$V(t)\underset{t\rightarrow\infty}{\longrightarrow}\frac{1}{2}\|{\bm{W}}^{\star}\|^{2}+\frac{1}{2}\Big{\|}\frac{d{\bm{W}}^{\star}}{d{\bm{\theta}}}\Big{\|}^{2}<\infty,$

where the last point follows from our assumption that
$\big{\|}d{\bm{W}}(t)/d{\bm{\theta}}\big{\|}<\infty$. We also have that

$V(0)=\frac{1}{2}\|{\bm{W}}_{0}\|^{2}+\frac{1}{2}\Big{\|}\frac{d{\bm{W}}(0)}{d{\bm{\theta}}}\Big{\|}^{2}<\infty.$

$V(t)$ being continuous as a composition of continuous functions, we conclude
that $V$ is bounded, and hence ${\bm{W}}(t)$ and ${\mathcal{S}}(t)$ are
bounded as well. ∎

##  Appendix E Experimental details

###  E.1 Choice of the numerical solver

In Section 5.2, we showed that the numerical solver had a significant impact
on the time to compute the meta-gradients (but not the memory requirements).
In Table 4, we show that using explicit Euler instead of an adaptive scheme
like Runge-Kutta leads to very similar results. In all our experiments we
therefore chose Runge-Kutta for its faster execution, and by convenience—it is
the default numerical solver available in JAX (Bradbury et al., 2018).

Table 4: The effect of the numerical solver on the performance of COMLN. The average accuracy (%) on $1,000$ held-out meta-test tasks is reported with $95\%$ confidence interval. Note that for a given setting, the same $1,000$ tasks are used for evaluation, making both methods directly comparable. RK: 4th-order Runge-Kutta with Dormand Prince adaptive step size. Euler: explicit Euler scheme. Method | miniImageNet 5-way
---|---
1-shot | 5-shot
COMLN (RK) | $53.01\pm 0.62$ | $70.54\pm 0.54$
COMLN (Euler) | $53.00\pm 0.83$ | $70.50\pm 0.72$

###  E.2 Details for measuring the efficiency

In Figures 3 & 6, we show the efficiency, both in terms of memory and runtime,
of COMLN compared to other meta-learning algorithms. The computational
efficiency (runtime) was measured as the average time taken to compute the
meta-gradients on a single 5-shot 5-way task from the miniImageNet dataset
over 100 runs.

[width=]figures/pgfplots/resnet12_memory_runtime

Figure 6: Empirical efficiency of COMLN on a single 5-shot 5-way task, with a
ResNet-12 backbone; this figure is similar to Fig. 3. (Left) Memory usage for
computing the meta-gradients as a function of the number of inner-gradient
steps. The extrapolated dashed lines correspond to the method reaching the
memory capacity of a Tesla V100 GPU with 32Gb of memory. (Right) Average time
taken (in ms) to compute the exact meta-gradients. The extrapolated dashed
lines correspond to the method taking over 3 seconds.

In order to ensure fair comparison between methods that rely on a discrete
number of gradient steps (MAML, ANIL, and iMAML), and COMLN which relies on a
continuous integration time $T$, we added a conversion between the number of
gradient steps and $T$. This corresponds to taking a learning rate of
$\alpha=0.01$ in Eq. 2 (which is standard practice for MAML and ANIL on
miniImageNet). This means that the number of gradient steps is $100\times$
larger than $T$. This can be formally justified by considering an explicit
Euler scheme for COMLN (see Section 3.3), with a constant step size
$\alpha=0.01$. The memory requirements is independent of the choice of the
numerical solver. For the computation time, this correspondence between $T$
and the number of gradient steps is no longer exact for _COMLN (RK)_ —the
comparison with _COMLN (Euler)_ is still valid though.

###  E.3 Analysis of the learned horizon

[width=]figures/pgfplots/evolution_and_fixed_T

Figure 7: (Left) Evolution of the meta-parameter $T$ controlling the amount of
adaptation necessary for all tasks during meta-training. Here, the backbone is
a ResNet-12. We normalized the duration of meta-training in $[0,1]$ to account
for early stopping; typically the model for miniImageNet requires an order of
magnitude fewer iterations. (Right) Comparison of COMLN (where $T$ is learned,
in red) to meta-learning with a fixed length of adaptation $T$ (in blue), on a
5-shot 5-way classification problem on the miniImageNet dataset (with a Conv-4
backbone).

Since one of the advantage of COMLN compared to other gradient-based methods
is its capacity to learn the amount of adaptation through the time horizon
$T$, Fig. 7 (left) shows the evolution of this meta-parameter during meta-
training. We can observe that for the more complex dataset tieredImageNet,
COMLN appropriately learns to use a longer sequence of adaptation. Similarly
within each dataset, it also learns to use shorter sequences of adaptation of
1-shot problems, possibly to allow for better generalization and to reduce
overfitting.

Besides adapting the amount of adaptation to the problem at hand, learning $T$
also has the advantage of saving computation while reaching high levels of
performance. If we were to fix $T$ ahead of meta-training (as is typically the
case in gradient-based meta-learning, where the number of gradient steps for
adaptation is a hyperparameter) to a large value in order to reach high
accuracy, as is shown in Fig. 7 (right), then it would induce larger
computational costs early on during meta-training compared to COMLN, which
achieves equal performance while tuning the value of $T$. In COMLN, the value
of $T$ is relatively small at the beginning of meta-training.

###  E.4 Additional Experiment: Preprocessed miniImageNet dataset

In addition to our experiments on the miniImageNet and tieredImageNet datasets
in Section 5, we want to evaluate the performance of continuous-time
adaptation in isolation from learning the embedding network. We evaluate this
using a preprocessed version of miniImageNet introduced in Rusu et al. (2018),
where the embeddings were trained using a Wide Residual Network (WRN) via
supervised classification on the meta-train set. For our purposes, we consider
these embeddings ${\bm{\phi}}\in{\mathbb{R}}^{640}$ as fixed, and only meta-
learn the initial conditions ${\bm{W}}_{0}$ as well as $T$.

The classification accuracies for COMLN and other baselines using pretrained
embeddings are shown in Table 5. COMLN achieves comparable or better
performance with a single linear classifier layer as other meta-learning
methods in both the 5-way 1-shot and 5-way 5-shot classification tasks. The
single exception is the 5-way 1-shot result of Meta-SGD from Rusu et al.
(2018), which exceeded the performance of COMLN. However, our implementation
of Meta-SGD achieved comparable performance to COMLN. This gap in performance
is likely due to the additional data used in Rusu et al. (2018) (meta-train
and meta-validation splits) during meta-training, as opposed to using only the
meta-training set for all other baselines. These results show that isolated
from representation learning, all these meta-learning algorithms (either
gradient-based or not) perform similarly, and COMLN is no exception.

One notable exception though is MetaOptNet (Lee et al., 2019), where the
performance is not as high as the other baselines when the backbone network is
not learned anymore—despite often being the best performing model in Table 2.
Our hypothesis is that this discrepancy is due to the accuracy of the QP
solver used in MetaOptNet, since learning individual SVMs on 1,000 held-out
meta-test tasks leads to performance matching all other methods (about 50% for
5-way 1-shot and about 70% for 5-way 5-shot).

Table 5: miniImageNet results using LEO embeddings and a single linear classifier layer. The average accuracy (%) on 1,000 held-out meta-test tasks is reported with $95\%$ confidence interval. * Results reported in (Rusu et al., 2018). ** Note that LEO uses more than a single linear classifier layer, but we add the numbers for completeness. Model | miniImageNet 5-way
---|---
1-shot | 5-shot
MAML (Finn et al., 2017) | $50.35\pm 0.63$ | $65.28\pm 0.54$
$\text{Meta-SGD}^{*}$ (Li et al., 2017) | $54.24\pm 0.03$ | $70.86\pm 0.04$
Meta-SGD (Li et al., 2017) | $50.57\pm 0.64$ | $69.09\pm 0.53$
iMAML (Rajeswaran et al., 2019) | $50.26\pm 0.61$ | $69.52\pm 0.51$
R2D2 (Bertinetto et al., 2018) | $50.33\pm 0.62$ | $70.38\pm 0.52$
LRD2 (Bertinetto et al., 2018) | $50.41\pm 0.62$ | $70.29\pm 0.52$
LEAP (Flennerhag et al., 2018) | $50.95\pm 0.62$ | $66.72\pm 0.55$
MetaOptNet (Lee et al., 2019) | $40.60\pm 0.60$ | $50.94\pm 0.62$
LEO** (Rusu et al., 2018) | $61.76\pm 0.08$ | $77.59\pm 0.12$
COMLN (Ours) | $50.39\pm 0.63$ | $70.06\pm 0.52$

###  E.5 Experimental details

For all methods and all datasets, we used SGD with momentum $0.9$ and Nesterov
acceleration, with a decreasing learning rate starting at $0.1$ and decreasing
according to the schedule provided by Lee et al. (2019). For meta-training, we
followed the standard procedure in gradient-based meta-learning, and meta-
trained with a fixed number of shots: for example in miniImageNet 5-shot
5-way, we only used tasks with $k=5$ training examples for each of the $N=5$
classes. This contrasts with Lee et al. (2019), which uses a larger number of
shots during meta-training than the one used for evaluation (e.g. meta-
training with $k=15$, and evaluating on $k=1$). This may explain the gap in
performance between COMLN and MetaOptNet, especially on 1-shot settings. We
opted to not follow this decision made by Lee et al. (2019) to ensure a fair
comparison with other gradient-based methods, which all used the process
described above.

#### Conv-4 backbone

We used a standard convolutional neural network with 4 convolutional blocks
(Finn et al., 2017). Each block consists of a convolutional layer with a
$3\times 3$ kernel and $64$ channels, followed by a batch normalization layer,
and a max-pooling layer with window size and stride $2\times 2$. The
activation function is a ReLU.

#### ResNet-12 backbone

We largely followed the architecture from Lee et al. (2019), which consists of
a 12-layer residual network. The neural network is composed of 4 blocks with
residual connections of 3 convolutional layers with a $3\times 3$ kernel. The
convolutional layers in the residual block have $k=64$, $160$, $320$ and $640$
channels respectively. The non-linearity functions are LeakyReLU$(0.1)$, and a
max-pooling layer with window size and stride $2\times 2$ is applied at the
end of each block. No global pooling is performed at the end of the embedding
network, meaning that the embedding dimension is $d=16,\\!000$. The only
notable difference with the architecture used by Lee et al. (2019) is the
absence of DropBlock (Ghiasi et al., 2018) regularization.