Adaptive Gradient-Based Meta-Learning Methods

Mikhail Khodak, Maria-Florina Balcan, Ameet Talwalkar

Introduction

Meta-learning, or learning-to-learn (LTL) , has recently re-emerged as an important direction for developing algorithms for multi-task learning, dynamic environments, and federated settings. By using the data of numerous training tasks, meta-learning methods seek to perform well on new, potentially related test tasks without using many samples. Successful modern approaches have also focused on exploiting the capabilities of deep neural networks, whether by learning multi-task embeddings passed to simple classifiers or by neural control of optimization algorithms .

Because of its simplicity and flexibility, a common approach is parameter-transfer, where all tasks use the same class of Θ\Theta-parameterized functions fθ:X↦Yf_{\theta}:\mathcal{X}\mapsto\mathcal{Y}; often a shared model ϕ∈Θ\phi\in\Theta is learned that is used to train within-task models. In gradient-based meta-learning (GBML) , ϕ\phi is a meta-initialization for a gradient descent method over samples from a new task. GBML is used in a variety of LTL domains such as vision , federated learning , and robotics . Its simplicity also raises many practical and theoretical questions about the task-relations it can exploit and the settings in which it can succeed. Addressing these issues has naturally led several authors to online convex optimization (OCO) , either directly or from online-to-batch conversion . These efforts study how to find a meta-initialization, either by proving algorithmic learnability or giving meta-test-time performance guarantees .

However, this recent line of work has so far considered a very restricted, if natural, notion of task-similarity – closeness to a single fixed point in the parameter space. We introduce a new theoretical framework, Average Regret-Upper-Bound Analysis (ARUBA), that enables the derivation of meta-learning algorithms that can provably take advantage of much more sophisticated structure. ARUBA treats meta-learning as the online learning of a sequence of losses that each upper bounds the regret on a single task. These bounds often have convenient functional forms that are (a) sufficiently nice, so that we can draw upon the existing OCO literature, and (b) strongly dependent on both the task-data and the meta-initialization, thus encoding task-similarity in a mathematically accessible way. Using ARUBA we introduce or dramatically improve upon GBML results in the following settings:

Adapting to the Task-Similarity: A major drawback of previous work is a reliance on knowing the task-similarity beforehand to set the learning rate or regularization , or the use of a sub-optimal guess-and-tune approach using the doubling trick . ARUBA yields a simple gradient-based algorithm that eliminates the need to guess the similarity by learning it on-the-fly.

Adapting to Dynamic Environments: While previous theoretical work has largely considered a fixed initialization , in many practical applications of GBML the optimal initialization varies over time due to a changing environment . We show how ARUBA reduces the problem of meta-learning in dynamic environments to a dynamic regret-minimization problem, for which there exists a vast array of online algorithms with provable guarantees that can be directly applied.

Adapting to the Inter-Task Geometry: A recurring notion in LTL is that certain model weights, such as feature extractors, are shared, whereas others, such as classification layers, vary between tasks. By only learning a fixed initialization we must re-learn this structure on every task. Using ARUBA we provide a method that adapts to this structure and determines which directions in Θ\Theta need to be updated by learning a Mahalanobis-norm regularizer for online mirror descent (OMD). We show how a variant of this can be used to meta-learn a per-coordinate learning-rate for certain GBML methods, such as MAML and Reptile , as well as for FedAvg, a popular federated learning algorithm . This leads to improved meta-test-time performance on few-shot learning and a simple, tuning-free approach to effectively add user-personalization to FedAvg.

Statistical Learning-to-Learn: ARUBA allows us to leverage powerful results in online-to-batch conversion to derive new bounds on the transfer risk when using GBML for statistical LTL , including fast rates in the number of tasks when the task-similarity is known and high-probability guarantees for a class of losses that includes linear regression. This improves upon the guarantees of Khodak et al. and Denevi et al. for similar or identical GBML methods.

Theoretical LTL: The statistical analysis of LTL was formalized by Baxter . Several works have built upon this theory for modern LTL, such as via a PAC-Bayesian perspective or by learning the kernel for the ridge regression . However, much effort has also been devoted to the online setting, often through the framework of lifelong learning . Alquier et al. consider a many-task notion of regret similar to the one we study in order to learn a shared data representation, although our algorithms are much more practical. Recently, Bullins et al. developed an efficient online approach to learning a linear data embedding, but such a setting is distinct from GBML and more closely related to popular shared-representation methods such as ProtoNets . Nevertheless, our approach does strongly rely on online learning through the study of data-dependent regret-upper-bounds, which has a long history of use in deriving adaptive single-task methods ; however, in meta-learning there is typically not enough data to adapt to without considering multi-task data. Analyzing regret-upper-bounds was done implicitly by Khodak et al. , but their approach is largely restricted to using Follow-the-Leader (FTL) as the meta-algorithm. Similarly, Finn et al. use FTL to show learnability of the MAML meta-initialization. In contrast, the ARUBA framework can handle general classes of meta-algorithms, which leads not only to new and improved results in static, dynamic, and statistical settings but also to significantly more practical LTL methods.

GBML: GBML stems from the Model-Agnostic Meta-Learning (MAML) algorithm and has been widely used in practice . An expressivity result was shown for MAML by Finn and Levine , proving that the meta-learner can approximate any permutation-invariant learner given enough data and a specific neural architecture. Under strong-convexity and smoothness assumptions and using a fixed learning rate, Finn et al. show that the MAML meta-initialization is learnable, albeit via an impractical FTL method. In contrast to these efforts, Khodak et al. and Denevi et al. focus on providing finite-sample meta-test-time performance guarantees in the convex setting, the former for the SGD-based Reptile algorithm of Nichol et al. and the latter for a regularized variant. Our work improves upon these analyses by considering the case when the learning rate, a proxy for the task-similarity, is not known beforehand as in Finn et al. and Denevi et al. but must be learned online; Khodak et al. do consider an unknown task-similarity but use a doubling-trick-based approach that considers the absolute deviation of the task-parameters from the meta-initialization and is thus average-case suboptimal and sensitive to outliers. Furthermore, ARUBA can handle more sophisticated and dynamic notions of task-similarity and in certain settings can provide better statistical guarantees than those of Khodak et al. and Denevi et al. .

Average Regret-Upper-Bound Analysis

Generality: Many algorithms of interest in meta-learning have regret guarantees U⁡t(x)\operatorname{\bf U}_{t}(x) with nice, e.g. smooth and convex, functional forms that depend strongly on both their parameterizations x∈Xx\in\mathcal{X} and the task-data. This data-dependence lets us adaptively set the parameterization xt∈Xx_{t}\in\mathcal{X}.

Consequences: By definition of U⁡t\operatorname{\bf U}_{t} we have that Uˉ⁡T\operatorname{\mathbf{\bar{U}}}_{T} bounds the task-averaged regret (TAR) Rˉ⁡T=1T∑t=1TR⁡t(xt)\operatorname{\bf\bar{R}}_{T}=\frac{1}{T}\sum_{t=1}^{T}\operatorname{\bf R}_{t}(x_{t}) . Thus if the average regret-upper-bound is small then the meta-learner will perform well on-average across tasks. In Section 5 we further show that a low average regret-upper-bound will also lead to strong statistical guarantees in the batch setting.

