summaries: The paper is concerned with meta reinforcement learning. The author(s) introduced a framework for learning high-order derivatives of value functions using off-policy evaluation. Specifically, the author(s) proposed to adopt existing off-policy evaluation methods to derive the value estimator and then take the derivative with respect to the target parameters to construct the estimate. This technique is further adopted to implement meta reinforcement learning.

summaries: This paper considers meta-RL with gradient-based adaptations, which relies heavily on the estimation of the Hessians of value functions. Though there are previous approaches that focus on unbiased/biased estimations of Hessian. This paper gives a unified view for estimating higher-order derivatives of value functions, through the lens of off-policy evaluation. This subsumes several previous works, also sheds some light on the bias-variance tradeoff of the estimate. Empirically, the new proposed second-order estimate is also incorporated in meta-RL algorithms to showcase the empirical efficiency.

summaries: This paper uses Huang and Jiang (2020) to talk about gradient-based meta learning in reinforcement learning using the language of off-policy evaluation. The logical connection behind those concepts stands from the fact that policy gradient methods in reinforcement learning are in fact instances of the likelihood ratio derivative estimator, which itself stems from important sampling. Proposition 3.2 and 3.3 are obtained from Huang and Jiang. Proposition 3.1 tries to extend their results to talk about higher order derivatives. The author's finally propose to use ideas from Taylor policy optimization from Tang et al. (2020) as a variance reduction method within the doubly robust policy gradient framework of Huang and Jiang.

summaries: The paper uses the recently developed framework of Taylor expansions for value functions (ref [29] in the paper), expands the results, and shows how they could be applied to Off-Policy Evaluation, and by a non-trivial extension to meta-learning. The main algorthmic contribution of the paper is the application of the ideas to meta-learning, but I feel the main contribution is a conceptual one, further developing the ideas from [29] and drawing theoretical connections between OPE and meta-learning.  The paper includes an experimental section demonstrating marginal improvement over baselines in several domains.

abstractText: Model-agnostic meta-reinforcement learning requires estimating the Hessian matrix of value functions. This is challenging from an implementation perspective, as repeatedly differentiating policy gradient estimates may lead to biased Hessian estimates. In this work, we provide a unifying framework for estimating higherorder derivatives of value functions, based on off-policy evaluation. Our framework interprets a number of prior approaches as special cases and elucidates the bias and variance trade-off of Hessian estimates. This framework also opens the door to a new family of estimates, which can be easily implemented with auto-differentiation libraries, and lead to performance gains in practice. We open source the code to reproduce our results1.

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Alex Graves
Ioannis Antonoglou
Daan Wierstra
Martin Riedmiller

title: Playing atari with deep reinforcement learning

venue: arXiv preprint arXiv:1312.5602,

authors:
David Silver
Aja Huang
Chris J Maddison
Arthur Guez
Laurent Sifre
George Van Den Driessche
Julian Schrittwieser
Ioannis Antonoglou
Veda Panneershelvam
Marc Lanctot

title: Mastering the game of go with deep neural networks and tree

venue: search. nature,

authors:
Sergey Levine
Chelsea Finn
Trevor Darrell
Pieter Abbeel

title: End-to-end training of deep visuomotor policies

venue: The Journal of Machine Learning Research,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: In International Conference on Machine Learning,

authors:
Ziyu Wang
Tom Schaul
Matteo Hessel
Hado Van Hasselt
Marc Lanctot
Nando De Freitas

title: Dueling network architectures for deep reinforcement learning

venue: arXiv preprint arXiv:1511.06581,

authors:
Yan Duan
John Schulman
Xi Chen
Peter L Bartlett
Ilya Sutskever
Pieter Abbeel

title: Rl2: Fast reinforcement learning via slow reinforcement learning

venue: arXiv preprint arXiv:1611.02779,

authors:
Rein Houthooft
Richard Y Chen
Phillip Isola
Bradly C Stadie
Filip Wolski
Jonathan Ho
Pieter Abbeel

title: Evolved policy gradients

venue: arXiv preprint arXiv:1802.04821,

authors:
Junhyuk Oh
Matteo Hessel
Wojciech M Czarnecki
Zhongwen Xu
Hado van Hasselt
Satinder Singh
David Silver

title: Discovering reinforcement learning algorithms

venue: arXiv preprint arXiv:2007.08794,

authors:
Zhongwen Xu
Hado van Hasselt
Matteo Hessel
Junhyuk Oh
Satinder Singh
David Silver

title: Meta-gradient reinforcement learning with an objective discovered online

venue: arXiv preprint arXiv:2007.08433,

authors:
Kate Rakelly
Aurick Zhou
Chelsea Finn
Sergey Levine
Deirdre Quillen

title: Efficient off-policy meta-reinforcement learning via probabilistic context variables

venue: In International conference on machine learning,

authors:
Luisa Zintgraf
Kyriacos Shiarlis
Maximilian Igl
Sebastian Schulze
Yarin Gal
Katja Hofmann
Shimon Whiteson

title: Varibad: A very good method for bayes-adaptive deep rl via metalearning

authors:
Zhongwen Xu
Hado P van Hasselt
David Silver

title: Meta-gradient reinforcement learning

venue: In Advances in neural information processing systems,

authors:
Tom Zahavy
Zhongwen Xu
Vivek Veeriah
Matteo Hessel
Junhyuk Oh
Hado van Hasselt
David Silver
Satinder Singh

title: A self-tuning actor-critic algorithm

venue: arXiv preprint arXiv:2002.12928,

authors:
Pedro A Ortega
Jane X Wang
Mark Rowland
Tim Genewein
Zeb Kurth-Nelson
Razvan Pascanu
Nicolas Heess
Joel Veness
Alex Pritzel
Pablo Sprechmann

title: Meta-learning of sequential strategies

authors:
Richard S Sutton
David A McAllester
Satinder P Singh
Yishay Mansour

title: Policy gradient methods for reinforcement learning with function approximation

authors:
Kavosh Asadi
Michael L Littman

title: An alternative softmax operator for reinforcement learning

authors:
James Bradbury
Roy Frostig
Peter Hawkins
Matthew James Johnson
Chris Leary
Dougal Maclaurin
George Necula
Adam Paszke
Jake VanderPlas
Skye Wanderman-Milne
Qiao Zhang

title: JAX: composable transformations of Python+NumPy programs, 2018

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in pytorch

authors:
Hao Liu
Richard Socher
Caiming Xiong

title: Taming maml: Efficient unbiased metareinforcement learning

authors:
Jakob Foerster
Gregory Farquhar
Maruan Al-Shedivat
Tim Rocktäschel
Eric Xing
Shimon Whiteson

title: Dice: The infinitely differentiable monte carlo estimator

authors:
Gregory Farquhar
Shimon Whiteson
Jakob Foerster

title: Loaded dice: Trading off bias and variance in any-order score function gradient estimators for reinforcement

authors:
Jingkai Mao
Jakob Foerster
Tim Rocktäschel
Maruan Al-Shedivat
Gregory Farquhar
Shimon Whiteson

title: A baseline for any order gradient estimation in stochastic computation graphs

authors:
Jonas Rothfuss
Dennis Lee
Ignasi Clavera
Tamim Asfour
Pieter Abbeel

title: Promp: Proximal meta-policy search

venue: arXiv preprint arXiv:1810.06784,

authors:
John Schulman
Nicolas Heess
Theophane Weber
Pieter Abbeel

title: Gradient estimation using stochastic computation graphs

venue: In Advances in Neural Information Processing Systems,

authors:
Doina Precup
Richard S Sutton
Sanjoy Dasgupta

title: Off-policy temporal-difference learning with function approximation

venue: In ICML,

authors:
Nan Jiang
Lihong Li

title: Doubly robust off-policy value evaluation for reinforcement learning

authors:
Yunhao Tang
Michal Valko
Rémi Munos

title: Taylor expansion policy optimization

venue: arXiv preprint arXiv:2003.06259,

authors:
Martín Abadi
Paul Barham
Jianmin Chen
Zhifeng Chen
Andy Davis
Jeffrey Dean
Matthieu Devin
Sanjay Ghemawat
Geoffrey Irving
Michael Isard

title: Tensorflow: A system for large-scale machine learning

