Task Weighting in Meta-learning with Trajectory Optimisation
Authors: Cuong C. Nguyen
Thanh-Toan Do
Gustavo Carneiro
Venue: Accepted by TMLR ()

Abstract: Developing meta-learning algorithms that are un-biased toward a subset of training tasks often requires hand-designed criteria to weight tasks, potentially resulting in sub-optimal solutions. In this paper, we introduce a new principled and fully-automated task-weighting algorithm for meta-learning methods. By considering the weights of tasks within the same mini-batch as an action, and the meta-parameter of interest as the system state, we cast the task-weighting meta-learning problem to a trajectory optimisation and employ the iterative linear quadratic regulator to determine the optimal action or weights of tasks. We theoretically show that the proposed algorithm converges to an $\epsilon_{0}$-stationary point, and empirically demonstrate that the proposed approach out-performs common hand-engineering weighting methods in two few-shot learning benchmarks.

Published in Transactions on Machine Learning Research (08/2023)
Task Weighting in Meta-learning with Trajectory Optimisa- tion
Cuong Nguyen
Australian Institute for Machine Learning The University of Adelaide
Thanh-Toan Do
edu
Department Data Science and AI Monash, University
Gustavo Carneiro
g.carnciro@surreeee
Centre for Vision, Speech and Signal Processing University of Surrey
Reviewed on OpenReview: https: openreview. net/ forum? id- SSkTBUyJ
Abstract
Developing meta-learning algorithms that are un-biased toward a subset of training tasks often requires hand-designed criteria to weight tasks, potentially resulting in sub-optimal solutions. Tn this paper, we introduce a new principled and fully-antomated task-weighting algorithm for meta-learning methods. Specifically, we frame the task-woighting problem as a trajectory optimization problem, where the weights of lasks within a mini-batch are treated as an action, and the meta- meta-parameter of interest viewed as the system state. Such modelling allows to employ the iterative linear quadratic regulator to determine the optimal task weights. We theoretically show that the proposed algorithı converges to eg-stationary point, and empirically demonstrate that the proposed approach out-performs common hand-engineering weighting methods on two few-shot learning benchmarks.
1 Introduction
Meta-learning has been studied from the carly 1990s (Schmidhuber, 1987; Naik Mammone, 1992; Thrun & Pratt, 1998) and recently gained renewed interest with the use of deep learning methods that achieves remarkable state-of-art results in several few-shot learning benchmarks (Vinyals al., 2016; Finn et al,, 2017; Snell al., 2017; Nichol et al,, 2018; Ravi & Beatson, 2019; Allen et al., 2019; Khodak al., 2019; Baik et al., 2020; Flennerhag et al., 2020). However, the majority of existing meta-learning algorithms simply minimise the average loss evaluated on validation subsets of training tasks, implicitly assuming that testing tasks come from the same distribution (over family of tasks) as training tasks (analogous to Mass-imbalance in single-task learning). This assumption hardly true in practice, and potentially biases the trained meta- learning models toward tasks observed more frequently during training, and consequently, resulting large variation of performance when evaluating on different subsets of testing tasks as shown in (Dhillon et al., 2019, Figure 1) and (Nguyen el al., 2021, Figure 1).
One way to address such issue is to exploit the diversity of training tasks, that the trained meta-learning models can generalise to a wider range of testing tasks. In fact, various studies in task relatedness or task similarity have shown that learning from certain tasks may facilitate the generalisation of meta-learning models (Thrun O'Sullivan, 1996; Zamir et al., 2018; Achille et al., 2019; Nguyen et al., 2021). This suggests the design of a re-weighting mechanism diversify the contribution of each training task when training a meta-learning model of interest. Existing re-weighting methods mostly rely on either hand-crafted criteria determine those weights (Collins et al., 2020), or additional validation tasks to learn the re- weightngg factors

---

Published in Transactions on Machine Learning Research (08/2023)
of interest (Xu et al., 2021). Such ad-hoc development of re-weighting mechanisms motivates us to design more principled approach to re-weight tasks for meta-learning- We note that there are also studies learning to balance the contribution of each training task, e.g., learning to balance (L2B) (Lee et al., 2020). However, L2B focuses on the task-adaptation step (also known as inner loop), while our interest is to explicitly weight the contribution of each lask at the mela-learning step (also known as outer-loop).
In this paper, we present a new principled and fully-automated weightiig algorithm,, called trajectory optimisation based task weighting for meta-learning (TOW). In particular, regard the iterative update of the maea--ea of interest as a trajectory, where the meta-parameet is the state and the task weighting vector is the corresponding action. This formulation casts the task-weighting problem in meta-learning as a trajectory optimisation problem, with the gradient-based update for the meta-parameter as the state transition dynamics and the weighted validation loss as the cost function. We then employ the techniques of trajectory optimisation determine the optimal action (i.e., weighting vector) w.r.t. the cost (i.e., weighted validation loss) to train the meta-parameter. We note that TOW is not a mota-learning method, but a lask weighting framework that can be integrated inlo existing meta-learning algorithms to circunvent the problematic assumption about the even distribution of training tasks. Our contributions can be summarised as follows:
We propose lo cast the lask-weighting problem in meta-learning to a finite-horizon discrele-time trajectory optimisation with its state denoted by the meta parammee and action denoted by the re-weight factors of tasks, and solve such problem using the itcrative linear quadratic regulator.
We prove that under the conditions of boundedness and smoothness of the loss function used, TOW converges to a particular st-ttaionnrry point.
We also demonstrate TOW's functionality by incorporating into two common meta-lcarning algorithms, namely MAML (Finn et al., 2017) and Prototypical Networks (Snell al., 2017), and showing that TOW enables meta-learning methods to converge with fewer number of iterations and achieves higher prediction accuracy than some common task re-weighting mechanisms in the literature.
2
Background
2.1
Trajectory optimisation
Given a continuous D-dimensional vector a state and an M-dimensional vector u as action, the objective of a trajectory optimisation is to find an optimal sequence of actions {ut as follows:
T E u.) s.l. Xe+1 = f(xe, u.),
(ut)== = arg min (ue)
(1)
where: c(x,, u.) and fxx,, uz) are the cost function and the state-transition dynamics evaluated at time step 6, respectively. These functions are assumed to be twice differentiable. In addition, the initial state x is assumed to be given, and trajectory optimisation means finding the optimal sequence of actions {u²)7 for a particular X1, not for all possible initial states.
In trajectory optimisation, the finite-horizon discrete-time problem shown in (1) can be solved approximately by iterative methods, such as differential dynamic programming (DDP) (Jacobson & Mayne, 1970) iter- ative linear quadratic regulator (iLQR) (Todorov & Li, 2005; Tassa et al., 2012). These methods rely a local approximation of the stato-transition dynamics and cost function using Taylor series about a nominal trajectory {x., (-) In DDP, both the stale-Iransition dynamics and cost function are approximated to the second order of their Taylor series, while in iLQR a "simplified" version of DDP, the state-transition dynamics is approximated up to the first order. In a loose sense, DDP is analogy to the Newton's method, while iL iLQR is analogous to a quasi-Newton's method. In this paper, we employ iLQR to solve for the task weights of interest due to its efficiency in terms of computation compared to DDP.
2

---

Published in Transactions on Machine Learning Research (08/2023)
The main idea of FLQR is to cast general non-lineat trajectory optimisation problem shown in (1) to linear quadratic problem (T.QP) by linearising the state-transition dynamics and quadraticising the cost function about a nominal trajectory {(8,, ú.)]== via Taylor's s series. In particular:
f(x,, u.) = f(x, ú.) V xx (x) - xx) + V Vuf(&,, (u ū.),
(2)
Exi
Fu,
where Fx, and Fu, denote the first gradients of the transition dynamies f f(,, Uz) w.r.t. x, and ue, respec- tively. Similarly, one can obtain the Taylor's scries for the cost function c(x,, Uz) with Cx, and cu, as the first gradients of the cost function, and Cx Cx, Cu,,x; and Cu,, as the second gradients w.r.t. the variables specified in their corresponding subscripts (see Appendix A for more details).
The approximate LQP can then be solved exactly by the linear quadratic regulator (LQR) (Anderson & Moore, 2007). Subsequently, the newly obtained trajectory is used as the nominal trajectory for the next iteration. This process is repeated until the cost function converged. The detailed derivation of iLQQ can be found in Appendix B. Further details of iLQR can be referred (Todorov & Li, 2005; Tassa et al., 2012). To our best knowledge, there are previous works that provide proof on the convergence of iLQQR. Therefore, for complete analysis, we provide the proof of convergence for iLQQ adopted from DDP (Yakowitz & Rutherford, 1984) in Appendix C.
2.2 Meta-learning
We follow the task environment (Baxter, 2000), where tasks are i.i.d. sampled from an unknown distribution P(7) over a family of tasks, to formulate the meta-learning problem. Each task T; is associated with two 1==1: where si denotes a training input and Vig denotes the corresponding training label, i E ( M). and validation (or query) subset s(9) which is similarly defined. For training tasks (Tr) both data subsets have labels, while for testing tasks TM+1, M+1 is labelled. The aim is to learn D-dimensional meta-parameter x shared across only the data in s(*)
data subsets: training (or support) subset s {(s]]) " (a))
(s)
all tasks, so that x can be efficiently fine-tuned on 5 to produce task-specific parameter for predicting the unlabelled data in s(9)) One of the simplest forms of meta-learning is analogous to an extension of parer-par optimisation in single-task learning, where the shared meta- parameere x is learnt from many tasks. The objective of meta-learning can generally be expressed as:
1 m
x" arg min u Lx(x) with Lx(x) x
Σ
l
(0) (0) (s M :
(q) m
j=1
le (s (2) (0(x)))()))]]])))]]] Vie (I....,
s.t.: of E arg min di
mi
ueu
(3)
where E(.) is the loss function that is non-negative and twice differentiable, u is an M-dimensional vector that re-weights the influence of M training tasks, U is the of feasible weights defined by some weighting critcrion, and ø(x) is the parameter fine-tuned task T, (also known as task-specific parameter).
The task adaptation in the lower-level optimisation of (3) varies depending on the type of meta-learning algorithm used. For example, it can be a gradient descent update with x being the initialisation of the neural network of interest (Finn et al., 2017). In metric- based meta-learning (Snell et al., 2017), that task adaptation step is slightly different where the class prototypes of training data are embedded into a latent space by the meta-model, and the validation loss based on the distance between the class prototypes to each data-point in s(9). There are also other extensions of (3) using probabilistic approaches (Yoon el al., 2018; Ravi & Beatson, 2019; Nguyen et al., 2020). Nevertheless, our approach proposed in Section 3 can generally be integrated into any of these meta-lcarning algorithms. In addition, our attention on weighting tasks, corresponding to the upper-level equation (3). As a result, we exclude the discussion of the lover-level optimisation for (x(x) in the rest of this paper to simplify the analysis.
3

---

Published in Transactions on Machine Learning Research (08/2023)
The task-weighting problem in (3) is corried out at the meta level (also known as "outer- ")op"). It. is, therefore, different. from some recent mela-learning methods (Khodak el al., 2019; Boik et al,, 2020; Flennerhag et al., 2020; Lee et al., 2020) that design different learning strategies for oz(x) at the task adaptation step (also known as "inner-loop") ") to estimate the meta-parameter with a uniform weighting u = (1/M)M.
Remark 1.
For the feasible weighting set U, the most widely-used weighting criterion is uniform: 2 = { }aa, which rosembles the conventional meta-learning. Another popular criterion is to select difficult tasks tasks that have the largest validation losses for training to optimise the performance on the wonst-case sce- narios (Collins et al., 2020). However, such difficult tasks may not always be preferred when outliers and noise are present. That leads to another weighting approach which favours the most familiar data-points in single-task learning (Kumar et al., 2010; Bengio et al., 2009; Wang et al., 2017) - often referred as curriculum learning. The two latter laig-weig approaches can be considered as the "exploration" and "exploita- tion" strategies used in reinforcement learning (RL), respectively. Similar to the exploration and exploitation dilemma in RL, we hypothesise that the optimality for task weighting is formed by a balance between these two approaches. In the following section, we propose principled approach to automate re-weighting lasks through an optimisation on sequence of many mini-batches rather than relying on manually-designed criteria as the previous studies.
3 Methodology
3.1 Task-weighting as a trajectory optimisation
The upper-level optimisation in (3) can also be written as: x* f(x, u), where f is a function denoting an optimiser, e.g., stochastic gradient descent (SGD) (Robbins & Monro, 1951). Such an update can be considered as state-transition dynamics where the meta-parameter x is the state and the weighting vector u is the action. In addition, the set of feasible re-weighting vectors, U, can be written result of an optimisation that minimises some cost function i.e.;:
s.t.: Xx+1 f(x,, u,) and x,is is given u u= (ut)) (ut)=_ arg min Σ c(x), {u) 1=1
where X1 is the initialisation of the meta-learning parameter, and the subscript denotes the time step.
Since our interest is the convergence speed and the generalisation of the learnt meta-model, we define the cost function as an un-discounted sum of mifformly-wiighted validation losses of sequence of T mini- batches, each has M tasks, plus a penalisation lisation on the action u, i.e., Gaussian prior with mean t4u and precision 8.. For the state-transition dynamics f. it is assumed to be the SGD to simplify the analysis. Such assumptions result in the following state- transsitanntiinn dynamics and cost function:
f(x,, u&) = X - a axx [u, [xx)] and c(x,, u.) = (xx)) + B4/2|lu, (h mrll,
(5)
where: L(x) is M-dimensional vector containing M validation losses Lx(x) defined in (3), 1M is an M-dimensional vector with all elements equal to 1 and ||.[| denotes the fy-norm.
Note that the action u, is not necessarily normalised to 1, as this constraint may not be effective in certain scenarios. For example, consider a mini-batch consisting entirely of familiar tasks and another consisting entirely of unfamiliar tasks. Our hypothesis is to assign small weights to the familiar tasks the former mini-batch and large weights to the unfamiliar tasks in the latter mini-batch to encourage diverse learning, However, normalizing u, to 1 would be disadvantageous it would make the contribution of tasks in both mini-batches equal, thereby further biasing the meta-learning model towards familiar tasks in the first mini- batch. To address this, we allow the weights to be automatically determined by the optimisation in equation (4) with a Gaussian prior. Nonetheless, one can choose to normalize u, to 1 by simply replacing it with softmax(u:).
In general, the solution in (4) cannot be solved exactly, but approximately using iterative methods such DDP iT.QR. Given the state-transition dynamics f(x:) follows the formulation of first-order gradient-
4

---

Published in Transactions on Machine Learning Research (08/2023)
based optimiser (see Eq. (5)), fxx) consists of the first gradient of the weighted loss u. L(xe) Hence, applying DDP will result. in an intractable solution since DDP requires the second gradient of f(x)); corresponding to the third gradient of the weighted loss u, L(xe). In contrast, iLQR needs only the first gradient of f(x)), which corresponds to the second gradient of the weighted loss u L(x). Although this means that iLQR no longer exhibits the quadratic convergence rale as DDP, in the context of meta-learning, the significant reduction in computation out-weights the speed of convergence for the task weighting vector u. In this paper, we use iLQR to solve for the re-weighting vectors {uī ) in (4).
w.r.t. X
3.2 Practical task-weig method based on trajectory optimisation
The approximation using Taylor's series on the state-transition dynamics and cost function is shown in Appendix A. This approximation leads to the calculation of two Hessian matrices: Fx, in the transition dynamics, and xx in the cost function (see their definitions in Eq. (2)). In addition, while performing recursive backward i.QQR (see Algorithm 2 in Appendix D), we need to calculate another intermediate matrix of the cost-to-go in (28) (please refer to Appendix B), denoted as Vt, which is analogous to an inverse Hessian matrix in Newton's method. Naively calculating these Hessian matrices comes at the quadratic complexity O(D)) (D is the dimension of x) in terms of running time and storage, resulting in an intractable solution for large-scaled models. To address such issue, the two Hessian matrices Fx, and Cx,x, may be approximatod by their diagonals which can be efficiently computed using the Hutchinson's 's method (Bekas et al., 2007). However, as the size of the model increases, using few samples from the uniform Rademacher distribution produces noisy estimations of the Hessian diagonals, resulting in a poor approximation (Yao et al., 2021). Instead of calculating the Hessian diagonals, we use the Gauss-Newton diagonals as replacements. As the Gauss-Newton matrix is known to be a good approximation of the Hessian matrix (Martens, 2010; Botev et al., 2017), this, therefore, results in a good approximation for the Hessian operator. In addition, Gauss- Newton diagonals can be efficiently calculated using a single backward pass (Dangel et al., 2020). For the matrix 1, we approximate it by its diagonal matrix. Since mairix V, is analogous to the inverse Hessian matrix in Newton's 's method, approximating it by its diagonal means performing Newton's method separately for each coordinate, which holds when the diagonal of V. is dominant. We also provide some additional results using full matrix V, in Appendix H. In general, we do not observe any significant difference in terms of accuracy evaluated on the validation set between the diagonal approximation and the one with full Gauss-Newton matrix. Nevertheless, these approximation increases the tractability of our proposed method, allowing to implement the proposed method for very large models, such as deep neural networks. The complete algorithm of the proposed task weighting meta-learning approach is outlined in Algorithm
To simplify the implementation and convergence analysis, we select the nominal actions that coincide with the uniform weighting, meaning that: ú. 1/M, AI E ".....,T,, E (1..... MJ. In addition, constrain that all elements of the weighting vector or action u are non-negative since each task would either contribute more or less or even not contribute to the learning for x. This constraint is incorporated into the stopping condition for iLQR shown in step 16 of Algorithm 1. If there is at least one element un,t E fl....,,Ty T), f1...., M} being negative, the backtracking line search will iterate one more time with è decaying toward 0, forcing u, stay close to the nominal ú, Thus, in the worst-case, is reduced to 0, making u, coincide with ú,, which is the uniform weighting. We also provide a complexity analysis of Algorithm 1 in Appendix E.
4 Convergence analysis
We prove that the training process using TOW to weight tasks converges to an co-stationary point where 6n is greater than some positive constant. Before analysing the convergence of TOW, we state lemma bounding the norm of the weighting vector (or action) u, obtained from iLQR:
Lemma If u, a stationary action of a nominal action ú, obtained from iLQR, then:
> >
0 S 8.
Proof. Please refer to Appendix F.2.1 for the detailed proof.
5