ARUBA’s applicability depends only on finding a low-regret algorithm over the functions U⁡t\operatorname{\bf U}_{t}; then by observation 2 we get a task-averaged regret bound where the first term vanishes as T→∞T\to\infty while by observation 1 the second term can be made small due to the data-dependent task-similarity:

Adapting to Similar Tasks and Dynamic Environments

Putting these together, we seek to define variants of Algorithm 1 for which as T→∞T\to\infty the average regret scales with VΨV_{\Psi}, where VΨ2=1T∑t=1TBR(θt∗∣∣ψt)V_{\Psi}^{2}=\frac{1}{T}\sum_{t=1}^{T}\mathcal{B}_{R}(\theta_{t}^{\ast}||\psi_{t}), without knowing this quantity in advance. Note for fixed ψt=θˉ∗=1Tθ1:T∗\psi_{t}=\bar{\theta}^{\ast}=\frac{1}{T}\theta_{1:T}^{\ast} this measures the empirical standard deviation of the optimal task-actions θt∗\theta_{t}^{\ast}. Thus achieving our goal implies that average performance improves with task-similarity.

On each task tt Algorithm 1 runs online mirror descent with regularizer 1ηtBR(⋅∣∣ϕt)\frac{1}{\eta_{t}}\mathcal{B}_{R}(\cdot||\phi_{t}) for initialization ϕt∈Θ\phi_{t}\in\Theta and learning rate ηt>0\eta_{t}>0. It is well-known that OMD and the related Follow-the-Regularized-Leader (FTRL), for which our results also hold, generalize many important online methods, e.g. OGD and multiplicative weights . For mtm_{t} convex losses with mean squared Lipschitz constant Gt2G_{t}^{2} they also share a convenient, data-dependent regret-upper-bound for any θt∗∈Θ\theta_{t}^{\ast}\in\Theta [48, Theorem 2.15]:

All that remains is to come up with update rules for the meta-initialization ϕt∈Θ\phi_{t}\in\Theta and the learning rate ηt>0\eta_{t}>0 in Algorithm 1 so that the average over TT of these upper-bounds U⁡t(ϕt,ηt)\operatorname{\bf U}_{t}(\phi_{t},\eta_{t}) is small. While this can be viewed as a single online learning problem to determine actions xt=(ϕt,ηt)∈Θ×(0,∞)x_{t}=(\phi_{t},\eta_{t})\in\Theta\times(0,\infty), it is easier to decouple ϕ\phi and η\eta by first defining two function sequences {ftinit}t\{f_{t}^{\textrm{init}}\}_{t} and {ftsim}t\{f^{\textrm{sim}}_{t}\}_{t}:

We show in Theorem 3.1 that to get an adaptive algorithm it suffices to specify two OCO algorithms, INIT⁡\operatorname{INIT} and SIM⁡\operatorname{SIM}, such that the actions ϕt=INIT⁡(t)\phi_{t}=\operatorname{INIT}(t) achieve good (dynamic) regret over ftinitf_{t}^{\textrm{init}} and the actions vt=SIM⁡(t)v_{t}=\operatorname{SIM}(t) achieve low (static) regret over ftsimf^{\textrm{sim}}_{t}; these actions then determine the update rules of ϕt\phi_{t} and ηt=vt/(Gtmt)\eta_{t}=v_{t}/(G_{t}\sqrt{m_{t}}). We will specialize Theorem 3.1 to derive algorithms that provably adapt to task similarity (Theorem 3.2) and to dynamic environments (Theorem 3.3).

To understand the formulation of ftinitf_{t}^{\textrm{init}} and ftsimf^{\textrm{sim}}_{t}, first note that ftsim(v)=U⁡t(ϕt,v/(Gtmt))f^{\textrm{sim}}_{t}(v)=\operatorname{\bf U}_{t}(\phi_{t},v/(G_{t}\sqrt{m_{t}})), so the online algorithm SIM⁡\operatorname{SIM} over ftsimf^{\textrm{sim}}_{t} corresponds to an online algorithm over the regret-upper-bounds U⁡t\operatorname{\bf U}_{t} when the sequence of initializations ϕt\phi_{t} is chosen adversarially. Once we have shown that SIM⁡\operatorname{SIM} is low-regret we can compare its losses ftsim(vt)f^{\textrm{sim}}_{t}(v_{t}) to those of an arbitrary fixed v>0v>0; this is the first line in the proof of Theorem 3.1 (below). For fixed vv, each ftinit(ϕt)f^{\textrm{init}}_{t}(\phi_{t}) is an affine transformation of ftsim(v)f^{\textrm{sim}}_{t}(v), so the algorithm INIT⁡\operatorname{INIT} with low dynamic regret over ftinitf^{\textrm{init}}_{t} corresponds to an algorithm with low dynamic regret over the regret-upper-bounds U⁡t\operatorname{\bf U}_{t} when ηt=v/(Gtmt) ∀ t\eta_{t}=v/(G_{t}\sqrt{m_{t}})~{}\forall~{}t. Thus once we have shown a dynamic regret guarantee for INIT⁡\operatorname{INIT} we can compare its losses ftinit(ϕt)f^{\textrm{init}}_{t}(\phi_{t}) to those of an arbitrary comparator sequence {ψt}t⊂Θ\{\psi_{t}\}_{t}\subset\Theta; this is the second line in the proof of Theorem 3.1.

Let INIT⁡\operatorname{INIT} be an algorithm whose dynamic regret over functions {ftinit}t\{f^{\textrm{init}}_{t}\}_{t} w.r.t. any reference sequence Ψ={ψt}t=1T⊂Θ\Psi=\{\psi_{t}\}_{t=1}^{T}\subset\Theta is upper-bounded by U⁡Tinit(Ψ)\operatorname{\bf U}^{\textrm{init}}_{T}(\Psi).

Let SIM⁡\operatorname{SIM} be an algorithm whose static regret over functions {ftsim}t\{f^{\textrm{sim}}_{t}\}_{t} w.r.t. any v>0v>0 is upper-bounded by a non-increasing function U⁡Tsim(v)\operatorname{\bf U}^{\textrm{sim}}_{T}(v) of vv.

If Algorithm 1 sets ϕt=INIT⁡(t)\phi_{t}=\operatorname{INIT}(t) and ηt=SIM⁡(t)Gtmt\eta_{t}=\frac{\operatorname{SIM}(t)}{G_{t}\sqrt{m_{t}}} then for VΨ2=∑t=1TBR(θt∗∣∣ψt)Gtmt∑t=1TGtmtV_{\Psi}^{2}=\frac{\sum_{t=1}^{T}\mathcal{B}_{R}(\theta_{t}^{\ast}||\psi_{t})G_{t}\sqrt{m_{t}}}{\sum_{t=1}^{T}G_{t}\sqrt{m_{t}}} it will achieve average regret