venue: In 12th {USENIX} symposium on operating systems design and implementation ({OSDI}

authors:
Richard S Sutton
Andrew G Barto

title: Reinforcement learning: An introduction, volume 1

venue: MIT press Cambridge,

authors:
Maruan Al-Shedivat
Trapit Bansal
Yuri Burda
Ilya Sutskever
Igor Mordatch
Pieter Abbeel

title: Continuous adaptation via meta-learning in nonstationary and competitive environments

venue: arXiv preprint arXiv:1710.03641,

authors:
Bradly C Stadie
Ge Yang
Rein Houthooft
Xi Chen
Yan Duan
Yuhuai Wu
Pieter Abbeel
Ilya Sutskever

title: Some considerations on learning to explore via meta-reinforcement learning

venue: arXiv preprint arXiv:1803.01118,

authors:
Vijay R Konda
John N Tsitsiklis

title: Actor-critic algorithms. In Advances in neural information processing systems, pages 1008–1014

authors:
Alekh Agarwal
Sham M Kakade
Jason D Lee
Gaurav Mahajan

title: Optimality and approximation with policy gradient methods in markov decision processes

venue: In Conference on Learning Theory,

authors:
Miroslav Dudík
Dumitru Erhan
John Langford
Lihong Li

title: Doubly robust policy evaluation and optimization

venue: Statistical Science,

authors:
Philip Thomas
Emma Brunskill

title: Data-efficient off-policy policy evaluation for reinforcement learning

authors:
Jiawei Huang
Nan Jiang

title: From importance sampling to doubly robust policy gradient

authors:
Hao Liu
Yihao Feng
Yi Mao
Dengyong Zhou
Jian Peng
Qiang Liu

title: Action-depedent control variates for policy optimization via stein’s identity

venue: arXiv preprint arXiv:1710.11198,

authors:
Cathy Wu
Aravind Rajeswaran
Yan Duan
Vikash Kumar
Alexandre M Bayen
Sham Kakade
Igor Mordatch
Pieter Abbeel

title: Variance reduction for policy gradient with action-dependent factorized baselines

venue: arXiv preprint arXiv:1803.07246,

authors:
George Tucker
Surya Bhupatiraju
Shixiang Gu
Richard Turner
Zoubin Ghahramani
Sergey Levine

title: The mirage of action-dependent baselines in reinforcement learning

authors:
Horia Mania
Aurelia Guy
Benjamin Recht

title: Simple random search provides a competitive approach to reinforcement learning

venue: arXiv preprint arXiv:1803.07055,

authors:
Mark Rowland
Will Dabney
Rémi Munos

title: Adaptive trade-offs in off-policy learning

venue: arXiv preprint arXiv:1910.07478,

authors:
Rémi Munos
Tom Stepleton
Anna Harutyunyan
Marc Bellemare

title: Safe and efficient off-policy reinforcement learning

authors:
John Schulman
Sergey Levine
Pieter Abbeel
Michael Jordan
Philipp Moritz

title: Trust region policy optimization

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
Lasse Espeholt
Hubert Soyer
Remi Munos
Karen Simonyan
Volodymir Mnih
Tom Ward
Yotam Doron
Vlad Firoiu
Tim Harley
Iain Dunning

title: Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures

venue: arXiv preprint arXiv:1802.01561,

authors:
Steven Kapturowski
Georg Ostrovski
John Quan
Remi Munos
Will Dabney

title: Recurrent experience replay in distributed reinforcement learning

authors:
Alireza Fallah
Aryan Mokhtari
Asuman Ozdaglar

title: On the convergence theory of gradientbased model-agnostic meta-learning algorithms

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Kaiyi Ji
Junjie Yang
Yingbin Liang

title: Multi-step model-agnostic meta-learning: Convergence and improved algorithms

venue: arXiv preprint arXiv:2002.07836,

authors:
Emanuel Todorov
Tom Erez
Yuval Tassa

title: Mujoco: A physics engine for model-based control

venue: In Intelligent Robots and Systems (IROS),

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

heading: 1 Introduction

text:
Recent years have witnessed the success of reinforcement learning (RL) in challenging domains such as game playing [1], board games [2], and robotics control [3]. However, despite such breakthroughs, state-of-the-art RL algorithms are still plagued by sample inefficiency, as training agents requires orders of magnitude more samples than humans would experience to reach similar levels of performance [1, 2]. One hypothesis on the source of such inefficiencies is that standard RL algorithms are not good at leveraging prior knowledge, which implies that whenever presented with a new task, the algorithms must learn from scratch. On the contrary, humans are much better at transferring prior skills to new scenarios, an innate ability arguably obtained through evolution over thousand of years.
Meta-reinforcement learning (meta-RL) formalizes the learning and transfer of prior knowledge in RL [4]. The high-level idea is to have an RL agent that interacts with a distribution of environments at meta-training time. The objective is that at meta-testing time, when the agent interacts with previously unseen environments, it can learn much faster than meta-training time. Here, faster learning is measured by the number of new samples needed to achieve a good level of performance. If an agent can achieve good performance at meta-testing time, it embodies the ability to transfer knowledge from prior experiences at meta-training time. meta-RL algorithms can be constructed in many ways, such as those based on recurrent memory [5, 6], gradient-based adaptations [4], learning loss functions [7–9], probabilistic inference of context variables [10–12], online adaptation of hyper-parameters within a single lifetime [13, 14] and so on. Some of these formulations have mathematical connections; see e.g. [15] for an in-depth discussion.
We focus on gradient-based adaptations [4], where the agent carries out policy gradient updates [16] at both meta-training and meta-testing time. Conceptually, since meta-RL seeks to optimize the
1https://github.com/robintyh1/neurips2021-meta-gradient-offpolicy-evaluation
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
way in which the agent adapts itself in face of new environments, it needs to differentiate through the policy gradient update itself. This effectively reduces the meta-RL problem into estimations of Hessian matrices of value functions.
Challenges for computing Hessian matrix of value functions. To calculate Hessian matrices (or unbiased estimates thereof) for supervised learning objectives, it suffices to differentiate through gradient estimates, which can be easily implemented with auto-differentiation packages [17–19]. However, it is not the case for value functions in RL. Intuitively, this is because value functions are defined via expectations with respect to distributions that themselves depend on the policy parameters of interest, whereas in supervised learning the expectations are defined with respect to a fixed data distribution. As a result, implementations that do not take this into account may lead to estimates [4] whose bias is not properly characterized and might have a negative impact on downstream applications [20].
Motivated by this observation, a number of prior works suggest implementation alternatives that lead to unbiased Hessian estimates [21], with potential variance reduction [22, 23, 20]; or biased estimates with small variance [24]. However, different algorithms in this space are motivated and derived in seemingly unrelated ways: for example, [21–23] derive code-level implementations within the general context of stochastic computation graphs [25]. On the other hand, [24] derives the estimates by explicitly analyzing certain terms with potentially high variance, which naturally produces bias in the final estimate. Due to apparently distinct ways of deriving estimates, it is not immediately clear how all such methods are related, and whether there could be other alternatives.
Central contribution. We present a unified framework for estimating higher-order derivatives of value functions, based on the concept of off-policy evaluation. The main insights are summarized in Table 1, where most aforementioned prior work can be interpreted as special cases of our framework. Our framework has a few advantages: (1) it conceptually unifies a few seemingly unrelated prior methods; (2) it elucidates the bias and variance trade-off of the estimates; (3) it naturally produces new methods based on Taylor expansions of value functions [29].
After a brief background introduction on meta-RL in Section 2, we will discuss the above aspects in detail in Section 3. From an implementation perspective, we will show in Section 4.1 that both the general framework and the new method can be conveniently implemented in auto-differentiation libraries [30, 19], making it amenable in a practical setup. Finally in Section 5, we validate important claims based on this framework with experimental insights.

heading: 2 Background

text: We begin with the notation and background on RL and meta-RL.

heading: 2.1 Task-based reinforcement learning

text:
Consider a Markov decision process (MDP) with state space X and action space A. At time t ≥ 0, the agent takes action at ∈ A in state xt ∈ X , receives a reward rt and transitions to a next state xt+1 ∼ p(·|xt, at). Without loss of generality, we assume a single starting state x0. Here, we assume the reward rt = r(xt, at, g) to be a deterministic function of state-action pair (xt, at) and the task variable g ∈ G. The task variable g ∼ pG is resampled for every episode. For example, xt ∈ X is the sensory space of a running robot, at ∈ A is the control, g is the episodic target direction in which to run and r(xt, at, g) is the speed in direction g. A policy π : X → P(A) specifies a distribution over actions at each state. For convenience, we define
the value function V π(x, g) := Eπ [ ∑∞ t=0 γ tr(xt, at, g) | x0 = x] and Q-function Qπ(x, a, g) :=
Eπ [ ∑∞ t=0 γ tr(xt, at, g) | x0 = x, a0 = a]. Without loss of generality, we assume the policy to be smoothly parameterized as πθ with parameter θ ∈ RD. We further assume that the MDPs terminate within a finite horizon of H <∞ under any policy.

heading: 2.2 Gradient-based meta-reinforcement learning

text:
The motivation of meta-RL is to identify a policy πθ such that given a task g, after updating the parameter from θ to θ′ with a parameter update computed under rewards r(x, a, g), the resulting policy πθ′ performs well. Formally, define U(θ, g) ∈ RD as a parameter update to θ, for example, one policy gradient ascent step. Model-agnostic meta-learning (MAML) [4] formulates meta-RL as optimizing the value function at the updated policy Eg∼pG [V πθ′ (x, g)] where θ′ = θ+U(θ, g) is the updated policy. The optimization is with respect to the initial policy parameter θ. The aim is to find θ such that it entails fast learning (or adaptation) to the environment, through the inner loop update operator U(θ, g), that leads to high-performing updated policy θ′. Consider the following problem,
max θ
F (θ) := Eg∼pG [V πθ′ (x0, g)] , θ′ = θ + U(θ, g). (1)
We can decompose the meta-gradient into two terms based on the chain rule,
∇θF (θ) = Eg∼pG [∇θF (θ, g)] := Eg∼pG [∇θθ′∇θ′V πθ′ (x0, g)] (2)
where∇θθ′ = I+∇θU(θ, g) ∈ RD×D is a matrix, and∇θ′V πθ′ (x0, g) is the vanilla policy gradient evaluated at the updated parameter θ′ [31]. A straightforward way to optimizing Eqn (1) is to carry out the outer loop update θ ← θ + α∇θF (θ) with learning rate α > 0.
Policy gradient update. Following MAML [4], we focus on policy gradient update where U(θ, g) = η∇θV πθ (x0, g) for a fixed step size η > 0. The matrix ∇θU(θ, g) = η∇2θV πθ (x0, g) is the Hessian of the value function with respect to policy parameters; henceforth, we define Hθ(x0, g) := ∇2θV πθ (x0, g) ∈ RD×D.
Estimating meta-gradients. Practical algorithms construct stochastic estimates of the metagradients using a finite number of samples. In Eqn (2), we decompose the meta-gradients into the multiplication of a Hessian matrix with a policy gradient vector. Given a fixed task variable g, a common practice of prior work is to construct the gradient estimate as the product of the Hessian estimate and the policy gradient estimate ∇̂F (θ, g) = (I + ηĤθ(x0, g))∇̂θ′V πθ′ (x0, g); see, e.g., [4, 32, 33, 24, 21–23, 20]. See Algorithm 1 for the full pseudocode for estimating meta-gradients by sampling multiple tasks. Since there is a large literature on constructing accurate estimates to policy gradient (e.g., actor-critic algorithms use baselines for variance reduction [34]), the main challenge consists in estimating the Hessian matrix accurately.
Bias of common plug-in estimators. Common practices in meta-RL algorithms rely on the premise that if both Hessian and policy gradient estimates are unbiased, then the meta-gradient estimate is unbiased too.
E [ Ĥθ(x0, g) ] = Hθ(x0, g),E [ ∇̂θ′V πθ′ (x0, g) ] = ∇θ′V πθ′ (x0, g)⇒ E [ ∇̂θF (θ, g) ] = ∇θF (θ, g).
Unfortunately, this is not true. This is because the two estimates are in general correlated when the sample size is finite, leading to the bias of the overall estimate. We provide further discussions in Appendix A. For the rest of the paper, we follow practices of prior work and focus on the properties of Hessian estimates, leaving a more proper treatment of this bias to future work.

heading: 3 Deriving Hessian estimates with off-policy evaluation

text: Since the meta-gradient estimates are computed by averaging over task variables g ∼ pG , in the following, we focus on Hessian estimates at a single state and task variable Hθ(x, g) with a fixed g. In this section, we also drop the dependency of the value function on g, such that, e.g., V πθ (x0) ≡ V πθ (x0, g) and Q(x, a, g) ≡ Q(x, a).

heading: 3.1 Off-policy evaluation: maintaining higher-order dependencies on parameters

text:
We assume access to data (xt, at, rt)∞t=0 generated under a behavior policy µ. Off-policy evaluation [26] consists in building estimators V̂ πθ (x, g) using the behavior data such that V̂ πθ (x) ≈ V πθ (x)
Algorithm 1 Pseudocode for computing meta-gradients for the MAML objective for i=1,2...n do
Sample task variable gi ∼ pG . Sample B trajectories under policy πθ; Compute B-trajectory policy gradient estimate ∇̂θV πθ (x0, gi), update parameter θ′ = θ + ∇̂θV πθ (x0, gi). Compute B-trajectory Hessian estimate Ĥθ(x0, gi) and an unbiased policy gradient estimate at θ′, i.e., ∇̂θ′V πθ (x0, g). Compute the i-th meta-gradient estimate ∇̂F (θ, gi) = ( I + ηĤθ(x0, gi) ) ∇̂θ′V πθ′ (x0, gi).
end for Output averaged meta-gradient estimate 1n ∑n i=1 ∇̂θF (θ, gi).
for a range of target policies πθ. Note that the estimate V̂ πθ (x) is a random variable depending on (xt, at, rt) ∞ t=0, it is also a function of θ. The approximation V̂
πθ (x) ≈ V πθ (x) implies that V̂ πθ (x) is indicative of how the value function V πθ (x) depends on θ, and hence maintains the higher-order dependencies on θ. Throughout, we assume πθ(a|x) > 0, µ(a|x) > 0 for all (x, a) ∈ X ×A.
Example: step-wise importance sampling (IS) estimate. As a concrete example, consider the unbiased step-wise IS estimate V̂ πθIS (x0) = ∑∞ t=0 γ t ( Πs≤tρ θ s ) rt where ρθs := πθ(as|xs)/µ(as|xs). Since the value function V πθ (x) is in general a highly non-linear function of θ (see discussions in [35]) we see that V̂ πθIS (x0) retains such dependencies via the sum of product of IS ratios.

heading: 3.2 Warming up: deriving unbiased estimates with variance reduction

text:
We start with a general result based on the intuition above: given an estimate V̂ πθ (x) to V πθ (x), we can directly use the mth-order derivative∇mθ V̂ πθ (x) ∈ RD K
as an estimate to∇Kθ V πθ (x). We introduce two assumptions: (A.1) V̂ πθ (x) is mth-order differentiable w.r.t. θ almost surely. (A.2)∥∥∥∇mθ V̂ πθ (x)∥∥∥∞ < M for some constant M for the order m of interest. These assumptions are fairly mild; see further details in Appendix C. The following result applies to general unbiased off-policy evaluation estimates.
Proposition 3.1. Assume (A.1) and (A.2) are satisfied. Further assume we have an estimator V̂ πθ (x) which is unbiased (Eµ [ V̂ πθ′ (x) ] = V πθ′ (x)) for all θ′ ∈ N(θ) where N(θ) is some
open set that contains θ. Under some additional mild conditions, the mth-order derivative of the estimate ∇mθ V̂ πθ (x0) are unbiased estimates to the mth-order derivative of the value function Eµ [ ∇mθ V̂ πθ (x) ] = ∇mθ V πθ (x) for m ≥ 1.
Doubly-robust estimates. As a special case, we describe the doubly-robust (DR) off-policy evaluation estimator [36, 27, 37]. Assume we have access to a state-action dependent critic Q(x, a, g), and we use the notation Q(x, π(x)) := ∑ a Q(x, a)π(a|x). The DR estimate is defined recursively as follows,
V̂ πθDR (xt) = Q(xt, πθ(xt)) + ρ θ t ( rt + γV̂ πθ DR (xt+1)−Q(xt, at) ) . (3)
The DR estimate is unbiased for all πθ and subsumes the step-wise IS estimate as a special case when Q(x, a) ≡ 0. If the critic is properly chosen, e.g., Q(x, a) ≈ Qπθ (x, a), it can lead to significant variance reduction compared to V̂ πθIS (x0). By directly differentiating the estimate ∇mθ V̂ πθ DR (x), we derive estimators for higher-order derivatives of the value function; the result for the gradient in Proposition 3.2 was shown in [38].
Proposition 3.2. Define πt := πθ(at|xt) and let δt := rt + γV̂ πθDR (xt+1)−Q(xt, at) be the sampled temporal difference error at time t. Note that∇θ log πt ∈ RD and∇2θ log πt ∈ RD×D. The estimates
of higher-order derivatives can be deduced recursively, and in particular for m = 1, 2,
∇θV̂ πθDR (xt) = ∇θQ(xt, πθ(xt)) + ρ θ t δt∇θ log πt + γρθt∇θV̂ πθ DR (xt+1), (4)
∇2θV̂ πθ DR (xt) = ρ θ t δt ( ∇2θ log πt +∇θ log πt∇θ log πTt ) + γρθt∇θV̂ πθ DR (xt)∇θ log π T t
+ γρθt∇θ log πt∇θV̂ πθ DR (xt) T +∇2θQ(xt, πθ(xt)) + γρθt∇2θV̂ πθ DR (xt+1). (5)
Bias and variance of Hessian estimates. Proposition 3.1 implies that∇θV̂πθ (x0) and∇2θV̂πθ (x0) are both unbiased. To analyze the variance, we start with m = 1: note that when on-policy µ = πt [16], ∇θV πθDR (x0) recovers a form of gradient estimates similar to actor-critic policy gradient with action-dependent baselines [39–41]; when Q(x, a) is only state dependent, ∇θV̂ πθDR (x) recovers the common policy gradient estimate with state-dependent baselines [34]. As such, the estimates are computed with potential variance reduction due to the critic. Previously, [38] started with the DR estimate and derived a more general result for the on-policy first-order case. For the Hessian estimate, we expect a similar effect of variance reduction as shown in experiments.
Recovering prior work on estimates to higher-order derivatives. When applied to meta-RL, DiCE [21] and its follow-up variants [22, 42] can be seen as special cases of∇2θV̂ πθ DR (x) with different choices of the critic Q when evaluated on-policy µ = πθ. See Table 1 for the correspondence between prior work and their equivalent formulations under the framework of off-policy evaluation. We will discuss detailed pseudocode in Section 4.1. See also Appendix F for more details.

heading: 3.3 Trading-off bias and variance with Taylor expansions

text:
Starting from unbiased off-policy evaluation estimates V̂ πθ (x), we can directly construct unbiased estimates to higher-order derivatives by differentiating the original estimate to obtain ∇mθ V̂ πθ (x). However, unbiased estimates can have large variance. Though it is possible to reduce variance through the critic, as we will show experimentally, this is not enough to counter the high variance due to products of IS ratios. This leads us to consider trading off bias with variance [43].
Since we postulate that the products of IS ratios lead to high variance, we might seek to control the number of IS ratios in the estimate. We briefly introduce Taylor expansion policy optimization (TayPO) [29], a natural framework to control for the number of IS ratios in the value estimate.
Taylor expansions of value functions. Consider the value function V πθ (x0) as a function of πθ. Using the idea of Taylor expansions, we can express V πθ (x0) as a sum of polynomials of πθ − µ. We start by defining the K th-order increment as Uπθ0 (x0) = V
µ(x0), which does not contain any IS ratio (zeroth order); and for K ≥ 1,
UπθK (x0) := Eµ [ ∞∑ t1=0 ∞∑ t2=t1+1 ... ∞∑ tK=tK−1+1 γtK ( ΠKi=1(ρ θ ti − 1) ) Qµ(xtK , atK )︸ ︷︷ ︸
Û πθ K (x0)
] . (6)
Intuitively, the K th-order increment only contains product of K IS ratios. Equation (6) also yields a natural sample-based estimate ÛπθK (x0), which we will discuss later. The K
th-order Taylor expansion is defined as the partial sum of increments V πθK (x0) := ∑K k=0 U πθ k (x0). Since V πθ K (x0) consists of products of up to K IS ratios, it is effectively the K th-order Taylor expansion of the value function. The properties are summarized as follows.
Proposition 3.3. (Adapted from Theorem 2 of [29].) Define ∥π − µ∥1 := maxx ∑
a |π(a|x) − µ(a|x)|. Let C be a constant and ε = 1−γγ . Then the following holds for all K ≥ 0,
V πθ (x0, g) = V πθ K (x0, g)︸ ︷︷ ︸
K-th order expansion
+C(∥πθ − µ∥1 /ε) K+1︸ ︷︷ ︸
residual
, (7)
If ∥πθ − µ∥1 < ε, then V πθ (x0) = limK→∞ V πθ K (x0) = Eµ [ ∑∞ k=0 U πθ k (x0)].
Sample-based estimates to Taylor expansions of value functions. As shown in Equation (6), ÛπθK (x0) is an unbiased estimate to U πθ K (x0). We can naturally define the sample-based estimate to the K th-order Taylor expansion, called the TayPO-K estimate,
V̂ πθK (x0) := K∑ k=0 Ûπθk (x0). (8)
The expression of Ûπθ (x0) contains O(TK) terms if the trajectory is of length T . Please refer to Appendix E for further details on computing the estimates in linear time O(T ) with sub-sampling. Note that V̂ πθK (x0) is a sample-based estimate whose bias against the value function is controlled by the residual term which decays exponentially when πθ and µ are close. Similar to how we derived the unbiased estimate∇mθ V̂ πθ DR (x), we can differentiate through the TayPO-K value estimate to produce estimates to higher-order derivatives∇mθ V̂ πθ K (x0).
Bias and variance of Hessian estimates. TayPO-K trades-off bias and variance with choices of K. To understand the variance, note that TayPO-K limits the number multiplicative IS ratios to be K. Though it is difficult to compute the variance, we argue that the variance generally increases with K as the number of IS ratios increase [26, 27, 44]. We characterize the bias of TayPO-K as follows.
Proposition 3.4. Assume (A.1) and (A.2) hold. Also assume ∥πθ − µ∥1 ≤ ε = (1− γ)/γ. For any tensor x, define ∥x∥∞ := maxi |x[i]|. The K th-order TayPO objective produces the following bias in estimating high-order derivatives,∥∥∥Eµ [∇mθ V̂ πθK ] (x0)−∇mθ V πθ (x0)∥∥∥∞ ≤ ∞∑ k=K+1 ∥∇mθ U πθ k (x0)∥∞ . (9)
Hence the upper bound for the bias decreases as K increases. Importantly, when on-policy µ = πθ, the K th-order TayPO objective preserves up to K th-order derivatives for any K ≥ 0,
Eµ [ ∇mθ V̂ πθ K (x0) ] = ∇mθ V πθ (x0),∀m ≤ K. (10)
Though IS ratios ρθt evaluate to 1 when on-policy, they maintain the parameter dependencies in differentiations. As such, higher-order expansions contains products of IS ratios of higher orders, and maintains the high-order dependencies on parameters more accurately. There is a clear trade-off between bias and variance mediated by K. When K increases, the higher-order derivatives are maintained more accurately in expectation, leading to less bias. However, the variance increases too.
Recovering prior work as special cases. Recently, [24] proposed a low variance curvature (LVC) Hessian estimate. This estimate is equivalent to the TayPO-K estimate with K = 1. As also noted by [24], their objective function bears similarities to first-order policy search algorithms [28, 45, 46], which have in fact been interpreted as first-order special cases of ∇θV̂ πθK (x) with K = 1 [29]. Importantly, based on Proposition 3.4, the LVC estimate only maintains the first-order dependency perfectly but introduces bias when approximating the Hessian, even when on-policy.
Limitations. Though the above framework interprets a large number of prior methods as special cases, it has some limitations. For example, the derivation of Hessian estimates based on the DR estimate (Proposition 3.2) involves estimates of the value function V̂ πθDR (xt). In practice, when near on-policy πθ ≈ µ, one might replace the DR estimate V̂ πθDR by other value function estimate, such as plain cumulative sum of returns or TD(λ), V̂ πθTD(λ) [22, 23] in Eqn (4). As such, the practical implementation might not strictly adhere to the conceptual framework. In addition, TMAML [20] is not incorporated as part of this framework: we show in Appendix B that the control variate introduced by TMAML in fact biases the overall estimate.

heading: 4 From Hessian estimates to meta-gradient estimates

text: A practical desiteratum for meta-gradient estimates is that it can be implemented in a scalable way using auto-differentiation frameworks [30, 19]. Below, we discuss how this can be achieved.

heading: 4.1 Auto-differentiating off-policy evaluation estimates for Hessian estimates

text:
In practice, we seek Hessian estimates that could be implemented with the help of an established framework, such as auto-differentiation libraries [30, 19]. Now we discuss how to conveniently implement ideas discussed in the previous section.
Algorithm 2 Example: an off-policy evaluation subroutine for computing the DR estimate Require: Inputs: Trajectory (xt, at, rt)Tt=0, target policy πθ, behavior policy µ, (optional) critic Q.
Initialize V̂ = Q(xT , πθ(xT ), g). for t = T − 1, . . . 0 do
Compute IS ratio ρθt = πθ(at|xt)/µ(at|xt). Recursion: V̂ ← Q(xt, πθ(at), g) + γρθt (rt + γQ(xt+1, πθ(xt+1), g)−Q(xt, at)) + γρθt V̂ .
end for Output V̂ as an estimate to V πθ (x0, g).
Auto-differentiating the estimates. We can abstract the off-policy evaluation as a function eval(D, θ) that takes in some data D and parameter θ, and outputs an estimate for V πθ (x0, g). In particular, D includes the trajectories and θ is input via the policy πθ. As an example, Algorithm 2 shows that for the doubly-robust estimator, the dependency of the estimator on θ is built through the recursive computation by eval(D, θ). In fact, if we implement Algorithm 2 with an auto-differentiation framework (for example, Tensorflow or PyTorch [17–19]), the higher-order dependency of V̂ on θ is maintained through the computation graph. We can compute the Hessian estimate by directly differentiating through the function output∇Kθ V̂ = ∇Kθ eval(D, θ) ≈ ∇Kθ V πθ (x0, g). In Appendix F we show how the estimates could be conveniently implemented, as in many deep RL agents (see, e.g., [44, 5, 47, 48]). We also show concrete ways to compute estimates with TayPO-K based on [29].

heading: 4.2 Practical implementations of meta-gradient estimates

text:
In Equation (2), we write the meta-gradient estimate as a product between an Hessian estimate and a policy gradient. In practice, the meta-gradient estimate is computed via Hessian-vector products to avoid explicitly computing the Hessian estimate of size D2. As such, the meta-gradient estimate could be computed by auto-differentiating through a scalar objective. See Appendix F for details.
Bias and variance of meta-gradient estimates. Intuitively, the bias and variance of the Hessian estimates translate into bias and variance of the downstream meta-gradient estimates. Prior work has showed that low variance of meta-gradient estimates lead to faster convergence [49, 50]. However, it is not clear how the bias (such as bias introduced by the Hessian estimates, or the bias due to correlated estimates) theoretically impacts the convergence. We will study empirically the effect of bias and variance in experiments, and leave further theoretical study to future work.

heading: 5 Experiments

text: We now carry out several empirical studies to complement the framework developed above. In Section 5.1, we use a tabular example to investigate the bias and variance trade-offs of various estimates, to assess the validity of our theoretical insights. We choose a tabular example because it is straightforward to compute exact higher-order derivatives of value functions and make comparison. In Section 5.2.1 and Section 5.2.2, we apply the new second-order estimate in high-dimensional meta-RL experiments, to assess the potential performance gains in a more practical setup. Though we can compute TayPO-K order estimates for general K, we focus on K ≤ 2 in experiments. Below, we also address TayPO-1 and TayPO-2 estimates as the first and second-order estimates respectively.

heading: 5.1 Investigating the bias and variance trade-off of different estimates

text:
We study the bias and variance trade-off of various estimates using a tabular exmaple. We consider random MDPs with |X | = 10, |A| = 5. The transition matrix of the MDPs are sampled from a Dirichlet distribution. See Appendix G for further details. The policy πθ is parameterized as πθ(a|x) = exp(θ(x, a))/ ∑ b exp(θ(x, b)). The behavior policy µ is uniform and θ is initialized so that θ(x, a) = log π(a|x) where π = (1− ε)µ+ επd for some deterministic policy πd and parameter ε ∈ [0, 1]. The hyper-parameter ε measures the off-policyness. In this example, there is no task variable. As performance metrics, we measure the component-wise correlation between the true derivatives∇Kθ V πθ (x0) (computed via an oracle) where x0 is a fixed starting state, and its estimates
∇Kθ V̂ πθ (x0), as commonly used in prior work [21–23]. The estimates are averaged across N samples. We study the effect of different choices of the off-policy estimate V̂ πθ , as a function of off-policyness ε and sample size N . We report results with mean± standard error over 10 seeds.
Effect of off-policyness. In Figure 1(a) and (b), we let N = 1000 and show the performance as a function of ε. When ε increases (more off-policy), the accuracy measures of most estimates decrease. This is because off-policyness generally increases both bias and variance (large IS ratios ρθ). At this level of sample size, the performance of the second-order estimate degrades more slowly than other estimates, making it the dominating estimate across all values of ε. We also include truncated DR estimate using min(ρθt , ρ) as a baseline inspired by V-trace [44, 47]. The truncation is motivated by controlling the variance of the overall estimate. However, the estimate is heavily biased and does not perform well unless πθ ≈ µ. See Appendix F for more.
Consider the case when ε = 0 and the setup is on-policy. When estimating the policy gradient, almost all estimates converge to the optimal level of accuracy, except for the step-wise IS estimate, where the variance still renders the performance sub-optimal. However, when estimating the Hessian matrix, the first-order estimate converges to a lower accuracy than both second-order estimate and DR estimate. This validates our theoretical analysis, as both the second-order estimate and DR estimate are unbiased when on-policy.
Effect of sample size. Figure 1(c) shows the accuracy measures as a function of sample size N when fixing ε = 0.5. Note that since sample sizes directly influence the variance, the results show the variance properties of different estimates. When N is small, the first-order estimate dominates due to smaller variance; however, when N increases, the first-order estimate is surpassed by the second-order estimate, due to higher bias. For more results and ablation on high-dimensional environments, see Appendix G.

heading: 5.2 High-dimensional meta-RL problems

text:
Next, we study the practical gains entailed by the second-order estimate in high-dimensional meta-RL problems. We first introduce a few important algorithmic baselines, and how the second-order estimate is incorporated into a meta-RL algorithm.
Baseline algorithms. All baseline algorithms use plain stochastic gradient ascent as the inner loop optimizer: θ′ = θ + η∇θV̂ πθ (x0, g) where g ∼ pG is a sampled goal and V̂ πθ (x0, g) is a sample-based estimate of policy gradients averaged over n trajectories. Different algorithms differ in how the inner loop loss is implemented, such that auto-differentiation produces different Hessian estimates: these include TRPO-DiCE [21–23], TRPO-MAML [4], TRPO-FMAML [4], TRPO-EMAML [32, 33]. Please refer to Appendix G for further details. Note despite the name, TRPO-DiCE baseline uses DR estimate to estimate Hessians. We implement the second-order
estimate using the proximal meta policy search (PROMP) [24] as the base algorithm. By default, PROMP uses the first-order estimate. Our new algorithm is named PROMP-TayPO-2.

heading: 5.2.1 Continuous control in 2D environments

text:
Environment. We consider a simple 2-D navigation task introduced in [24]. The state xt is the coordinate of a ball placed inside a room, the action at is the direction in which to push the ball. The goal g ∈ R4 is an one-hot encoding of which corner of the room contains positive rewards. With 3 adaptation steps, the agent should ideally navigate to the desired corner indicated by g.
Training performance. In Figure 1(d), we show the performance curves of various baseline algorithms. Though MAML and EMAML learns quickly during the initial phase of training, they ultimately become unstable. TRPO-DiCE generally underperforms other methods potentially due to bias. On the other hand, FMAML and PROMP both reduce variance at the cost of bias, but they both achieve a slightly lower level of performance compared to PROMP-TayPO-2. Overall, PROMP-TAyPO-2 achieves much more stable training curves compared to others, potentially owing to the better bias and variance trade-off in the Hessian estimates.

heading: 5.2.2 Large scale locomotion experiments

text:
Environments. We consider the set of meta-RL tasks based on simulated locomotion in MuJoCo [51]. Across these tasks, the states xt consist of robotic sensory inputs, the actions at are torque controls applied to the robots. The task g is defined per environment: for example, in random goal environment, g ∈ R2 is a random 2-d goal location that the robot should aim to reach. Experiment setup. We adapt the open source code base by [24] and adopt exactly the same experimental setup as [24]. At each iteration, the agent samples n = 40 task variables. For each task, the agent carries out K = 1 adaptation computed based on B = 20 trajectories sampled from the environment, each of length T = 100. See Appendix G for further details on the architecture and other hyper-parameters. We report averaged results with mean± std over 10 seeds.
Results. The training performance of different algorithmic baselines are shown in Figure 2. Comparing TRPO-DiCE, TRPO-MAML and PROMP: we see that the results are compatible with those reported in [24], where PROMP outperforms TRPO-MAML, while TRPO-DiCE generally performs the worst potentially due to high variance in the gradient estimates. As a side observation, TRPO-FMAML generally underperforms TRPO-MAML, which implies the necessity of carrying out approximations to the Hessian matrix beyond the identity matrix. PROMP-TayPO-2 slightly outperforms PROMP in a few occasions, where the new algorithm achieves slightly faster learning speed and sometimes higher final performance. However, overall, we see that the empirical gains are marginal. This implies that under the default setup of these meta-RL experiments, the variance might be a major factor in gradient estimates, and the first-order estimate is near optimal compared to other estimates. See Appendix G for additional experiments.

heading: 6 Conclusion

text:
We have unified a number of important prior work on meta-gradient estimations for model-agnostic meta-RL. Our analysis entails the derivations of prior methods based on the unifying framework of
differentiating through off-policy evaluation estimates. This framework provides a principled way to reason about the bias and variance in the higher-order derivative estimates of value functions, and opens the door to a new family of estimates based on novel off-policy evaluation estimates. As an important example, we have theoretically and empirically studied the properties of the family of TayPO-based estimates. It is worth noting that this framework further suggests any future advances in off-policy evaluation could be conveniently imported into potential improvements in meta-gradient estimates. As future work, we hope to see the applications of such principled estimates in broader meta-RL applications.
Acknowledgement. The authors thank David Abel for reviewing an early draft of this work and providing very useful feedback. Yunhao and Tadashi are thankful for the Scientific Computation and Data Analysis section at the Okinawa Institute of Science and Technology (OIST), which maintains a cluster we used for many of our experiments.

title: Unifying Gradient Estimators for Meta-Reinforcement Learning via Off-Policy Evaluation


summaries: The paper proposes MCM ~ algorithm to perform bidirectional compression in distributed setting. The authors claim similar convergence guarantees as vanilla setting.  They introduce the notion of sending different models to different clients while keeping the global model preserved.

summaries: This paper proposed a new method: MCM for bidirectional compression in federated learning. Convergence analysis showed that linear convergence can be achieved with the proposed method. Empirical results also show the method could achieve non-trivial performance with high compression rate.

summaries: The paper proposes a bidirectional compression algorithm for parameter-server-based distributed training. Two variants, MCM and Rand-MCM, are proposed to achieve the same improvement as the uplink-compression-only compression algorithm, where the downlink-compression-only impacts local models while the global model is preserved. The Rand-MCM also introduces the concept of (simulated) random smoothing via model compression.

summaries: In this paper, the authors propose MCM, a new algorithm achieves bidirectional compression with better convergence rate. Theoretical analysis is provided. The empirical results show good performance.

abstractText: We develop a new approach to tackle communication constraints in a distributed learning problem with a central server. We propose and analyze a new algorithm that performs bidirectional compression and achieves the same convergence rate as algorithms using only uplink (from the local workers to the central server) compression. To obtain this improvement, we design MCM, an algorithm such that the downlink compression only impacts local models, while the global model is preserved. As a result, and contrary to previous works, the gradients on local servers are computed on perturbed models. Consequently, convergence proofs are more challenging and require a precise control of this perturbation. To ensure it, MCM additionally combines model compression with a memory mechanism. This analysis opens new doors, e.g. incorporating worker dependent randomized-models and partial participation.

authors:
Martín Abadi
Paul Barham
Jianmin Chen
Zhifeng Chen
Andy Davis
Jeffrey Dean
Matthieu Devin
Sanjay Ghemawat
Geoffrey Irving
Michael Isard
Manjunath Kudlur
Josh Levenberg
Rajat Monga
Sherry Moore
Derek G. Murray
Benoit Steiner
Paul Tucker
Vijay Vasudevan
Pete Warden
Martin Wicke
Yuan Yu
Xiaoqiang Zheng

title: TensorFlow: a system for large-scale machine learning

venue: In Proceedings of the 12th USENIX conference on Operating Systems Design and Implementation,

authors:
Naman Agarwal
Ananda Theertha Suresh
Felix Xinnan X Yu
Sanjiv Kumar
Brendan McMahan

title: cpSGD: Communication-efficient and differentially-private distributed SGD

venue: Advances in Neural Information Processing Systems

authors:
Dan Alistarh
Demjan Grubic
Jerry Li
Ryota Tomioka
Milan Vojnovic

title: QSGD: Communication-Efficient SGD via Gradient Quantization and Encoding

venue: Advances in Neural Information Processing Systems,

title: Self-concordant analysis for logistic regression

venue: Electronic Journal of Statistics,

title: Large-Scale Machine Learning with Stochastic Gradient Descent

venue: Proceedings of COMPSTAT’2010,

authors:
Sebastian Caldas
Sai Meher Karthik Duddu
Peter Wu
Tian Li
Jakub Konečný
H. Brendan McMahan
Virginia Smith
Ameet Talwalkar

title: LEAF: A Benchmark for Federated Settings

authors:
Rich Caruana
Thorsten Joachims
Lars Backstrom

title: KDD-Cup 2004: results and analysis

venue: ACM SIGKDD Explorations Newsletter,

authors:
Chih-Chung Chang
Chih-Jen Lin

title: LIBSVM: A library for support vector machines

venue: ACM Transactions on Intelligent Systems and Technology,

authors:
Igor Colin
Aurelien Bellet
Joseph Salmon
Stéphan Clémençon

title: Gossip Dual Averaging for Decentralized Optimization of Pairwise Functions

venue: In International Conference on Machine Learning,

authors:
Yann Fraboni
Richard Vidal
Marco Lorenzi

title: Free-rider attacks on model aggregation in federated learning

authors:
Eduard Gorbunov
Dmitry Kovalev
Dmitry Makarenko
Peter Richtarik

title: Linearly Converging Error Compensated SGD

title: Communication Energy efficiency and economic fallacies

venue: Energy Policy,

title: A data-driven statistical model for predicting the critical temperature of a superconductor

venue: Computational Materials Science,

authors:
Samuel Horváth
Peter Richtárik

title: A Better Alternative to Error Feedback for Communication- Efficient Distributed Learning

authors:
Samuel Horváth
Dmitry Kovalev
Konstantin Mishchenko
Sebastian Stich
Peter Richtárik

title: Stochastic Distributed Learning with Gradient Quantization and Variance Reduction

venue: [math],

authors:
Peter Kairouz
H. Brendan McMahan
Brendan Avent
Aurélien Bellet
Mehdi Bennis
Arjun Nitin Bhagoji
Keith Bonawitz
Zachary Charles
Graham Cormode
Rachel Cummings
Rafael G.L. D’Oliveira
Salim El Rouayheb
David Evans
Josh Gardner
Zachary Garrett
Adrià Gascón
Badih Ghazi
Phillip B. Gibbons
Marco Gruteser
Zaid Harchaoui
Chaoyang He
Lie He
Zhouyuan Huo
Ben Hutchinson
Justin Hsu
Martin Jaggi
Tara Javidi
Gauri Joshi
Mikhail Khodak
Jakub Konečný
Aleksandra Korolova
Farinaz Koushanfar
Sanmi Koyejo
Tancrède Lepoint
Yang Liu
Prateek Mittal
Mehryar Mohri
Richard Nock
Ayfer Özgür
Rasmus Pagh
Mariana Raykova
Hang Qi
Daniel Ramage
Ramesh Raskar
Dawn Song
Weikang Song
Sebastian U. Stich
Ziteng Sun
Ananda Theertha Suresh
Florian Tramèr
Praneeth Vepakomma
Jianyu Wang
Li Xiong
Zheng Xu
Qiang Yang
Felix X. Yu
Han Yu
Sen Zhao

title: Advances and Open Problems in Federated Learning

authors:
Jakub Konečný
H. Brendan McMahan
Daniel Ramage
Peter Richtárik

title: Federated Optimization: Distributed Machine Learning for On-Device Intelligence

authors:
Zhize Li
Dmitry Kovalev
Xun Qian
Peter Richtarik

title: Acceleration for Compressed Gradient Descent in Distributed and Federated Optimization

authors:
Xiaorui Liu
Yao Li
Jiliang Tang
Ming Yan

title: A Double Residual Compression Algorithm for Efficient Distributed Learning

authors:
Laurens van der Maaten
Geoffrey Hinton

title: Visualizing Data using t-SNE

venue: Journal of Machine Learning Research,

authors:
Horia Mania
Xinghao Pan
Dimitris Papailiopoulos
Benjamin Recht
Kannan Ramchandran
Michael I. Jordan

title: Perturbed Iterate Analysis for Asynchronous Stochastic Optimization

authors:
Brendan McMahan
Eider Moore
Daniel Ramage
Seth Hampson
Blaise Aguera y Arcas

title: Communication-Efficient Learning of Deep Networks from Decentralized Data

venue: In Artificial Intelligence and Statistics,

authors:
Konstantin Mishchenko
Eduard Gorbunov
Martin Takáč
Peter Richtárik

title: Distributed Learning with Compressed Gradient Differences

authors:
Constantin Philippenko
Aymeric Dieuleveut

title: Artemis: tight convergence guarantees for bidirectional compression in Federated Learning

authors:
Herbert Robbins
Sutton Monro

title: A Stochastic Approximation Method

venue: Annals of Mathematical Statistics,

authors:
Felix Sattler
Simon Wiedemann
Klaus-Robert Müller
Wojciech Samek

title: Robust and Communication-Efficient Federated Learning From Non-i.i.d

venue: Data. IEEE Transactions on Neural Networks and Learning Systems,

authors:
Kevin Scaman
Francis Bach
Sebastien Bubeck
Laurent Massoulié
Yin Tat Lee

title: Optimal Algorithms for Non-Smooth Distributed Optimization in Networks

authors:
Frank Seide
Amit Agarwal

title: CNTK: Microsoft’s Open-Source Deep-Learning Toolkit

venue: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD

authors:
Frank Seide
Hao Fu
Jasha Droppo
Gang Li
Dong Yu

title: 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns

venue: In Fifteenth Annual Conference of the International Speech Communication Association. Citeseer,

authors:
Sebastian U Stich
Sai Praneeth Karimireddy

title: The error-feedback framework: Better rates for SGD with delayed gradients and compressed updates

authors:
Sebastian U Stich
Jean-Baptiste Cordonnier
Martin Jaggi

title: Sparsified SGD with Memory

authors:
Hanlin Tang
Xiangru Lian
Ming Yan
Ce Zhang
Ji Liu

title: $D^2$: Decentralized Training over Decentralized Data

authors:
Hanlin Tang
Chen Yu
Xiangru Lian
Tong Zhang
Ji Liu

title: DoubleSqueeze: Parallel Stochastic Gradient Descent with Double-pass Error-Compensated Compression

authors:
Paul Vanhaesebrouck
Aurélien Bellet
Marc Tommasi

title: Decentralized Collaborative Learning of Personalized Models over Networks

authors:
Jianqiao Wangni
Jialei Wang
Ji Liu
Tong Zhang

title: Gradient Sparsification for Communication-Efficient Distributed Optimization

authors:
Jiaxiang Wu
Weidong Huang
Junzhou Huang
Tong Zhang

title: Error Compensated Quantized SGD and its Applications to Large-scale Distributed Optimization

authors:
An Xu
Zhouyuan Huo
Heng Huang

title: Training Faster with Compressed Gradient

authors:
Shuai Zheng
Ziyue Huang
James Kwok

title: Communication-Efficient Distributed Blockwise Momentum SGD with Error-Feedback

venue: In Advances in Neural Information Processing Systems,

heading: 1 Introduction

text:
Large scale distributed machine learning is widely used in many modern applications [1, 6, 28]. The training is distributed over a potentially large number N of workers that communicate either with a central server [see 17, 22, on federated learning], or using peer-to-peer communication [9, 34, 32].
In this work, we consider a setting using a central server that aggregates updates from remote nodes. Formally, we have a number of features d ∈ N∗, and a convex cost function F : Rd → R. We want to solve the following distributed convex optimization problem using stochastic gradient algorithms [25, 5]: minw∈Rd F (w) with F (w) = 1N ∑N i=1 Fi(w), where (Fi) N i=1 is a local risk function (empirical risk or expected risk in a streaming framework). This applies to both instances of distributed and federated learning.
An important issue of those frameworks is the high communication cost between the workers and the central server [16, Sec. 3.5]. This cost is a concern from several points of view. First, exchanging information can be the bottleneck in terms of speed. Second, the data consumption and the bandwidth usage of training large distributed models can be problematic; and furthermore, the energetic and environmental impact of those exchanges is a growing concern. Over the last few years, new algorithms were introduced, compressing messages in the upload communications (i.e., from remote devices to the central server) in order to reduce the size of those exchanges [29, 3, 36, 2, 35, 31, 30, 23, 18]. More recently, a new trend has emerged to also compress the downlink communication: this is bidirectional compression.
The necessity for bidirectional compression can depend on the situation. For example, a single uplink compression could be sufficient in asymmetric regimes in which broadcasting a message to N workers (“one to N”) is faster than aggregating the information coming from each node (“N to one”). However, in other regimes, e.g. with few machines, where the bottleneck is the transfer time of a heavy model (up to several GB in modern Deep Learning architectures) the downlink communication cannot be disregarded, as the upload and download speed are of the same order [24].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Furthermore, in a situation in which participants have to systematically download an update (e.g., on their smartphones) to participate in the training, participants would prefer to receive a small size update (compressed) rather than a heavier one. To encompass all situations, we consider algorithms for which the information exchanged is compressed in both directions.
To perform downlink communication, existing bidirectional algorithms [33, 38, 26, 19, 24, 14, 37, 11] first aggregate all the information they have received, compress them and then carry out the broadcast. Both the main “global” model and the “local” ones perform the same update with this compressed information. Consequently, the model hold on the central server and the one used on the local workers (to query the gradient oracle) are identical. However, this means that the model on the central server has been artificially degraded: instead of using all the information it has received, it is updated with the compressed information.
Here, we focus on preserving (instead of degrading) the central model: the update made on its side does not depend on the downlink compression. This implies that the local models are different from the central model. The local gradients are thus measured on a “perturbed model” (or “perturbed iterate”): such an approach requires a more involved analysis and the algorithm must be carefully designed to control the deviation between the local and global models [21]. For example, algorithms directly compressing the model or the update would simply not converge.
We propose MCM - Model Compression with Memory - a new algorithm that 1) preserves the central model, and 2) uses a memory scheme to reduce the variance of the local model. We prove that the convergence of this method is similar to the one of algorithms using only unidirectional compression. Potential Impact. Proposing an analysis that handles perturbed iterates is the key to unlock three major challenges of distributed learning run with bidirectionally compressed gradients. First, we show that it is possible to improve the convergence rate by sending different randomized models to the different workers, this is Rand-MCM. Secondly, this analysis also paves the way to deal with partially participating machines: the adaptation of Rand-MCM to this framework is straightforward; while adapting existing algorithms [26] to partial participation is not practical. Thirdly, this framework is also promising in terms of business applications, e.g., in the situation of learning with privacy guarantees and with a trusted central server. We detail those three possible extensions in Subsection 4.1.
Broader impact. This work is aligned with a global effort to make the usage of large scale Federated Learning sustainable by minimizing its environmental impact. Though the impact of such algorithms is expected to be positive, at least on environmental concerns, cautiousness is still required, as a rebound effect may be observed [12]: having energetically cheaper and faster algorithms may result in an increase of such applications, annihilating the gain made by algorithmic progress.
Contributions. We make the following contributions: 1. We propose a new algorithm MCM, combining a memory process to the “preserved” update. To
convey the key steps of the proof, we also introduce an auxiliary hypothetical algorithm, Ghost. 2. For those algorithms, we carefully control the variance of the local models w.r.t. the global one.
We provide a contraction equation involving the control on the local model’s variance and show that MCM achieves the same rate of convergence as single compression in strongly-convex, convex and non-convex regimes. We give a comparisons of MCM’s rates with existing algorithms in Table 2. 3. We propose a variant, Rand-MCM incorporating diversity into models shared with the local workers and show that it improves convergence for quadratic functions.
This is the first algorithm for double compression to focus on a preserved central model. We underline, both theoretically and in practice, that we get the same asymptotic convergence rate for simple and double compression - which is a major improvement. Our approach is one of the first to allow for worker dependent model, and to naturally adapt to worker dependent compression levels.
The rest of the paper is organized as follows: in Section 2 we present the problem statement and introduce MCM and Rand-MCM. Theoretical results on these algorithms are successively presented in Sections 3 and 4. Finally, we present experiments supporting the theory in Section 5.