---

Published in Transactions on Machine Learning Research (08/2023)
Algorithm Task- weighting for meta-learning
procedure TOW TRAINING(X1, RLLQU)
2: D. X11 initial mela-parametee
3: > TILQRA number of iterations in iLQR
4: while x is not converged do
5: get mini-batches, each consists of M tasks
ù, + 1/M1Mr Vt E (T....:7) 8,41 6 fx,, ûr), Vi E (
for n= 1,..., M.QQR do
KK 01 + ILQREACKWARD((&.. úz ( 6 2
Algorithm 2 (Appendia D)
10:
11:
repeat
b Backdracking line search
12:
+ }
13:
for t T do
iLQRR forward pass
14:
u, = (x; - x2) + ck, ú,
15:
X++1 ((x,;
16:
until ((,NN) -((i)) S f50 and u, 2 0
Eg. (40)
17:
{xt] t [x11==
b Update nominal state
18:
19:
X1 f Xy
Update meta-porameter
20:
return X
4.1
Assumptions the boundedness and smoothness of the loss function and its gradients
To analyse the convergence of general non-convex function, one typically assumes that the loss function, its first derivative and second derivative are bounded and Lipschitz-contimuuus (Collins et al., 2020; Fallah et al., 2020) as shown in the following Assumptions 1 to 3, respectively.
4.1.1 Boundedness and smoothness of the loss function
Assumption 1. The loss function of interest, i, mentioned in (3) is B-bounded and L-Iipschitz.
Formally, Assumption means that the loss function é has the following properties:
Boundedness: 3B 0: Vx E R.P [8(sg Meri x)1 S B,
Lipschitz continuily: 3L
vx. DRD, [((s Wy: x) ((sij, Miys x) S L||x x||.
The boundedness assumption is to bound the second moment of the loss function, while the Lipschitz- continuity assumption of the loss function l implics that the gradient norm of w.r.t. x is bounded above (see Lemma 14 in Appendix F.4):
| (= S L.
(6)
4.1.2 Smoothness of the Jacobian and Hessian
We then state the assumptions about the gradient and Hessian of the loss function as follows: Assumption 2. The gradient of the loss function l(s, x) w.r.l. x is S-Lipschitz.
Assumption means that:
V (S::) (el(s S| x||.
Assumption 3. The Hessian matrir Vllls, x) is p- Tápsschitz.

---

Published in Transactions on Machine Learning Research (08/2023)
Assumption implies that:
3p 0: VX, R.P
V (ligs Miyi x)
V ( Wifi x)| S pl|x x||.
These assumptions are used bound the gradient of the "true" validation loss of task Ti (see Lemma 12 in Appendix F.2.3), which is defined as follows:
(x(x) = B (s(5))) Mij (1]) o(x))]
(7)
(9) where g indicates the expectation over all data pairs {(s]]) vij دزر sampled from the true (validation) probability distribution D.
Given the above assumptions and lemmas, the convergence of TOW can be shown in Theorem 2. We also provide examples of loss functions satisfying Assumptions 1 to 3 in Appendix I.
4.2 The convergence of the proposed method
Theorem (Main result). If Assumptions 1 to hold, the learning rale a < <//3(6/M-1), and z is randomly Thee returned by Algorithm 1, then:
sampled from {x, J
(xx) Tieer Brootey M [1Vau, LIM [2]][]]
K
E Bx-
S 60 +
Titer
where:
46B M + 0 ( 11M++))
2u, LI.M (x1)
cD
0,
a 2 22- aड (6VM+) +1)
2 as 88 88VM +
a
with Titer as the number of grudient-update for the meta-panameter, and B (D(to) dr as the expectation taken over all query dala sampled from I. mini-bolches.
Proof. Please refer to Appendix F.3 for the detailed proof.
Theorem shows that the expectation of squared gradient norm of the weighted validation loss upper- bounded by a monotonically reducing function w..tt. the number of iterations Titer. This implies that Algorithm 1 converges in expectation to an eg-stationary point.
Remark 2. The result in Theorem 2 agrees with previous studies on task-weighting for meta-learning, e.g-, (Collins el al., 2020, Ineq. (75)) where the gradienl norm is bounded above by some posilive constant. The tightness of the bound in Theorem mostly depends that limp-~o 60 0. Thus, to ensure that eo is small, 8 needs to be small. We a strong prior of u, i.e., setting large value of Bu in Eq. (5), e.g., Bu = 10 in our experiments presented in Section 5. In addition, the integrated the backtracking line search in Algorithm 1 forces u to slay close to the uniform weighting ú, making å very small. Another factor contributes to the small value of eo the inverse of the number of tasks in a mini-balch, 1/M, as seen in (2). In practice, M cannot be too large and often in the range of 5 to 10 to fit into the memory of a processing unit, e.g., GPU. And since we can impose constraints to make 8 tiny, can guarantee to obtain small cn to make the bound in Theorem 2 tight.
how small the value of to is. In fact, we observe.
make ö small by imposing
5
Experiments
5.1 N-way k-shot classification
In this section, we evaluate the performance of the proposed trajectory optimisation task weighting (TOW) approach with three baselines: one with uniform weighting, denoted uniform, one with higher weights
7

---

Published in Transactions on Machine Learning Research (08/2023)
98
Accuracy (%)
Accuracy (%)
96
90
94
uniform
exploration
exploitation
92
TOW
80
90
0
5
10
15
20
4
6
8
Ns of iterations (x 1,000
Nt of iterations (x 1,000)
(a) MAML on Omniglot
(b) Protonet Omniglot
Accuracy (%6
42
40
10
20
30
40
10
20
30
10
20
30
40
No of iterations (x1,000)
No of iterations (x1,000)
Xt of iterations (x1,000)
(c) MAML on mini-ImageNet
(d) Protonet on mini mini--mageNet
(e) MAML with Resnet-10 10
Figure 1: Validation accuracy exponential moving average (with smoothing factor 0.1) of different task-weighting strategies evaluated on: (a) and (b) Omniglot, and (c), (d) and (e) mini- Maini--mageNet.
difficult tasks (or tasks with higher losses), denoted as exploration, and the other one with higher weights easier tasks (or tasks with lower losses), denoted exploitation. The experiments based n-way k-shot classification setting used in few-shot learning with tasks formed from Omniglot (Take et al., 2015) and mini-Imageel (Vinyals et al., 2016) the two most widely-used datasets in meta-learning.
Baselines Naively implementing the two baselines, exploration and exploitation, will easily lead to trivial solutions where only the task with largest or smallest loss within mini-batch is selected. Thus, only one task in each mini-batch is used for learning, and consequently, making the learning noisy and unstable. We, therefore, introduce a prior, dcnoted p(u), as regularisation to prevent many tasks within the same mini-batch from being discarded. Formally:
arg min -u T ((x) - In p(u) for exploration
u'
arg aig T [(x) Inp(u) for exploitation.
In general, the prior p(u) be any distribution that has support in (0, +o0) such as Beta, Gamma Cauchy distribution. For simplicity, p(u) is selected as Dirichlet distribution with concentration к > I to constrain the weight vector within probability simplex. One can then use a non-lincar optimisation solver to solve (8) to oblain an optimal u² for one the two baselines. In the implementation, we use Sequential Least SQuares Programming (SLSQP) to obtain u*. Note that the definition of the ezploration bascline above resembles TR-MAMI, (Collins al., 2020), but applicable for common few-shot lcarning benchmarks where the number of tasks is large. Similarly, the exploitation an analogy to robust Bayesian data ingghttting (Wang et al., 2017) or curriculum learning in single-task learning.
8

---

Published in Transactions on Machine Learning Research (08/2023)
Table The classification accuracy averaged on 1,000 random testing tasks generated from Omniglot and 600 tasks from mini-ImageNNe with 95 percent confident interval; the bold numbers denote their statistically differences from the ones in the same column.
Weighting method
Mini-ImageNet
Omniglot
4 layer
CNN
Resnet-10
Uniform
94.86 +0.43
48.70 +1.84
49.12 #00.6
Exploration
92.64 #0.52
48.80 10.72
48.72 10.74
MAML
Exploitation
95.34 #0.42
49.22
#0.74
48.44 #0.76
TOW
95.94 +0.40
51.55 :0.75
52.32 =0.80
Uniform
95.21 +0.37
49.42
+0078
Exploration
94.57 +0.38
48.56
#0.77
Protonet
Exploitation
95.78 +0.40
48.39
+0.79
TOW
96.84 =0.37
51.05 =0.80
Table 2: Rumning time (in GPU-houu) of different weightiig methods based on MAML.
mini-Image
Omniglot
CNN
Resnet-10
Exploration
1.55
5.24
7.18
Exploitation
1.55
5.24
7.18
Uniform
1.35
5.03
7.16
TOW
7.50
38.12
67.78
Datasets Omniglot and mini-ImageNNet are used in the evaluation. For Omniglot, we follow the original train-test split (Lake al., 2015) by using 30 alphabets for training and 20 alphabets for testing. We also utilise the hierarchy structure of alphabet-character to form finer-grained classification tasks to make the classification more difficult than the random train-test split. For mini- ImageNet, we follow the standard train-test split that uses 64 classes for training, 16 classes for validation and 20 for testing (Ravi & Larochelle, 2017) in our evaluation.
Models used The base model used across the experiments is the CNN module network that is widoly used in few-shot image classification (Vinyals et al., 2016; Finn el al., 2017). Two common meta-learning algorithms considered in this section include MAML (Finn al., 2017) and Prototypical Networks (Snell et al., 2017). In MAML, the flattened features are passed to linear fully-connected layer to classify, while in Prototypical Networks, the classification is based Euclidean distances to the prototypes of each class.
Hyper-paramcters Please refer to Appendix G for the details of hyper-paramoters used.
Results Figures la to 1d plot the testing accuracy evaluated on 100 validation tasks drawn from Omniglot and mini-ImageNett following the 5-way 1-shot setting. The validation accuracy curves along the training process show that TOW can achieve higher performance comparing to the three baselines on various datasets and meta-learning methods. We also carry out an experiment using Resnet-10 (He et al., 2016) on mini- ImageNet to demonstrate the scalability of TOW. The results Resnet-10 in Figure 1e shows similar observation that TOW out-performs other task-wwighting methods. We note that the validation accuracy curves of Resnet-10 fluct uates due lo our injected dropout to regularise the network from overfitting since it is known that larger networks, such as Resnet-10 or Resnet-18, severely overfit in the few-shot setting (Nguyen et al., 2020). For the ovaluation on testing sets, we follow the standard sctting in few-shot lcarning by measuring the prediction accuracy on 1,000 and 600 lesting tasks formed from Omniglot and mini- mini-Image. Net,
reported in (Finn al., 2017) 2 Prepoted in (Snell al., 2017)
9

---

Published in Transactions on Machine Learning Research (08/2023)
70
70
65
Accuracy (%)
Accurscy (%)
65
60
uniform
55
exploration
60
exploitation
TOW
50
55
20
40
60
80
0
5
10
15
20
25
Nt of training tasks (x10,000)
No of training tasks (x 10,000)
(a) MAML
(b) Protonet
Figure 2: Exponential moving average with smoothing factor 0.1 of the prediction accuracy evaluated on validation tasks formed from the any-shot setting of mini-TmageNet dataset mentioned in Section 5.2 where the base model is 4-module CNN.
respectively (Vinyals et al., 2016; Finn et al., 2017). The results in Table show that TOW can be at least 2 percent more accurate than the best baseline among Uniform, Exploration, and Exploitation. Note that there a difference between the results shown in Figure and Table due to their differences in terms of (i) the tasks form: one from validation set, while the other from testing set, and (ii) the number tasks evaluated. Despite the promising results, the downside of TOW is the overhead caused by approximating the cost and state-trane transtrransition dynamics over mini- mini-batches of tasks to determine the locally-optimal {ut J As shown in Table 2, TOW is about 7 to 9 times slower than the three baselines. We also provide visualisation of the weights in Section 5.3.
5.2 Any-shot classification
We also follow the realistic task distribution (Lee et al., 2020) to evaluate further the performance of TOW. The new setting is mostly similar N-way k-shot, except k not fixed and might be different for each class within task. Specifically, with probability of 0.5, the number of shots for each class is sampled from a uniform distribution: l ~ Uniform(1, 50) to simulate class imbalance. With the other 0.5 probability, the same number of shots h ~ Uniform(1, 50) is used for all classes within that task. The number of validation (or query) samples is kept at 15 samples per class.
Similar to the experiments carried out in Section 5.1, TOW demonstrates a higher performance compared to the three basclines: exploitation, exploration and uniform along the training process, as shown in Figure In general, TOW can achieve the state-of-the-art results when evaluating on 3,000 testing tasks formed from mini-ImageNNt,,
shown in Table 3, compared to common meta-learning methods. To further evaluate TOW, we follow the same setting and use the models trained on mini- ImageNet to test on 50 classes of bird images split from CUB dataset. The results in the last column of Table 3 show that TOW can also work well on out-of-distribution tasks formed from CUB compared to most of the methods in the literature. The main reason that explains the worse performance of TOW, compared with Bayesian TAML, that the meta-learning based methods used by TOW are MAML and Protonet, which have a smaller number of perameersss to model tasks than Bayesian TAML.
5.3 Weight visualisation
To further understand how the weight u, varies along the training process, we conduct a study to monitor the weights u, of a set of pre-defined Omniglot tasks. In the first sctting, we fix T 5 mini-batches of tasks, each consisting of M = 5 tasks formed from the same alphabet. In the second case, the configuration is similar, but each mini-batch contains M 5 tasks formed from different alphabets. For each case, we
10

---

Published in Transactions on Machine Learning Research (08/2023)
Table Prediction results on any-shot classification evaluated on 3,000 testing tasks with percent confident interval; the bold numbers denote the results that are statistically significant. The results of previous methods are reported in (Lee et al., 2020).
Training set
mini- mmag-lmageNeet
Testing set
mini-ImageNet
CUB
MAML (Finn et al., 2017)
66.64 #0.22
65.77 +0.24
Meta-SGD (Li et al., 2017)
69.95 #0.20
65.94 +0.22
MT-net (Lee & Choi, 2018)
67.63 #0.23
66.09 +0.23
ABML (Ravi & Beatson, 2019)
56.91 #0.19
57.88 #0.20
Protonet (Snell et al., 2017)
69.11 +0.19
60.80 +0.19
Proto-MAML (Triantafillou et al., 2020)
68.96 #0.18
61.77 #0.19
Bayesian TAML (Lee et al., 2020)
71.46 10.19
71.71 #0.21
ROW-MAML
70.02 00.24
68.34 #0.25
TOW-Protones
72.12 =0.21
64.79 +0.25
0.22
0.24
Weight value
Weight value
0.22
0.2
0.2
train-same trffn-different
train test
0.18
test-same
diff-different
0.18
10
15
20
10
20
30
40
N of meta-updates (x 10,000)
Ns of meta-updates (x 10,0000
(a) Omniglot
(b) Mini-ImageNet
Figure 3: Visualisation of the weight values associated with (a) Omniglot and (b) mini- NeageNN... For the Omniglot dataset, the tasks are drawn from the same Omniglot alphabet (either training or testing set), and the notation same means that all tasks mini-batch are formed from one alphabet, while different indicates the mini-batch consists of tasks formed from different alphabets.
also run with tasks drawn from training and testing sets. Our desire is to observe how the weight changes and whether there is a difference between training and testing tasks. We plot the weight for a single task belonging to the "controlled" 5-task mini-batch of interest in Figure 3. The results show the variation of the weights of the tasks of interest with a common trend among all the tasks considered, in which the weights are large at the boginning and gradually reduce when training progresses. This is, indced, expected since most tasks are unfamiliar the model the early stale, and gradually becomes more familiar. In addition, observe that the weights for testing tasks are slightly larger than the ones for training tasks.
We also conduct the same ablation study for mini-ImageNNe tasks. Since in mini-emageNet, we do not have any information regarding to the hierarchical structure of classes, we carry out the experiment on training and testing tasks only. One could, of course, utilise the word-net structure to categorise classes into "alphabet" as Omniglot, but for the sake simplicity, we proceed without including such information. Figure 3b shows the weight values of some tasks each training checkpoint, which also has similar but noisier trend as the ones in Omniglot.
11

---

Published in Transactions on Machine Learning Research (08/2023)
45
Accuracy (%)
40
Accuracy (%)
40
35
35
2
10
B+LQR=2
T = 10
30
30
20
40
60
80
100
20
40
60
80
100
Nt of iterations (x 1,000)
No of iterations (x 1,000)
(a) Number of iLQR iterations TSLQR
(b) Length of trajectory T
Figure 4: Ablation studies on 5-way 1-shot Imageett with different values of (a) number of iterations TanQR in iLQR (with smnotthig factor 0.1), and (b) the length of trajectory T (without smoothing).
5.4 Ablation studies
We carry out ablation studies to analysc the effect of each hyper-paramettr to the performance of the proposed method, TOW. In particular, we follow the experiment setting of 5-way 1-shot tasks formed from the mini- Imag-lmaeNNet dataset as shown in Section 5.1. Specifically, the model is the 4 convolutional neural network, where the mini-batch size M is set to 5, the number of iLQR iterations niQQR is set to 2, T = 2, Bu = 10 and fuu = 1/M.
5.4.1 Number of iLQR iterations
We carry out an ablation study about the effect the number iterations used in iLQR, PALQR. by analysing the validation accuracy on two different settings: one with 2 iterations and the other with iterations. The qualitative result between the two settings in Figure 4a agrees with our intuition that the larger the number of iterations in iLQR, the more accurate the weighting and hence, the larger the validation accuracy. However, the trade-off is the significant increase in terms of training time, which is approximately 6 times larger than the one with 2 iterations.
5.4.2 The length of the trajectory T
As the proposed method, TOW, is based on iLQQ, it also shares some of the propertics of iLQR. According to the iLQR paper (Tassa et al., 2012, Section IV. A), the value of the horizon T is a problem-depeudent quantity which must be found by trial-and-error. The optimal value T is, therefore, chosen by further fine- tuning. However, as our aim is connect the trajectory optimisation and teighteig in meta-learning and analyse the convergence of the proposed method theoretically, we select the value of T' specified in the paper for demonstration purposes. As shown in Figure 4b, increasing T might not always result in the best performance. It is due to the approximation nature of iLQR where the transition dynamics f and the cost function e are lincarised and quadraticised via Taylor series. Such approximations are, however, only accurate for short trajectory (small T). For a long trajectory, the future states predicted by iLQR would be deviated further away, making the optimisation sub-optimal. Even if we assume that we have enough computational power to treat the entire training set of tasks as one trajectory and optimise that by iLQR, the result might not be better than short trajectory.
5.4.3 papper-prammters of the prior of u
We also present two ablation studies to understand the impact of the Gaussian prior of u on the performance of TOW. In this case, we vary either Bn pa introduced in Eq. (5), while fixing other hyper-parameters.
12

---

Published in Transactions on Machine Learning Research (08/2023)
40
Accuracy (%)
Accuracy (%)
35
35
fv. 0.05
Bu
fa 0.2
By
10
30
0.5
0
20
40
60
80
100
10
20
30
40
50
Nt of iterations (x 1,000)
No of iterations (x 1,000)
(a) Prior precision Bu
(b) Prior mean
Figure Ablation studies on the prior the weighting vector u introduced in Eq. (5) (without smoothing).
40
Accuracy (%)
35
uniform
exploration
30
exploitation
TOW
10
iterations (x 1,000)
Figure 6: Comparison between TOW and the three baselines where each the baselines trained times more than TOW, where the additional tasks used to train the baselines randomly sampled from the training task distribution.
For Bu, run two different values: 8n 1 and 3. 10 and show the result in Figure 5a. For You run experiment with Pa 0.05 and μa = 0.5 with mini-batch size of 5 tasks (corresponding to Duniform = 0.2 x 1) and show the result in Figure 5b.
According to this ablation study, smaller 8n results in slightly better performance compared to one with a larger Su, while the signal in Figure 5b is unclear. the case of Su. we hypothesise that larger value would place stronger regularisation, forcing the weighting value to be close a uniform one. Reducing 3, would allow the weighting values to fluctuate and easily overfit. In case of pas TOW seems work well with different values of pa.
5.4.4 Train the baselines more
To provide a fairer comparison against competing approaches, instead of performing one iteration for each baseline, we increase the number of iterations when training the baselines. In particular, each of the baselines is trained on five times more tasks compared with the number tasks for training TOW. To compensato for this larger number iterations, we also reduce by five folds the learning rate, the comparison is fair. The results in Figure 6 show that TOW qualitatively out-performs the three baselines in this fairer setting.
13

---

Published in Transactions on Machine Learning Research (08/2023)
6 Related work
Our study directly relates to re-weighting tasks in meta-learning. One recent work is TR-MAML (Collins et al., 2020) which places higher weights on tasks with larger validation losses to optimise performance for worst-case scenarios. However, when the number of training lasks is large, e.g. (1%00) & 8.25 x 10 5-way classification tasks formed from 1000 characters in Omniglot dataset (Lake et al., 2015), learning weight for each training task is intractable. TR-MAML. circumvents such issue by clustoring tasks into a small number of clusters based on some ad-hoc intuition and learn the weight for each cluster. This, however, reduces the practicability of TR-MAML. Another work, a-MAML (Cai al., 2020), provides an upper- bound on the distance between the weighted risk evaluated on training tasks to the expected risk on testing tasks. The re-weight factors can then be obtained to minimise that upper-bound, reducing the variance between training and testing tasks. Note that a-MAMT. follows a transductive setting where labelled data of testing tasks are available in the training phase, while we follow the inductive setting where no information about testing is available during training. In reinforcement learning (RL), MWL-MAML (Xu et al., 2021) is recently proposed to employ meta-learning to learn the local optimal re-weight factor of each trajectory using a few gradient descent steps. The downside of MWT.-MAML, is the need of validation trajectories (or validation tasks in meta-learning) that are representative enough to learn those weights. In addition, MWL-MAML also requires to store the weight value of each task as in TR-MAML, making the approach intractable in setting, Furthermore, TR-MAML, a-MAML and MWL-MAML rely single mini-batch of tasks to determinc the weights without considering the effect of soquence of mini-batches when training a mela-model, potentially rendering sub-optinal solutions. In contrast, our proposed method does not need to cluster tasks nor require additional set of validation tasks. In addition, our proposed method automates the calculation of task-woighting through an optimisation over a sequence of mini-batches, allowing to obtain better local-optimal solutions outside of a single mini-batch of tasks. There are also other studies about task balancing, such as Learn to Balance (I2BB) (Lee et al., 2020). However, T.2B introduces additional parameters in the task adaptation step (inner-loop), while our method explicitly introduces a weighting vector at the meta- parameter update step (outer-loop).
Our study is also similar to task-weighting in multi-task learning (Chen et al., 2018; Sener & Koltun, 2018; Guo et al., 2018; Liu et al., 2021) where the goal is to obtain an optimal re-weighting vector u for all tasks. Such modelling can, therefore, work well with a small number of tasks, but potentially fall short when the number of tasks is very large, e.g. in the magnitude of 10 training tasks for 5-way Omniglot classification, due to the poor scalability of the computational and storage complexities of that modelling. In comparison, our proposed approach does not explicitly learn the weighting vector for all training tasks, but determines the weighting vector for tasks in current and some following mini-batches via a trajectory optimisation technique. In a loose sense, the multi-task learning approaches can be considered as an analogy to a "batch" learning setting w.r.t. the weighting vector u, while ours is analogous to an "online" learning setting which can scale well to the number of training lasks. There is also a concurrent work Auto-Tambda (Liu el al., 2022) that is designed to use meta-leariing to learn how to weight tasks in the multi-task setting. However, Auto- Lambda is similar to TR-MAML, which is designed for fixed number of tasks in the multi-task learning, while being intractable when the number of tasks is large. Furthermore, Auto-Lambda is also similar to MAML since Auto-Tambda employs validation subsets of training tasks to meta-learn the weighting of tasks.
This paper is motivated from the observation of large variation in terms of prediction performance made by meta-learning algorithms on various testing tasks (Dhillon et al., 2019, Figure 1), implying that the trained meta-model may be biased toward certain training tasks. Such observation may be rooted in task relatedness or task similarity which is growing research topic in the field of transfer learning. Existing works include task-clustering using k-nearest neighbours (Thrun & O'Sullivan, 1996) using convex optimisation (Jacob et al., 2009), learning task relationship through task covariance matrices (Zhang & Young, 2012), or the- oretical guarantees to learn similarity between tasks (Shui et al., 2019). Recently a large-scale empirical study, known as Taskonomy (Zamir et al., 2018), investigated the relationship between 26 computer vision tasks. Another promising direction to quantify task similarity to employ task representation, notably Task2Vec (Achille et al., 2019), which is based on Fisher Information matrix to embed tasks into a latent space. One commonality among those studies is that learning from certain training tasks may be beneficial
14

---

Published in Transactions on Machine Learning Research (08/2023)
to generalise to unseen tasks. This suggests the design of mechanism to re-weight the contribution of each training task to improve the performance of the meta-model of interest.
7 Discussion and conclusion
We propose principled approach based trajectory optimisation to mitigate the issue of non-uniform distribution of training tasks in meta-learning. The idea is to model the training process in meta-learning by trajectory optimisation with state as parameer and action as the weights of training tasks. The local optimal weights obtained from iLQR - a trajectory optimiser are then used to re-weight tasks to train the meta-parameter of interest. We demonstrate that the proposed approach converges with less number of training lasks and has a final prediction accuracy that out-performs some common hand-crafted task- weighting baselines.
Our proposed method also has some limitations that could be addressed in future work. TOW relies iLQR which is not ideal for large-scale systems with high dimensional stale space such as deep neural networks. Despite the approximation of Hessian matrices to use only diagonals as mentioned in Section 3.1, the lineari- sation of the state-transition dynamics and quadraticisation of the cost function are still time-consuming, and consequently, reduce TOW's efficiency. Future work might find a faster approximation to optimise the running time for TOW as well as evaluate on large-scaled datasots, such as meta-dataset (Triantafillou al., 2020).
Furthermore, our method is local in nature due to the Taylor's series approximation about nominal tra- jectory used in iLQQR. One way to improve further is to define a "global" or "stationary" policy To(x;, u.), which is similar to Guided Policy Search (Levine & Koltun, 2013; 2014). This policy can then be trained on multiple local optimal trajectories obtained from iLQR. While this approach may offer a superior generali- sation for the policy, scalability an issue since the policy needs to process the high-dimensional state x,. As result, a very large model may be required implement such policy.
Acknowledgments
This work was funded by the Australian Research Council through grant FT1101005255. The computation performed in this paper was managed and supported by the Deep Purple service at the Australian Institute for Machine Learning and the Phoenix HPC service at the University of Adelaide.
References
Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless C Fowlkes, Stefano Soatto, and Pictro Perona. TASK2VEC: Task embedding for meta-learning. In In- ternational Conference on Computer Vision, pp. 6430 6439, 2019.
Kelsey Allen, Evan Shelhamer, Hamul Shin, and Joshua Tenenbaum. Infinite mixture prototypes for few-shot learning. In International Conference on Machine Leurning, pp. 232 241. PMLR, 2019.
Brian DO Anderson and John B Moore. Optimal control: linear quadratic methods. Courier Corporation, 2007.
Sungyong Baik, Myungsub Choi, Janghoon Choi, Heewon Kim, and Kyoung Mu Lee. Meta-learning with adaptive hyperparameters. In Advances in Neural Information Processing Systems, 2020.
Jonathan Baxter. model of inductive bias learning. Journal of Artificial Intelligence Research, 12:149-198, 198, 2000.
Costas Bokas, Effrosyni Kokiopoulou, and Yousef Saad. An estimator for the diagonal of a matrix. Applied numericca mathemalics, 1211 1229, 2007.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Interna- tional Conference on Machine Learning, pp. 41 41-48, 2009.
15

---

Published in Transactions on Machine Learning Research (08/2023)
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. In International Conference on Machine Learning, pp. 557-565. PMLR, 2017.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends& in Machine Learning, 8(3-4):231 357, 2015.
Diana Cai, Rishit Sheth, Lester Mackey, and Nicolo Fusi. Weighted meta-learning. In ICML Workshop on Automated Machine Learning, 2020.
Zhao Chen, Vijay Badrinarayanan, Chen-Yi Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normal- ization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pp. 794 803. PMLR, 2018.
Liam Collins, Aryan Mokhtari, and Sanjay Shakkottai. Task-robust model-agnostic meta-learning. In Ad- vanced Neural Information System Process, 2020.
Felix Dangel, Frederik Kunstner, and Philipp Hennig. BackPACK: Packing more into backprop. In Inter- national Conference on Learning Representations, 2020.
Guneet $ Dhillon, Pratik Chaudhari, Avinash Ravichandran, and Stefano Soatto. A baseline for few-shot image classification. In International Conference on Learning Representations, 2019.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. On the convergence theory of gradient- gradient-based model- agnostic meta-learning algorithus.. In International Conference on Artificial Intelligence and Slatislics, pp. 1092-1092. PMLR, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, pp. 1126-1135, 2017.
Scbastian Flennerhag, Andrei A Rusu, Razvan Pascanu, Francesco Visin, Hujun Yin, and Rain Hadsell. Meta-learning with warped gradient descent. In International Conference Learning Representation, 2020.
Michelle Guo, Albert Haque, De-An Huang, Serena Yeung, and Fei-Fei. Dynamic task prioritization for multitask learning. In European Conference on Computer Vision (ECCV), 9 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference Computer Vision and Paltern Recognation, pp. 770 778, 2016.
Laurent Jacob, Jean-philippe Vert, and Francis R Bach. Clustered multi-task learning: A convex formulation. In Advances in Neural Information Processing Systems, pp. 745-752, 2009.
David Jacobson and David Q Mayne. Differential dynamic programming. Elsevier, 1970.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in Neural Information Processing Systems, 26:315-323, 2013.
M Khodak, M Balcan, and Talwalkar. Adaptive gradient-based meta-learning methods. In Advances in Neural Information Processing Systems, 2019.
M Pawan Kumar, Benjamin Packer, and Daphne Koller. Self-paced learning for latent variable models. In Advances in Neural Information Processing Systems, volume 1, pp- 2, 2010.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332 1338, 2015.
Hae Boom Lee, Haycon Lee, Donghyun Na, Sachoon Kim, Minscop Park, Eunho Yang, and Sung Ju Hwang. Learning to balance: Bayesian meta-learning for imbalanced and distrriiistiibutionn tasks. In International Conference Learning Represendations, 2020.
16

---

Published in Transactions on Machine Learning Research (08/2023)
Yoonho Lee and Scungjin Choi. Gradient-based meta-learning with learned layerwise metric and subspace. In International Conference on Machine Learning, pp.
PMLR, 2018.
Sergey Levine and Vladlen Koltun. Guided policy search. In International conference on machine learning. pp. 1-9. PMI.R, 2013.
Sergey Levine and Vladlen Koltun. Learning complex neural network policies with trajectory optimization. In International Conference on Machine Learning, pp. 829-837. PMLR, 2014.
Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few-shot learning. arXiv preprint arv: 1.f7:1700.000 000.0, 2017.
Liyang Liu, Yi Li, Zhanghui Kuang, Jing- Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Towards impartial multi-task learning. In International Conference on Learning Representations, 2021.
Shikun Liu, Stephen James, Andrew Davison, and Edward Johns. Auto-lambda: Disentangling dynamic task relationships. Transactions Machine Learning Research, 2022. URL forum?id-KKeCMinnSW.
/ net/
James Martens. Deep learning via hessian-free optimization. In International Conference on Machine Learning, volume 27, pp. 735 742, 2010.
Devang K Naik and RJ Mammone. Meta-neural networks that learn by learning. In International Joint Conference on Neural Networks, volume 1, pp. 437 442. IEEE, 1992.
Yurii Nesterov. Introductory lectures on conver optimization: basic course, volume 87. Springer Science Business Media, 2003.
Cuong Nguyen, Thanh- Toan Do, and Gustavo Carnciro. Uncertainty in model-agnostic meta-learning using variational inference. In Winter Conference on Applications of Computer Vision, pp. 3090-3100, 2020.
Cuong C Nguyen, Thanh-Toar Do, and Gustavo Carnciro. Probabilistic task modelling for meta-learning. In Conference Uncertainty in Artificial Intelligence, pp. 781-791 PMLR, 2021.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms, arXiv preprint arXiv:1803.0200 2018.
Elijah Polak.
Computational methods in optimization: a unified approach, volume 77. Academic press, 1971.
Sachin Ravi and Alex Beatson. Amortized bayesian meta-learning. In International Conference on Learning Representations, 2019.
Sachin Ravi and Hugo Larochelle. Optimization as model for few-shot learning, In International Conference on Learning Representations, 2017.
Herbert Robbins and Sutton Monro. stochastic approximation method. The unnals of mathemalicul statistics, pp. 400-407, 1951.
Jürgen Schmidhuber. Evolutionary principles in self-referential learning (On learning how to learn: the meta-meta-... hook). Diploma thesis, Technische Universität München, 1987.
Ozan Sener and Vladlen Koltun. Multi-task learning as multi-objective optimization. In Advances in Neural Information Processing Systems, 2018.
Changjian Shui, Mahdieh Abbasi, Louis-Emile Robitaille, Boyu Wang, and Christian Gagné. A principled approach for learning task similarity in multitask learning. In International Joint Conference on Artificial Intelligence, pp. 3446 3452, 2019.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems, pp. 4077-4087, 2017.
17

---

Published in Transactions on Machine Learning Research (08/2023)
Yuval Tassa, Tom Erez, and Emanuel Todorov. Synthesis and stabilization of complex behaviors through online trajectory optimization. In International Conference on Intelligent Robols and Systems, pp. 4906- 4913. IEEE, 2012.
Sebastian Thrun and Joseph O'Sullivan. Discovering structure in multiple learning tasks: The TC algorithm. In International Conference on Machine Learning, volume 96, pp. 489-497, 1996.
Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science Business Media, 1998.
Emanuel Todorov and Weiwei Li. A generalized iterative lag method for locally-optimal feedback control of constrained nonlinear stochastic systems. In American Control Conference, pp. 300 306. IEEE, 2005.
Eleni Triintafiilou,, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Utku Evci, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, al. Meta-datasct: A dataset of datasets for learning to learn from few examples. In International Conference on Learning Regresenlations, 2020.
Oriol Vinyals, Charles Blundell, Tim Lillicrap, Daan Wierstra, el al. Matching networks for one shot. learning. 2016.
In Advances in Neural Information Processing Systems, pp.
Yixin Wang, Alp Kucukelbir, and David M Blei. Robust probabilistic modeling with bayesian data reweight- ing. In International Conference on Machine Learning, pp. 3646 3655. PMLR, 2017.
Zhixiong Xu, Xiliang Chen, Wei Tang, Jun Lai, and Lei Cao. Meta weight learning via model-agnostic meta-learning. Neurocomputing, 432:124-132, 2021.
Sidney Yakowitz and Brian Rutherford. Computational aspects of discrete-time optimal control. Applied Mathematics and Computation, 15(1):29-44 45, 1984.
Zhewei Yao, Amir Gholami, Shong Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoncy. ADA- HESSIAN: an adaptive second order optimizer for machine learning. AAAI Conference on Artificial Intelligence, pp. 10665-10673, 2021.
Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungjin Ahn. Bayesian model-agnostic meta-learning. Advances in Neural Information Processing Systems, pp. 7332-7342, 2018.
Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik, and Silvio Savaresc. Taskonomy: Disentangling task transfer learning. In Conference on Computer Vision and Pottern Recog- nilion, pp. 3712-3722, 2018.
Yu Zhang and Ran Yeung. A convex formulation for learning task relationships in multi-tase task learning. In Conference on Uncertointy in Artificial Intelligence, 2012.
18

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix A Approximation of state-transition dynamics and cost function
A.1 Linearise the state-transition dynamics
In this section, we present the explicit form of the linearisation for the state-transition dynamics f. In particular, the main purpose is to obtain the first gradient of the state-transition dynamics f w.r.t. x and u defined in Eq. (2). We consider two cases of f: stochastic gradient descent (SGD) and Adam.
A.1.1 Stochastic gradient descent (SGD)
The transition dynamics is given as:
X++1 f(x,, uz) = X2 -a xx [u. L(xx)]
(9)
Applying Taylor's expansion to the first order around a state-action pair (X, (2)) gives:
X++1 = 84+1 + (TD - aV [ (xx))) (x+ - x,) - a [L(x))] (u, - ú.)
(10)
If ōx,- = x, -x, & (similar for ōu,), then:
8x = (ID - V a L(x:))) 8x: + (-avi (Lxx)) őu.
(11)
Hence, the coefficient matrices of the Taylor's series for the stato-transition dynamics following the SGD update can be expressed as:
Fx, In V i L(x,)
(12a)
Fu. L(x)))
(12b)
A.1.2 Adam
The gradient-based update for the parameter of interest using Adam keeps track of the mean and variance:
m, = + [1)3, u,
(13)
Vt = -2v,-11 + (1-12) (J. u.) o (J,T))
where:
= 0, = 0.
J, the Jacobian matrix of the validation losses for tasks minibatch t-th:
J, V [L(xx))
(14)
④ is the elementwise multiplication.
The corrected-bias estimators are then defined as:
m11
m
1 g
(15)
Vt
V
The update or state-transition dynamics is then given as:
m
X f(x; u.) = X2
(16)
19

---

Published in Transactions on Machine Learning Research (08/2023)
The update for new state of the model parameter can also be written by substitution:
m
1
X
X
3i
Vi
+
(17)
1-
The approximation using Taylor's expansion up to the first order will result with the following matrices:
Fx,
ID
(18a)
+
1-00
Fu, 1 $
OF
m;
(18b)
1-85
In the following, we determine the gradient terms in each of the expressions for Fx, and F Fu,- For simplicity, assume that m and Vt-1 constant w.r.t. both X;-1 and u1-1-
The gradient w.r.t. on m, and can be expressed as:
Vn, m [m.] (1 -11-810 B. [v 2(1 - 32 uJ)
(19)
Note that given functions f (a notation for general function, not the state-transition dynamic) and g
f
g
9
u [f
V
U. [g]
(20)
Hence:
m,
[m
m
(21)
U.
+
2
13)v,
To calculate the gradient w.r.t. xt, the update of Adam is rewritten as:
m. -8mm,-+ (1- x In [xx)]
(22)
Vi 82 (1 32) V Ju L(xi) Vx, [u, Lxx))
The gradient w.r.t. x; on m, and v: can be expressed as:
[m:] (1- 81) V | L(xx)]
[v]] = 2(1 B2) V [n L(x:)
v2 [u, [.,])
(23)
Hence:
m
[mg]
m
Xx [V]]
(24)
X:
V.
+ c
2
83)
I- b
1-3
(
1-
Vi
20

---

Published in Transactions on Machine Learning Research (08/2023)
A.2 Quadraticise cost function
The cost. function consists of two terms: validation loss on the query subset and penalisation of the action uy. Since the cost function itself is already in the second-order function w.r.t. the re-weighting vector u, we only need to quadraticisc the function using Taylor's scries w.r.t. the meta-parameter x.
A.2.1 Quadraticise validation loss The loss term in (5) can be approximated to second order
(xx)) R ((x)) + (x+ 2,)" V Vx. [[))( + (xi -2) V2 [X))(())] (xi (x--))). &().
(25)
A.2.2 Quadraticise the penalisation of the action u,
The penalisation term is already in the quadratic form. Here, it is rewritten following the Taylor's series follows:
Bu
Bu
(u )uls) (u: HulM) = [uu - û.) + u MulMl [(ue ú.) + ûr - pulMs
2
2
1
2 (u: û.) (B.IM) (u: ú.) + (u: á.)) Bu (ū, - uular)
(26)
Given the Docaly--uadrrate approximation of the cost function in (25) and (26), the coefficient matrices and vector of the quadratic form of the cost function can be writlen as:
Cx
V2 x; [xx)))
(27a)
Cx, Cu..x =
(27b)
Cu.u BuIar
(27c)
Cx. Vx 1
(27d)
Cu. = Bu (ú, - pnlM)
(27e)
21

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix B Derivation of iLQR
This section presents the derivation for one itcration of iLQR adopted from the original iLQR papers (Todorov & Li, 2005; Tassa et al., 2012). One difference is at the backtracking line-search where we employ similar acceptance criterion as in DDP.
For brevity, the finite-horizon discrete-time optimal control problem in (1) is restated as follows:
T cxx,, ur) s.t. Xc+1 = f(xe, uz),
[ui]a - arg min (ue)ez
Σ
(1)
Let Qs be the cost-to-go defined as:
T
Qx(TET uer) = > uy),
(28)
and Vi be the value function at time step In other words, Vi is the cost-to-go given the local optimal action sequence:
(x(x,) = min (x uer).
(29)
According to the Dynamic Programming Principle, the objective function that minimises the cost function over an entire sequence of actions uer can be reduced to a sequence of minimisation over a single action up proceeding backwards in time:
Vi(x,) = min Σ
cxxx uy)
(according to the definition of Q defined in Eq. (28))
u j=t
min
c(x,, u.) + min
c(x,, uy)
Us
min
[c(x;, u²) + V (X++1)]
U.
= -min [c(x;, u.) + V (f(x;,
(30)
us
By applying the Taylor's series to the first order for the state-transition dynamics f and to the second order for the cost function about a nominal trajectory (X,, úz), the term inside the minimisation in (30) can be approximated to:
1
qx
qu
1
1
c(x;, u.) V (f(x;,
8x, Ju.
qx,
Qx
Qx
dx:
(31)
2
qu:
Qui
Que
du:
where:
8x. = Xt.
Et,
Su.
(32)
and
qx. Cx2 F Vi+1
qu. Cu, + FT FT. Vi+1
Qx Cx,x, + F..
( F Fx,
Qu, Cu.,u, F VetiFu.
Qx QU,x, Cx,,, + F V FI.I
(33)
with Fx, and Fu, are the first gradients of the state-transtition dynamics f w.r.t. the variable in the subscripts, C() and C()) are the first and second gradients of the cost function e w.r.t. the variables in the subscripts (please refer Eq. (2) for details), and Vi+1 and Ve+i are the first and second gradients of the value function V w.r.t. the state X
22

---

Published in Transactions on Machine Learning Research (08/2023)
Minimising the quadratic problem in (31) w.r.t. ou, results in:
ou; K,dx, + k,,
(34)
where:
K2 Qui,u. Qui
k
-Qu quo
(35)
Substituting the action ou, obtained in (34) into (31) gives a quadratic model:
Vi((a) ~ 2 ox, Voox, 5x, + const.,
(36)
where V, and are the second and first gradients of the value function V. w.r.t. the state x,, which can be expressed as:
V, = Qx Qx K,
Ve = 9x, + Qx.,u, kz.
(37)
Hence, one can recursively calculate the local quadratic model {V,, of the value function V2, and the linear controller (K, k backward through time. Once is completed, a new trajectory can be calculated using the forward pass with the general non-linear state-transition dynamics f as follows:
&1 X1
u, Kr (xt 82) + k. + ûr
X++1 = =x(,, ut).
(38)
Although the trajectory obtained in (38) converges to a local minimum for the approximate model of the value function V,, it does not guarantee the convergence for general non-linear models such as the one in (1). The reason is that the new trajectory might deviate too far from the nominal trajectory, resulting in a poor Taylor approximation of the true model. To overcome, backtracking linear-search with parameter E (0, 1] is introduced:
u K2 (x, x2) ck₃ + ú,
(39)
The criterion to accept the trajectory produced in the iteration with backtracking line-search is similar to the one in DDP:
(uir)) J(ūrT)) < 5eh.
(40)
where
T Juir)) = Σ c(xe, uz) 1=1
(41)
0- 0,+1 qu, Que que
(42)
Hence, in the worst case when the new Irajectory strays too far from the model's region of validity, then + 0 and the trajectory is the same the nominal trajectory. The procedure of one iL.QR iteration is outlined in Algorithm 2 in Appendix D. In addition, the convergence proof for iLQR adapted from DDP (Yakowitz & Rutherford, 1984) is provided in Appendix C to complete the analysis.
23

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix C
Convergence of iLQR
To prove the convergence of iLQR algorithm, we rely on Theorem 3 in (Polak, 1971, p. 14), which is re-stated in Theorem 3 in Appendix C.1
C.1 Auxiliary to prove convergence
Definition 1 (Algorithm model). Given a: T is algorithm, and c:: T - R. is function used slopping crilerion, the algoritim model can be described us:
1. Compute a TT.
Set 0.
3. Compute alzz).
4. Set 2++1 = a(2i).
5. If (f((2))) 2 c(zi), stop, else set i= i 1 and go to step 3.
Theorem 3 will show what such an algorithm will compute. Theorem ((Polak, 1971, Theorem 3, p. 14)). Suppose that:
(1) cl)) is either continsuuus at all non-desirable points E TT or else c(z) is bounded from below for 2 ZET;
(i) for every 2 T which is not desirable, there erist on 5(x) and a 8(2) such that:
c(a(z')) <(=') s 5(2) <, B(z,e(2)) {a E T || - 2|1s S <(z)).
(43)
Then, either the sequence {2i} constructed by algorithom defined in Definition is finite and next last element is desirable, or else it is infinite and every accumulation point of {zi is desirable.
C.2 Proof of iLQR convergence
Since iLQR is an algorithm model as described in Definition 1, prove its convergence property, one needs to assert that the 2 conditions in Theorem are met.
First condition is satisfied since the cost J(ur)) is continuous, and twice differentiable. We will prove that iQR satisfies the second condition of Theorem 3.
From (39) and (42), the sequence actions obtained from iLQR u1.7 and 8 depend on the nominal trajectory. Hence, we explicitly denote them as u(e, û) and (a). In addition, since the loss function {(.) is assumned to be twice differentiable, the cost function and state-transition dynamics are continuous. Hence, both u(s, û) and 0, (ŭ) are also continuous.
The convergence proof for iLQR is presented in Appendix C.2.2. The proof requires some auxiliary lemmas in Appendix C.2.1.
C.2.1 Auxiliary lemmas
Lemma 4. If u(E, u) is non-stationary sequence of actions, then 0, < 0.
3 A direct for determining if 2; desirable may substituted for the test c(2(+1) c(2)).
24

---

Published in Transactions on Machine Learning Research (08/2023)
Proof. It is apparent from the derivation of HLQR that qu, = =0, E (1...., T) if and only if û is stationary sequence of actions. Hence, by negation, if á is a non-stationary sequence of actions, qu, o 0 for some time steps & E {1, {I....,T}.
Also note that, if u(E, u) is non-stationary, then also non-stationary. In addition, the Hessian matrix w.r.t. action is assumed be positive-definit (PSD), then from the construction, Quaay and its inverse are also PSD.
Th PSD of Qui.a, combining with (42) leads to the fact that 01 (ū) < 0 for any non-stationary sequence of actions ú.
Lemma 5. Ifu(e, u) is non-slationary sequence of actions, then:
J(u(s, u)) (() = ( (
(44)
Proof. If we define AQ, as the resulting incremental change in the cost-to-go Q from the perturbation Au; using Taylor expansion, then:
AQ: = Qc(xe, u, + Au)) - Qx un) = qu. Au, + +OIlAu,
(45)
Note that:
J(u(s, ú)) - J(ú) = J(uIN + AuJ:N) J(UINN
AQ
qu. Au, + +0 (IAuI.N(1)).
(46)
1=1
1=1
Taking the derivative w.r.t. gives:
aJ(((e))
T du,
qu: de
(47)
de
From (39), u, is lincar function w.r.t. E, resulting in:
du
k; Qua qu.-
(48)
de
Hence:
aJ(u(s))
qu
Qui,u qu, (1)).
(49)
de
which implies (44).
Lemma If (((ee ū) is a non-stationary sequence of actions, then there erists e such that:
Juue,, ū)) - J(û) s (u)) Ve E [0, Vn E (0.5, 1).
(50)
Proof. The result in Lemma can be re-written as:
Jul(, 0)) (() n00, (û) + (1)
E (0.5, 1).
(51)
Hence, there must exist value E > 0 such that: (1 (() O < 0. This leads to:
J(u(e, a)) (() ne 0 (f) < n (i). Ye [0,=]]-
(52)
Note that the second inequality holds due to the negativity of 02 proved in Lemma 4. This completes the proof.
25

---

Published in Transactions on Machine Learning Research (08/2023)
Lemma 7. Given that 81 and any action sequence continuous, and for any 6 >> 0 such that |jû-ul 5, there evists 62 > O such that:
(i(u) - 52 < <
Proof. Since 81 is assumed to be continuous, and the variable u also continuous, we can employ the definition continuity of function 0, u to obtain the following:
Vsz > 0 = 38 > 0: (ū) - 0 [u)l < E2, vú {u u || - ul < o}.
(53)
The negation form the statement (53) can be expressed as:
AA - 3e > 0: 0, () 0, < e2, va E {u Eu: : : ul < <d}.
(54)
Solving the inequation above completes the proof.
C.2.2 Convergence of iLQR
Theorem Let the state-transition dynamies f and the cost function have continuous second portial derivalives w.r.l the continuous slate and action If u(e, û) denoles the successive application of iLQR, then any accumulation point of u(s, û) is stationary w.r.t the finite-horizon cost J (uir)).
Proof. We first determine the condition that u(e, ú) caleulated in (39) is the successor of iLQ at ú. Ac- cording to iLQR algorithm, u(s, ûi iLQR(()) only if:
Jule, û)) ((ū) s & [0)(0).
(55)
2
Note that from Lemma 6:
Jul(, - J(a) < nedl (i) < 8 (0). Vn E (0.5, 1), Ve E [0,si]
(56)
which indicates that u(e, û) is successor of iLQR when E [0, =1].
If uec ú) is successor of iLQQR, the acceptance criterion combining with the result in Lemma 7 leads to:
E (a) < E (m) + 02), vû E {ú e u: | ull < 6}. (iLQQ(()) - (() <
(57)
2
2
Note that if 5 is set to be small enough, then is also very small, resulting in (n) + 82 < 0. II we set d() = ö and e e(ū) (u) + E2], then iLQR satisfies the second condition in Theorem 3.
26

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix D
Trajectory optimisation algorithm(s)
Algorithm 2 Implementation of iLQR backward to determine the controller of interest.
procedure I..QRBACKWARD 0x)=1)
Vr+1 0, and VI+1 = 0
Quadratic matrie and vector of cost-to-go Stopping criterion for nominal trajectory
0T+I =0 for t = : do
Backward through time
5:
Fx, Fu, = linearise dynamics
see Appendic A.1 D see Appendiz A.2
6:
Xx Cx.,u,, Cua Cu. Cx. Cu. = quadraticise cost function
7:
xx = Cx.x, + Xt Fx.
2nd derivatives of cost-to-go go
Qx, F X1
V IttiFu.
Qu, F., V V4IFF Fx,
10:
Qui.ur Cur.ue F. V Fu.
11:
9x. = Cx. + FT Vi+1
1st derivatives of cost-to-go
12:
qu, = Cu. F u, Ve+1
13:
K; Qui.. Qu.:.
Linear controller
14:
k; Qui qu.
15:
V, xx + Qx K,
2nd derivatives of value function
16:
= 9x, + Qx
> 1st derivatives of value function
17:
8, = 0141 - qn Qu qu.
18:
return (Kr, kr);=1,0
27

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix E
Complexity analysis
Before analysing the complexity of the proposed method presented in Algorithm 1, we provide a detail description of the annotations used for brevity in Table 4.
Table 4: Detailed description of the notations used in the paper.
Notation
Description
O((T))
Time complexity train meta-learning method fol-
lowing uniform weighting strategy
D
Dimension of meta-parameter
M
Number of tasks within "meta" mini-batch
mg = max.e (1..... m
(4)
Total number of validation samples within task Number of iterations used in iLQR to calculate the weighting vector
n.QQR
Tis
Number of iterations performing line-search Number of arithmetic operations in the model of interest.
The downside of TOW is the overhead due to the linearisation and quadraticisation for the state-transition dynamics and cost function, and the calculation to obtain the controller K, and kG shown in Algorithm 1. If O(T)) the time complexity to train a meta-learning method following a uniform weighting strategy, then the time complexity required by TOW will consist of the following:
nominal trajectory to obtain (&,, ú.): O(To)
linearisation and quadraticisation using Gauss-Newton matrices: O(LLoommnDD)
iLQR backward: O(LLQRAM
iLQR forward with back-tracking line search: OmLLQR")TTT)
Thus, the final complexity of TOW is: O( + 1)To + (...QR(mo)) + M)D))
comparing to O(T))
the conventional meta- meta-learning
28

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix F Convergence analysis for trajectory optimisation based task weighting meta-learning
F.1 Notations
The following notations are used throughout the paper:
|||| is the 1,2 norm of vector RD, e.g. Vx x
||| the matrix norm of matrix A E RM induced by vector norm:
|x
|
sup
Vx E R
||x + 0.
Given the vector and matrix norm, two common inequalities used in this section are:
Triangle inequality:
A B|| < ||||| + BE
Sub-multiplication:
Ax S | [
F.2
Auxiliary lemmas
F.2.1
Boundedness of action (or weighting vector)
Lemma 1. If U, is stationary action of nominal action ú, obtained from iLQR, then:
38 [u, à.l S 6.
Proof. According the procedure of iLQR shown in (39):
ú. K, (x+ &,) + ckc.
(58)
Since the matrix K, vectors k, x and are well-defined, the norm of á is also well-defined.
In addition, 8 is implicitly related to the Gaussian prior N( Mal, IM) A larger value of 3r in (5) would result in smaller value for 8.
Corollary 9. If Ut and ut, stationary actions of two nominal actions ú,, út Duniform 1/M 1 oblained from iLQR at lime sleps iı and l2, respectively, then the followings hold:
[us u,al $ 28
(bounded La norm of difference between 2 actions)
(59a)
[4,
+ VM
(bounded L2 norm)
(59b)
u, < <8VM + 1.
(bounded L1 norm)
(39c)
Proof. The first incquality (59a) can be proved by simply applying triangle incquality on the & norm and employing the result in Lemma
Un
|l
= [1( úmuiform)
+ (uuniform uz)il
< [n Anniforma + |/u unniforml
<
[m ún + [u; ût
25.
(60)
The second inequality (59b) can similarly be proved using triangle inequality:
juu = [u, üuniform) + üuniforml < jus auniform + [
5
(61)
v M
29

---

Published in Transactions on Machine Learning Research (08/2023)
The last inequality (59c) can be proved by the Sauchy-Schwart inequality:
1 1J u, $ [1 un
S V |u, |l
(Cauchy-Schwarz inequality)
< LVVM M + 1.
(due to Ineq. (59b))
(62)
F.2.2 Bounded variances of the loss' gradient and the weighted validation loss
The bounded gradient norm in (6) (see Assumption 1) also implics that the variance of gradient of the loss function w.r.t. training samples bounded as shown in Lemuna 10.
Lemma 10 (Bounded variance of loss' gradient). If Assumption I holds, then the variance of gradient of the loss function is bouuddedd:
3o > 0: E (BI
V ((igs Bijix) (( V (( Mayi x) x)]||| & o0, Vx R
Proof. The term inside the expectation in the left-hand side can be written as:
x (sij, V)jix) - B(s x ((( Viyi Viji x)]]
x ((1)) Viyi x)|1 (B V [
Myj x)]]
(triangle inequality)
x May x]] + B))
(Jensen's inequality) (Boundedness of gradient in (6), which is proved in Lemma 14).
(S Byy x)ll
S 2L
(63)
Selecting a 2T. completes the proof.
The result in Lemma also leads to the boundedness of the variance of the weighted validation loss as shown in Lemma 11.
Lemma 11 (Bounded variance of weighted loss). Given Lemma 10, the variance of V [u] L(x:)] is bounded above by 52 = o² (6 10++M-
Proof. We use the following well-known inequality for variance as a part of the proof.
If X, E (I,...,m) n}, are random variables with finite variance: Var (X,) < +00, then:
Var
Σ
Xx
Var (Xi).
Note that the validation loss Lx(x;) defined in (3) is the empirical expected values of loss é evaluated task i-th. Hence, applying the above inequality for variance gives:
my
m,
Var vxL((x))) Var
xl
(0)
Vij
x
Sip
Vikk
x
mg
m
j=l
1
m
Var
Sij
(g) (4)
Vij
x
m,
Vx
(a)
Bikk
xxx
my
m
/=1
k=
2 (due to Lemma 10)
(64)
30

---

Published in Transactions on Machine Learning Research (08/2023)
The variance of the weighted validation loss x L(xx) can be written
M
Var (Vxu, L(xx))
-aa Σ us VxL((Xe))
Σ
i Var (VL((xx))
M
S 0 2 u = o²||
(due to the result in (64))
i=1
S 0² (6 + M m0
(Corollary 9).
(65)
F.2.3 Smoothness of validation loss
In this subsubsection, we prove Lemma 12 about the smoothness of validation loss. To make the subsub- section self-contained, we re-state the definition of the task-specific parameter (I(x) and the true validation loss L(xx) follows:
m
0x(x) =
Vx l
Sik Mek xx)
(8) (x)
(66)
[(x) = E
( (
(0)
~D(1)
(9 Mig : [e (s(g
(7)
Note that in (3), the task-specific parameter is defined as:
le (5/2 Ui ox(x))
(x) E arg min di
Here, we approximate that lower-level optimisation by one step of gradient descont as shown in Eq. (66).
Lemma 12 and proof are shown as follows:
Lemma 12. If the conditions in Assumptions to 3 hold, then the gradient of the true volidation loss [((x)
defined in Eq. (7) is S-iiscchitz, where: = S(1 + +S8 + TpL.
Proof. Before starting the proof, exploit the notation of gradient of the loss function at point follows:
vxl(s,g:v) V xx((s,
677)
x=V.
Given the definition of the true validation loss in Eq. (7), its gradient w.r.t. x can be calculated using chain rule as follows:
vxLi(x) V ( D!
[e (s (g) (0)
s Me ; dx(())]
= B of >DD(
V (x)((x)
(sl :
(68)
And since de(x) obtained by gradient descent does not depends validation (or query) samples, we can, therefore, rewrite the above gradient as:
(xL( Vxúx(x) x E
V
(g) (g) [x(x)) (
(() (4) -D(()
.u.
I -
I
7I (el)) (/ v55 s(s))))
4
(5(2)
Ui x }}}
x E (e w wD(())
V (s(())) y/ d1(x)) (69)
31

---

Published in Transactions on Machine Learning Research (08/2023)
Note that we abuse the notation and use in set S
to indicate the average evaluated on all data points
In the following, we omit sample (s. y) from the expectation to simplify the notations. In particular, the above gradient can be ro-written as:
(x) = { 11- B s(
8ik "ikk [2;;x Vxl
x
]] R
(5)) Uij (x)x
(70)
Thus, we can calculate the difference of the gradient evaluated on the same task Tj but with two different mela-parameters:
V x
(X)
B S(-))
((-))
Sik Vik :x)
(e)
E D(o)
yik
I bIF 1/
V l
1i Mik
(9)
B BD((
Sik yik
9 x
}
18 D
5ik
(4)
I E
(2/2 ui x
+ T
E S(D)
(912) . x
E D
Dik
(adding and substract the green terms)
B S(-))
g())
V 1 Mik
E D(1))
Mikk
(x)
xx (5 with :
-7B s(())
v
Pik :x
2
Vik
xx)
Vx
(s/2 212 di (5x )]
(71)
Applying the triangle inequality gives:
V x (x) VLL.
(x)
Vill
( Bikk
(x) :xx }} Epn
Vxl
(n)
Yikk
(0)
r xx(x))
- -xxl (s yek ;
+
s(o))
Sike
(a)
Vik xx) &el
sike Vik
(o)
(a),
x
E
V
(4)
v (x) (x)
(72)
Next, we upper-bound the two terms in the right-hand side of Ineq. (72). The first term can be upper- bounded as:
First term
I
BB g(+)
Sik Mik
(s) (4)
x
E pp"0)
V
5 Vik
(9) (g)
:e: (.()))
v
(g)
<
I
(5)
Vik
(s)
E E
V l
(g) (g) sik yik (x
V
s Uik
(n)
(73)
Applying Jensen's 's inequality on the L2 norm of the expectation in the right-hand side of the above inequality to bring the expectation outside of the L2 norm, then employing the smoothness of in Assumption 2 to obtain the following:
First term <
I
7B
g(o)
(0)
E
p(()) V2 (s U (1(()))) Vxl [x)
(g) (5 Vikk (2
2E S(C)
x
(x) (V is S-Lipschitz Assumption 2)
(74)
32

---

Published in Transactions on Machine Learning Research (08/2023)
Given the definition of Φ(x) in Eq. (66), we can obtain the following:
[x)(()) (x)]1
(x x) L1 s/rr
Big
xx
Vx
x
Sik Mij ;
| x|| Y
IE s(
Vx l
iik Vif (5)
vx
Sik Vij""
(5) (8).
(triangle incquality)
x X|| + + )(-))
vx
Sik Dij
(a)
x
V x
(5/x) Mig
(e) (e),
(Jensen's 's inequality)
S x x|| S x
X|| (Assumption 2)
(75)
Thus, can upper-bound further Inoq. (74) as follows:
First term S(1 + 7:5 I ylB
- s(n)
Vil
(iik -ikk'' ; )]][| || ||8-x|] x|
(8) (a)
(76)
If (Aa)2=1 are the eigenvalues of the Hessian matrix E Bi(())
V
Sit Vikk x
(s) (5)
then due to Lemma 13,
the eigenvalues of BBg()) V (() :ik x)]
In addition, since I is symmetric and positive semi-definite, its norm equals to the largest eigen-
are {1 - yda )
gg [
(0)
ik ; xx)
value (refer to Lemma 15):
ylE s(o) g(
V
Bit Wikk (2); ;
max [1 >dal S + y max (da]-
(77)
According to Assumption 2 and the result in (Bubeck et al., 2015, section 3.2, page 266), the eigenvalues of are smaller than S, implying:
the Hessian matrix g V (skk Vin (5), xx)
- TlB g(0)
Vill
sikk Bikk x
]]
< 1 +SS.
(78)
Therefore, the first term on the right-hand side of (72) is upper-bounded by:
First term I S1+ ||x x||-
(79)
The second term in the side of (72) can be upper- bounded as:
Second term
71 s(*))
Mik
x
Mik
(x)
E E
V
5ik yikk d
(q)
BB g(o)
V
Sik Vik
(a)
(e)
x
Va
E Bio
V
(x
7 3
(s) (e)
(s(2
Epp.o)
Vxl (s(2)
(Jenson's 's incquality)
s YpL |X x|| (Eq. (6) and Assumption 3).
(80)
Combining the results in (72), (79) and (80) gives:
V (x) - V.L. (x)
< [S(1 + + +xpD]] ||x x||
(81)
This completes the proof.
33

---

Published in Transactions on Machine Learning Research (08/2023)
F.3
Convergence of TOW
Theorem 2 (Main result). If Assumplions 1 to 3 hold, the learning rale < <2/8(6vM-1)) and z is randomly
sampled from {xa]emt Tiner returned by Algorithm 1, them:
B ( M || L1: [2)][]
2~ (xe) (xx)
Thear
S 6o Titer
where:
48BVM M + a (iv (6VM+1) +1)
2u LI.M (x1)
co with Titer the number of grudient-update for the meta-parumeter, and E (D(t)) " over all query data sampled from t mini-batches.
a 2 as (6VM+ +1]
> 0,
==
a [ 2 2 a$ (6VM M + +1)]
the expectation taken
Proof. According to Eq. (7), [x(x E R is the expected validation loss of task i-th using meta-paramcter M
x. In addition, the notation (D(-)
indicates the data probability that generates query data pairs for M
tasks in mini-batch at time step 1.
For convenience, we also denote L(x) E RM is the vector consisting of expected validation losses on M tasks in a mini-batch of tasks:
[(x) = [1)))
L2(x)
Lar(z))
(82)
From Lemma 12, the gradient of the validation loss L, (x) is S-Lipschitz continuous. Hence, applying Taylor's theorem for the function (x(+11) at point x, gives:
3 (x++1) L (xx) + V Li (X++1 - xx) + 2 xx+1 x,11?
(83)
Note that un is constrained be non-negative as mentioned in Section 3.1 or in step 16 of Algorithm 1. Hence, multiply both sides by us 2 0,4 (1....;7)) T) E {I,..., M) and sum side-by-side to obtain:
3 | - 1 1 u.
u L (x++1) s un (x.)) +
[u,L (x:)] (x++1 - xx) + 2
(84)
For simplicity, we assume that the state-transiting dynamics is SGD as mentioned in (5). Thus:
Xe+1 - X, = - Vx [u, L Lxx))
(85)
Thus, one can simplify to: n [ S [x)((x)) - a a (xi) Vx [u, L [xx)] + " Vx [u,L [x:))x:]]ll 1
(86)
The above inequality contains the true validation loss [(x) and the empirical validation loss L(x). The relation between them can be written as: [(x) 1 )Dw))ar [L(x)]. Thus, we take the expectation of validation data (expectation over the f-th mini-batch, each consists of M tasks) on both sides to obtain:
Vx u [x,)) + $
E (D!D))* or [x++(( s E ( u [T(())) a V L (x:)
Vx [u,T [x)][]]] u
(87)
where (D(
M is the probability distribution of task data in the mini-batch of M tasks at time step L.
34

---

Published in Transactions on Machine Learning Research (08/2023)
Note that one can use Eq. (7) to imply that:
Vx x [u, (xx)] = B (D)o)) [Vx [u. L [xx)]]
which also implies:
V 1uL (xx)] =R (D!O))
" V [u L [xx)]]
(88)
Thus, the above inequality can be written as:
E (oooa) " (xx+1)] S (xx) -xx)) Vx [u. (xx)]1
a²$
Var V [u. [x+)] ]] + V| [u,[ (xx)]][] Th
+ 2
(since E[X²| Var[X] + (B[X])*)
1 a$ 1T | |urL (xx)]1]
S u [x,)
+ a²S
Var
V [u. L [xxx)] Th
(89)
Since:
Thu
$ S8VM M +1
(Corollary 9 in Appendix F.2.1)
Var Vx [x:)]]
S /2
(Temma 11)
then:
E (ofjo) " [u, (x(()x)+)) S (xx)) - -a 1 a5 2
(5VM+ + :)) | [u.[(x))]]++
2²525
(VM+1)
2
(90)
Re-arranging the gradient norm of the weighted validation loss to the left-hand side gives:
- a$ (OVM+1) | [x)][( s u, L [xx)) - E B(Doo0 [u. L (xx+1]] + a²o25
(VM+1))
(91)
The
excluding the constant term at the end, can be written as:
u, L (xx) - E F(D(o) or [u,
(x(+1)]
= uu L (x,)) - B (ofo)) or [u (xe+1)] + (oooo d [uTtiL (xx+1)] (oojo ar [u] L (xe+1)]
(92)
The second part in the right-hand side of the above expression can be upper-bounded as:
[ (TTI(((e+)) - EE (D/W)) ar [u. L [xx+1)]
(D(w) ar
= -E (pfo)) x
(ui+1 ue) L (xx+1)
< E (p(o) xr
|1411 - uell welly
Σ La (x4+1)
(Cauchy-Schwaal inequality)
< BBM M E (D. M ||uet1 -ell
(E is B-bounded, and hence, L, is B-bounded)
< 2 (Corollary 9).
(93)
35

---

Published in Transactions on Machine Learning Research (08/2023)
Hence, upper-bound further the hanh side of (91), resulting in:
a$ (DVM +1) || [x))(
[xx) -E (D((o))
" [ult ( + 28B
2²o2 (5vM+1).
(94)
2
We take the expectation over all the bati-batches used at time step t, I..... on both sides of Ineq. (94)
to obtain:
as 1 2 (oVM + +1 (Df [IVV [u,L (xx)) [1"]
2²o25
$ (D(2) m [u. L (x+)] - (D/O) m [uftz (xx+1)) xi] + 28BVA
(DVM +1),
(95)
2
where: E (D(Lo) x [-] = E (Df(o)) xr (D)0) a ... (DfO)arl))
Summing (95) from 1= to Titer gives:
a$ (6VM+1) +
Tiser Σ
E (D( x || |urL (x:) |1]
a
2
a²025 (OVM + +1)
$ u [ (x1) - E ( (D/t) yar
[unnert (XTre++1)] + Titer 28BV +
(96)
2
Note that Uni 20 (1..... Titor E (1..... M}, and L; 2 0 due to the non-negativity of the loss function assumed in (3). Hence, we can upper- upper-bound bound further the right-hand side to:
1 a$ (6VM+1) +1
Ther Σ 1=1
B(1s)
[[|V> |u, [x.[(())] s (x,)
Thtr 28BVM + 2025
(DVM++)
(97)
If the learning rate a is selected such that: - 03/2 (6VM+1) > O, then dividing both sides by
oTTierr [ [1-a3 a5/2 (6vM+1 +1)] gives:
2u L [(xxxx Ther 455V/M + 2222 [1M++]]]]]]
Tisar
E (Dfo) m | [u,L(xx))] (x:)]
(98)
Titer
a'liter
2 2-
(6VM+ +
Next, we use a similar trick in the SVRG paper (Johnson & Zhang, 2013) to make the left-hand side term useful. The idea. is to output some randomly chosen x, rather than outputing XT For simplicity, let z = X with uniform probability for t {l...., J.ter.. The expectation of the gradient norm in this case can be expressed as:
Titor
Σ
E (D(O) x
[| || [urL (x:)(x:)]l1 = E (x:) Titer
E
V [n. (x:)]1])
(99)
Titer
This combining with (98) leads to:
48BVM M + 2225 (8VM + +1]
2u L (xx)) + Ther
E (x. Tizer [<(o:e)" or [I V [u,[ (2)) [1]
(100)
aTter 2-aš a (6VM+ +1)
which concludes the proof.
36

---

Published in Transactions on Machine Learning Research (08/2023)
F.4
Miscellaneous lemmas
Lemma 13. If X is an eigenvalue of matrix A, then À 1 is eigenvalue of matrix A I, where I is the identity matria.
Proof. According the definition of eigenvalue, is an eigenvalue of A if:
det (A - XI) =
Rewriting the above equation gives:
det ((A - - -A- -A)) =
=0 =0
(101)
Hence, À- 1 is eigenvalue of A I.
Lemma 14 (Adpated from https: htap::::/maat.sssceeecchhe stackexchaagg....... http://m..ccoo///33333////779990
If a function R - Rm, where n, m N, is differentiable and L-Lipschitz, then gradient norm is bounded by L.
Proof. According to the definition of vector or matrix norm:
Vf(x) wll
sup
Ya, v E R": :|o||| 0
|0||
V (z()
ssu lim
A
X-0
Alllll
f()) (Av)
sup
(102)
A 0
||xl||
Applying the triangle inequality gives:
[[(x + 2u)
((x) f(x) (Av)l
[[ + Au) - s(x)
V f(x)| < sup lim
A-0
Mevlll
Ilf(x + Av)
f(x)
V fxx) (Av)ll
L|l + du - æll
sup lim
(ʃ is L-Lipschitz)
X
|eel
Avel
< L.
(103)
The first term equals to 0 since it corresponds to the definition of Fréchet derivative.
Lemma 15 (Adapted from https://math. stackexchange. wtt. com////119991//74777) If puxn is a positive definite symmetric matriz, then its largest eigenvalue
|A |
Amax max
+
(104)
Proof. According the spectral theorem, if A is a real and symmetric matrix, then there exists an orthonor- mal basis consisting of eigenvectors of A, denoted vn, U2...., Un Without loss of generality, let's assume that the corresponding eigenvalues are: X1, M......., An.
For any non-zero vector I, we can represent it in this orthonormal basis as:
C Ci E R.
(105)
37

---

Published in Transactions on Machine Learning Research (08/2023)
Then, the function of interest can be calculated follows:
|| A
EE== i= GA
x
n2 i=1 &
(tu)) is an orthonormal basis)
i==
< max A;. k
(106)
kE
The oquality occurs when the vector x equals to the corresponding eigenvoctor of the cigenvalue (I..,( di.
38

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix G
Experimental settings
For all experiments, GD is used to obtain (x(x) with 5 itcrations and a learning rate of 0.1 for Omniglot and 0.01 for mini-ImageNet. The learning rate for the meta-parameters, a, is set at 10 for all the setting, The mini-batch size is M = 10 tasks for Omniglot and M = 5 tasks for mini-ImageNet.. For the Dirichlet concentration of the prior in the exploration and exploitation baselines, we try three values x E {0.2, 1.2, 5), and found that too small value of Ieads noisy learning since only the easiest or hardest task is selected, while loo large value of makes both the baselines identical to uniform weighting. Hence, we select that balances between these two strategies. Note that 1 results in random prior, leading to trivial solution. For the trajectory optimiser iLQR, the state-transition dynamics f follows the formula of Adam optimiser since Adam provides a less noisy training as in SGD, The nominal trajectory is, as mentioned in Section 3.1, selected with uniform actions: û, = 1/Mr, Vj E iterations used iLQR is to speed up the training, although higher number of iterations can be used to achieve belter performance by trading off the rumning time. We also provide an ablation study with two different numbers of iterations in iLQR in Section 5.4, where the larger number of iterations in iLQR slightly improves the prodiction accuracy on the validation set of mini-ImageNer The mumber of time steps (or number of mini-batches) is T 10 for Omniglot and 5 for mini-ImageNet. The parameters of the prior on the action u, are Pa. = 1/M and 8, = 10. As we do not observe any major difference between different configuration of M and 7" used in this experiment, we report the result for the case M = 10 and = 5. Each experiment is carried out on a single NVIDIA Tesla V100 GPU with 32 GB memory following the configuration of NVIDIA DGX-1. The experiments mentioned in Section 5.4 is performed on a single NIVIDIA A100 GPU with 40 GB memory.
1.2
t E (1....,7) 7}. The number of
39

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix H Additional results when calculating with full matrix V.
In this Appendix, we provide additional results of prediction accuracy tasks formed from the evaluation sets of Omniglot and mini-lmageNet. These results are based on the naive implementation of iLQR where the auxiliary Hessian matrix V, of the cost-lo-go in (28) is exact without any approximation. We note that due to the quadratic complexity of running time, we cannot train the one on mini- mmag-ee until convergence. The running time using the full matrix Ve is 160 GPU-hour for Omniglot and 184 GPU-hour for mini-ImageNet..
Testing accuracy (%)
94
Testing accuracy (%)
92
35
96
90
uniform
88
30
exploration
exploitation
94
TOW
25
10
30
15
20
Number of tasks (x10,000)
Number training tasks (x10,000)
(a) MAML on Omniglot
(b) MAML on mini-ImageNNt
Figure 7: Additional results of prediction accuracy 100 random validation tasks using MAML when the matrix V, exact without any approximation.
40

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix Examples of loss functions satisfying Assumptions 1 to 3
In this section, analyse some common loss functions inchuding mean square error (MSE) and cross-entropy (CE) to see if they satisfy Assumptions to 3 made in Section 4.
Table 5:
Examples of some common loss functions used with the assumption on the linearity in Eq. (107) that satisfy Assumptions specified in Section 4.
Loss (Assumption 1)
Lipschitz gradient
Lipschitz Hessian
Bounded Lipschitz
(Assumption 2)
(Assumption 3)
MSE
CE
Before analysing the loss functions of interest, it is important note that the objoctive function mostly depends on the model used. For example, g(s; x) denotes the output of a model parameterised by x for the input s, then the objective function of interest {(g(s;)),)) Depending the model g used, the objective function is different and would be very complicated if g represents deep neural network since it would result in high-dimensional, non-linear, non-convex function. Analysing such general function is, however, beyond the scope of this paper. To simplify, we assume that is linear function of x:
g(s; Sxx
(107)
Boundedness in Assumption both the loss functions of interest, MSE and CE, are theoretically unbounded, and thus, does not satisfy this assumption. However, we can simply clip the loss to ensure the satisfaction of the boundedness assumption.
Mean square error
The loss function is defined as:
((x) = |Sx y122
(108)
Lipschitz continuity In general, MSE does not satisfy the Lipschitz contimuity property. However, when the vector norm of x is bounded: |x|] I Xo, then MSE Lipschitz-contimmuus.
Proof. The Lipschitz continuity means:
|Sx -yll
-1|xx
const. |x] -x2ll.
(109)
when the loss clipped Swhen the of parameter hounded
41

---

Published in Transactions on Machine Learning Research (08/2023)
The left hand side term can be expanded as follows:
[x - v112 ||5x; - y112
x Sx1 - -((xx)) ) y - -x x Sx2 (Sx2)) 81
||x Sxi - S Sx2 - 2(xi - x2) x
[|(x -x2) x T
Sxi +x, S S(x1 x2) - -2(x1 -x²) -x²) TTyl
(Iriangle inequality)
[[xx - x²) " Sx1 | S(x1 x2)o| + +2|1(xx - ST yl
(submultipliccatiie itc.uuuutiiiiimuuuutttytttyl inequality)
+
+ ||ssT|] |||| x2ll
< [||sTs]) [||"s|||i>] + ||s ||||| + 2 T x1) ||x| - x2
s 2 (Xo ST s|| + y|) ||x
-x2ll-
(110)
where the norm is the operator norm the argument is a matrix.
Smoothness
MSE is smooth, or its gradient is Lipschitz-conttmmuus.
Proof. The gradient of MSE can be written as:
Vl(x) = TSx y).
(111)
Therefore, for any and X2:
Vi(xx) V((x)) = = (Sx: - y) ST SSxx² - y)
2 S(x1 x2)
2 S S || X2
(112)
Thus, MSE in this setting is Lipschitz--ooniimous.
Lipstnos Hessian Since the Hessian in this case is constant, it also satisfies the Lipschitz- continuity.
1.2 Cross-entropy loss
The loss function be defined
((x) -y In |softmax (Sx)]
(113)
where:
exp (Sx)
softmax (Sx) 1 exp (Sx)
(114)
is the softmax function.
Lipschitz continuity The cross-entropy loss Lipschitzz-oominuons.
Proof. To prove, we employ the mean-value theorem and then find the upper-bound of the gradient norm.
42

---

Published in Transactions on Machine Learning Research (08/2023)
First, calculate the gradient of the softmax function w.r.t. each element z; in x:
@softmax (Sx); θ exp (Sx);
dr,
0z, 1 exp (Sx)
exp (Sx); 0(Sx);
(Sx), 1 exp (Sx) Ox,
= softmax(Sxx) (Sx) [1 (j = softmax(Sx)) Sij.
(115)
where the subscript denotes an element in the corresponding vector and (-) denotes the indicator function.
Thus, the derivative of the cross-entropy loss w.r.t. the paramcter of interest x can be written following the chain rule:
d((x) dx.
y In softmax (Sx),
0r:
no yj
(softmax(Sx);
Σ
softmax(Sx).
Or:
Re
yı [1(j i) softmax(Sx). Sij.
(116)
where ne the number of classes. Therefore, apply the triangle incquality obtain the following:
(l(x)
y [1[ = i) - softmax(Sxx) S.yl
Ox.
<
y
x [[ = i) softmax x S
f=1
22
$ 2 Ssyl
(117)
Since each element of the gradient is bounded, the gradient norm is also bounded.
By applying the mean-value theorem, can obtain the following:
[[((x)) - E(x2)l sup xE[x
V x ||x x2l1.
(118)
where: z E [x1,x)) denotes vector contained in the set of points between X1, E R.D
Since the gradient norm bounded, the mean-value theorem results in the Lipschitz contimuity for the cross-entropy loss.
Smoothness
The cross-entropy loss this case is smooth, or other words, its gradient is Lipschitz- continuous.
Proof. According (116), the difference of derivative between x and X2 can be written as:
dl(x1)
d6(x2)
y, [Softmax(Sxx); s0
0x:
Or:
no Si
[softmax(Sxg). softmax(Sxi))
Σ /=1
(119)
43

---

Published in Transactions on Machine Learning Research (08/2023)
Therefore:
2
nhe
[x()) vx((xe)] =
|softmax - softmax((xx ):
Σ Si
j=1
ne
max i (l,.., D)
S (Sxtmax(SXx)): softmax
max ie(
y Si
|softmax(Sxii) - softmax(Sx:)ll
(120)
And since softmax function is Lipschitz-contimtuus the gradient of cross-entropy loss, in this case, also Lipschitz-coutinnous
Lipschitz continous Hessian The Hessian matrix of the cross-entropy loss Tipschitz-cootiimous..
Proof. To prove the Lipschitz continuity of the Hessian matrix for the cross-entropy loss, we will prove that the norm of its third order derivative tensor is bounded. According to (Nesterov, 2003, Page 175), the norm of the third order tensor, also known as Terssian, is bounded, then the Hessian matrix will satisfies the Lipschitz contimuity. To do this, calculate each element in the Terssian tensor as follows:
The second derivative of the cross-entropy loss w.r.t. an element 2; can be written as:
02((x) = drk 1=1
[j = 1) softmax(Sx)
dr,0xx
@softmax (SS))
Tle
Sij
Ozk
nee
--sotmax(SSx)) [L(k =) softmax(Sx) Sik
S
(121)
Thus, the third order derivative can also be derived
(2x() -Sk 02,0xz011
ne S.y
Σn
[softmax (Sx); [1: = =) softmax(Sx).]]
(122)
Oxi
To this point, we can see that the third order derivative involves the derivative of softmax functions and some constants relating to input data and its label y. Also, as the derivative of soft max function shown in (115) consists of softmax and indicator functions, the derivative of softmax function is bounded. Thus, can conclude that the third order derivative in (122) is also bounded with the bound depending the norm of the input max;,; [Syy and label max, y,|-
When the norm of each element in the Terssian is bounded, its norm is also bounded. Thus, according to the result in (Nesterov, 2003, Page 175), we can conclude that the Hossian is Lipschitz-contimoous.
1.3 Values of the Lipschitz continuity constants
The values of the constants used in Assumptions 1 to 3 heavily depend on some other parameters, including the maximum value of norm of the inputs and labels as well as the clipping norm value of the parameter x
see Proposition 4 *On the properties of the softmax function with application game theory and reinforcement learning".
arXiv preprint xrX1v:1704.00800
44

---

Published in Transactions on Machine Learning Research (08/2023)
(in case of MSE). Thus, cannot give exact values in numbers for those constants, but only provide few expressions to show the dependency of their values on the input and output.
Table 6: The values of some Lipschitz-coonimmity constants assumed in Assumptions 1 to 3.
Loss function
L
s
MSE
sups,y 2
[Xol|sTs s + ||ST>|]) sups 2|
sups 2| s
CE
45

---

Published in Transactions on Machine Learning Research (08/2023)
Appendix J
Additional visualisation
We provide additional visualisation for the qualitative results reported in Section 5 with uncertainty in the evaluation. Conventionally, meta-learning methods are evaluated on a certain number of tasks where the accuracy reported is averaged over those testing tasks (Vinyals et al., 2016; Finn et al., 2017: Ravi & Larochelle, 2017; Snell et al., 2017; Ravi Beatson, 2019; Nguyen et al., 2020). In this case, the uncertainty, represented by the confidence intervals, is calculated basod on the accuracy values obtained on those testing lasks. In Section 5, do not plot this uncertainly associated with the empirical results, e-g-, Figures 1 and 2, to clarify the visualisation. Hence, provide additional qualitative results with 95 percent confidence intervals in Figure 8, where Figure 8a is the accuracy on 100 validation 5-way 1-shot tasks formed from mini- ImageNet, and Figure 8b is the "zoom-in" version.
45
Accuracy (%)
Accuracy (%)
40
uniform
40
exploration
exploitation
30
tow
80
95
100
Xt iterations (x 1,000)
of iterations (x 1,000)
(a) Zoomed-out version
(b) Zoomed-in version
Figure 8: Additional visualisation of the prediction accuracy evaluated on 100 5-way 1-shot mini-ImageNet with 95 percent confidence intervals.
46


Reviews:
[{"review_id": "yqDjs6nOJQ", "reviewer": "Reviewer_SZgw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n- The motivation for the trajectory weighting is relatively clear from the perspective of prior work.\n- The paper seems to do a good job going through the related work and clarifying how exactly their method/approach relates. Specifically, the introduction and section 4 have a lot of really nice explanations on prior work and their detriments.\n- \n\n## Weaknesses:\n1. **BASELINES**: Given the rich set of potential baselines provided in the introduction and section 4, it is not clear why the only baselines provided were naive ones proposed by the authors. For instance, the potential for sub-optimal solutions is not a valid reason to not include the baseline in a scientific setting. \n    - (1.1) I think $\\alpha$-MAML and MWL-MAML (TR-MAML likely would be\n  intractable) are great options for incorporating another baseline from the literature. While they are approaching the problem differently than here, this can be contextualized in the presented results more fully and discuss their trade-offs. This would also probably help understand what the computational trade-offs of your method over these would be. MWL does seem to only be tested w/ RL, but their approach should be applicable to the classification setting, no? Although, I'm not sure how the validation trajectories are a hurdle, could the authors go into more detail?\n2. **The Study**: There are a few questions about the empirical study and some suggestions given the format of TMLR. \n    - (a) I'm not as familiar with meta learning, but I'm not sure exactly why the results are presented without error shading in figure 1 and 2. This calls into question how the confidence intervals were calculated in table 2. Was this from multiple runs? If so how many? If not, is this standard in this literature?\n    - (b) How does the TOW method compare computationally to the methods in Table 3? Is it still a large trade-off like the other baselines, or more reasonable?\n    - (c) I think the studies in Appendix H and I have a place in the main paper to further understand the proposed algorithm. This also relates to the hypothesis discussed in section 5. These provide more evidence for your approach and suggest further study, imo.\n    - (d) I think Appendix G should be included in the main paper, but that is more a style choice so the authors can decide.\n3. **Writing**: I think the paper is very densly written and not very friendly to an audience who is not familiar with trajectory optimisation. There are also several parts of the text which could do with a re-organization to make things much clearer to the reader.\n    - (a) I think section 3.1 should more directly relate back to algorithm 1 in a precise manner. Currently, the algorithm is a bit obfuscated on the precise details because of the writing in 3.1. For example, It is unclear if this approach uses Hessian diagonals or the Gauss-Newton diagonals. I believe the approach is using the Gauss-newton diagonals, but this is really unclear on first read.\n    - (b) This also relates to section 4. While there is a lot of great information, I think the organization could have been done in such a way that it is not overly dense.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new method for task weighting through trajectory optimization. They propose a method using iLQR which finds the optimal weightings over a trajectory of mini-batches. The paper proves convergence through this method to a stationary point of the policy. Finally, they provide several empirical examples using hand-crafted baselines. These baselines provide extreme versions of the exploration and exploitation their method should balance, and a uniform method which just uniformly weights the tasks. Finally, the paper compares favorably to the baseline performance found in (Lee et al, 2020). ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The weaknesses are ordered by importance. 1 and 2 are most critical to address while 3 is much more related to the clarity and writing of the paper. So the authors should focus on 1 and 2, using 3 to improve the paper in how they see fit. ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "92tVoCXez3", "reviewer": "Reviewer_YUvV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n - The paper motivates the problem setting well and provides the necessary background in Section 2.\n - The authors provide a novel approach (TOW) for solving a non-trivial problem (optimizing task weightings in meta-learning)\n - The authors provide a solid theoretical convergence analysis of the proposed procedure\n - The benefit of TOW over simpler baselines is convincingly demonstrated in the experiments.\n - The authors also discuss limitations of ROW, such as increased computational runtime.\n\nWeaknesses:\n - Evaluation is limited to few-shot classification on relatively simple problems (OmniGlot, Mini-Imagenet). Evaluating also on other meta-learning problems would strengthen the empirical evidence.\n - TOW is not tested in conjunction with state-of-the-art meta-learning methods, like e.g. Bayesian TAML (Lee et al., 2020), but only with MAML and Protonet.\n - Runtime of TOW is relatively high (7-9x slower than baselines). It is not clear if this is a constant factor or increases with size of network of number of tasks. An analysis on the asymptotic behaviour of TOW in terms of large networks or large number of tasks would be helpful.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose the novel approach TOW for optimizing task weighting in the outer-loop of meta-learning based on a trajectory optimization approach. For this, they cast the task-weighting problem to a finite-horizon discrete-time trajectory optimisation and apply iterative linear quadratic regulator for optimizing task weighting. They also prove that under certain conditions the procedure converges. They show benefits of TOW compared to simpler baselines on few-shot classification on OmniGlot and Mini-ImageNet.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": " - Having the related works section between methods and experiment sections is suboptimal as it disconnects these two parts. I would recommend to have related works either after the introduction or before the conclusion.\n - Section 3.1 is somewhat lengthy and not easy to digest. I think splitting it into subsubsection could help. Also, on page 5 it is not always clear which problems the authors want to address in the respective paragraphs.\n - Equation 5 has two hyperparameters ($\\beta$ and $\\mu_u$). How important is tuning these parameters in terms of preventing overfitting? It would be helpful to study this empirically.\n\n", "broader_impact_concerns": "There are no concerns from the reviewer", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Tr7VWLgytY", "reviewer": "Reviewer_YmYi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper proposes an interesting connection by optimizing a \"hyperparameter\" like task weights via optimal control, and to my knowledge, the use of iLQR within the optimization procedure like this is an interesting novel application. It does not appear to me that the proposed trajectory optimization is specific to the metalearning setup with task weighting action; it could be interesting to consider applying this idea to generic optimization problems where the action can be any other hyperparameter for learning (for example, weighting of individual data examples or scaling factors for different coordinates of the learning rate).\n- The paper shows promising improvements in benchmarks (albeit at the cost of much increased compute) and proof of convergence is provided under some typical assumptions.\n\nWeaknesses:\n- The iLQR loop seems quite computationally expensive to run, needing to reevaluate the trajectory of iterates many times for each update. Authors show that runtime for the proposed method can be 7-9x slower than the compared baselines. I suggest expanding the comparisons to include variants of the baselines with additional training (with the same # of tasks) to analyze how much benefit the optimized task reweighting over naively increasing the computation used to comparable amounts.\n- I did find the presentation of the method a bit difficult to follow, in particular trying to understand where exactly the trajectory optimization algorithm fit in. In particular, it was not immediately clear that there were now three nested levels of iteration, one outer loop, one for the iLQR iterations, and one for the inner loop of the metalearning algorithm.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes to use iLQR as part of an optimization algorithm for meta-learning. They frame the iterates of the outer-loop optimization problem as the states, the action space as weights on different sampled tasks in each minibatch, with the dynamics of the system as gradient descent on the weighted loss induced by the selected action and cost given by the unweighted loss over tasks. \n\nThe authors show improved results on several few shot learning benchmarks compared to not using any task weighting or simple heuristic task weighing schemes.They also provide convergence proofs for the introduced algorithm, as well as for iLQR, which apparently did not appear in prior works introducing it. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "To major concerns regarding acceptance of the paper\n- I do not see any analysis on the horizon $T$ of the trajectory optimization problem in the paper, and appears to be arbitrarily chosen to be 5 or 10 based on appendix G. Naturally, the proposed algorithm would have no impact with the trivial case of horizon 1, and it would be important to see how the quality benefits materialize as we increase the horizon. What would happen if we treat our *entire* training set of tasks as one trajectory optimized with iLQR?\n- As mentioned above in the weaknesses section, comparisons against the baselines with extra compute. How does the method compare if we simply run through the minibatches with the baseline algorithm multiple times to compensate for the drastically increased compute costs?\n\nSome other suggested changes to strengthen the work.\n- For clarity: I would suggest introducing the structure of the proposed TOW algorithm earlier, in the end of the introduction. I think it would be very helpful to explicitly state early on we consider a sequence of meta-parameter iterates as a trajectory, and we optimize the average loss over the iterates as the cost.\n", "broader_impact_concerns": "no concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes an approach to perform task-weighting in the outer loop of the meta-learning process based on the trajectory optimization process.
The paper is clearly motivated and well-placed with respect to the state of the art. The problem the authors face is interesting, and the proposed approach is novel. The authors provide both theoretical analysis about the convergence of the proposed method and show empirically its effectiveness over simple baselines. The paper also highlights some limitations of the proposed approach.
During the discussion, the authors updated their paper to address some of the weaknesses pointed out by the reviewers. In particular, they have improved the clarity of the presentation, and the empirical analysis by adding some baselines and ablation studies.
The reviewers reached a consensus on accepting this paper.
Decision: Accept (poster)