For σt=Gtmt\sigma_{t}=G_{t}\sqrt{m_{t}} we have by the regret bound on OMD/FTRL (2) that

where the last line follows by substituting v=max⁡{VΨ,U⁡Tinit(Ψ)/σ1:T}v=\max\left\{V_{\Psi},\sqrt{\operatorname{\bf U}^{\textrm{init}}_{T}(\Psi)/\sigma_{1:T}}\right\}. ∎

By Theorem 3.1, if we can specify algorithms INIT⁡\operatorname{INIT} and SIM⁡\operatorname{SIM} with sublinear regret over ftinitf^{\textrm{init}}_{t} and ftsimf^{\textrm{sim}}_{t} (3), respectively, then the average regret will converge to O(VΨm)\mathcal{O}(V_{\Psi}\sqrt{m}) as desired. We first show an approach in the case when the optimal actions θt∗\theta_{t}^{\ast} are close to a fixed point in Θ\Theta, i.e. for fixed ψt=θˉ∗=1Tθ1:T∗\psi_{t}=\bar{\theta}^{\ast}=\frac{1}{T}\theta_{1:T}^{\ast}. Henceforth we assume the Lipschitz constant GG and number of rounds mm are the same across tasks; detailed statements are in the supplement.

Under the assumptions of Theorem 3.1 and boundedness of BR\mathcal{B}_{R} over Θ\Theta, if INIT⁡\operatorname{INIT} plays ϕt+1=1tθ1:t∗\phi_{t+1}=\frac{1}{t}\theta_{1:t}^{\ast} and SIM⁡\operatorname{SIM} uses ε\varepsilon-EWOO (4) with ε=1/T\varepsilon=1/\sqrt{T} then Algorithm 1 achieves average regret

Related Tasks in Changing Environments:

In many settings we have a changing environment and so it is natural to study dynamic regret. This has been widely analyzed by the online learning community , often by showing a dynamic regret bound consisting of a sublinear term plus a bound on the variation in the action or function space. Using Theorem 3.1 we can show dynamic guarantees for GBML via reduction to such bounds. We provide an example in the Euclidean geometry using the popular path-length-bound PΨ=∑t=2T∥ψt−ψt−1∥2P_{\Psi}=\sum_{t=2}^{T}\|\psi_{t}-\psi_{t-1}\|_{2} for reference actions Ψ={ψt}t=1T\Psi=\{\psi_{t}\}_{t=1}^{T} . We use a result showing that OGD with learning rate η≤1/β\eta\leq 1/\beta over α\alpha-strongly-convex, β\beta-strongly-smooth, and LL-Lipschitz functions has a bound of O(L(1+PΨ))\mathcal{O}(L(1+P_{\Psi})) on its dynamic regret [42, Corollary 1]. Observe that in the case of R(⋅)=12∥⋅∥22R(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2} the sequence ftinitf^{\textrm{init}}_{t} in Theorem 3.1 consists of DGmDG\sqrt{m}-Lipschitz quadratic functions. Thus using Theorem 3.1 we achieve the following:

Under Theorem 3.1 assumptions, bounded Θ\Theta, and R(⋅)=12∥⋅∥22R(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2}, if INIT⁡\operatorname{INIT} is OGD with learning rate 1Gm\frac{1}{G\sqrt{m}} and SIM⁡\operatorname{SIM} uses ε\varepsilon-EWOO (4) with ε=1/T\varepsilon=1/\sqrt{T} then by using OGD within-task Algorithm 1 will achieve for any fixed comparator sequence Ψ={ψt}t∈[T]⊂Θ\Psi=\{\psi_{t}\}_{t\in[T]}\subset\Theta the average regret

for VΨ2=12T∑t=1T∥θt∗−ψt∥22V_{\Psi}^{2}=\frac{1}{2T}\sum_{t=1}^{T}\|\theta_{t}^{\ast}-\psi_{t}\|_{2}^{2} and PΨ=∑t=2T∥ψt−ψt−1∥2P_{\Psi}=\sum_{t=2}^{T}\|\psi_{t}-\psi_{t-1}\|_{2}.

This bound controls the average regret across tasks using the deviation VΦV_{\Phi} of the optimal task parameters θt∗\theta_{t}^{\ast} from some reference sequence Φ\Phi, which is assumed to vary slowly or sparsely so that the path length PΦP_{\Phi} is small. Figure 2 illustrates when such a guarantee improves over Theorem 3.2. Note also that Theorem 3.3 specifies OGD as the meta-update algorithm INIT⁡\operatorname{INIT}, so under the approximation that each task tt’s last iterate is close to θt∗\theta_{t}^{\ast} this suggests that simple GBML methods such as Reptile or FedAvg are adaptive. The generality of ARUBA also allows for the incorporation of other dynamic regret bounds and other non-static notions of regret .

Adapting to the Inter-Task Geometry

Observe the similarity between this update AdaGrad , which is also inversely related to the sum of the element-wise squares of all gradients seen so far. Our method adds multi-task information by setting the numerator to depend on the sum of squared distances between the initializations ϕt\phi_{t} set by the algorithm and that task’s optimal action θt∗\theta_{t}^{\ast}. This algorithm has the following guarantee:

As T→∞T\to\infty the average regret converges to the minimum over ϕ,H\phi,H of the last two terms, which corresponds to running OMD with the optimal initialization and per-coordinate learning rate on every task. The rate of convergence of T−2/5T^{-2/5} is slightly slower than the usual 1/T1/\sqrt{T} achieved in the previous section; this is due to the algorithm’s adaptivity to within-task gradients, whereas previously we simply assumed a known Lipschitz bound GtG_{t} when setting ηt\eta_{t}. This adaptivity makes the algorithm much more practical, leading to a method for adaptively learning a within-task learning rate using multi-task information; this is outlined in Algorithm 2 and shown to significantly improve GBML performance in Section 6. Note also the per-coordinate separation of the left term, which shows that the algorithm converges more quickly on non-degenerate coordinates. The per-coordinate specification of ηt\eta_{t} (9) can be further generalized to learning a full-matrix adaptive regularizer, for which we show guarantees in Theorem 4.2. However, the rate is much slower, and without further assumptions such methods will have Ω(d2)\Omega(d^{2}) computation and memory requirements.

for ε=1/T\varepsilon=1/\sqrt{T} and ζ=m/T\zeta=\sqrt{m}/\sqrt{T}. Then for λj\lambda_{j} corresponding to the jjth largest eigenvalue we have

Fast Rates and High Probability Bounds for Statistical Learning-to-Learn