heading: 2 Problem statement

text:
We consider the minimization problem described in Section 1. In the convex case, we assume there exists an optimal parameter w∗, and denote F∗ = F (w∗). We use ‖·‖ to denote the Euclidean norm. To solve this problem, we rely on a stochastic gradient descent (SGD) algorithm. A stochastic gradient gik+1 is provided at iteration k in N to the device i in J1, NK. This gradient oracle can be
computed on a mini-batch of size b. This function is then evaluated at point wk. In the classical centralized framework (without compression), for a learning rate γ, SGD corresponds to:
wk+1 = wk − γ 1
N N∑ i=1 gik+1(wk) . (1)
We now describe the framework used for compression.

heading: 2.1 Bidirectional compression framework

text:
Bidirectional compression consists in compressing communications in both directions between the central server and remote devices. We use two different compression operators, respectively Cup and Cdwn to compress the message in each direction. Roughly speaking, the update in eq. (1) becomes:
wk+1 = wk − γCdwn ( 1
N N∑ i=1 Cup(gik+1(wk)) ) .
However, this approach has a major drawback. The central server receives and aggregates information 1 N ∑N i=1 Cup(gik+1(wk)). But in order to be able to broadcast it back, it compresses it, before applying the update. We refer to this strategy as the “degraded update” approach. Its major advantage is simplicity, and it was used in all previous papers performing double compression. Yet, it appears to be a waste of valuable information. In this paper, we update the global model wk+1 independently of the downlink compression:{
wk+1 = wk − γ 1N ∑N i=1 Cup ( gik+1(ŵk) ) .
ŵk+1 = Cdwn(wk+1) (2)
However, bluntly compressing wk+1 in eq. (2) hinders convergence, thus the second part of the update needs to be refined by adding a memory mechanism. We now describe both communication stages of the real MCM, which is entirely defined by the following uplink and downlink equations.
Downlink Uplink{ Ωk+1 = wk+1 −Hk , ŵk+1 = Hk + Cdwn(Ωk+1) Hk+1 = Hk + αdwnCdwn(Ωk+1).  ∀i ∈ J1, NK,∆ik = gik+1(ŵk)− hik wk+1 = wk − γN ∑N i=1 Cup(∆ik) + hik
hik+1 = h i k + αupCup(∆ik).
(3)
Downlink Communication. We introduce a downlink memory term (Hk)k, which is available on both workers and central server. The difference Ωk+1 between the model and this memory is compressed and exchanged, then the local model is reconstructed from this information. The memory is then updated as defined on left part of eq. (3), with a learning rate αdwn.
Introducing this memory mechanism is crucial to control the variance of the local model ŵk+1. To the best of our knowledge MCM is the first algorithm that uses such a memory mechanism for downlink compression. This mechanism was introduced by Mishchenko et al. [23] for the uplink compression but with the other purpose of mitigating the impact of heterogeneity, while we use it here to avoid divergence of the local model’s variance.
Uplink Communication. The motivation to introduce an uplink memory term hik for each device i ∈ J1, NK is different, and better understood. Indeed, for the uplink direction, this mechanism is only necessary (and then crucial) to handle heterogeneous workers [i.e., with different data distributions, see e.g. 24]. Here, the difference ∆ik between the stochastic gradient g i k+1 at the local model ŵk (as defined in eq. (3)) and the memory term is compressed and exchanged. The memory is then updated as defined on right part of eq. (3) with a rate αdwn. Remark 1 (Rate αdwn). It is necessary to use αdwn < 1. Otherwise, the compression noise tends to propagate and is amplified, because of the multiplicative nature of the compression. In Figure 1 we compare MCM, with 3 other strategies: compressing only the update, compressing wk − ŵk−1, (i.e., αdwn = 1), and compressing the model (i.e., Hk = 0), showing that only MCM converges. Remark 2 (Memory vs Error Feedback). Error feedback is another technique, introduced by Seide et al. [29]. In the context of double compression, it has been shown to improve convergence for a restrictive class of contracting compression operators (which are generally biased) by Zheng et al. [38], Tang et al. [33]. However, we note several differences to our approach. (1) For unbiased operators - as considered in Dore, it did not lead to any theoretical improvement [Remark 2 in Sec. 4.1., 19]. (2) Moreover, only a fraction (namely (1 + ωdwn)−1) of the “error” wk+1 − ŵk+1 can be preserved in the EF term (see line 18 in algo 1 in Liu et al.). It is thus impossible to recover the central preserved model as a function of the degraded model and the EF term. (3) [38] consider a biased operator and the same compression level for uplink and downlink compression. They also rely on stronger assumptions on the gradient (uniformly bounded) and only tackle the homogeneous case.
In Table 1 we summarize the main algorithms for compression in distributed training. As downlink communication can be more efficient than uplink, we consider distinct operators Cdwn, Cup and allow the corresponding compressions levels to be distinct: those quantities are defined in Assumption 1. Assumption 1. There exists constants ωup , ωdwn ∈ R∗+, such that the compression operators Cup and Cdwn satisfy the two following properties for all w in Rd: E[Cup/dwn(w)] = w, and E[‖Cup/dwn(w)− w‖2] ≤ ωup/dwn‖w‖2. The higher is ω, the more aggressive the compression is. We only consider unbiased operators, that encompass sparsification, quantization and sketching. References and a discussion on those operators, and possible extensions of our results to biased operators are provided in Appendix A.1. Remark 3 (Related work on Perturbed iterate analysis). The theory of perturbed iterate analysis was introduced by Mania et al. [21] to deal with asynchronous SGD. More recently, it was used by Stich and Karimireddy [30], Gorbunov et al. [11] to analyze the convergence of algorithms with uplink compressions, error feedback and asynchrony. Using gradients at randomly perturbed points can also be seen as a form of randomized smoothing [27], a point we discuss in Appendix A.2.
2.2 The randomization mechanism, Rand-MCM In this subsection, we describe the key feature introduced in Rand-MCM: randomization. It consists in performing an independent compression for each device instead of performing a single one for all of them. As a consequence, each worker holds a different model centered around the global one. This introduces some supplementary randomness that stabilizes the algorithm. Formally, we will consider N mutually independent compression operators Cdwn,i instead of a single one Cdwn, and the central server will send to the device i at iteration k + 1 the compression of the difference between its model and the local memory on worker i: Cdwn,i(wk+1 −Hik). The tradeoffs associated with this modification are discussed in Section 4.
The pseudocode of Rand-MCM is given in Algorithm 1 in Appendix A. It incorporates all components described above: 1) the bidirectional compression, 2) the model update using the non-degraded point, 3) the two memories, 4) the up and down compression operators, 5) the randomization mechanism.