In the general case, Theorem 5.1 provides bounds on the excess transfer risk decreasing with Uˉ⁡/m\operatorname{\mathbf{\bar{U}}}/m and 1/mT1/\sqrt{mT}. Thus if Uˉ⁡\operatorname{\mathbf{\bar{U}}} improves with task-similarity so will the transfer risk as T→∞T\to\infty. Note that the second term is 1/mT1/\sqrt{mT} rather than 1/T1/\sqrt{T} as in most-analyses ; this is because regret is mm-bounded but the OMD regret-upper-bound is O(m)\mathcal{O}(\sqrt{m})-bounded. The results also demonstrate ARUBA’s ability to utilize specialized results from the online-to-batch conversion literature. This is witnessed by the guarantee for self-bounded losses, a class which Zhang shows includes linear regression; we use a result by the same author to obtain high-probability bounds, whereas previous GBML bounds are in-expectation . We also apply a result due to Kakade and Tewari for the case of strongly-convex regret-upper-bounds, enabling fast rates in the number of tasks TT. The strongly-convex case is especially relevant for GBML since it holds for OGD with fixed learning rate.

In the setting of Theorems 3.2 & 5.1, if δ≤1/e\delta\leq 1/e and Algorithm 1 uses within-task OGD with initialization ϕt+1=1tθ1:t∗\phi_{t+1}=\frac{1}{t}\theta_{1:t}^{\ast} and step-size ηt=VQ+1/TGm\eta_{t}=\frac{V_{\mathcal{Q}}+1/\sqrt{T}}{G\sqrt{m}} for VQV_{\mathcal{Q}} as above, then w.p. 1−δ1-\delta

If ηt\eta_{t} is set adaptively using ε\varepsilon-EWOO as in Theorem 3.2 for ε=1/mT+1/m\varepsilon=1/\sqrt{mT}+1/\sqrt{m} then w.p. 1−δ1-\delta

Empirical Results: Adaptive Methods for Few-Shot & Federated Learning

ARUBA++: starting with ηT,1=ηT\eta_{T,1}=\eta_{T} and gT,1=gTg_{T,1}=g_{T}, adaptively reset the learning rate by setting g^T,i+1←g^T,i+c∇i2\hat{g}_{T,i+1}\leftarrow\hat{g}_{T,i}+c\nabla_{i}^{2} for some c>0c>0 and then updating ηT,i+1←bT/gT,i+1\eta_{T,i+1}\leftarrow\sqrt{b_{T}/g_{T,i+1}}. Isotropic: btb_{t} and gtg_{t} are scalars tracking the sum of squared distances and sum of squared gradient norms, respectively.

We first examine if Algorithm 2 can improve performance on Omniglot and Mini-ImageNet , two standard few-shot learning benchmarks, when used to modify Reptile, a simple meta-learning method . In its serial form Reptile is roughly the algorithm we study in Section 3 when OGD is used within-task and η\eta is fixed. Thus we can set Reptile+ARUBA to be Algorithm 2 with θ^t\hat{\theta}_{t} the last iterate of OGD and the meta-update a weighted sum of θ^t\hat{\theta}_{t} and ϕt\phi_{t}. In practice, however, Reptile uses Adam to exploit multi-task gradient information. As shown in Table 1, ARUBA matches or exceeds this baseline on Mini-ImageNet, although on Omniglot it requires the additional within-task updating of ARUBA++ to show improvement.

It is less clear how ARUBA can be applied to MAML , as by only taking one step the distance traveled will be proportional to the gradient, so η\eta will stay fixed. We also do not find that ARUBA improves multi-step MAML – perhaps not surprising as it is further removed from our theory due to its use of held-out data. In Table 1 we compare to Meta-SGD , which does learn a per-coordinate learning rate for MAML by automatic differentiation. This requires more computation but does lead to consistent improvement. As with the original Reptile, our modification performs better on Mini-ImageNet but worse on Omniglot compared to MAML and its modification Meta-SGD.

Federated Learning:

A main goal in this setting is to use data on heterogeneous nodes to learn a global model without much communication; leveraging this to get a personalized model is an auxiliary goal , with a common application being next-character prediction on mobile devices. A popular method is FedAvg , where at each communication round rr the server sends a global model ϕr\phi_{r} to a batch of nodes, which then run local OGD; the server then sets ϕr+1\phi_{r+1} to the average of the returned models. This can be seen as a GBML method with each node a task, making it easy to apply ARUBA: each node simply sends its accumulated squared gradients to the server together with its model. The server can use this information and the squared difference between ϕr\phi_{r} and ϕr+1\phi_{r+1} to compute a learning rate ηr+1\eta_{r+1} via Algorithm 2 and send it to each node in the next round. We use FedAvg with ARUBA to train a character LSTM on the Shakespeare dataset, a standard benchmark of a thousand users with varying amounts of non-i.i.d. data . Figure 3 shows that ARUBA significantly improves over non-tuned FedAvg and matches the performance of FedAvg with a tuned learning rate schedule. Unlike both baselines we also do not require step-size tuning when refining the global model for personalization. This reduced need for hyperparameter optimization is crucial in federated settings, where the number of user-data accesses are extremely limited.

Conclusion

In this paper we introduced ARUBA, a framework for analyzing GBML that is both flexible and consequential, yielding new guarantees for adaptive, dynamic, and statistical LTL via online learning. As a result we devised a novel per-coordinate learning rate applicable to generic GBML procedures, improving their training and meta-test-time performance on few-shot and federated learning. We see great potential for applying ARUBA to derive many other new LTL methods in a similar manner.

Acknowledgments

We thank Jeremy Cohen, Travis Dick, Nikunj Saunshi, Dravyansh Sharma, Ellen Vitercik, and our three anonymous reviewers for helpful feedback. This work was supported in part by DARPA FA875017C0141, National Science Foundation grants CCF-1535967, CCF-1910321, IIS-1618714, IIS-1705121, IIS-1838017, and IIS-1901403, a Microsoft Research Faculty Fellowship, a Bloomberg Data Science research grant, an Amazon Research Award, an Amazon Web Services Award, an Okawa Grant, a Google Faculty Award, a JP Morgan AI Research Faculty Award, and a Carnegie Bosch Institute Research Award. Any opinions, findings and conclusions, or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of DARPA, the National Science Foundation, or any other funding agency.

References

Appendix A Background and Results for Online Convex Optimization

We first state the related definitions of strong convexity and strong smoothness:

Finally, we will also consider functions that are exp-concave :

We now turn to the Bregman divergence and a discussion of several useful properties :

The definition directly implies that Bf(⋅∣∣y)\mathcal{B}_{f}(\cdot||y) preserves the (strong or strict) convexity of ff for any fixed y∈Sy\in S. Strict convexity further implies Bf(x∣∣y)≥0 ∀ x,y∈S\mathcal{B}_{f}(x||y)\geq 0~{}\forall~{}x,y\in S, with equality iff x=yx=y. Finally, if ff is α\alpha-strongly-convex, or β\beta-strongly-smooth, w.r.t. ∥⋅∥\|\cdot\| then Definitions A.1 and A.2 imply Bf(x∣∣y)≥α2∥x−y∥2\mathcal{B}_{f}(x||y)\geq\frac{\alpha}{2}\|x-y\|^{2} or Bf(x∣∣y)≤β2∥x−y∥2\mathcal{B}_{f}(x||y)\leq\frac{\beta}{2}\|x-y\|^{2}, respectively.

By Definition A.4 the last expression has a unique minimum at y=xˉy=\bar{x}. ∎

A.2 Online Algorithms

Here we provide a review of the online algorithms we use. Recall that in this setting our goal is minimizing regret:

Within-task our focus is on two closely related meta-algorithms, Follow-the-Regularized-Leader (FTRL) and (linearized lazy) Online Mirror Descent (OMD).

for all θ∗∈Θ\theta^{\ast}\in\Theta and G2≥1T∑t=1TGt2G^{2}\geq\frac{1}{T}\sum_{t=1}^{T}G_{t}^{2}.

We next review the online algorithms we use for the meta-update. The main requirement here is logarithmic regret guarantees for the case of strongly convex loss functions, which is satisfied by two well-known algorithms:

Kakade and Shalev-Shwartz [32, Theorem 2] and Bartlett et al. [7, Theorem 2.1] provide for FTL and AOGD, respectively, the following regret bound:

Finally, we state the EWOO algorithm due to Hazan et al. . While difficult to run in high-dimensions, we will be running this method in single dimensions, when computing it requires only one integral.

Hazan et al. [28, Theorem 7] provide the following guarantee for EWOO, which is notable for its lack of explicit dependence on the Lipschitz constant.

A.3 Online-to-Batch Conversion

Finally, as we are also interested in distributional meta-learning, we discuss some techniques for converting regret guarantees into generalization bounds, which are usually named online-to-batch conversions. We first state some standard results.

for θˉ=1Tθ1:T\bar{\theta}=\frac{1}{T}\theta_{1:T} and any θ∗∈Θ\theta^{\ast}\in\Theta.

For nonnegative bounded losses we have the following fact [14, Proposition 1]:

Note that Cesa-Bianchi et al. only prove the first inequality; the second follows via the same argument but applying the symmetric version of the Azuma-Hoeffding inequality . The inequalities above can be easily used to derive the following competitive bounds:

for any θ∗∈Θ\theta^{\ast}\in\Theta. If the losses are also convex then for θˉ=1Tθ1:T\bar{\theta}=\frac{1}{T}\theta_{1:T} we have

Apply linearity of expectations to get the first inequality and Jensen’s inequality to get the second. ∎

We now discuss some stronger guarantees for certain classes of loss functions. The first, due to Kakade and Tewari [33, Theorem 2], yields faster rates for strongly convex losses:

We can also obtain a data-dependent bound using a result of Zhang under a self-bounding property. Cesa-Bianchi and Gentile [13, Proposition 2] show a similar but less general result.

Apply Jensen’s inequality and Zhang [54, Theorem 4]. ∎

Note that nonnegative 1-bounded convex losses satisfy the conditions of Theorem A.5 with ρ=1\rho=1. However, we are interested in a different result that can yield a data-dependent competitive bound:

Zhang [54, Lemma 7] shows that the conditions are satisfied for ρ=4\rho=4 by least-squares regression.

A.4 Dynamic Regret Guarantees

Here we review several results for optimizing dynamic regret. We first define this quantity:

Mokhtari et al. [42, Corollary 1] show the following guarantee for OGD over strongly convex functions:

Appendix B Strongly Convex Coupling

Our first result is a simple trick that we believe may be of independent interest. It allows us to bound the regret of FTL on any (possibly non-convex) sequence of Lipschitz functions so long as the actions played are identical to those played on a different strongly-convex sequence of Lipschitz functions. The result is formalized in Theorem B.1.

We start with some standard facts about convex functions.

Next we state some technical results, starting with the well-known be-the-leader lemma [48, Lemma 2.1].

The final result depends on a stability argument for FTL on strongly-convex functions adapted from Saha et al. :

Adding these two inequalities and applying Claim B.1 yields

Dividing by ∥θt−θt+1∥\|\theta_{t}-\theta_{t+1}\| yields the result. ∎

In the convex case we instead apply Claim B.1 and Lemma B.2 to get

B.2 Applications

We now show two applications of strongly convex coupling. The first shows logarithmic regret for FTL run on a sequence of Bregman regularizers. Note that these functions are nonconvex in general.

Note that αtBR(θt∣∣⋅)\alpha_{t}\mathcal{B}_{R}(\theta_{t}||\cdot) is αtGt\alpha_{t}G_{t}-Lipschitz w.r.t. ∥⋅∥\|\cdot\|. Let R′(⋅)=12∥⋅∥22R^{\prime}(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2}, so BR′(θt∣∣ϕ)=12∥θt−ϕ∥22 ∀ ϕ∈Θ,t∈[T]\mathcal{B}_{R^{\prime}}(\theta_{t}||\phi)=\frac{1}{2}\|\theta_{t}-\phi\|_{2}^{2}~{}\forall~{}\phi\in\Theta,t\in[T]. The function αtBR′(θt∣∣⋅)\alpha_{t}\mathcal{B}_{R^{\prime}}(\theta_{t}||\cdot) is thus αt\alpha_{t}-strongly-convex and DD-Lipschitz w.r.t. ∥⋅∥2\|\cdot\|_{2}. Now by Claim A.1 FTL run on this new sequence plays the same actions as FTL run on the original sequence. Applying Theorem B.1 yields the result. ∎

Appendix C Adaptive and Dynamic Guarantees

INIT⁡\operatorname{INIT}: a method that has dynamic regret U⁡Tinit(Ψ)=∑t=1Tftinit(ϕt)−ftinit(ψt)\operatorname{\bf U}^{\textrm{init}}_{T}(\Psi)=\sum_{t=1}^{T}f^{\textrm{init}}_{t}(\phi_{t})-f^{\textrm{init}}_{t}(\psi_{t}) w.r.t. reference actions Ψ={ψt}t=1T⊂Θ\Psi=\{\psi_{t}\}_{t=1}^{T}\subset\Theta over the sequence ftinit(⋅)=BR(θt∗∣∣⋅)Gtmtf^{\textrm{init}}_{t}(\cdot)=\mathcal{B}_{R}(\theta_{t}^{\ast}||\cdot)G_{t}\sqrt{m_{t}} .

SIM⁡\operatorname{SIM}: a method that has (static) regret U⁡Tsim(x)\operatorname{\bf U}^{\textrm{sim}}_{T}(x) decreasing in x>0x>0 over the sequence of functions ftsim(x)=(BR(θt∗∣∣ϕt)x+x)Gtmtf^{\textrm{sim}}_{t}(x)=\left(\frac{\mathcal{B}_{R}(\theta_{t}^{\ast}||\phi_{t})}{x}+x\right)G_{t}\sqrt{m_{t}}.

Then if Algorithm 1 sets ϕt=INIT⁡(t)\phi_{t}=\operatorname{INIT}(t) and ηt=SIM⁡(t)Gtmt\eta_{t}=\frac{\operatorname{SIM}(t)}{G_{t}\sqrt{m_{t}}} it will achieve