heading: 3 Assumptions and Theoretical analysis

text:
We make standard assumptions on F : Rd → R. We first assume that the loss function F is smooth. Assumption 2 (Smoothness). F is twice continuously differentiable, and is L-smooth, that is for all vectors w1, w2 in Rd: ‖∇F (w1)−∇F (w2)‖ ≤ L‖w1 − w2‖. Results in Section 3 are provided in a convex, strongly-convex and non-convex setting. Assumption 3 (Strong convexity). F is µ-strongly convex (or convex if µ = 0), that is for all vectors w1, w2 in Rd: F (w2) ≥ F (w1) + (w2 − w1)T∇F (w1) + µ2 ‖w2 − w1‖22 .
Next, we present the assumption on the stochastic gradients.
Assumption 4 (Noise over stochastic gradients computation). The noise over stochastic gradients for a mini-batch of size b, is uniformly bounded: there exists a constant σ ∈ R+, such that for all k in N, for all i in J1, NK and for all w in Rd we have: E[‖gik(w)−∇F (w)‖2] ≤ σ2/b. We here provide guarantees of convergence for MCM. MCM incorporates an uplink memory term, designed to handle heterogeneous workers. To highlight our main contributions, that concerns the downlink compression, we present the results in the homogeneous setting, that is with Fi = Fj and αup = 0. Similar results (almost identical, up to constant numerical factors) in to the heterogeneous setting are described in Appendix G. Experiments are also performed on heterogeneous workers. We provide here convergence results in the strongly-convex, then convex case. Notations and settings. For k in N, we denote Υk = ‖wk −Hk−1‖2, and define Vk = E[‖wk − w∗‖2] + 32γLω2dwnE[Υk], which serves as Lyapunov function. Vk is composed of two terms: the first one controls the quadratic distance to the optimal model, and the second controls the variance of the local models ŵk. For both theorems, we choose αdwn = (8ωdwn)−1. We denote Φ(γ) := (1 + ωup) ( 1 + 64γLω2dwn ) .
Limit learning rate: There exists a maximal learning rate to ensure convergence. More specifically, we define γmax := min(γupmax, γ dwn max , γ Υ max), where γ up max := (2L (1 + ωup/N))
−1 corresponds to the classical constraint on the learning rate in the unidirectional regime [see 23, 24], γdwnmax := (8Lωdwn)
−1 is a similar constraint coming from the downlink compression, and γΥmax := ( 8 √ 2Lωdwn √ 8ωdwn + ωup/N )−1
is a combined constraint that arises when controlling the variance term Υ.1 Overall, this constraints are weaker than in the “degraded” framework [19, 24], in which γDoremax ≤ ( 8L(1 + ωdwn)(1 + ωup/N) )−1 . Especially, in the regime in which ωup,dwn → ∞ and ωdwn ' ωup ': ω, the maximal learning rate for MCM is (Lω3/2)−1, while it is (Lω2)−1 in [19, 24]. Our γmax is thus larger by a factor √ ω, see Table 2. We define L̃ such that γmax = (2L̃) −1.
Theorem 1 (Convergence of MCM in the homogeneous and strongly-convex case). Under Assumptions 1 to 4 with µ > 0, for k in N, for any sequence (γk)k≥0 ≤ γmax we have:
Vk ≤ (1− γkµ)Vk−1 − γkE [F (ŵk−1)− F (w∗)] + γ2kσ 2Φ(γk)
Nb , (4)
Consequently, (1) if σ2 = 0 (noiseless case), for γk ≡ γmax we recover a linear convergence rate: E[‖wk − w∗‖2] ≤ (1− γmaxµ)kV0; (2) if σ2 > 0, taking for all K in N, γK = 2/(µ(K + 1) + L̃), for the weighted Polyak-Ruppert average w̄K = ∑K k=1 λkwk−1/ ∑K k=1 λk, with λk := (γk−1) −1,
E [F (w̄K)− F (w∗)] ≤ µ+ 2L̃
4µK2 ‖w0 − w∗‖2 +
4σ2(1 + ωup)
µKNb
( 1 +
64Lω2dwn µK ln(µK + L̃)
) .
(5)
Limit Variance (Equation (4)). For a constant γ, the variance term (i.e., term proportional to σ2) in Equation (4) is upper bounded by γ 2σ2
Nb (1 + ωup)(1 + 64γLω 2 dwn). The impact of the downlink
compression is attenuated by a factor γ. As γ decreases, this makes the limit variance similar to the one of Diana, i.e., without downlink compression [23, Eq. 16 in Th. 2] and much lower than the variance for previous algorithms using double compression for which the variance scales quadratically with the compression constants as γ2σ2(1 + ωup)(1 + ωdwn)/N : (1) for Dore, see Corollary 1 in Liu et al. [19] (who indicate (1− ρ)−1 ≥ (1 +ωup/N)(1 +ωdwn)), (2) for Artemis see Table 2 and Th. 3 point 2 in [24], (3) for [11], see Theorem I.1. (with γD′1 ∝ γ 2σ2(1 + ωup)(1 + ωdwn)/N ).
Bound 5 has a quadratic dependence on ωdwn, but the corresponding term is divided by an extra factor K, the number of iterations. For example in experiments, for w8a using quantization with s = 20, we have ωdwn ' 17, and after only 50 epoch with a batch size b = 12, we have K ' 2500. Hence, the term ω2/K is vanishing through iterations and we asymptotically recover a rate of convergence equivalent to algorithms using unidirectional compression.
1The dependency in ω3/2 is similar to the one obtained by Horváth et al. [15] in unidirectional compression in the non-convex case (Theorem 4).
Convergence and complexity: With a decaying sequence of steps, we obtain a convergence rate scaling as O(K−1) in Equation (5), without dependency on the ωdwn in the dominating term, which only appears in faster decaying terms scaling as K−2. The iteration complexity (i.e., number of iterations to achieve expected error) is thus at first order O →0( σ2(1+ωup) µ Nb ). Again, this matches the complexity of Diana [15, see Theorem 1 and Corollary 1] and is smaller by a factor 1 + ωdwn than the one of Artemis, Dore, DIANAsr-DQ (see Corollary I.1. in [11]). Next, we give a convergence result in the convex case.
Theorem 2 (Convergence of MCM, convex case). Under Assumptions 1 to 4 with µ = 0. For all k > 0, for any γ ≤ γmax, we have, for w̄k = 1k ∑k−1 i=0 wi,
γE [F (wk−1)− F (w∗)] ≤ Vk−1 − Vk + γ2σ2Φ(γ)
Nb =⇒ E[F (w̄k)− F∗] ≤ V0 γk + γσ2Φ(γ) Nb . (6)
Consequently, for K in N large enough, a step-size γ = √ ‖w0−w∗‖2Nb (1+ωup)σ2K , we have:
E[F (w̄K)− F∗] ≤ 2
√ ‖w0 − w∗‖2 (1 + ωup)σ2
NbK +O(K−1). (7)
Moreover if σ2 = 0 (noiseless case), we recover a faster convergence: E[F (w̄K)− F∗] = O(K−1). Limit Variance (Eq. (6)). The variance term is identical to the strongly-convex case.
Convergence and complexity (Equation (7)). The downlink compression constant only appears in the second-order term, scaling as 1/K. In other words, the convergence rate is equivalent to the convergence rate of Diana, in the non-strongly-convex. As K increases, this complexity scales as (1+ωup) n 2 independently of the downlink compression. Again, for previous algorithms with double
compression the complexity is at least O (
(1+ωup)(1+ωdwn) n 2
) (see Corollary I.2 in [11]).
Control of the variance of the local model.
0 150 300 450 600 Number of passes on data
−2
0
2
lo g 1
0( F (w
k ) − F (w
∗) )
Artemis-ND MCM - α = 0 MCM - α = 1 MCM
0 150 300 450 600 Number of passes on data
−2
0 2 lo g 1 0( F (w k ) − F (w ∗) ) Artemis-ND MCM - α = 0 MCM - α = 1
MCM
Figure 1: Comparing MCM on two datasets with three other algorithms using a non-degraded update, γ = 1/L. Artemis-ND stands for Artemis with a non-degraded update.
We here present the backbone Lemma of MCM’s proof. It allows to control the variance of the local model E[‖ŵk − wk‖2 |wk] (which is upperbounded by ωdwnE[‖Υk‖2 |wk]) and to build the Lyapunov function defined in Theorems 1 and 2.
This result highlights the impact of the downlink memory term. Without memory, i.e., with αdwn = 0, the variance of the local model ‖ŵk − wk‖2 increases with the number of iterations. On the other hand, if αdwn is too large (close to 1), this variance diverges. This behavior is illustrated on
two real datasets on Figure 1. This phenomenon is similar to the divergence observed in frameworks involving error feedback, when the compression operator is not contractive.
Theorem 3. Consider the MCM update as in eq. (2). Under Assumptions 1, 2 and 4 with µ = 0, if γ ≤ (8ωdwnL)−1 and α ≤ (4ωdwn)−1, then for all k in N:
E [Υk] ≤ (
1− αdwn 2
) E [Υk−1] + 2γ2 ( 1
αdwn + ωup N
) E [ ‖∇F (ŵk−1)‖2 ] + 2γ2σ2(1 + ωup)
Nb .
This bound provides a recursive control on Υk. Beyond the (1 − αdwn) contraction, the bound comprises the squared-norm of the gradient at the previous perturbed iterate, and a noise term.
Summary of rates. In Table 2, we summarize the rates and complexities, and maximal learning rate for Diana, Artemis, Dore and MCM. For simplicity, we ignore absolute constants, and provide asymptotic values for large ωup, ωdwn, and complexities for → 0. Proof in the heterogeneous case. To extend Theorems 1 to 3 in the heterogeneous setting for a convex objective (Appendix G), we assume that there exists a constant B in R+, s.t.:
1 N ∑N i=0 ‖∇Fi(w∗)‖2 = B2 . We further define Ξk = 1N2 ∑N i=1 ∥∥hik −∇Fi(w∗)∥∥2, where for all i in J1, NK. This term is recursively controled [23, 24] and combined into the Lyapunov function. Proofs. To convey the best understanding of the theorems and the spirit of the proof, we introduce a Ghost algorithm (impossible to implement) in Appendix D.1. A sketch of the proof describes the main steps in the case of Ghost, those steps are similar for MCM. Fundamentally, our proof relies on a tight analysis, related to perturbed iterate analysis [21]. Proofs of Theorems 1 to 3 are given in Appendix E. Th. S11 in Appendix E.4 ensures convergence for a non-convex F . Note that the proof for non-convex follows a different approach than the one in Theorems 1 and 2.
As mentioned in the introduction, our analysis of perturbed iterate in the context of double compression opens new directions: in particular, it opens the door to handling a different model for each worker. In the next section, we detail those possibilities, and provide theoretical guarantees for Rand-MCM, the variant of MCM in which instead of sending the same model to all workers, the compression noises are mutually independent. Remark 4 (Communication budget). How to split a given communication budget between uplink and downlink to optimize the convergence is an open question which is intrinsically related to the situation. Indeed it depends on many factors like the selected operators of compression, the upload/downlink speed or the number of participating workers at each iteration. However, our approach provides some insights on this question. Because asymptotically the impact of double compression is marginal, for a fixed budget, Theorem 2 suggests to strongly compress on the downlink direction (which leads to a large ωdwn), but to perform a weaker compression in the uplink direction.
4 Extension to Rand-MCM

heading: 4.1 Communication and convergence trade-offs

text:
In Rand-MCM, we leverage the fact that the compressions used for each worker need not to be identical. On the contrary, it is possible to consider independent compressions. By doing so, we reduce the impact of the downlink compression.
The relevance of such a modification depends on the framework: while the convergence rate will be improved, the computational time can be slightly increased. Indeed, N compressions need to be computed instead of one: however, this computational time is typically not a bottleneck w.r.t. the communication time. A more important aspect is the communication cost. While the size of each message will remain identical, a different message needs to be sent to each worker. That is, we go from a “one to N” configuration to N “one to one” communications. While this is a drawback, it is not an issue when the bandwidth/transfer time are the bottlenecks, as Rand-MCM will result in a better convergence with almost no cost. Furthermore, we argue that handling worker dependent models is essential for several major applications. Rand-MCM can directly be adapted to those frameworks.
1. Worker dependent compression. A first simple situation is the case in which workers are allowed to choose the size (or equivalently the compression level) of their updates.
2. Partial participation (PP). Similarly, having N different messages to send to each worker may be unavoidable in the case of partial participation of the workers. This is a key feature in Federated Learning frameworks [22]. In the classical distributed framework (without downlink constraints) it is easy to deal with it, as each available worker just queries the global model to compute its gradient on it [see for example 14]. On the other hand, for bidirectional compression, to ensure that all the local models match the central model, the adaptation to partial participation relies on a synchronization step. During this step, each worker that has not participated in the last S steps receives the last S corresponding messages as long as it costs less to send this sequence than a full uncompressed model. This is described in the description of the adaptation to partial participation in [24], in the remark preceding Eq. (20) in [26] and by Tang et al. [33, v2 on arxiv for the distributed case], who use
a buffer. On the contrary, Rand-MCM naturally handles a different model, memory and update per worker. The adaptation to partial participation is thus straightforward. Though theoretical results are out of the scope of this paper, we provide experiments on PP in Appendix B.1.1 and fig. 4.
One drawback is the necessity to store the N memories (Hik)i∈[N ] instead of one, which results in an additional memory cost. To circumvent this issue we propose two independent solutions. 1) Keep and use a single memory H̄k = N−1 ∑N i=1H i k (as suggested in [24]). It is then necessary to periodically reset the local memories Hik on all workers to the averaged value H̄k (rarely enough not to impact the communication budget). This is illustrated in fig. 4. 2) Use Rand-MCM with an arbitrary number of groups G N of workers. In each group Gg , g ∈ [G], all workers share the same memory (Hgk ) and receive the same update Cdwn,g(wk+1 −Hgk ). We call this algorithm Rand-MCM-G. Remark 5 (Protecting the global model from honest-but-curious clients). Another business advantage of MCM and Rand-MCM is that providing degraded models to the participants can be used to guarantee privacy, or to ensure the workers participate in good faith, and not only to obtain the model. This issue of detecting ill-intentioned clients (free-riders) that want to obtain the model without actually contributing has been studied by Fraboni et al. [10].