for VΨ2=1∑t=1TGtmt∑t=1TBR(θt∗∣∣ψt)GtmtV_{\Psi}^{2}=\frac{1}{\sum_{t=1}^{T}G_{t}\sqrt{m_{t}}}\sum_{t=1}^{T}\mathcal{B}_{R}(\theta_{t}^{\ast}||\psi_{t})G_{t}\sqrt{m_{t}}.

Letting xt=SIM⁡(t)x_{t}=\operatorname{SIM}(t) be the output of SIM⁡\operatorname{SIM} at time tt, defining σt=Gtmt\sigma_{t}=G_{t}\sqrt{m_{t}} and σ1:T=∑t=1Tσt\sigma_{1:T}=\sum_{t=1}^{T}\sigma_{t}, and substituting into the regret-upper-bound of OMD/FTRL (2), we have that

where the last line follows by substituting x=max⁡{VΨ,U⁡Tinit(Ψ)σ1:T}x=\max\left\{V_{\Psi},\sqrt{\frac{\operatorname{\bf U}^{\textrm{init}}_{T}(\Psi)}{\sigma_{1:T}}}\right\}. ∎

Under the assumptions of Theorem C.1 and boundedness of BR\mathcal{B}_{R} over Θ\Theta, if INIT⁡\operatorname{INIT} uses FTL, or AOGD in the case of R(⋅)=12∥⋅∣22R(\cdot)=\frac{1}{2}\|\cdot|_{2}^{2}, and SIM⁡\operatorname{SIM} uses ε\varepsilon-FTL as defined in Proposition B.2, then Algorithm 1 achieves

for V2=min⁡ϕ∈Θ∑t=1TσtBR(θt∗∣∣ϕ)V^{2}=\min_{\phi\in\Theta}\sum_{t=1}^{T}\sigma_{t}\mathcal{B}_{R}(\theta_{t}^{\ast}||\phi) and constant CC the product of the constant CC from Proposition B.1 and the bound on the gradient of the Bregman divergence. Assuming σt=Gm ∀ t\sigma_{t}=G\sqrt{m}~{}\forall~{}t and substituting ε=1T\varepsilon=\frac{1}{\sqrt{T}} yields

Substitute Propositions B.1 and B.2 into Theorem C.1. ∎

Under the assumptions of Theorem C.1 and boundedness of BR\mathcal{B}_{R} over Θ\Theta, if INIT⁡\operatorname{INIT} uses FTL, or AOGD in the case of R(⋅)=12∥⋅∥22R(\cdot)=\frac{1}{2}\|\cdot\|_{2}^{2}, and SIM⁡\operatorname{SIM} uses ε\varepsilon-EWOO as defined in Proposition C.2, then Algorithm 1 achieves

for V2=min⁡ϕ∈Θ∑t=1TσtBR(θt∗∣∣ϕ)V^{2}=\min_{\phi\in\Theta}\sum_{t=1}^{T}\sigma_{t}\mathcal{B}_{R}(\theta_{t}^{\ast}||\phi) and constant CC the product of the constant CC from Proposition B.1 and the bound on the gradient of the Bregman divergence. Assuming σt=Gm ∀ t\sigma_{t}=G\sqrt{m}~{}\forall~{}t and substituting ε=1T\varepsilon=\frac{1}{\sqrt{T}} yields

Substitute Proposition B.1 and Corollary C.2 into Theorem C.1. ∎

Under the assumptions of Theorem 3.1 and boundedness of Θ\Theta, if INIT⁡\operatorname{INIT} is OGD with learning rate 1σmax⁡\frac{1}{\sigma_{\max}} and SIM⁡\operatorname{SIM} uses ε\varepsilon-EWOO as defined in Proposition C.2 then Algorithm 1 achieves

for PT(Ψ)=∑t=2T∥ψt−ψt−1∥2P_{T}(\Psi)=\sum_{t=2}^{T}\|\psi_{t}-\psi_{t-1}\|_{2}. Assuming σt=Gm ∀ t\sigma_{t}=G\sqrt{m}~{}\forall~{}t and substituting ε=1T\varepsilon=\frac{1}{\sqrt{T}} yields

Substitute Theorem 3.3 and Corollary C.2 into Theorem C.1. ∎

Appendix D Adapting to the Inter-Task Geometry

for c‾p=1−(23)1−p1−p\underline{c}_{p}=\frac{1-\left(\frac{2}{3}\right)^{1-p}}{1-p} and c‾p=11−p\overline{c}_{p}=\frac{1}{1-p}.

We will denote the spectral norm by ∥⋅∥2\|\cdot\|_{2} and the Frobenius norm by ∥⋅∥F\|\cdot\|_{F}.

[43, Theorem 3.1] The function f(X)=−log⁡det⁡Xf(\bm{X})=-\log\det\bm{X} is 1σ2\frac{1}{\sigma^{2}}-strongly-convex w.r.t. ∥⋅∥2\|\cdot\|_{2} over the set of symmetric positive-definite matrices with spectral norm bounded by σ\sigma.

By the Löwner-Heinz theorem , x−1,x,x^{-1},x, and x2x^{2} are operator convex. The result follows by applying Claim D.5. ∎

for any x>0\bm{x}>0 and some constant CpC_{p} depending only on pp.

Taking the summation over the coordinates yields

for Cp,1=4c‾1−32p2(1+1c‾p)/c‾p3/2C_{p,1}=4\overline{c}_{1-\frac{3}{2}p}\sqrt{2}\left(1+\frac{1}{\underline{c}_{p}}\right)/\underline{c}_{p}^{3/2} and Cp,2=42(1+1c‾p)∑t=1∞1t1+p2/c‾p3/2C_{p,2}=4\sqrt{2}\left(1+\frac{1}{\underline{c}_{p}}\right)\sum_{t=1}^{\infty}\frac{1}{t^{1+\frac{p}{2}}}/\underline{c}_{p}^{3/2}. Thus we have

Separating again per-coordinate we have that

Define bt2=12(θt∗−ϕt)2\bm{b}_{t}^{2}=\frac{1}{2}(\bm{\theta}_{t}^{\ast}-\bm{\phi}_{t})^{2} and gt2=∇1:m2\bm{g}_{t}^{2}=\bm{\nabla}_{1:m}^{2}. Then applying Proposition D.1 yields

Substituting η+1dmT\bm{\eta}+\frac{\bm{1}_{d}}{\sqrt{mT}} for the optimum and the values of ε,ζ,p\varepsilon,\zeta,p completes the proof. ∎

for constant CσC_{\sigma} depending only on σB,σG\sigma_{B},\sigma_{G}.

Since by Claim D.4 −log⁡det⁡∣X∣-\log\det|\bm{X}| is ζ2σB2+ε2\frac{\zeta^{2}}{\sigma_{B}^{2}+\varepsilon^{2}}-strongly-convex we have by Theorem B.1 that