heading: 4.2 Theoretical results

text:
In this Section, we provide two main theoretical results for Rand-MCM. First Theorem 4 ensures that the theoretical guarantees are at least as good for Rand-MCM as for MCM. Then, in Theorem 5, we provide convergence result for both MCM and Rand-MCM in the case of quadratic functions.
Theorem 4. Theorems 1 to 3 are valid for Rand-MCM and Rand-MCM-G. The improvement in Rand-MCM comes from the fact that we are ultimately averaging the gradients at several random points, reducing the variance coming from this aspect. The goal is obviously to reduce the impact of ωdwn. Keeping in mind that the dominating term in the rate is independent of ωdwn, we can thus only expect to reduce the second-order term. Next, the uplink compression noise increases with the variance of the randomized model, which will not be directly reduced by Rand-MCM. As a consequence, we only expect the improvement to be visible in the part of the second-order term that does not depend on ωup (that is, the effect would be the most significant if ωup is small or 0).
This intuition is corroborated by the following result, in which we show that the convergence is improved when adding the randomization process for a quadratic function. Extending the proof beyond quadratic functions is possible, though it requires an assumption on third or higher order derivatives of F (e.g., using self-concordance [4]) to control of E [ ||∇F (ŵk−1)− E[∇F (ŵk−1)]||2
∣∣ wk−1]. Theorem 5 (Convergence in the quadratic case). Under Assumptions 1 to 4 with µ = 0, if the function is quadratic, after running K > 0 iterations, for any γ ≤ γmax, and we have
E[F (w̄K)− F∗] ≤ V0 γK + γσ2ΦRd(γ) Nb ,
with ΦRd(γ) = (1 + ωup) ( 1 + 4γ 2L2ωdwn K ( 1 C + ωup N ) )
and C = N for Rand-MCM, C = G Rand-MCM-G, and C = 1 for MCM.
This result is derived in Appendix F. We can make the following comments: (1) The convergence rate for quadratic functions is slightly better than for smooth functions. More specifically, the right hand term in Φ is multiplied by an additional γ ( 1 C + ωup N ) (w.r.t. Theorem 2), which is decaying at the same rate as γ. Besides, the proof for Rand-MCM is substantially modified, as E[∇F (ŵk−1)] is an unbiased estimator of∇F (wk−1). (2) Moreover, the randomization in Rand-MCM (resp. Rand-MCM-G) further reduces by a factor N (resp. G) this term. Depending on the relative sizes of ωup and N , this can lead to a significant improvement up to a factor of N . In practice the impact of Rand-MCM is noticeable, as illustrated in the following experiments.

heading: 5 Experiments

text:
In this section, we illustrate the validity of the theoretical results given in the previous section on both synthetic and real datasets, on (1) least-squares linear regression (LSR), (2) logistic regression (LR), and (3) non-convex deep learning. We compare MCM with classical algorithms used in distributed settings: Diana, Artemis, Dore and of course the simplest setting - SGD, which is the baseline.
In these experiments, we provide results on the log of the excess loss F (wk)− F∗, averaged on 5 runs (resp. 2) in convex settings (resp. deep learning), with errors bars displayed on each figure (but not in the “zoom square”), corresponding to the standard deviation of log10(F (wk)− F∗). On Figure 3, the X-axis is respectively the number of iterations and the number of bits exchanged.
Each experiment has been run with N = 20 workers using stochastic scalar quantization [3], w.r.t. 2-norm. To maximize compression, we always quantize on a single level (s = 20), unless for PP (s = 21) and neural network (the value of s depends on the dataset).
We used 9 different datasets.
• One toy dataset devoted to linear regression in an homogeneous setting. This toy dataset allows to illustrate MCM properties in a simple framework, and in particular to ilustrate that when σ2 = 0, we recover a linear convergence2, see Figure 2b. • Five datasets commonly used in convex optimization (a9a, quantum, phishing, superconduct and w8a); see Table S1 for more details. Experiments were conducted with heterogeneous workers obtained by clustering (using TSNE [20]) the input points. • Four dataset in a non-convex settings (CIFAR10, Fashion-MNIST, FE-MNIST, MNIST); see Table S2 for more details.
All experiments are performed without any tuning of the algorithms, (e.g., with the same learning rate for all algorithms and without reducing it after a certain number of epochs). Indeed, our goal is to show that our method achieves a performance close to the unidirectional-compression framework (Diana), while performing an important downlink compression. More details about experiments can be found in Appendix B.
On Figure 3, we display the excess loss for quantum and a9a w.r.t. the number of iteration and number of communicated bits. The plots of phising, superconduct and w8a are not provided but can be found on our github repository. We only report their excess loss after 450 iterations in Table 3.
Saturation level. All experiments are performed with a constant learning rate γ to observe the bias (initial reduction) and the variance (saturation level) independently. Stochastic gradient descent results in a fast convergence during the first iterations, and then reaches a saturation at a given level proportional to σ2. Theorem 2 states that the variance of MCM is proportional to ωup, this is experimentally observed on Tables 3 and 4 and figs. 2 and 3: MCM meets Diana while Artemis and Dore saturate at a higher level (scaling as ωup×ωdwn). These trade-offs are preserved with optimized learning rates.
2Even stronger, we show in experiments that we recover a linear rate if we have σ∗ = 0 (the noise over stochastic gradient computation at the optimum point w?).
Linear convergence when σ2 = 0. The six algorithms present a linear convergence when σ2 = 0. This is illustrated by Figure 2b: we ran experiments with a full gradient descent. Note that in these settings MCM has a slightly worse performance than other methods; however, this slow-down is compensated by Rand-MCM.
Impact of randomization. The impact of randomization is noticeable on Figures 2b and S5b. Randomization helps to stabilise convergence of it reduces the variance of the runs and when σ2 = 0, it performs identically to SGD. Figure 4 illustrates the impact of using a single memory, instead of N , to alleviate the memory cost in the PP setting (Subsection 4.1), with or without periodic reset. Without reset, performance are slightly degraded, but with it, we recover previous results.
Deep learning. Table 4 and figs. 2c and 2d illustrate experiments with neural networks, details on dataset settings and networks architecture are given in Appendix B.2. Again, MCM meets Diana rates as stated by Theorem S11 (theorem in the non-convex case).
Overall, these experiments show the benefits of MCM and Rand-MCM, that reach the saturation level of Diana while exchanging at 10x to 100x fewer bits. More experiments with partial participation for Rand-MCM are given in Appendix B.1.1. All the code is provided on our github repository.

heading: 6 Conclusion

text: In this work, we propose a new algorithm to perform bidirectional compression while achieving the convergence rate of algorithms using compression in a single direction. One of the main application of this framework is Federated Learning. With MCM we stress the importance of not degrading the global model. In addition, we add the concept of randomization which allows to reduce the variance associated with the downlink compression. The analysis of MCM is challenging as the algorithm involves perturbed iterates. Proposing such an analysis is the key to unlocking numerous challenges in distributed learning, e.g., proposing practical algorithms for partial participation, incorporating privacy-preserving schemes after the global update is performed, dealing with local steps, etc. This approach could also be pivotal in non-smooth frameworks, as it can be considered as a weak form of randomized smoothing.

heading: Acknowledgments

text: We would like to thank Richard Vidal, Laeticia Kameni from Accenture Labs (Sophia Antipolis, France) and Eric Moulines from École Polytechnique for insightful discussions. This research was supported by the SCAI: Statistics and Computation for AI ANR Chair of research and teaching in artificial intelligence, by Hi!Paris, and by Accenture Labs (Sophia Antipolis, France).

title: Preserved central model for faster bidirectional compression in distributed settings