for some CσC_{\sigma} depending on σB2,σG2\sigma_{B}^{2},\sigma_{G}^{2}. Therefore

for ε=1/T\varepsilon=1/\sqrt{T} and ζ=m/T\zeta=\sqrt{m}/\sqrt{T}. Then we achieve

Let DD and GG be the diameter of Θ\Theta and Lipschitz bound on the losses, respectively. Then applying Proposition D.2 yields

Appendix E Online-to-Batch Conversion for Task-Averaged Regret

where θˉ=1mθ1:m\bar{\theta}=\frac{1}{m}\theta_{1:m} is generated by randomly sampling t∈U[T]t\in\mathcal{U}[T], running the online algorithm with state sts_{t}, and averaging the actions {θi}i∈[m]\{\theta_{i}\}_{i\in[m]}. If on each task the meta-learning algorithm runs an online algorithm with regret upper bound U⁡m(st)\operatorname{\bf U}_{m}(s_{t}) a convex, nonnegative, and BmB\sqrt{m}-bounded function of the state st∈Xs_{t}\in\mathcal{X}, where X\mathcal{X} is a convex Euclidean subset, and the total regret upper bound is Uˉ⁡T\operatorname{\mathbf{\bar{U}}}_{T}, then we also have the bound

where θˉ=1mθ1:m\bar{\theta}=\frac{1}{m}\theta_{1:m} is generated by running the online algorithm with state sˉ=1Ts1:T\bar{s}=\frac{1}{T}s_{1:T} and averaging the actions {θi}i∈[m]\{\theta_{i}\}_{i\in[m]}.

For the second inequality, applying Proposition A.1, Jensen’s inequality, and Proposition A.2 yields

The first inequality follows similarly except using R⁡m\operatorname{\bf R}_{m} instead of U⁡m\operatorname{\bf U}_{m}, linearity of expectation instead of Jensen’s inequality, 1 instead of BB, and Rˉ⁡T\operatorname{\bf\bar{R}}_{T} instead of Uˉ⁡T\operatorname{\mathbf{\bar{U}}}_{T}. ∎

Note that since regret-upper-bounds are nonnegative one can easily replace 8 by 2 in the second inequality by simply multiplying and dividing by BmB\sqrt{m} in the third line of the above proof.

Under the assumptions of Theorems 3.2 and 5.1, if the loss functions are Lipschitz and we use Algorithm 1 with ηt\eta_{t} also learned, using ε\varepsilon-EWOO as in Theorem 3.2 for ε=1/mT+1/m\varepsilon=1/\sqrt{mT}+1/\sqrt{m}, and set the initialization using ϕt+1=1t∑s≤tθs∗\phi_{t+1}=\frac{1}{t}\sum_{s\leq t}\theta_{s}^{\ast}, then w.p. 1−δ1-\delta we have

Substitute Corollary C.3 into Theorem E.1 using the fact the the regret-upper-bounds are O(mε)\mathcal{O}(\frac{\sqrt{m}}{\varepsilon})-bounded. Conclude by applying Claim E.1. ∎

where θˉ=1mθ1:m\bar{\theta}=\frac{1}{m}\theta_{1:m} is generated by randomly sampling t∈U[T]t\in\mathcal{U}[T], running the online algorithm with state sts_{t}, and averaging the actions {θi}i∈[m]\{\theta_{i}\}_{i\in[m]}. If on each task the meta-learning algorithm runs an online algorithm with regret upper bound U⁡m(st)\operatorname{\bf U}_{m}(s_{t}) a convex, nonnegative, and BmB\sqrt{m}-bounded function of the state st∈Xs_{t}\in\mathcal{X}, where X\mathcal{X} is a convex Euclidean subset, and the total regret upper bound is Uˉ⁡T\operatorname{\mathbf{\bar{U}}}_{T}, then we also have the bound

where θˉ=1mθ1:m\bar{\theta}=\frac{1}{m}\theta_{1:m} is generated by running the online algorithm with state sˉ=1Ts1:T\bar{s}=\frac{1}{T}s_{1:T} and averaging the actions {θi}i∈[m]\{\theta_{i}\}_{i\in[m]}.

By Corollary A.2 and Jensen’s inequality we have w.p. 1−δ21-\frac{\delta}{2} that

As in the proof of Theorem E.1, by Proposition A.2 we further have w.p. 1−δ21-\frac{\delta}{2} that

Substituting the second inequality into the first yields the second bound. The first bound follows similarly except using R⁡m\operatorname{\bf R}_{m} instead of U⁡m\operatorname{\bf U}_{m}, linearity of expectation instead of Jensen’s inequality, 1 instead of BB, and Rˉ⁡T\operatorname{\bf\bar{R}}_{T} instead of Uˉ⁡T\operatorname{\mathbf{\bar{U}}}_{T}. ∎

Applying Proposition A.1 and Theorem A.4 we have w.p. 1−δ21-\frac{\delta}{2} that

This yields the first bound since. The second bound follows similarly except for the application of Corollary A.2 in the second step w.p. 1−δ21-\frac{\delta}{2}. ∎

Appendix F Adapting to Task-Similarity under Parameter Growth

In this appendix we cast the problem of adaptively learning the task-similarity in the framework of Khodak et al. . We do this specifically to show that our basic results extend to approximate meta-updates under quadratic growth. We first provide a generalized version of their Ephemeral method in Algorithm 3. We then state the relevant approximation assumptions and proceed to prove guarantees on the average regret-upper-bound for the case of a fixed task-similarity in Theorem F.1 and for adaptively learning it in Theorem F.2. Then the quadratic-growth results of Khodak et al. , specifically Propositions B.1, B.2, and B.3, can be applied directly to show average regret-upper-bound guarantees of the same order as those in the main paper but with additional om(1)o_{m}(1) terms inside the parentheses. Note that our results, especially in the batch-within-online setting, will in general be stronger because we do not incur the Δmax⁡\Delta_{\max}-error term that is needed to account for the doubling trick in Khodak et al. .

Assume the data given to Algorithm 3 and define the following quantities:

convenience coefficients σt=Gtmt\sigma_{t}=G_{t}\sqrt{m_{t}}

sequence of update parameters {θ^t∈Θ}t∈[T]\{\hat{\theta}_{t}\in\Theta\}_{t\in[T]} with average update ϕ^=1σ1:T∑t=1Tσtθ^\hat{\phi}=\frac{1}{\sigma_{1:T}}\sum_{t=1}^{T}\sigma_{t}\hat{\theta}

a sequence of reference parameters {θt′∈Θ}t∈[T]\{\theta_{t}^{\prime}\in\Theta\}_{t\in[T]} with average reference parameter ϕ′=1σ1:T∑t=1Tσtθt′\phi^{\prime}=\frac{1}{\sigma_{1:T}}\sum_{t=1}^{T}\sigma_{t}\theta_{t}^{\prime}

a sequence {θt∗∈Θ}t∈[T]\{\theta_{t}^{\ast}\in\Theta\}_{t\in[T]} of optimal parameters in hindsight

we will say we are in the “Exact" case if θ^t=θt′=θt∗ ∀ t\hat{\theta}_{t}=\theta_{t}^{\prime}=\theta_{t}^{\ast}~{}\forall~{}t and the “Approx" case otherwise

κ≥1,Δt∗≥0\kappa\geq 1,\Delta_{t}^{\ast}\geq 0 s.t. ∑t=1TαtBR(θt∗∣∣ϕt)≤Δ1:T∗+κ∑t=1TαtBR(θ^t∣∣ϕt)\sum_{t=1}^{T}\alpha_{t}\mathcal{B}_{R}(\theta_{t}^{\ast}||\phi_{t})\leq\Delta_{1:T}^{\ast}+\kappa\sum_{t=1}^{T}\alpha_{t}\mathcal{B}_{R}(\hat{\theta}_{t}||\phi_{t}) for some αt≥0\alpha_{t}\geq 0

ν≥1,Δ′≥0\nu\geq 1,\Delta^{\prime}\geq 0 s.t. ∑t=1TσtBR(θ^t∣∣ϕ^)≤Δ′+ν∑t=1TσtBR(θt′∣∣ϕ′)\sum_{t=1}^{T}\sigma_{t}\mathcal{B}_{R}(\hat{\theta}_{t}||\hat{\phi})\leq\Delta^{\prime}+\nu\sum_{t=1}^{T}\sigma_{t}\mathcal{B}_{R}(\theta_{t}^{\prime}||\phi^{\prime})

average deviation V2=1σ1:T∑t=1TσtBR(θt′∣∣ϕ′)V^{2}=\frac{1}{\sigma_{1:T}}\sum_{t=1}^{T}\sigma_{t}\mathcal{B}_{R}(\theta_{t}^{\prime}||\phi^{\prime}) of the reference parameters

action diameter D2=max⁡{D∗2,max⁡θ∈ΘBR(θ∣∣ϕ1)}D^{2}=\max\{{D^{\ast}}^{2},\max_{\theta\in\Theta}\mathcal{B}_{R}(\theta||\phi_{1})\} in the Exact case or max⁡θ,ϕ∈ΘBR(θ∣∣ϕ)\max_{\theta,\phi\in\Theta}\mathcal{B}_{R}(\theta||\phi) in the Approx case

effective action space Θ^=Conv⁡({θ^t}t∈[T])\hat{\Theta}=\operatorname{Conv}(\{\hat{\theta}_{t}\}_{t\in[T]}) if INIT⁡\operatorname{INIT} is FTL or Θ\Theta if INIT⁡\operatorname{INIT} is AOGD

upper bound G′G^{\prime} on the Lipschitz constants of the functions {BR(θ^t∣∣⋅)}t∈[T]\{\mathcal{B}_{R}(\hat{\theta}_{t}||\cdot)\}_{t\in[T]} over Θ^\hat{\Theta}

we will say we are in the “Nice" case if BR(θ∣∣⋅)\mathcal{B}_{R}(\theta||\cdot) is 1-strongly-convex and β\beta-strongly-smooth w.r.t. ∥⋅∥ ∀ θ∈Θ\|\cdot\|~{}\forall~{}\theta\in\Theta

in the general case INIT⁡\operatorname{INIT} is FTL; in the Nice case INIT⁡\operatorname{INIT} may instead be AOGD

convenience indicator ι=1INIT⁡=FTL⁡\iota=1_{\operatorname{INIT}=\operatorname{FTL}}

TASK⁡η,ϕ=FTRL⁡η,ϕ(R)\operatorname{TASK}_{\eta,\phi}=\operatorname{FTRL}_{\eta,\phi}^{(R)} or OMD⁡η,ϕ(R)\operatorname{OMD}_{\eta,\phi}^{(R)}

at t=1t=1 the update algorithm INIT⁡\operatorname{INIT} plays ϕ1∈Θ\phi_{1}\in\Theta satisfying max⁡θ∈ΘBR(θ∣∣ϕ1)<∞\max_{\theta\in\Theta}\mathcal{B}_{R}(\theta||\phi_{1})<\infty

in the Approx case RR is β\beta-strongly-smooth for some β≥1\beta\geq 1

The following theorem does not appear in the main paper but is used in discussion. It shows guarantees for the case when the task-similarity is known in advance and so SIM⁡\operatorname{SIM} always returns a constant.

Make Assumption F.1 and suppose SIM⁡\operatorname{SIM} always plays Dt=εD_{t}=\varepsilon. Then Algorithm 3 has a regret upper-bound of

for C=G′22C=\frac{{G^{\prime}}^{2}}{2} in the Nice case or otherwise C=2C′D′G′C=2C^{\prime}D^{\prime}G^{\prime}.

F.2 Average Regret when Learning Task Similarity

for C=G′22C=\frac{{G^{\prime}}^{2}}{2} in the Nice case or otherwise C=2C′D′G′C=2C^{\prime}D^{\prime}G^{\prime}.

F.3 Statistical Task-Similarity under Quadratic Growth

In this section we relate our task-similarity measure to that of Denevi et al. under α\alpha-QG.

Following the argument of Shalev-Shwartz et al. [49, Theorem 2] but applying α\alpha-QG instead of strong-convexity in Equation 8, which holds by definition of α\alpha-QG, we obtain

Appendix G Experimental Details

Code is available at https://github.com/mkhodak/ARUBA.

For our Reptile experiments we use the code and default settings provided by Nichol et al. , except we tune the learning rate, which for ARUBA corresponds to ε/ζ\varepsilon/\zeta, and the coefficient cc in ARUBA++. In addition to the the parameters listed in the above tables, we set ζ=p=1.0\zeta=p=1.0 for all experiments. All evaluations are averages of three runs.

G.2 FedAvg

For FedAvg we train a 2-layer stacked LSTM model with 256 hidden units, 8-dimensional trained character embeddings, with a maximum input string size of 80 characters; these settings are used to match those of McMahan et al. . Similarly, we take their approach of only removing those actors from the Shakespeare dataset with fewer than two lines and split each user temporally into train/test sets with a training fraction of 0.8. Unlike McMahan et al. , we also split the users into meta-training and meta-testing sets, also with a fraction of 0.8, in order to evaluate meta-test performance. We run both algorithms for 500 rounds with a batch of 10 users per round and a within-task batch-size of 10, as in Caldas et al. . For unmodified FedAvg we found that an initial learning rate of η=1.0\eta=1.0 worked well – this is similar to those reported in McMahan et al. and Caldas et al. – and for the tuned variant we found that a multiplicative decay of 0.990.99. At meta-test-time we tuned the refinement learning rate over {10−3,10−2,10−1}\{10^{-3},10^{-2},10^{-1}\}. For ARUBA and its isotropic variant we set ε=ζ=0.05\varepsilon=\zeta=0.05 and p=1.0p=1.0, so that η=ε/ζ=1.0\eta=\varepsilon/\zeta=1.0 in our setting as well.