# Adaptive Client Sampling in Federated Learning via Online Learning with
Bandit Feedback

Boxin Zhao The work was done when Boxin Zhao was a research intern at Ant
Financial Group. University of Chicago Booth School of Business, USA  Ant
Financial Group, China  Ziqi Liu  Ant Financial Group, China  Chaochao Chen
Zhejiang University, China  Mladen Kolar  University of Chicago Booth School
of Business, USA  Zhiqiang Zhang  Ant Financial Group, China  Jun Zhou  Ant
Financial Group, China

###### Abstract

Due to the high cost of communication, federated learning (FL) systems need to
sample a subset of clients that are involved in each round of training. As a
result, client sampling plays an important role in FL systems as it affects
the convergence rate of optimization algorithms used to train machine learning
models. Despite its importance, there is limited work on how to sample clients
effectively. In this paper, we cast client sampling as an online learning task
with bandit feedback, which we solve with an online stochastic mirror descent
(OSMD) algorithm designed to minimize the sampling variance. To handle the
tuning parameters in OSMD that depend on the unknown problem parameters, we
use the online ensemble method and doubling trick. We prove a dynamic regret
bound relative to the theoretically optimal sampling sequence. The regret
bound depends on the total variation of the comparator sequence, which
naturally captures the intrinsic difficulty of the problem. To the best of our
knowledge, these theoretical contributions are new and the proof technique is
of independent interest. Through both synthetic and real data experiments, we
illustrate advantages of the proposed client sampling algorithm over the
widely used uniform sampling and existing online learning based sampling
strategies. The proposed adaptive sampling procedure is applicable beyond the
FL problem studied here and can be used to improve the performance of
stochastic optimization procedures such as stochastic gradient descent and
stochastic coordinate descent.

##  1 Introduction

Modern edge devices, such as personal mobile phones, wearable devices, and
sensor systems in vehicles, collect large amounts of data that are valuable
for training of machine learning models. If each device only uses its local
data to train a model, the resulting generalization performance will be
limited due to the number of available samples on each device. Traditional
approaches where data are transferred to a central server, which trains a
model based on all available data, have fallen out of fashion due to privacy
concerns and high communication costs. Federated Learning (FL) has emerged as
a paradigm that allows for collaboration between different devices (clients)
to train a global model while keeping data locally and only exchanging model
updates [27].

In a typical FL process, we have clients that contain data and a central
server that orchestrates the training process [22]. The following process is
repeated until the model is trained: (i) the server selects a subset of
available clients; (ii) the server broadcasts the current model parameters and
sometimes also a training program (e.g., a Tensorflow graph [1]); (iii) the
selected clients make updates to the model parameters based on their local
data; (iv) the local model updates are uploaded to the server; (v) the server
aggregates the local updates and makes a global update of the shared model. In
this paper, we focus on the first step and develop a practical strategy for
selecting clients with provable guarantees.

To train a machine learning model in a FL setting with $M$ clients, we would
like to minimize the following objective:111We use $[M]$ to denote the set
$\\{1,\dots,M\\}.$

$\min_{w}F(w)\coloneqq\sum_{m\in[M]}\lambda_{m}\phi\left(w;\mathcal{D}_{m}\right),$
(1)

where $\phi(w;\mathcal{D}_{m})$ is the loss function used to assess the
quality of a machine learning model parameterized by the vector $w$ based on
the local data $\mathcal{D}_{m}$ on the client $m\in[M]$. The parameter
$\lambda_{m}$ denotes the weight for client $m$. Typically, we have
$\lambda_{m}=n_{m}/n$, where $n_{m}=|\mathcal{D}_{m}|$ is the number of
samples on the client $m$, and the total number of samples is
$n=\sum^{M}_{m=1}n_{m}$. At the beginning of the $t$-th communication round,
the server uses the sampling distribution
$p^{t}=(p^{t}_{1},\dots,p^{t}_{M})^{\top}$ to choose $K$ clients by sampling
with replacement from $[M]$. Let $S^{t}\subseteq[M]$ denote the set of chosen
clients with $|S^{t}|=K$. The server transmits the current model parameter
vector $w^{t}$ to each client $m\in S^{t}$. The client $m$ computes the local
update $g^{t}_{m}$ and sends it back to the server.222Throughout the paper we
do not discuss how $g^{t}_{m}$ is obtained. One possibility that the reader
could keep in mind for concreteness is the LocalUpdate algorithm [8], which
covers well-known algorithms such as minibatch SGD and FedAvg [27]. After
receiving local updates from clients in $S^{t}$, the server constructs a
stochastic estimate of the global gradient as

$g^{t}=\frac{1}{K}\sum_{m\in S^{t}}\frac{\lambda_{m}}{p^{t}_{m}}g^{t}_{m},$
(2)

and makes the global update of the parameter $w^{t}$ using $g^{t}$. For
example, $w^{t+1}=w^{t}-\mu_{t}g^{t}$, if the server is using stochastic
gradient descent (SGD) with the stepsize sequence $\\{\mu_{t}\\}_{t\geq 1}$
[4]. However, the global update can be obtained using other procedures as
well.

The sampling distribution in FL is typically uniform over clients,
$p^{t}=p^{\rm unif}=(1/M,\dots,1/M)^{\top}$. However, nonuniform sampling
(also called importance sampling) can lead to faster convergence, both in
theory and practice, as has been illustrated in stochastic optimization [45,
29]. While the sampling distribution can be designed based on prior knowledge
[45, 21, 29, 36], we cast the problem of choosing the sampling distribution as
an online learning task and need no prior knowledge about (1).

Existing approaches to designing a sampling distribution using online learning
focus on estimation of the best sampling distribution under the assumption
that it does not change during the training process. However, the best
sampling distribution changes with iterations during the training process, and
the target stationary distribution does not capture the best sampling
distribution in each round. In the existing literature, the best _fixed_
distribution in hindsight is used as the comparator to measure the performance
of the algorithm used to design the sampling distribution. Here, we focus on
measuring the performance of the proposed algorithm against the best _dynamic_
sampling distribution. We use an online stochastic mirror descent (OSMD)
algorithm to generate a sequence of sampling distributions and prove a regret
bound relative to the dynamic comparator that involves a total variation term
that characterizes the intrinsic difficulty of the problem. To the best of our
knowledge, this is the first bound on the dynamic regret with intrinsic
difficulty characterization in importance sampling.

###  1.1 Contributions

We develop an algorithm based on OSMD that generates a sequence of sampling
distributions $\\{p^{t}\\}_{t\geq 1}$ based on the partial feedback available
to the server from the sampled clients. The learning rate of the proposed
algorithm depends on the total variation of the best sequence of sampling
distributions, which is generally unknown. Furthermore, it also depends on the
total number of iterations that can also be unknown. We make the algorithm
practical by addressing a few technical challenges. In particular, we adapt
the follow-the-regularized-leader algorithm and doubling trick to
_automatically_ choose the learning rate that performs as well as the best
learning rate asymptotically. We prove a bound on regret relative to the best
_dynamic_ comparator, which allows us to consider the best sequence of
sampling distributions as they change over iterations. The bound includes a
_total variation_ term that characterizes the _intrinsic difficulty_ of the
problem by capturing the difficulty of following the best sequence of
distributions. Such a characterization of problem difficulty is novel. We
demonstrate the empirical superiority of the proposed algorithm through
synthetic and real data experiments. In addition to client sampling in FL, our
proposed algorithm will have a broad impact on stochastic optimization.
Adapting our algorithm to any stochastic optimization procedure that chooses
samples, such as SGD, or coordinates, such as Stochastic Coordinate Descent,
can improve their performance. For example, [44] illustrated the practical
benefits of our algorithm to speed up L-SVRG and L-Katyusha.

###  1.2 Related work

Our paper is related to client sampling in FL, importance sampling in
stochastic optimization, and online convex optimization. We summarize only the
most relevant literature, without attempting to provide an extensive survey.

For client sampling, [9] proposed to use the theoretically optimal sampling
distribution to choose clients. However, their method requires all clients to
compute local updates in each round, which is impractical due to stragglers.
[33] modelled the parameters of the model during training by an Ornstein-
Uhlenbeck process, which was then used to derive an optimal sampling
distribution. [11] developed a biased client selection strategy and analyzed
its convergence property. [23, 10, 41] considered client sampling as a multi-
armed bandit problem, but provided only limited theoretical results. [39] used
reinforcement learning for client sampling with the objective of maximizing
accuracy, while minimizing the number of communication rounds.

Our paper is also closely related to importance sampling in stochastic
optimization. [45, 29] illustrated that by sampling observations from a
nonuniform distribution when using a gradient-based stochastic optimization
method, one can achieve faster convergence. They designed a fixed sampling
distribution using prior knowledge on the upper bounds of gradient norms. [13]
extended the importance sampling to minibatches. [36, 21, 16] developed
adaptive sampling strategies that allow the sampling distribution to change
over time. [30, 31, 46, 35] discussed importance sampling in stochastic
coordinate descent methods. [28, 34, 3, 2, 18] illustrated how to design the
sampling distribution by solving an online learning task with bandit feedback.
[28, 34] designed the sampling distribution by solving a multi-armed bandit
problem with the EXP3 algorithm [25, Chapter 11]. [3] used the follow-the-
regularized-leader algorithm [25, Chapter 28] to solve an online convex
optimization problem and make updates to the sampling distribution. [2]
restricted the sampling distribution to be a linear combination of
distributions in a predefined set and used an online Newton step to make
updates to the mixture weights. The above approaches estimate a stationary
distribution, while the best distribution is changing with iterations and,
therefore, is intrinsically dynamic. In addition to having suboptimal
empirical performance, these papers provide theoretical results that only
establish a regret relative to a fixed sampling distribution in hindsight. To
address this problem, [18] took a non-stationary approach where the most
recent information for each client was kept. A decreasing stepsize sequence is
required to establish a regret bound. Furthermore, the regret bound does not
capture the intrinsic difficulty of the problem. In comparison, we establish a
regret bound relative to a dynamic comparator—a sequence of sampling
distributions—without imposing assumptions on the stepsize sequence, and this
bound includes the dependence on the total variation term characterizing the
intrinsic difficulty of the problem.

Our paper also contributes to the literature on online convex optimization. We
cast the client sampling problem as an online learning problem [20] and adapt
algorithms from the dynamic online convex optimization literature to solve it.
[17, 42, 14] proposed methods that achieve sublinear dynamic regret relative
to dynamic comparator sequences. In particular, [17] used a dynamic mirror
descent algorithm to achieve sublinear dynamic regret with total variation
characterizing the intrinsic difficulty of the environment. However, the
optimal tuning parameters depend on the unknown total variation. On the other
hand, [37, 43] proposed different online ensemble approaches to automatically
choose the tuning parameters for online gradient descent. Compared with the
problem settings in the above studies, there are two key new challenges that
we need to address. First, we only have partial information—bandit
feedback—instead of the full information about the loss functions. Second, the
loss functions in our case are unbounded, which violates the common
boundedness assumption in the online learning literature. To overcome the
first difficulty, we construct an unbiased estimator of the loss function and
its gradient, which are then used to make an update to the sampling
distribution. We address the second challenge by first bounding the regret of
our algorithm when the sampling distributions in the comparator sequence lie
in a region of the simplex for which the loss is bounded, and subsequently
analyze the additional regret introduced by projecting the elements of the
comparator sequence to this region.

###  1.3 Notation

Let $\mathbb{R}^{M}_{+}=[0,\infty)^{M}$ and
$\mathbb{R}^{M}_{++}=(0,\infty)^{M}$. For $M\in\mathbb{N}^{+}$, let
$\mathcal{P}_{M-1}\coloneqq\\{x\in\mathbb{R}^{M}:\sum^{M}_{i=1}x_{i}=1\\}$ be
the $(M-1)$-dimensional probability simplex. We use
$p=(p_{1},\dots,p_{M})^{\top}$ to denote a sampling distribution with support
on $[M]\coloneqq\\{1,\dots,M\\}$. We use $p^{1:T}$ to denote a sequence of
sampling distributions $\\{p^{t}\\}_{t=1}^{T}$. Let
$\Phi:\mathcal{D}\subseteq\mathbb{R}^{M}\mapsto\mathbb{R}$ be a differentiable
convex function defined on $\mathcal{D}$, where $\mathcal{D}$ is a convex open
set, and we use $\bar{\mathcal{D}}$ to denote the closure of $\mathcal{D}$.
The Bregman divergence between any $x,y\in\mathcal{D}$ with respect to the
function $\Phi$ is given as
$D_{\Phi}\left(x\,\|\,y\right)=\Phi(x)-\Phi(y)-\langle\nabla\Phi(y),x-y\rangle$.
The unnormalized negative entropy is denoted as
$\Phi_{e}(x)=\sum^{M}_{m=1}x_{m}\log x_{m}-\sum^{M}_{m=1}x_{m}$,
$x=(x_{1},\dots,x_{M})^{\top}\in\mathcal{D}=\mathbb{R}^{M}_{+}$, with $0\log
0$ defined as $0$. We use $\|\cdot\|_{p}$ to denote the $L_{p}$-norm for
$1\leq p\leq\infty$. For $x\in\mathbb{R}^{n}$, we have
$\|x\|_{p}=(\sum^{n}_{i=1}x^{p}_{i})^{1/p}$ when $1<p<\infty$,
$\|x\|_{1}=\sum^{n}_{i=1}|x_{i}|$, and $\|x\|_{\infty}=\max_{1\leq i\leq
n}|x_{i}|$. Given any $L_{p}$-norm $\|\cdot\|$, we define its dual norm as
$\|z\|_{\star}\coloneqq\sup\\{z^{\top}x:\|x\|\leq 1\\}$.

For two sequences $\\{a_{n}\\}$ and $\\{b_{n}\\}$, $a_{n}=O(b_{n})$ if there
exists $C>0$ such that $|a_{n}/b_{n}|\leq C$ for all $n$ large enough;
$a_{n}=\Theta(b_{n})$ if $a_{n}=O(b_{n})$ and $b_{n}=O(a_{n})$ simultaneously.
Similarly, $a_{n}=\tilde{O}(b_{n})$ if $a_{n}=O(b_{n}\log^{k}b_{n})$ for some
$k\geq 0$; $a_{n}=\tilde{\Theta}(b_{n})$ if $a_{n}=\tilde{O}(b_{n})$ and
$b_{n}=\tilde{O}(a_{n})$ simultaneously.

###  1.4 Organization of the paper

We motivate importance sampling in FL, introduce an adaptive client sampling
algorithm, and establish a bound on the dynamic regret in Section 2. We design
two extensions to the sampling algorithm that make it adaptive to the unknown
problem parameters in Section 3. Sampling without replacement is discussed in
Section 4. We provide the experimental results in Section 5 and Section 6.
Conclusions and discussion of interesting future directions are given in
Section 7. The technical proofs are deferred to appendix. Code to replicate
the results in this paper is available at
……………..https://github.com/boxinz17/FL-Client-Sampling.

##  2 Adaptive client sampling

We show how to cast the client sampling problem as an online learning task.
Subsequently, we solve the online learning problem using OSMD algorithm and
provide a regret analysis for it.

###  2.1 Client sampling as an online learning problem

Recall that at the beginning of the $t$-th communication round, the server
uses a sampling distribution $p^{t}$ to choose a set of clients $S^{t}$, by
sampling with replacement $K$ clients from $[M]$, to update the parameter
vector $w^{t}$. For a chosen client $m\in S^{t}$, the local update is denoted
as $g^{t}_{m}$. For example, the local update
$g^{t}_{m}=\nabla\phi(w^{t};\mathcal{D}_{m})$ may be the full gradient; when
minibatch SGD/FedSGD is used, then
$g^{t}_{m}=(1/B)\sum^{B}_{b=1}\nabla\phi(w^{t};\xi^{t}_{m,b})$, where
$\xi^{t}_{m,b}\overset{i.i.d.}{\sim}\mathcal{D}_{m}$ and $B$ is the batch
size; when FedAvg is used, then
$g^{t}_{m}=(w^{t}_{m,B}-w^{t})/\eta_{\text{local}}$, where
$w^{t}_{m,b}=w^{t}_{m,b-1}-\eta_{\text{local}}\nabla
f\left(w^{t}_{m,b-1};\xi^{t}_{m,b-1}\right)$, $b\in[B]$, $w^{t}_{m,0}=w^{t}$,
$\xi^{t}_{m,b}\overset{i.i.d.}{\sim}\mathcal{D}_{m}$, and
$\eta_{\text{local}}$ is the local stepsize. We define the aggregated oracle
update at the $t$-th communication round as

$J^{t}=\sum^{M}_{m=1}\lambda_{m}g^{t}_{m}.$

The oracle update $J^{t}$ is constructed only for theoretical purposes and is
not computed in practice. The stochastic estimate $g^{t}$, defined in (2), is
an unbiased estimate of $J^{t}$, that is,
$\mathbb{E}_{S^{t}}\left[g^{t}\right]=J^{t}$. Note that we only consider the
randomness of $S^{t}$ and treat $g^{t}_{m}$ as given. The variance of $g^{t}$
is

$\displaystyle\mathbb{V}_{S^{t}}[g^{t}]=\frac{1}{K}\left(\sum^{M}_{m=1}\frac{\lambda_{m}^{2}\|g^{t}_{m}\|^{2}}{p^{t}_{m}}-\left\|J^{t}\right\|^{2}\right).$
(3)

Our goal is to design the sampling distribution $p^{t}$, used to sample
$S^{t}$, so to minimize the variance in (3). In doing so, we can ignore the
second term as it is independent of $p^{t}$.

Let $a^{t}_{m}=\lambda^{2}_{m}\|g^{t}_{m}\|^{2}$. For any sampling
distribution $q=(q_{1},\dots,q_{M})^{\top}$, the _variance reduction loss_
333The variance reduction loss $l_{t}(\cdot)$ should be distinguished from the
training loss $\phi(\cdot)$. is defined as

$l_{t}(q)=\frac{1}{K}\sum^{M}_{m=1}\frac{a^{t}_{m}}{q_{m}}.$ (4)

Given a sequence of sampling distributions $q^{1:T}$, the cumulative variance
reduction loss is defined as $L(q^{1:T})\coloneqq\sum^{T}_{t=1}l_{t}(q^{t})$.
When the choice of $q^{1:T}$ depends on
$S^{1:T}\coloneqq\\{S^{t}\\}_{t=1}^{T}$, the expected cumulative variance
reduction loss is defined as
$\bar{L}(q^{1:T})\coloneqq\sum^{T}_{t=1}\mathbb{E}_{S^{0:(t-1)}}[l_{t}(q^{t})]$,
where $S^{0}$ is empty set.

The variance reduction loss appears in a bound on the suboptimality of a
stochastic optimization algorithm. Suppose $F(\cdot)$ in (1) is
$\sigma$-strongly convex. Furthermore, suppose the local update
$g^{t}_{m}=\nabla\phi(w^{t};\mathcal{D}_{m})$ is the full gradient of the
local loss and the global update is made by SGD with stepsize
$\mu_{t}=2/(\sigma t)$. Theorem 3 of [34] then states that for any $T\geq 1$:

$\mathbb{E}_{S^{1:T}}\left[F\left(\frac{2}{T(T+1)}\sum^{T}_{t=1}t\cdot
w^{t}\right)\right]-F(w^{\star})\leq\frac{2}{\sigma T(T+1)}\bar{L}(p^{1:T}),$
(5)

where $w^{\star}$ is the minimizer of the objective in (1). Therefore, by
choosing the sequence of sampling distributions $p^{1:T}$ to make the
$\bar{L}(p^{1:T})$ small, one can achieve faster convergence. While we have
provided a specific example here, choosing the sequence of sampling
distributions to minimize $\bar{L}(p^{1:T})$ improves the convergence of a
number of algorithms [45, 21, 29, 36, 2, 3, 28, 18, 34, 44]. We develop an
algorithm that creates a sequence of sampling distributions $p^{1:T}$ to
minimize $\bar{L}(p^{1:T})$ using only the norm of local updates, and without
imposing assumptions on the loss functions or how the local and global updates
are made. As a result, the algorithm can be applied to design sampling
distributions for essentially any stochastic optimization procedure.

Suppose that at the beginning of the $t$-th communication round we know all
$\\{a^{t}_{m}\\}^{M}_{m=1}$. Then the optimal sampling distribution

$p^{t}_{\star}=(p^{t}_{\star,1},\dots,p^{t}_{\star,M})^{\top}=\arg\min_{p\in\mathcal{P}_{M-1}}l_{t}(p)$

is obtained as
$p^{t}_{\star,m}=\sqrt{a^{t}_{m}}/(\sum^{M}_{m=1}\sqrt{a^{t}_{m}})$. Computing
the distribution $p^{t}_{\star}$ is impractical as it requires local updates
of all clients, which eradicates the need for client sampling. From the form
of $p^{t}_{\star}$, we observe that clients with a large $a^{t}_{m}$ are more
“important” and should have a higher probability of being selected. Since we
do not know $\\{a^{t}_{m}\\}^{M}_{m=1}$, we will need to explore the
environment to learn about the importance of clients before we can exploit the
best strategy. Finally, we note that the relative importance of clients will
change over time, which makes the environment dynamic and challenging.

Based on the above discussion, we cast the problem of creating a sequence of
sampling distributions as an online learning task with bandit feedback, where
a game is played between the server and environment. Let $p^{1}$ be the
initial sampling distribution. At the beginning of iteration $t$, the server
samples with replacement $K$ clients from $[M]$, denoted $S^{t}$, using
$p^{t}$. The environment reveals $\\{a^{t}_{m}\\}_{m\in S^{t}}$ to the server,
where $a^{t}_{m}=\lambda^{2}_{m}\|g^{t}_{m}\|^{2}$. The environment also
computes $l_{t}(p^{t})$; however, this loss is not revealed to the server. The
server then updates $p^{t+1}$ based on the feedback $\\{\\{a^{u}_{m}\\}_{m\in
S^{u}}\\}^{t}_{u=1}$ and sampling distributions $\\{p^{u}\\}^{t}_{u=1}$. Note
that in this game, the server only gets information about the chosen clients
and, based on this partial information, or bandit feedback, needs to update
the sampling distribution. On the other hand, we would like to be competitive
with an oracle that can calculate the cumulative variance reduction loss. We
will design $p^{t}$ in a way that is agnostic to the generation mechanism of
$\\{a^{t}\\}_{t\geq 1}$, and will treat the environment as deterministic, with
randomness coming only from $\\{S^{t}\\}_{t\geq 1}$. We describe an OSMD-based
approach to solve this online learning problem.

###  2.2 OSMD Sampler

The variance-reduction loss function $l_{t}$ is a convex function on
$\mathcal{P}_{M-1}$ and

$\nabla
l_{t}(q)=-\frac{1}{K}\left(\frac{a^{t}_{1}}{(q_{1})^{2}},\dots,\frac{a^{t}_{M}}{(q_{M})^{2}}\right)^{\top}\in\mathbb{R}^{M}\quad\text{for
all}\quad q=(q_{1},\dots,q_{M})^{\top}\in\mathbb{R}^{M}_{++}.$

Since we do not observe $a^{t}$, we cannot compute $l_{t}(\cdot)$ or $\nabla
l_{t}(\cdot)$. Instead, we can construct unbiased estimates of them. For any
$q\in\mathcal{P}_{M-1}$, let $\hat{l}_{t}(q;p^{t})$ be an estimate of
$l_{t}(q)$ defined as

$\hat{l}_{t}(q;p^{t})=\frac{1}{K^{2}}\sum^{M}_{m=1}\frac{a^{t}_{m}}{q_{m}p^{t}_{m}}\mathcal{N}\left\\{m\in
S^{t}\right\\},$ (6)

and $\nabla\hat{l}_{t}(q;p^{t})\in\mathbb{R}^{M}$ has the $m$-th entry defined
as

$\left[\nabla\hat{l}_{t}(q;p^{t})\right]_{m}=-\frac{1}{K^{2}}\cdot\frac{a^{t}_{m}}{q^{2}_{m}p^{t}_{m}}\mathcal{N}\left\\{m\in
S^{t}\right\\}.$ (7)

The set $S^{t}$ is sampled with replacement from $[M]$ using $p^{t}$ and
$\mathcal{N}\left\\{m\in S^{t}\right\\}$ denote the number of times that a
client $m$ is chosen in $S^{t}$. Thus, $0\leq\mathcal{N}\left\\{m\in
S^{t}\right\\}\leq K$. Given $q$ and $p^{t}$, $\hat{l}_{t}(q;p^{t})$ and
$\nabla\hat{l}_{t}(q;p^{t})$ are random variables in $\mathbb{R}$ and
$\mathbb{R}^{M}$ that satisfy

$\displaystyle\mathbb{E}_{S^{t}}\left[\hat{l}_{t}(q;p^{t})\mid
p^{t}\right]=l_{t}(q),\qquad\mathbb{E}_{S^{t}}\left[\nabla\hat{l}_{t}(q;p^{t})\mid
p^{t}\right]=\nabla l_{t}(q).$

When $S^{t}$ and $p^{t}\in\mathbb{R}^{M}_{++}$ are given,
$\hat{l}_{t}(q;p^{t})$ is a convex function with respect to $q$ on
$\mathbb{R}^{M}_{++}$ and satisfies
$\hat{l}_{t}(q;p^{t})-\hat{l}_{t}(q^{\prime};p^{t})\leq\langle\nabla\hat{l}_{t}(q;p^{t}),q-q^{\prime}\rangle$,
for $q,q^{\prime}\in\mathbb{R}^{M}_{++}$. The constructed estimates
$\hat{l}_{t}(q;p^{t})$ and $\nabla\hat{l}_{t}(q;p^{t})$ are crucial for
designing updates to the sampling distribution.

Algorithm 1 OSMD Sampler

1: Input: A sequence of learning rates $\\{\eta_{t}\\}_{t\geq 1}$; parameter
$\alpha\in(0,1]$, $\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha/M,\infty)^{M}$;
number of iterations $T$.

2: Output: $\hat{p}^{1:T}$.

3: Initialize: $\hat{p}^{1}=p^{\rm unif}$.

4: for $t=1,2,\dots,T-1$ do

5: Sample $S^{t}$ by $\hat{p}^{t}$.

6: Compute $\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})$ via (7).

7: $\displaystyle\hat{p}^{t+1}=\arg\min_{p\in\mathcal{A}}\,\eta_{t}\langle
p,\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\rangle+D_{\Phi}\left(p\,\|\,\hat{p}^{t}\right)$.

8: end for

OSMD Sampler is an online stochastic mirror descent algorithm for updating the
sampling distribution, detailed in Algorithm 1. The sampling distribution is
restricted to lie in the space
$\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha/M,\infty)^{M}$, $\alpha\in(0,1]$, to
prevent the server from assigning small probabilities to devices. Let
$\Phi:\mathcal{D}\subseteq\mathbb{R}^{M}\mapsto\mathbb{R}$ be a continuously
differentiable convex function defined on $\mathcal{D}$, with
$\mathcal{A}\subseteq\bar{\mathcal{D}}$. The learning rates
$\\{\eta_{t}\\}_{t\geq 1}$ are positive and nonincreasing.444We use the term
_learning rate_ when discussing an online algorithm that learns a sampling
distribution, while the term _stepsize_ is used in the context of an
optimization algorithm. Line 7 of Algorithm 1 provides an update to the
sampling distribution using the mirror descent update. The available feedback
is used to construct an estimate of the loss, while the Bregman divergence
between the current and next sampling distribution is used as a regularizer,
ensuring that the updated sampling distribution does not change too much. The
update only uses the most recent information, while forgetting the history,
which results in nonstationarity of the sequence of sampling distributions. In
Line 5 of Algorithm 1, we choose $S^{t}$ by sampling with replacement. In
Section 4, we discuss how to extend the results to sampling without
replacement.

The mirror descent update in Line 7 is not available in a closed form in
general and an iterative solver may be needed. However, when $\Phi(\cdot)$ is
chosen as the negative entropy $\Phi_{e}(\cdot)$, a closed-form efficient
solution can be obtained as shown in Proposition 1.

######  Proposition 1.

Suppose $\Phi=\Phi_{e}$ is the unnormalized negative entropy in Algorithm 1.
Let

$\tilde{p}^{t+1}_{m}=p^{t}_{m}\exp\left\\{\mathcal{N}\left\\{m\in
S^{t}\right\\}\eta_{t}a^{t}_{m}/(K^{2}(p^{t}_{m})^{3})\right\\},\qquad
m\in[M].$

Let $\pi:[M]\mapsto[M]$ be a permutation such that
$\tilde{p}^{t+1}_{\pi(1)}\leq\tilde{p}^{t+1}_{\pi(2)}\leq\dots\leq\tilde{p}^{t+1}_{\pi(M)}$.
Let $m^{t}_{\star}$ be the smallest integer $m$ such that

$\tilde{p}^{t+1}_{\pi(m)}\left(1-\frac{m-1}{M}\alpha\right)>\frac{\alpha}{M}\sum^{M}_{j=m}\tilde{p}^{t+1}_{\pi(j)}.$

Then

$\hat{p}^{t+1}_{m}=\begin{cases}\alpha/M&\text{if }\pi(m)<m^{t}_{\star}\\\
\left((1-((m^{t}_{\star}-1)/M)\alpha)\tilde{p}^{t+1}_{m}\right)/\left(\sum^{M}_{j=m^{t}_{\star}}\tilde{p}^{t+1}_{\pi(j)}\right)&\text{otherwise}.\end{cases}$

###### Proof.

See Appendix A.1. ∎

Algorithm 2 Solver of Step 7 of Algorithm 1 when $\Phi$ is unnormalized
negative entropy

1: Input: $\hat{p}^{t}$, $S^{t}$, $\\{a^{t}_{m}\\}_{m\in S^{t}}$, and
$\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha/M,\infty)^{M}$.

2: Output: $\hat{p}^{t+1}$.

3: Let $\tilde{p}^{t+1}_{m}=p^{t}_{m}\exp\left\\{\mathcal{N}\left\\{m\in
S^{t}\right\\}\eta_{t}a^{t}_{m}/(K^{2}(p^{t}_{m})^{3})\right\\}$ for
$m\in[M]$.

4: Sort $\\{\tilde{p}^{t+1}_{m}\\}^{M}_{m=1}$ in a non-decreasing order:
$\tilde{p}^{t+1}_{\pi(1)}\leq\tilde{p}^{t+1}_{\pi(2)}\leq\dots\leq\tilde{p}^{t+1}_{\pi(M)}$.

5: Let $v_{m}=\tilde{p}^{t+1}_{\pi(m)}\left(1-\frac{m-1}{M}\alpha\right)$ for
$m\in[M]$.

6: Let $u_{m}=\frac{\alpha}{M}\sum^{M}_{j=m}\tilde{p}^{t+1}_{\pi(j)}$ for
$m\in[M]$.

7: Find the smallest $m$ such that $v_{m}>u_{m}$, denoted as $m^{t}_{\star}$.

8: Let $\hat{p}^{t+1}_{m}=\begin{cases}\alpha/M&\text{if
}\pi(m)<m^{t}_{\star}\\\
\left((1-((m^{t}_{\star}-1)/M)\alpha)\tilde{p}^{t+1}_{m}\right)/\left(\sum^{M}_{j=m^{t}_{\star}}\tilde{p}^{t+1}_{\pi(j)}\right)&\text{otherwise}.\end{cases}$

An efficient implementation is shown in Algorithm 2. The main cost comes from
sorting the sequence $\\{\tilde{p}^{t+1}_{m}\\}^{M}_{m=1}$, which can be done
with the computational complexity of $O(M\log M)$. However, note that we only
update a few entries of $p^{t}$ to get $\tilde{p}^{t+1}$ and $p^{t}$ is
sorted. Therefore, most entries of $\tilde{p}^{t+1}$ are also sorted. Using
this observation, we can usually achieve a much faster running time, for
example, by using an adaptive sorting algorithm [15].

We provide a bound on the dynamic regret for OSMD Sampler next.

###  2.3 Dynamic regret of OSMD Sampler

We first describe the dynamic regret used to measure the performance of an
online algorithm that generates a sequence of sampling distributions
$\\{\hat{p}\\}_{t\geq 1}$ in a non-stationary environment. Given any
comparator sequence $q^{1:T}\in\mathcal{P}_{M-1}^{T}$, the dynamic regret is
defined as

$\text{D-Regret}_{T}(q^{1:T})=\bar{L}\left(\hat{p}^{1:T}\right)-L\left(q^{1:T}\right).$
(8)

In contrast, the static regret measures the performance of an algorithm
relative to the best fixed sampling distribution, that is, it restricts
$q^{1}=\dots=q^{T}$ [28, 34, 3, 2]. Throughout the paper, we primarily focus
on the optimal sampling distribution sequence $\\{p^{t}_{\star}\\}_{t\geq 1}$,
and write $\text{D-Regret}_{T}$ to denote
$\text{D-Regret}_{T}(p^{1:T}_{\star})$.

The following quantity describes the dynamic complexity of a comparator
sequence and will appear in the regret bound below.

######  Definition 2.1 (Total Variation).

The total variation of a comparator sequence $q^{1:T}$ with respect to the
norm $\|\cdot\|$ on $\mathbb{R}^{M}$ is
$\text{TV}\left(q^{1:T}\right)=\sum^{T-1}_{t=1}\|q^{t+1}-q^{t}\|$.

The total variation measures how variable a sequence is. The larger the total
variation $\text{TV}(q^{1:T})$, the more variable $q^{1:T}$ is, and such a
comparator sequence is harder to match. To give an upper bound on the dynamic
regret of OSMD Sampler, we need the following assumptions.

######  Assumption 1.

The function $\Phi(\cdot)$ is $\rho$-strongly convex, $\rho>0$ with respect to
$\|\cdot\|$:

$\Phi(x)\geq\Phi(y)+\langle\nabla\Phi(y),x-y\rangle+\frac{\rho}{2}\|x-y\|^{2},\
\ x,y\in\mathcal{D}.$

######  Assumption 2.

There exist positive functions $H(M,\alpha)$, $\\{Q_{t}(M,\alpha)\\}_{t\geq
1}$, and $D_{\max}(M,\alpha)$, such that $\|\nabla\Phi(p)\|_{\star}\leq
H(M,\alpha)$, $\|\nabla\hat{l}_{t}(p;q)\|_{\star}\leq Q_{t}(M,\alpha)$ almost
surely, and $D_{\Phi}(q\,\|\,p^{\rm unif})\leq D_{\max}(M,\alpha)$ for all
$p,q\in\mathcal{A}$.

These assumptions are standard in the literature [17]. Under Assumption 1, we
have $D_{\Phi}\left(x\,\|\,y\right)\geq(\rho^{2}/2)\|x-y\|^{2}$. To simplify
the notation, we will omit $M$ and $\alpha$ from $H$, $\\{Q_{t}\\}_{t\geq 1}$,
and $D_{\max}$. We also need the following quantities that quantify how far
$q^{t}$ is from $\mathcal{A}$. Given $q^{t}\in\mathcal{P}_{M-1}$ and
$\alpha\in(0,1]$, let

$\begin{gathered}\psi(q^{t},\alpha)\coloneqq\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\},\qquad\omega(q^{t},\alpha)\coloneqq\frac{\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}}{\sum^{M}_{m=1}\left(q^{t}_{m}-\frac{\alpha}{M}\right)\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}},\\\
\phi(q^{t},\alpha)\coloneqq\frac{\omega(q^{t},\alpha)}{1-\omega(q^{t},\alpha)\left(1-\frac{\alpha}{M}\right)}.\end{gathered}$
(9)

We will use these quantities to characterize the projection error in the
following theorem, which is the main result of this section.

######  Theorem 1.

Suppose Assumptions 1-2 hold and we use $\|\cdot\|=\|\cdot\|_{1}$ to define
the total variation. Assume that $\\{\eta_{t}\\}_{t\geq 1}$ is a nonincreasing
sequence. Let $\hat{p}^{1:T}$ be a sequence generated by Algorithm 1. For any
comparator sequence $q^{1:T}$, we have

$\text{D-Regret}_{T}(q^{1:T})\leq\underbrace{\frac{D_{\max}}{\eta_{1}}+\frac{2H}{\eta_{T}}\text{TV}\left(q^{1:T}\right)+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t}}_{\text{Intrinsic
Regret}}+\underbrace{\frac{8H}{\eta_{T}}\sum^{T}_{t=1}\psi(q^{t},\alpha)+\sum^{T}_{t=1}\phi(q^{t},\alpha)l_{t}(q^{t})}_{\text{Projection
Error}}.$

###### Proof.

See Appendix A.2. ∎

From Theorem 1, we see that the bound on the dynamic regret consists of two
parts. The first part is the intrinsic regret, quantifying the difficulty of
tracking a comparator sequence in $\mathcal{A}^{T}$; the second part is the
projection error, arising from projecting the comparator sequence onto
$\mathcal{A}^{T}$. Note that the intrinsic regret depends on $\alpha$ through
$D_{\max}$, $H$, and $\\{Q_{t}\\}_{t\geq 1}$. As shown in Appendix A.2, we
have $0\leq\omega(q^{t},\alpha)\leq 1$ for all $\alpha\in[0,1]$, which implies
that $\phi(q^{t},\alpha)\leq M/\alpha$. Furthermore,

$\psi(q^{t},\alpha)\leq\sum^{M}_{m=1}\frac{\alpha}{M}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}\leq\alpha.$

Therefore, the projection error can be upper bounded by
$(8H\alpha)/\eta_{T}+(M/\alpha)\sum^{T}_{t=1}l_{t}(q^{t})$. More importantly,
when $q^{t}_{m}\in\mathcal{A}$, we have
$\psi(q^{t},\alpha)=\omega(q^{t},\alpha)=\phi(q^{t},\alpha)=0$. Thus, when the
comparator sequence belongs to $\mathcal{A}^{T}$, the projection error
vanishes and we only have the intrinsic regret. As $\alpha$ decreases from one
to zero, the intrinsic regret gets larger (it often tends to infinity as shown
in Corollary 1), while we are allowing a larger class of comparator sequences;
on the other hand, the projection error decrease to zero, since the gap
between $\mathcal{A}$ and $\mathcal{P}_{M-1}$ vanishes with $\alpha$. An
optimal choice of $\alpha$ balances the two sources of regret.

When $\|\cdot\|=\|\cdot\|_{1}$ and $\Phi$ is the unnormalized negative
entropy, the Step 7 of Algorithm 1 has a closed-form solution (see Proposition
1). By Pinsker’s inequality, the unnormalized negative entropy $\Phi_{e}(x)$
is $1$-strongly convex on $\mathcal{P}_{M-1}$ with respect to $\|\cdot\|_{1}$.
Applying Theorem 1 with the comparator sequence $\\{p^{t}_{\star}\\}_{t\geq
1}$, we have the following corollary.

######  Corollary 1.

Suppose conditions of Theorem 1 hold and $\Phi=\Phi_{e}$. Then

$\text{D-Regret}_{T}\leq\frac{\log
M}{\eta_{1}}+\frac{2\log(M/\alpha)}{\eta_{T}}\text{TV}\left(p^{1:T}_{\star}\right)+\frac{2M^{6}}{K^{4}\alpha^{6}}\sum^{T}_{t=1}\eta_{t}\left(\bar{a}^{t}\right)^{2}\\\
+\frac{8H}{\eta_{T}}\sum^{T}_{t=1}\psi(p^{t}_{\star},\alpha)+\sum^{T}_{t=1}\phi(p^{t}_{\star},\alpha)l_{t}(p^{t}_{\star}),$

where $\bar{a}^{t}\coloneqq\max_{1\leq m\leq M}a^{t}_{m}=\max_{1\leq m\leq
M}\lambda^{2}_{m}\|g^{t}_{m}\|^{2}$.

###### Proof.

See Appendix A.3. ∎

The tuning parameters $\\{\eta_{t}\\}_{t\geq 1}$ and $\alpha$ affect the
performance of OSMD Sampler. The following corollary provides one way to
choose these parameters.

######  Corollary 2.

Suppose conditions of Corollary 1 hold. Let $\alpha=\alpha_{\star}$, where
$\alpha_{\star}$ is a small enough constant so that
$p^{t}_{\star}\in\mathcal{A}$ for all $t\in[T]$. Let $\eta_{t}=\eta^{\star}$,
where

$\eta^{\star}=\frac{K^{2}\alpha_{\star}^{3}}{M^{3}\bar{a}^{1}}\sqrt{\frac{\log
M+2\log(M/\alpha_{\star})\text{TV}\left(p^{1:T}_{\star}\right)}{2T}}.$ (10)

Then

$\text{D-Regret}_{T}\leq\frac{2\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\sqrt{T\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]}.$
(11)

###### Proof.

The proof follows directly from Corollary 1. ∎

Note that as the training proceeds, the norms of local updates are decreasing
and $\sum^{T}_{t=1}(\bar{a}^{t})^{2}\leq T(\bar{a}^{1})^{2}$. When
$\lambda_{m}=1/M$, $m\in[M]$, we have
$\sum^{T}_{t=1}\left(\bar{a}^{t}\right)^{2}=O(T/M^{2})$. Thus,

$\text{D-Regret}_{T}=\tilde{O}\left(M^{2}/(K^{2}\alpha_{\star}^{3})\sqrt{1+\text{TV}\left(p^{1:T}_{\star}\right)}\sqrt{T}\right).$

When $\text{TV}\left(p^{1:T}_{\star}\right)=o(T)$, we achieve a sublinear
cumulative regret. In general, $\text{TV}\left(p^{1:T}_{\star}\right)\leq
2(T-1)$ and a large value indicates a more difficult problem.

##  3 Adaptive-OSMD and Adaptive-Doubling-OSMD Sampler

The choice of the sequence of learning rates $\\{\eta_{t}\\}_{t\geq 1}$ has a
large effect on the performance of OSMD Sampler. The choice in (10) depends on
unknown quantities such as $\bar{a}^{1}$, $\text{TV}(p^{1:T}_{\star})$, and
$T$. To get $\bar{a}^{1}$, we can add a pre-training phase where we broadcast
the initial model parameter $w^{0}$ to all devices before the start of the
training, and collect the returned $\|g^{0}_{m}\|^{2}$ from all responsive
devices, which we denote as $S^{0}$. Then
$\hat{\bar{a}}^{1}\coloneqq\max_{m\in S^{0}}\lambda_{m}\|g^{0}_{m}\|^{2}$. On
the other hand, $\text{TV}(p^{1:T}_{\star})$ and $T$ are hard to estimate in
advance of the training. We discuss how to use an online ensemble method to
choose the learning rate without the knowledge of
$\text{TV}(p^{1:T}_{\star})$, and how to get rid of the dependence on $T$
using the doubling trick.

Algorithm 3 Adaptive-OSMD Sampler

1: Input: Meta learning rate $\gamma$; the set of expert learning rates
$\mathcal{E}=\\{\eta_{1}\leq\eta_{2}\leq\dots\leq\eta_{E}\\}$ with
$E=|\mathcal{E}|$; parameter $\alpha_{\star}\in(0,1]$,
$\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha_{\star}/M,\infty)^{M}$; number of
iterations $T$; initial distribution $p^{\text{init}}$.

2: Output: $\hat{p}^{1:T}$.

3: Set $\theta^{1}_{e}=(1+1/E)/(e(e+1))$ and
$\hat{p}^{1}_{e}=p^{\text{init}}$, $e\in[E]$.

4: for $t=1,2,\dots,T-1$ do

5: Compute $\hat{p}^{t}=\sum^{E}_{e=1}\theta^{t}_{e}\hat{p}^{t}_{e}$.

6: Sample $S^{t}$ by $\hat{p}^{t}$.

7: for $e=1,2,\dots,E$ do

8: Compute $\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})$ via (6) and
$\nabla\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})$ via (7).

9: Solve $\hat{p}^{t+1}_{e}=\arg\min_{p\in\mathcal{A}}\,\eta_{e}\langle
p,\nabla\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\rangle+D_{\Phi}\left(p\,\|\,\hat{p}^{t}_{e}\right)$
via Algorithm 2.

10: end for

11: Update the weight of each expert:
$\theta^{t+1}_{e}=\frac{\theta^{t}_{e}\exp\left\\{-\gamma\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\right\\}}{\sum^{E}_{e=1}\theta^{t}_{e}\exp\left\\{-\gamma\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\right\\}},\quad
e\in[E].$

12: end for

The main idea is to run a set of expert algorithms, each with a different
learning rate for Algorithm 1. We then use a prediction-with-expert-advice
algorithm to track the best performing expert algorithm.555We refer the reader
to Chapter 2 of [7] for an overview of prediction-with-expert-advice
algorithms. More specifically, we define the set of expert learning rates as

$\mathcal{E}\coloneqq\left\\{\left.2^{e-1}\cdot\frac{K^{2}\alpha^{3}}{M^{3}\bar{a}^{1}}\sqrt{\frac{\log
M}{2T}}\,\right|\,e=1,2,\dots,E\right\\},$ (12)

where

$E=\left\lfloor\frac{1}{2}\log_{2}\left(1+\frac{4\log(M/\alpha)}{\log
M}(T-1)\right)\right\rfloor+1.$ (13)

Then for each $\eta_{e}\in\mathcal{E}$, Adaptive-OSMD Sampler algorithm runs
an expert algorithm to generate a sequence of sampling distributions
$\hat{p}^{1:T}_{e}$. Meanwhile, it also runs a meta-algorithm that uses
exponentially-weighted-average strategy to aggregate
$\\{\hat{p}^{1:T}_{e}\\}^{E}_{e=1}$ into a single output $\hat{p}^{1:T}$,
which achieves performance close to the best expert. Algorithm 3 details
Adaptive-OSMD Sampler. Note that since we can compute
$\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})$ and
$\nabla\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})$ directly, there is no need to
use a surrogate loss as in [37] and [43].

From the computational perspective, the major cost comes from solving step 9
of Algorithm 3, which needs to be run for a total number of
$T|\mathcal{E}|=T\left\lfloor\log_{2}T\right\rfloor$ times. Compared with
Algorithm 1, the computational complexity only increases by a $\log(T)$
factor. We have the following result on Algorithm 3.

######  Theorem 2.

Let $\Phi=\Phi_{e}$ and we use $\|\cdot\|=\|\cdot\|_{1}$ to define total
variation. Let $\alpha=\alpha_{\star}$, where $\alpha_{\star}$ is a small
enough constant so that $p^{t}_{\star}\in\mathcal{A}$ for all $t\in[T]$. Let
$\hat{p}^{1:T}$ be the output of Algorithm 3 with
$\gamma=\frac{\alpha_{\star}}{M}\sqrt{\frac{8K}{T\bar{a}^{1}}}$,
$p^{\text{init}}=p^{\rm unif}$ and $\mathcal{E}$ as in (12). Then

$\text{D-Regret}_{T}\leq\frac{3\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\left\\{T\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]\right\\}^{1/2}\\\
+\frac{M}{\alpha_{\star}}\sqrt{\frac{T\bar{a}^{1}}{8K}}\left(1+2\log
E\right).$

###### Proof.

See Appendix A.4. ∎

Since the additional regret term is $\tilde{O}((M/\alpha_{\star})\sqrt{T/K})$,
which is no larger than the first term asymptotically, the bound on the regret
is of the same order as in (11). However, we do not need to know the total
variation $\text{TV}(p^{1:T}_{\star})$ to set the learning rate.

Based on Theorem 2, the choice of $\alpha_{\star}$ relies on prior knowledge
about $\\{p^{t}_{\star}\\}_{t\geq 1}$. Specifically, we need $\alpha_{\star}$
to be small enough so that $p^{t}_{\star}\in\mathcal{A}$ for all $t\in[T]$.
While this prior knowledge is not generally available, in Section 5.3 we
experimentally show that the proposed algorithm is robust to the choice of
$\alpha$. As long as the chosen $\alpha$ is not too small or too large, we
obtain a reasonable solution. We always set $\alpha=0.4$ in experiments.

We now discuss how to use the doubling trick to implement Adaptive-OSMD
Sampler without the knowledge of the total number of iterations $T$ needed to
optimize the objective (1). The idea is motivated by Section 2.3 of [7]. The
basic idea is to restart Adaptive-OSMD Sampler at exponentially increasing
time points $T_{b}=2^{b-1}$, $b\geq 1$. The learning rates of experts in
Algorithm 3 are reset at the beginning of each time interval, and the meta-
algorithm learning rate $\gamma$ is chosen optimally for the interval length.
We set $\bar{a}^{T_{b}}$ at the beginning of each time interval using the
maximum environment feedback from the previous interval.

More specifically, at the time point $T_{b}$, we let

$\mathcal{E}_{b}\coloneqq\left\\{\left.2^{e-\frac{b}{2}-1}\cdot\frac{K^{2}\alpha^{3}\sqrt{\log
M}}{M^{3}\hat{a}^{b}}\,\right|\,e=1,2,\dots,E_{b}\right\\},$ (14)

where

$E_{b}=\left\lfloor\frac{1}{2}\log_{2}\left(1+\frac{4\log(M/\alpha)}{\log
M}(2^{b-1}-1)\right)\right\rfloor+1,$ (15)

and $\gamma_{b}=\frac{\alpha}{M}\sqrt{\frac{8K}{2^{b-1}\hat{a}^{b}}}$, $b\geq
1$. We set $\hat{\bar{a}}^{T_{b}}=\max_{m\in S^{T_{b}-1}}a^{T_{b}-1}_{m}$. In
a practical implementation, at the time point $t=T_{b}$, instead of
initializing all expert algorithms using uniform distribution, we can
initialize them with the output of the meta-algorithm for $t=T_{b}-1$. To get
$a^{0}$, the server uses a pre-training phase where the initial model
parameter $w^{0}$ is broadcast to all devices before the start of the
training. Subsequently, the server collects the returned $\|g^{0}_{m}\|^{2}$
from all responsive devices, which are denoted as $S^{0}$. Then
$\hat{\bar{a}}^{1}\coloneqq\max_{m\in S^{0}}a^{0}_{m}$ where
$a^{0}_{m}=\lambda_{m}\|g^{0}_{m}\|^{2}$. Adaptive-Doubling-OSMD Sampler is
detailed in Algorithm 4.

Algorithm 4 Adaptive-OSMD Sampler with Doubling Trick (Adaptive-Doubling-OSMD)

1: Input: Paramter $\alpha_{\star}$.

2: Output: $\hat{p}^{t}$ for $t=1,\dots,T$.

3: Use $w^{0}$ to get $\\{a^{0}_{m}\\}_{m\in S^{0}}$, where $S^{0}$ is the set
of responsive clients in the pre-training phase.

4: Initialize $\hat{\bar{a}}^{1}=\max_{m\in S^{0}}a^{0}_{m}$ and set $b=1$.

5: while True do

6: Set $\mathcal{E}_{b}$ as in (14).

7: Let
$\gamma_{b}=\frac{\alpha_{\star}}{M}\sqrt{\frac{8K}{2^{b-1}\hat{\bar{a}}^{b}}}$.

8: Obtain $\\{\hat{p}^{t}\\}^{2^{b}-1}_{t=2^{b-1}}$ from Algorithm 3 with
parameters: $\gamma_{b}$, $\mathcal{E}_{b}$, $\alpha_{\star}$, the number of
iterations $2^{b-1}$, and the initial distribution $p^{\rm unif}$ or
$\hat{p}^{2^{b-1}-1}$ (when $b>1$).

9: if Training Process is Converged then

10: Break.

11: end if

12: Let $b\leftarrow b+1$.

13: Let $\hat{\bar{a}}^{b}=\max_{m\in[M]}a^{T_{b}}_{m}$.

14: end while

From the computational perspective, by the proof of Theorem 3, Algorithm 4
needs to run Step 9 of Algorithm 3 for a total number of
$O(T|\mathcal{E}|^{2})=O(T\left\lfloor\log_{2}T\right\rfloor)$ times.
Therefore, the computational complexity of Adaptive-Doubling-OSMD Sampler is
asymptotically the same as that of Adaptive-OSMD Sampler, while it increases
by only a $\log(T)$ factor compared to OSMD Sampler. The following theorem
provides a bound on the dynamic regret for Adaptive-Doubling-OSMD Sampler.

######  Theorem 3.

Let $\Phi=\Phi_{e}$ and we use $\|\cdot\|=\|\cdot\|_{1}$ to define the total
variation. Let $\alpha=\alpha_{\star}$, where $\alpha_{\star}$ is a small
enough constant so that $p^{1:T}_{\star}\in\mathcal{A}^{T}$ and the training
is stopped after $T$ iterations. Suppose that there exists a constant $C>1$
such that $\hat{\bar{a}}^{T_{b}}\leq C\bar{a}^{T_{b}}$ for all
$b=1,2,\dots,B$, where $B=\left\lfloor\log_{2}(T+1)\right\rfloor$. Let
$\hat{p}^{1:T}$ be the output of Algorithm 4, where $p^{\rm unif}$ is used in
Step 8. Then

$\text{D-Regret}_{T}\\\
\leq\frac{\sqrt{2(T+1)}}{\sqrt{2}-1}\left\\{\frac{(2C+1)\sqrt{2}M^{3}}{K^{2}\alpha_{\star}^{3}}\sqrt{\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)}+\frac{M}{\alpha_{\star}}\sqrt{\frac{\bar{a}^{1}}{8K}}(C+2\log
E)\right\\}.$

###### Proof.

See Appendix A.5. ∎

From Theorem 3 we observe that the asymptotic regret bound has the same order
as that of OSMD Sampler and Adaptive-OSMD Sampler. However, Adaptive-Doubling-
OSMD Sampler does not need to know $\text{TV}(p^{1:T}_{\star})$ or $T$ in
advance.

##  4 Adaptive sampling without replacement

Algorithm 5 Adaptive sampling without replacement

1: Input: $w^{1}$ and $\hat{p}^{1}$.

2: for $t=1,2,\dots,T-1$ do

3: Let $\hat{p}^{t}_{(1)}=\hat{p}^{t}$ and sample $m^{t}_{1}$ from $[M]$ by
$\hat{p}^{t}_{(1)}$.

4: for $k=2,\cdots,K$ do

5: /* Design the sampling distribution for sampling the $k$-th client in the
$t$-th round */

6: Construct $\hat{p}^{t}_{(k)}$ by letting
$\hat{p}^{t}_{(k),m}=\left\\{\begin{array}[]{ll}\left(1-\sum^{k-1}_{l=1}\hat{p}^{t}_{m^{t}_{l}}\right)^{-1}\hat{p}^{t}_{m}&\text{if
}m\in[M]\backslash\\{m^{t}_{1},\dots,m^{t}_{k-1}\\}\\\
0&\text{otherwise.}\end{array}\right.$

7: /* Sample the $k$-th client */

8: Sample $m^{t}_{k}$ from $[M]\backslash\\{m^{t}_{1},\dots,m^{t}_{k-1}\\}$ by
$\hat{p}^{t}_{(k)}$.

9: end for

10: Let $S^{t}=\\{m^{t}_{1},\cdots,m^{t}_{K}\\}$.

11: The server broadcasts the model parameter $w^{t}$ to clients in $S^{t}$.

12: The clients in $S^{t}$ compute and upload the set of local gradients
$\left\\{g^{t}_{m^{t}_{1}},\cdots,g^{t}_{m^{t}_{K}}\right\\}$.

13: /* Construct global gradient estimate */

14: Let
$g^{t}_{(1)}=\lambda^{t}_{m^{t}_{1}}g^{t}_{m^{t}_{1}}/\hat{p}^{t}_{(1),m^{t}_{1}}$.

15: for $k=2,\cdots,K$ do

16: Let
$g^{t}_{(k)}=\lambda^{t}_{m^{t}_{k}}g^{t}_{m^{t}_{k}}/\hat{p}^{t}_{(k),m^{t}_{k}}+\sum^{k-1}_{l=1}\lambda^{t}_{m^{t}_{l}}g^{t}_{m^{t}_{l}}$.

17: end for

18: Let $\tilde{g}^{t}=K^{-1}\sum^{K}_{k=1}g^{t}_{(k)}$.

19: /* Update the model weight based on the global gradient estimate */

20: Obtain the updated model parameter $w^{t+1}$ using $w^{t}$ and
$\tilde{g}^{t}$.

21: /* Update sampling distribution */

22: Let $a^{t}_{m}=\lambda^{2}_{m}\|g^{t}_{m}\|^{2}$ for $m\in S^{t}$.

23: Input $\\{a^{t}_{m}\\}_{m\in S^{t}}$ into Adaptive-OSMD Sampler to get
$\hat{p}^{t+1}$.

24: end for

In the discussion so far, we have assumed that the set $S^{t}$ is obtained by
sampling with replacement from $p^{t}$. However, when $K$ is relatively large
compared to $M$ and $p^{t}$ is far from uniform distribution, sampling without
replacement can be more efficient than sampling with replacement. However,
when sampling without replacement using $p^{t}$, the variance reduction loss
does not have a clean form as in (4). As a result, an online design of the
sampling distribution is more challenging. In this section, we discuss how to
use the sampling distribution obtained by Adaptive-OSMD Sampler to sample
clients without replacement, following the approach taken in [18].

The detailed sampling procedure is described in Algorithm 5. We still use
Adaptive-OSMD Sampler to update the sampling distribution. However, we use the
designed sampling distribution in a way that no client is chosen twice.
Furthermore, Step 18 of Algorithm 5 constructs the gradient estimate with the
following properties.

######  Proposition 2 (Proposition 3 [18]).

Let $\hat{p}^{t}=p$ and let $\tilde{g}^{t}$ be as in Step 18 of Algorithm 5.
Note that $\tilde{g}^{t}=\tilde{g}^{t}(p)$ depends on $p$. Recall that
$J^{t}=\sum^{M}_{m=1}\lambda_{m}g^{t}_{m}$. We have

$\mathbb{E}_{S^{t}}\left[\tilde{g}^{t}\right]=J^{t}\qquad\text{and}\qquad\arg\min_{p\in\mathcal{P}_{M-1}}\mathbb{E}_{S^{t}}\left[\|\tilde{g}^{t}-J^{t}\|^{2}_{2}\right]=\arg\min_{p\in\mathcal{P}_{M-1}}l_{t}(p),$

where $l_{t}(\cdot)$ is defined in (4) and the expectation is taken over
$S^{t}$.

From Proposition 2, we see that $\tilde{g}^{t}$ is an unbiased stochastic
gradient. Furthermore, the variance of $\tilde{g}^{t}$ is minimized by the
same sampling distribution that minimizes the variance reduction loss in (4).
Therefore, it is reasonable to use the sampling distribution generated by
Adaptive-OSMD Sampler to design $\tilde{g}^{t}$.

##  5 Synthetic experiments

We use synthetic data to demonstrate the performance of Adaptive-OSMD Sampler.
In Section 5.1, we compare it with uniform sampling, which is used in majority
of the literature on FL. In Section 5.2, we compare it with existing bandit-
feedback online learning samplers: MABS [34] and VRB [3]. We also include the
theoretical optimal sampling distribution as the golden standard for
comparison in the above experiments. In addition, we examine the robustness of
Adaptive-OSMD Sampler to the choice of $\alpha$ in Section 5.3, while in
Section 5.4, we compare it with the sampling without replacement variant
discussed in Section 4.

In all experiments, we set $\alpha$ in Adaptive-OSMD Sampler as $\alpha=0.4$.
The tuning parameters for MABS and VRB are set as in their original papers.
Specifically, the parameter in MABS and VRB that prevents them from assigning
zero probability to any client is denoted as $\alpha$, and set to $0.4$.

We generate data as follows. We set the number of clients as $M=100$, and each
client has $n_{m}=100$ samples, $m\in[M]$. Samples on each client are
generated as

$y_{m,i}=\langle w_{\star},x_{m,i}\rangle+N(0,0.1),\qquad i\in[n_{m}],$

where the coefficient vector $w_{\star}\in\mathbb{R}^{d}$ has elements
generated as i.i.d. $N(10,3)$, and the feature vector
$x_{m,i}\in\mathbb{R}^{d}$ is generated as $x_{m,i}\sim N(0,\Sigma_{m})$,
where $\Sigma_{m}=s_{m}\cdot\Sigma$, $\Sigma$ is a diagonal matrix with
$\Sigma_{jj}=\kappa^{(j-1)/(d-1)-1}$ and $\kappa>0$ is the condition number of
$\Sigma$. We generate $\\{s_{m}\\}^{M}_{m=1}$ i.i.d. from
$e^{N(0,\sigma^{2})}$ and rescale them as
$s_{m}\leftarrow(s_{m}/\max_{m\in[M]}s_{m})\times 10$ so that $s_{m}\leq 10$
for all $m\in[M]$. In this setting, $\kappa$ controls the difficulty of each
problem when solved separately, while $\sigma$ controls the level of
heterogeneity across clients. In all experiments, we fix $\kappa=25$, which
corresponds to a hard problem, and change $\sigma$ to simulate different
heterogeneity levels. We expect that uniform sampling suffers when the
heterogeneity level is high. The dimension $d$ of the problem is set as
$d=10$. The results are averaged over $10$ independent runs.

We use the mean squared error loss defined as

$L(w)=\frac{1}{M}\sum^{M}_{m=1}L_{m}(w),\qquad\text{where}\qquad
L_{m}(w)=\frac{1}{2n_{m}}(y_{m,i}-\langle w_{\star},x_{m,i}\rangle)^{2}.$

We use the stochastic gradient descent to make global updates. At each round
$t$, we choose a subset of $K=5$ clients, denoted as $S^{t}$. For each client
$m\in S^{t}$, we choose a mini-batch of samples, $\mathcal{B}^{t}_{m}$, of
size $\bar{B}=10$, and compute the mini-batch stochastic gradient. The
parameter $w$ is updated as

$w^{t+1}=w^{t}+\frac{\mu_{\text{SGD}}}{MK\bar{B}}\sum_{m\in
S^{t}}\frac{1}{p^{t}_{m}}\sum_{i\in\mathcal{B}^{t}_{m}}\left(y_{m,i}-\langle
w_{\star},x_{m,i}\rangle\right)\cdot x_{m,i},$

where $\mu_{\text{SGD}}$ is the learning rate, set as $\mu_{\text{SGD}}=0.1$
in simulations.

###  5.1 Adaptive-OSMD Sampler vs Uniform Sampling

Figure 1: The training loss (top row) and the cumulative regret (bottom row)
of Adaptive-OSMD Sampler vs Uniform vs Optimal with $\sigma=1.0$,
$\sigma=3.0$, and $\sigma=10.0$. The solid line denotes the mean and the
shaded region covers $\text{mean}\pm\text{standard deviation}$ across
independent runs. Figure 2: The training loss (top row) and the cumulative
regret (bottom row) of Adaptive-OSMD Sampler vs MABS vs VRB with $\sigma=1.0$,
$\sigma=3.0$ and $\sigma=10.0$. The solid line denotes the mean and the shaded
region covers $\text{mean}\pm\text{standard deviation}$ across independent
runs.

The results of the training process and the cumulative regret are shown in
Figure 1. For the training loss, we see that when the heterogeneity level is
low ($\sigma=1.0$), the uniform sampling performs as well as Adaptive-OSMD
Sampler and theoretically optimal sampling; however, as the heterogeneity
level increases, the performance of uniform sampling gradually suffers; when
$\sigma=10.0$, uniform sampling performs poorly. On the other hand, Adaptive-
OSMD Sampler performs well across all levels of heterogeneity and is very
close to the theoretically optimal sampling. Similarly, for the cumulative
regret, when the heterogeneity level is low, the cumulative regret of uniform
sampling is close to Adaptive-OSMD Sampler; however, when the heterogeneity
level increases, the cumulative regret of uniform sampling gets much larger
than Adaptive-OSMD Sampler. Based on the above results, we can conclude that
while the widely used choice of uniform sampling may be reasonable when
heterogeneity is low, our proposed sampling strategy is robust across
different levels of heterogeneity, and thus should be considered as the
default option.

###  5.2 Adaptive-OSMD Sampler vs MABS vs VRB

We compare Adaptive-OSMD Sampler to other bandit feedback online learning
samplers: MABS [34] and VRB [3]. Training loss and cumulative regret are shown
in Figure 2. We see that while VRB performs better when the heterogeneity
level is low and MABS performs better when the heterogeneity level is high,
Adaptive-OSMD Sampler always achieves the best in both training loss and
cumulative regret across all different levels of heterogeneity. Thus, we
conclude that Adaptive-OSMD is a better choice than MABS and VRB.

###  5.3 Robustness of Adaptive-OSMD Sampler to the choice of $\alpha$

We examine the robustness of Adaptive-OSMD Sampler to the choice of $\alpha$.
We run Adaptive-OSMD Sampler separately for each
$\alpha\in\\{0.01,0.1,0.4,0.7,0.9,1.0\\}$. Note that when $\alpha=1.0$, the
Adaptive-OSMD Sampler outputs a uniform distribution. Training loss and
cumulative regret are shown in Figure 3. We observe that Adaptive-OSMD Sampler
is robust to the choice of $\alpha$, and performs well as long as $\alpha$ is
not too close to zero or too close to one.

###  5.4 Experiments on sampling with replacement vs without replacement

We compare sampling with replacement and sampling without replacement when
used together with Adaptive-OSMD sampler. Sampling without replacement is
described in Section 4. Training loss and cumulative regret are shown in
Figure 4. We observe that using sampling with replacement results in a
slightly smaller cumulative regret and a slightly better training loss.
However, these differences are not significant.

Figure 3: The training loss (top row) and the cumulative regret (bottom row)
of Adaptive-OSMD Sampler with different choices of $\alpha$ under
$\sigma=1.0$, $\sigma=3.0$ and $\sigma=10.0$. The solid line denotes the mean
and the shaded region covers $\text{mean}\pm\text{standard deviation}$ across
independent runs. Figure 4: The training loss (top row) and the cumulative
regret (bottom row) of Adaptive-OSMD Sampler with replacement vs without
replacement across $\sigma=1.0$, $\sigma=3.0$ and $\sigma=10.0$. The solid
line denotes the mean and the shaded region covers
$\text{mean}\pm\text{standard deviation}$ across independent runs.

##  6 Real data experiment

We compare Adaptive-OSMD Sampler with uniform sampling on real data. We use
three commonly used computer vision data sets: MNIST [26], KMINIST [12], and
FMINST [40]. We set the number of devices to be $M=500$. To better simulate
the situation where our method brings significant convergence speed
improvement, we create a highly skewed sample size distribution of the
training set among clients: $65$% of clients have only one training sample,
$20$% of clients have $5$ training samples, $10$% of clients have $30$
training samples, and $5$% of clients have $100$ training samples. This
setting tries to illustrate a real-life situation where most of the data come
from a small fraction of users, while most of the users have only a small
number of samples. The skewed sample size distribution is common in other FL
data sets, such as LEAF [6]. The sample size distribution in the training set
is shown in Figure 5. In addition, each client has $10$ validation samples
used to measure the prediction accuracy of the model over the training
process.

Figure 5: The sample size distribution in the training set across clients.
Figure 6: Comparison between Adaptive-OSMD Sampler and uniform sampler on real
data in terms of training loss (top row) and validation accuracy (bottom row).
Different columns correspond to different data sets. Adaptive-OSMD Sampler is
both faster and more stable. The result is the average performance over $5$
independent runs.

We use a multi-class logistic regression model. For a given gray scale picture
with the label $y\in\\{1,2,\dots,C\\}$, we unroll its pixel matrix into a
vector $x\in\mathbb{R}^{p}$. Given a parameter matrix $W\in\mathbb{R}^{C\times
p}$, the training loss function defined in (1) is

$\phi(W;x,y)\coloneqq l_{\text{CE}}\left(\varsigma(Wx)\,;y\right),$

where $\varsigma(\cdot):\mathbb{R}^{C}\rightarrow\mathbb{R}^{C}$ is the
softmax function defined as

$\left[\varsigma(x)\right]_{i}=\frac{\exp(x_{i})}{\sum^{K}_{j=1}\exp(x_{j})},\quad\text{
for all }x\in\mathbb{R}^{C},$

and $l_{\text{CE}}(x\,;y)=\sum^{C}_{i=1}\mathds{1}(y=i)\log x_{i}$,
$x\in\mathbb{R}^{C}$, $y\in\\{1,\dots,C\\}$, is the cross-entropy function.

We use the same algorithms and tuning parameters as in Section 5. Learning
rate in SGD is set to $0.075$ for MNIST and KMNIST, and is set to $0.03$ for
FMNIST. The total number of communication rounds is to $1,000$. In each round
of communication, we choose $K=10$ clients to participate ($2$% of total
number of clients). For a chosen client $m$, we compute its local mini-batch
gradient with the batch size equal to $\min\\{5,n_{m}\\}$, where $n_{m}$ is
the training sample size on the client $m$.

Figure 6 shows both the training loss and validation accuracy. Each figure
shows the average performance over 5 independent runs. We use the same random
seed for both Adaptive-OSMD Sampler and uniform sampling in each run, and
change random seeds across different runs. The main focus is on minimizing the
training loss, and the validation accuracy is only included for completeness.
We observe that Adaptive-OSMD Sampler performs better than uniform sampling
across all data sets. Given the cheap computational cost and the significant
practical advantage, we recommend using Adaptive-OSMD Sampler as the default
option in practice.

##  7 Discussion

We studied the client sampling problem in FL. We proposed an online learning
with bandit feedback approach to tackle client sampling. We used online
stochastic mirror descent to solve the online learning problem and applied the
online ensemble method with doubling trick to choose the tuning parameters. We
established an upper bound on the dynamic regret relative to the theoretically
optimal sequence of sampling distributions. The total variation of the
comparator is explicitly included in the upper bound as a measure of the
difficulty of the problem. Extensive numerical experiments demonstrated the
benefits of our approach over both widely used uniform sampling and other
competitors.

In this paper, we have focused on sampling with replacement. However, sampling
without replacement would ideally be a more efficient approach. In Section 4,
we discussed a natural extension of Adaptive-OSMD Sampler to a setting where
sampling without replacement is used. However, this approach does not directly
minimize the variance of the gradient $g^{t}$. When sampling without
replacement is used, the variance function becomes more complicated and the
design of an algorithm to directly minimize the variance is an interesting
future direction.

Another fruitful future direction is the design of sampling algorithms for
minimizing personalized FL objectives, which we discuss next.

###  7.1 Client sampling with personalized FL objective

Data distributions across clients are often heterogeneous. Personalized FL has
emerged as one effective way to handle such heterogeneity [24]. [19]
illustrated how many existing approaches to personalization can be studied
through a unified framework, and, in this section, we discuss a natural
extension of Adaptive-OSMD Sampler to this personalized objective.
Specifically, we study the following optimization problem

$\min_{w,\beta}F(w,\beta)\coloneqq\sum^{M}_{m=1}\lambda_{m}\phi(w,\beta_{m};\mathcal{D}_{m}),$
(16)

where $w\in\mathbb{R}^{d_{0}}$ corresponds to the shared parameter and
$\beta=(\beta_{1},\dots,\beta_{M})$ with $\beta_{m}\in\mathbb{R}^{d_{m}}$
corresponds to the local parameters. The objective in (16) coves a wide range
of personalized federated learning problems [19]. We further generalize the
approach and study the following bilevel optimization problem:

$\displaystyle\min_{w}$ $\displaystyle
h(w)\coloneqq\sum^{M}_{m=1}\lambda_{m}F_{m}(w,\hat{\beta}_{m}(w))\coloneqq\sum^{M}_{m=1}\lambda_{m}\phi(w,\hat{\beta}_{m}(w);\mathcal{D}_{m})$
(17) $\displaystyle\text{subject
to}\quad\hat{\beta}_{m}(w)=\arg\min_{\beta_{m}}G_{m}(w,\beta_{m})\coloneqq\phi(w,\beta_{m};\bar{\mathcal{D}}_{m}).$

When $\mathcal{D}_{m}=\bar{\mathcal{D}}_{m}$, then (17) recovers (16). When
$\bar{\mathcal{D}}_{m}\neq\mathcal{D}_{m}$, we then optimize the shared and
local parameters on different datasets, which may prevent overfitting. The
formulation in (17) is closely related to the implicit MAML [32].

In the following, we use $\nabla_{w}$ to denote a partial derivative with
respect to $w$ with $\beta_{m}$ fixed, $\nabla_{\beta_{m}}$ to denote a
partial derivative with respect to $\beta_{m}$ with $w$ fixed, and $\nabla$ to
denote a derivative with respect to $w$ where $\beta_{m}(w)$ is treated as a
function of $w$. Let
$\nabla^{2}_{\beta_{m}\beta_{m}}G_{m}(w,\beta_{m})\in\mathbb{R}^{d_{m}\times
d_{m}}$ be the Hessian matrix of $G_{m}$ with respect to $\beta_{m}$ where $w$
is fixed, and
$\nabla^{2}_{w\beta_{m}}G_{m}(w,\beta_{m})\in\mathbb{R}^{d_{0}\times d_{m}}$
be the Hessian matrix of $G_{m}$ with respect to $w$ and $\beta_{m}$, that is,

$\displaystyle\left[\nabla^{2}_{\beta_{m}\beta_{m}}G_{m}(w,\beta_{m})\right]_{i,j}$
$\displaystyle=\frac{\partial
G_{m}(w,\beta_{m})}{\partial\beta_{m,i}\beta_{m,j}}\quad\text{for
all}\,i,j=1,2,\dots,d_{m},$
$\displaystyle\left[\nabla^{2}_{w\beta_{m}}G_{m}(w,\beta_{m})\right]_{i,j}$
$\displaystyle=\frac{\partial G_{m}(w,\beta_{m})}{\partial
w_{i}\beta_{m,j}}\quad\text{for all}\,i=1,2,\dots,d_{0},j=1,2,\dots,d_{m}.$

By the implicit function theorem, we have

$\nabla
h(w)=\underbrace{\frac{1}{M}\sum^{M}_{m=1}\lambda_{m}\nabla_{1}F_{m}(w,\hat{\beta}_{m}(w))}_{\nabla_{1}h(w)}+\underbrace{\frac{1}{M}\sum^{M}_{m=1}\lambda_{m}\nabla_{2}F_{m}(w,\hat{\beta}_{m}(w))}_{\nabla_{2}h(w)}$
(18)

where

$\displaystyle\nabla_{1}F_{m}(w,\hat{\beta}_{m}(w))$
$\displaystyle\coloneqq\nabla_{w}F_{m}(w,\hat{\beta}_{m}(w)),$
$\displaystyle\nabla_{2}F_{m}(w,\hat{\beta}_{m}(w))$
$\displaystyle\coloneqq-\nabla^{2}_{w\beta_{m}}F_{m}(w,\hat{\beta}_{m}(w))\left[\nabla^{2}_{\beta_{m}\beta_{m}}F_{m}(w,\hat{\beta}_{m}(w))\right]^{-1}\nabla_{\beta_{m}}F_{m}(w,\hat{\beta}_{m}(w)).$

There are two parts to $\nabla h(w^{t})$ and, therefore, instead of choosing a
single subset of clients for computing both parts, we decouple $S^{t}$ into
two subsets $S^{t}_{1}$ and $S^{t}_{2}$, $S^{t}=S^{t}_{1}\cup S^{t}_{2}$. We
use clients in $S^{t}_{1}$ to compute local updates of the first part, and
clients in $S^{t}_{2}$ to compute the local updates of the second part. To get
an estimate of $\nabla h(w)$, we can estimate $\nabla_{1}h(w)$ and
$\nabla_{2}h(w)$ separably and then combine. Assume that $g^{t}_{1,m}$ is an
estimate of $\nabla_{1}F_{m}(w,\hat{\beta}_{m}(w))$ and $g^{t}_{2,m}$ is an
estimate of $\nabla_{2}F_{m}(w,\hat{\beta}_{m}(w))$, we can then construct
estimates of $\nabla_{1}h(w)$ and $\nabla_{2}h(w)$ as

$\displaystyle g^{t}_{1}=\frac{1}{K_{1}}\sum_{m\in
S^{t}_{1}}\lambda_{m}\frac{g^{t}_{1,m}}{p^{t}_{1,m}},\qquad
g^{t}_{2}=\frac{1}{K_{2}}\sum_{m\in
S^{t}_{2}}\lambda_{m}\frac{g^{t}_{2,m}}{p^{t}_{2,m}},$

where $K_{1}=|S^{t}_{1}|$ and $K_{2}=|S^{t}_{2}|$. Then
$g^{t}=g^{t}_{1}+g^{t}_{2}$ is an estimate of $\nabla h(w)$.

We design $p^{t}_{1}$ and $p^{t}_{2}$ to choose $S^{t}_{1}$ and $S^{t}_{2}$ by
minimizing the variance of the gradients. Note that

$\min_{p^{t}_{1}}\mathbb{E}_{1}\left[\mathbb{E}_{S^{t}_{1}}\left[\left\|g^{t}_{1}-\nabla_{1}h(w^{t})\right\|^{2}\right]\right]+\min_{p^{t}_{2}}\mathbb{E}_{2}\left[\mathbb{E}_{S^{t}_{2}}\left[\left\|g^{t}_{2}-\nabla_{2}h(w^{t})\right\|^{2}\right]\right]\\\
\leq\min_{p^{t}_{1}=p^{t}_{2}=p^{t}}\mathbb{E}\left[\mathbb{E}_{S^{t}}\left[\left\|g^{t}_{1}-\nabla_{1}h(w^{t})+g^{t}_{2}-\nabla_{2}h(w^{t})\right\|^{2}\right]\right],$

so that the decomposition allows us to better minimize the variance. We term
this approach as _doubly variance reduction for personalized Federated
Learning_. The first part minimizes the variance of updates to the shared
global parameter, when the best local parameters are fixed; and the second
part minimizes the variance of updates to local parameters, when the global
part is fixed. While these two parts are related, any given machine will have
different contributions to these two tasks.

Adaptive-OSMD Sampler can be used to minimize the variance for both parts of
the gradient. We note that this is a heuristic approach to solving the client
sampling problem when minimizing a personalized FL objective. Personalized FL
objectives have additional structures that should be used to design more
efficient sampling strategies. Furthermore, designing sampling strategies that
improve the statistical performance of trained models, rather than improving
computational speed, is important in the heterogeneous setting. Addressing
these questions is an important area for future research.

## References

  * Abadi et al., [2016] Abadi, M., Agarwal, A., Barham, P., Brevdo, E., Chen, Z., Citro, C., Corrado, G. S., Davis, A., Dean, J., Devin, M., Ghemawat, S., Goodfellow, I. J., Harp, A., Irving, G., Isard, M., Jia, Y., Józefowicz, R., Kaiser, L., Kudlur, M., Levenberg, J., Mané, D., Monga, R., Moore, S., Murray, D. G., Olah, C., Schuster, M., Shlens, J., Steiner, B., Sutskever, I., Talwar, K., Tucker, P. A., Vanhoucke, V., Vasudevan, V., Viégas, F. B., Vinyals, O., Warden, P., Wattenberg, M., Wicke, M., Yu, Y., and Zheng, X. (2016).  Tensorflow: Large-scale machine learning on heterogeneous distributed systems.  CoRR, abs/1603.04467.
  * Borsos et al., [2019] Borsos, Z., Curi, S., Levy, K. Y., and Krause, A. (2019).  Online variance reduction with mixtures.  In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 705–714. PMLR.
  * Borsos et al., [2018] Borsos, Z., Krause, A., and Levy, K. Y. (2018).  Online variance reduction for stochastic optimization.  In Bubeck, S., Perchet, V., and Rigollet, P., editors, Conference On Learning Theory, COLT 2018, Stockholm, Sweden, 6-9 July 2018, volume 75 of Proceedings of Machine Learning Research, pages 324–357. PMLR.
  * Bottou et al., [2018] Bottou, L., Curtis, F. E., and Nocedal, J. (2018).  Optimization methods for large-scale machine learning.  Siam Review, 60(2):223–311.
  * Boyd et al., [2004] Boyd, S., Boyd, S. P., and Vandenberghe, L. (2004).  Convex optimization.  Cambridge university press.
  * Caldas et al., [2018] Caldas, S., Wu, P., Li, T., Konečný, J., McMahan, H. B., Smith, V., and Talwalkar, A. (2018).  LEAF: A benchmark for federated settings.  CoRR, abs/1812.01097.
  * Cesa-Bianchi and Lugosi, [2006] Cesa-Bianchi, N. and Lugosi, G. (2006).  Prediction, learning, and games.  Cambridge university press.
  * Charles and Konečný, [2020] Charles, Z. and Konečný, J. (2020).  On the outsized importance of learning rates in local update methods.  CoRR, abs/2007.00878.
  * Chen et al., [2020] Chen, W., Horvath, S., and Richtárik, P. (2020).  Optimal client sampling for federated learning.  CoRR, abs/2010.13723.
  * [10] Cho, Y. J., Gupta, S., Joshi, G., and Yagan, O. (2020a).  Bandit-based communication-efficient client selection strategies for federated learning.  In Matthews, M. B., editor, 54th Asilomar Conference on Signals, Systems, and Computers, ACSCC 2020, Pacific Grove, CA, USA, November 1-4, 2020, pages 1066–1069. IEEE.
  * [11] Cho, Y. J., Wang, J., and Joshi, G. (2020b).  Client selection in federated learning: Convergence analysis and power-of-choice selection strategies.  CoRR, abs/2010.01243.
  * Clanuwat et al., [2018] Clanuwat, T., Bober-Irizar, M., Kitamoto, A., Lamb, A., Yamamoto, K., and Ha, D. (2018).  Deep learning for classical japanese literature.  CoRR, abs/1812.01718.
  * Csiba and Richtárik, [2018] Csiba, D. and Richtárik, P. (2018).  Importance sampling for minibatches.  J. Mach. Learn. Res., 19:27:1–27:21.
  * Daniely et al., [2015] Daniely, A., Gonen, A., and Shalev-Shwartz, S. (2015).  Strongly adaptive online learning.  In Bach, F. R. and Blei, D. M., editors, Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 1405–1411. JMLR.org.
  * Estivill-Castro and Wood, [1992] Estivill-Castro, V. and Wood, D. (1992).  A survey of adaptive sorting algorithms.  ACM Comput. Surv., 24(4):441–476.
  * Gopal, [2016] Gopal, S. (2016).  Adaptive sampling for SGD by exploiting side information.  In Balcan, M. and Weinberger, K. Q., editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 364–372. JMLR.org.
  * Hall and Willett, [2015] Hall, E. C. and Willett, R. M. (2015).  Online convex optimization in dynamic environments.  IEEE J. Sel. Top. Signal Process., 9(4):647–662.
  * Hanchi and Stephens, [2020] Hanchi, A. E. and Stephens, D. A. (2020).  Adaptive importance sampling for finite-sum optimization and sampling with decreasing step-sizes.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.
  * Hanzely et al., [2021] Hanzely, F., Zhao, B., and Kolar, M. (2021).  Personalized federated learning: A unified framework and universal optimization techniques.  arXiv preprint arXiv:2102.09743.
  * Hazan, [2016] Hazan, E. (2016).  Introduction to online convex optimization.  Found. Trends Optim., 2(3-4):157–325.
  * Johnson and Guestrin, [2018] Johnson, T. B. and Guestrin, C. (2018).  Training deep models faster with robust, approximate importance sampling.  In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 7276–7286.
  * Kairouz et al., [2021] Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K. A., Charles, Z., Cormode, G., Cummings, R., D’Oliveira, R. G. L., Eichner, H., Rouayheb, S. E., Evans, D., Gardner, J., Garrett, Z., Gascón, A., Ghazi, B., Gibbons, P. B., Gruteser, M., Harchaoui, Z., He, C., He, L., Huo, Z., Hutchinson, B., Hsu, J., Jaggi, M., Javidi, T., Joshi, G., Khodak, M., Konečný, J., Korolova, A., Koushanfar, F., Koyejo, S., Lepoint, T., Liu, Y., Mittal, P., Mohri, M., Nock, R., Özgür, A., Pagh, R., Qi, H., Ramage, D., Raskar, R., Raykova, M., Song, D., Song, W., Stich, S. U., Sun, Z., Suresh, A. T., Tramèr, F., Vepakomma, P., Wang, J., Xiong, L., Xu, Z., Yang, Q., Yu, F. X., Yu, H., and Zhao, S. (2021).  Advances and open problems in federated learning.  Found. Trends Mach. Learn., 14(1-2):1–210.
  * Kim et al., [2020] Kim, T., Bae, S., Lee, J., and Yun, S. (2020).  Accurate and fast federated learning via combinatorial multi-armed bandits.  CoRR, abs/2012.03270.
  * Kulkarni et al., [2020] Kulkarni, V., Kulkarni, M., and Pant, A. (2020).  Survey of personalization techniques for federated learning.  In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pages 794–797. IEEE.
  * Lattimore and Szepesvári, [2020] Lattimore, T. and Szepesvári, C. (2020).  Bandit algorithms.  Cambridge University Press.
  * LeCun and Cortes, [2010] LeCun, Y. and Cortes, C. (2010).  MNIST handwritten digit database.  CoRR.
  * McMahan et al., [2017] McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. (2017).  Communication-efficient learning of deep networks from decentralized data.  In Singh, A. and Zhu, X. J., editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA, volume 54 of Proceedings of Machine Learning Research, pages 1273–1282. PMLR.
  * Namkoong et al., [2017] Namkoong, H., Sinha, A., Yadlowsky, S., and Duchi, J. C. (2017).  Adaptive sampling probabilities for non-smooth optimization.  In Precup, D. and Teh, Y. W., editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 2574–2583. PMLR.
  * Needell et al., [2016] Needell, D., Srebro, N., and Ward, R. (2016).  Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm.  Math. Program., 155(1-2):549–573.
  * Nesterov, [2012] Nesterov, Y. E. (2012).  Efficiency of coordinate descent methods on huge-scale optimization problems.  SIAM J. Optim., 22(2):341–362.
  * Perekrestenko et al., [2017] Perekrestenko, D., Cevher, V., and Jaggi, M. (2017).  Faster coordinate descent via adaptive importance sampling.  In Singh, A. and Zhu, X. J., editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA, volume 54 of Proceedings of Machine Learning Research, pages 869–877. PMLR.
  * Rajeswaran et al., [2019] Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. (2019).  Meta-learning with implicit gradients.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R., editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 113–124.
  * Ribero and Vikalo, [2020] Ribero, M. and Vikalo, H. (2020).  Communication-efficient federated learning via optimal client sampling.  CoRR, abs/2007.15197.
  * Salehi et al., [2017] Salehi, F., Celis, L. E., and Thiran, P. (2017).  Stochastic optimization with bandit sampling.  arXiv preprint arXiv:1708.02544.
  * Salehi et al., [2018] Salehi, F., Thiran, P., and Celis, L. E. (2018).  Coordinate descent with bandit sampling.  In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 9267–9277.
  * Stich et al., [2017] Stich, S. U., Raj, A., and Jaggi, M. (2017).  Safe adaptive importance sampling.  In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R., editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4381–4391.
  * van Erven and Koolen, [2016] van Erven, T. and Koolen, W. M. (2016).  Metagrad: Multiple learning rates in online learning.  In Lee, D. D., Sugiyama, M., von Luxburg, U., Guyon, I., and Garnett, R., editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 3666–3674.
  * Wainwright, [2019] Wainwright, M. J. (2019).  High-dimensional statistics: A non-asymptotic viewpoint, volume 48.  Cambridge University Press.
  * Wang et al., [2020] Wang, H., Kaplan, Z., Niu, D., and Li, B. (2020).  Optimizing federated learning on non-iid data with reinforcement learning.  In 39th IEEE Conference on Computer Communications, INFOCOM 2020, Toronto, ON, Canada, July 6-9, 2020, pages 1698–1707. IEEE.
  * Xiao et al., [2017] Xiao, H., Rasul, K., and Vollgraf, R. (2017).  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  CoRR, abs/1708.07747.
  * Yang et al., [2020] Yang, M., Wong, A., Zhu, H., Wang, H., and Qian, H. (2020).  Federated learning with class imbalance reduction.  CoRR, abs/2011.11266.
  * Yang et al., [2016] Yang, T., Zhang, L., Jin, R., and Yi, J. (2016).  Tracking slowly moving clairvoyant: Optimal dynamic regret of online learning with true and noisy gradient.  In Balcan, M. and Weinberger, K. Q., editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 449–457. JMLR.org.
  * Zhang et al., [2018] Zhang, L., Lu, S., and Zhou, Z. (2018).  Adaptive online learning in dynamic environments.  In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 1330–1340.
  * Zhao et al., [2022] Zhao, B., Lyu, B., and Kolar, M. (2022).  L-svrg and l-katyusha with adaptive sampling.  arXiv preprint arXiv:2201.13387.
  * Zhao and Zhang, [2015] Zhao, P. and Zhang, T. (2015).  Stochastic optimization with importance sampling for regularized loss minimization.  In Bach, F. R. and Blei, D. M., editors, Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 1–9. JMLR.org.
  * Zhu et al., [2016] Zhu, Z. A., Qu, Z., Richtárik, P., and Yuan, Y. (2016).  Even faster accelerated coordinate descent using non-uniform sampling.  In Balcan, M. and Weinberger, K. Q., editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 1110–1119. JMLR.org.

##  Appendix A Technical proofs

###  A.1 Proof of Proposition 1

First, we show that the solution $\hat{p}^{t+1}$ in Step 7 of Algorithm 1 can
be found as

$\displaystyle\tilde{p}^{t+1}$
$\displaystyle=\arg\min_{p\in\mathcal{D}}\,\eta_{t}\langle
p,\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\rangle+D_{\Phi_{e}}\left(p\,\|\,\hat{p}^{t}\right),$
$\displaystyle\hat{p}^{t+1}$
$\displaystyle=\arg\min_{p\in\mathcal{A}}\,D_{\Phi_{e}}\left(p\,\|\,\tilde{p}^{t+1}\right).$

The optimality condition for $\tilde{p}^{t+1}$ implies that

$\eta_{t}\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})+\nabla\Phi_{e}(\tilde{p}^{t+1})-\nabla\Phi_{e}(\hat{p}^{t})=0.$
(19)

By Lemma 1, the optimality condition for $\hat{p}^{t+1}$ implies that

$\langle
p-\hat{p}^{t+1},\nabla\Phi_{e}(\hat{p}^{t+1})-\nabla\Phi_{e}(\tilde{p}^{t+1})\rangle\geq
0,\quad\text{for all }p\in\mathcal{A}.$

Combining the last two displays, we have

$\langle
p-\hat{p}^{t+1},\eta_{t}\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})+\nabla\Phi_{e}(\hat{p}^{t+1})-\nabla\Phi_{e}(\hat{p}^{t})\rangle\geq
0,\quad\text{for all }p\in\mathcal{A}.$

By Lemma 1, this is the optimality condition for $\hat{p}^{t+1}$ to be the
solution in Step 7 of Algorithm 1.

Note that (19) implies that

$-\frac{\eta_{t}}{K^{2}}\cdot\frac{a^{t}_{m}}{(p^{t}_{m})^{3}}\mathcal{N}\left\\{m\in
S^{t}\right\\}+\log(\tilde{p}^{t+1}_{m})-\log(\hat{p}^{t}_{m})=0,\qquad
m\in[M].$

Therefore,

$\tilde{p}^{t+1}_{m}=\hat{p}^{t}_{m}\exp\left(\frac{\eta_{t}a^{t}_{m}}{K^{2}\left(\hat{p}^{t}_{m}\right)^{3}}\mathcal{N}\left\\{m\in
S^{t}\right\\}\right),\qquad m\in[M],$

and the final result follows from Lemma 4.

###  A.2 Proof of Theorem 1

We first state a proposition that will be used to prove Theorem 1. The key
difference between Theorem 1 and Proposition 3 is that in Proposition 3 the
comparator sequence lies in $\mathcal{A}$, and, as a result, there is no
projection error.

######  Proposition 3.

Suppose the conditions of Theorem 1 hold. For any comparator sequence
$q^{1:T}$ with $q^{t}\in\mathcal{A}$, $t\in[T]$, we have

$\text{D-Regret}_{T}(q^{1:T})\leq\frac{D_{\max}}{\eta_{1}}+\frac{2H}{\eta_{T}}\text{TV}\left(q^{1:T}\right)+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t}.$

###### Proof.

By Lemma 1 and the definition of $\hat{p}^{t+1}$ in Step 7 of Algorithm 1, we
have

$\langle\hat{p}^{t+1}-q^{t},\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\rangle\leq\frac{1}{\eta_{t}}\langle\nabla\Phi(\hat{p}^{t})-\nabla\Phi(\hat{p}^{t+1}),\hat{p}^{t+1}-q^{t}\rangle.$
(20)

By the convexity of $\hat{l}_{t}(\cdot;\hat{p}^{t})$, we have

$\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\hat{l}_{t}(q^{t};\hat{p}^{t})\leq\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-q^{t}\rangle=\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t+1}-q^{t}\rangle+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle.$

Then, by (20), we further have

$\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\hat{l}_{t}(q^{t};\hat{p}^{t})\leq\frac{1}{\eta_{t}}\langle\nabla\Phi(\hat{p}^{t})-\nabla\Phi(\hat{p}^{t+1}),\hat{p}^{t+1}-q^{t}\rangle+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle.$

From the definition of $\mathcal{D}$, we have

$D_{\Phi}\left(x_{1}\,\|\,x_{2}\right)=D_{\Phi}\left(x_{3}\,\|\,x_{2}\right)+D_{\Phi}\left(x_{1}\,\|\,x_{3}\right)+\langle\nabla\Phi(x_{2})-\nabla\Phi(x_{3}),x_{3}-x_{1}\rangle,\qquad
x_{1},x_{2},x_{3}\in\mathcal{D}.$

Then

$\displaystyle\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})$
$\displaystyle-\hat{l}_{t}(q^{t};\hat{p}^{t})$
$\displaystyle\leq\frac{1}{\eta_{t}}\left[D_{\Phi}\left(q^{t}\|\hat{p}^{t}\right)-D_{\Phi}\left(q^{t}\|\hat{p}^{t+1}\right)-D_{\Phi}\left(\hat{p}^{t+1}\|\hat{p}^{t}\right)\right]+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle$
$\displaystyle=\frac{1}{\eta_{t}}\left[D_{\Phi}\left(q^{t}\|\hat{p}^{t}\right)-D_{\Phi}\left(q^{t+1}\|\hat{p}^{t+1}\right)\right]+\frac{1}{\eta_{t}}\left[D_{\Phi}\left(q^{t+1}\|\hat{p}^{t+1}\right)-D_{\Phi}\left(q^{t}\|\hat{p}^{t+1}\right)\right]$
$\displaystyle\qquad-\frac{1}{\eta_{t}}D_{\Phi}\left(\hat{p}^{t+1}\|\hat{p}^{t}\right)+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle.$
(21)

We bound the second term in (21) as

$\displaystyle
D_{\Phi}\left(q^{t+1}\|\hat{p}^{t+1}\right)-D_{\Phi}\left(q^{t}\|\hat{p}^{t+1}\right)$
$\displaystyle=\Phi(q^{t+1})-\Phi(q^{t})-\langle\nabla\Phi(\hat{p}^{t+1}),q^{t+1}-q^{t}\rangle$
$\displaystyle\overset{(a)}{\leq}\langle\nabla\Phi(q^{t+1})-\nabla\Phi(\hat{p}^{t+1}),q^{t+1}-q^{t}\rangle$
$\displaystyle\overset{(b)}{\leq}\|\nabla\Phi(q^{t+1})-\nabla\Phi(\hat{p}^{t+1})\|_{\star}\|q^{t+1}-q^{t}\|$
$\displaystyle\overset{(c)}{\leq}2H\|q^{t+1}-q^{t}\|,$ (22)

where (a) follows from the convexity of $\Phi(\cdot)$, (b) follows from the
definition of the dual norm, and (c) follows from the definition of $H$. Since
$\Phi(\cdot)$ is $\rho$-strongly convex, we can bound the third and fourth
term in (21) as

$-\frac{1}{\eta_{t}}D_{\Phi}\left(\hat{p}^{t+1}\|\hat{p}^{t}\right)+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle\leq-\frac{\rho}{2\eta_{t}}\|\hat{p}^{t+1}-\hat{p}^{t}\|^{2}+\|\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\|_{\star}\|\hat{p}^{t}-\hat{p}^{t+1}\|.$

Since $ab\leq a^{2}/(2\epsilon)+b^{2}\epsilon/2$, $a,b,\epsilon>0$, we further
have

$\displaystyle-\frac{1}{\eta_{t}}D_{\Phi}\left(\hat{p}^{t+1}\|\hat{p}^{t}\right)+\langle\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t}),\hat{p}^{t}-\hat{p}^{t+1}\rangle\leq\frac{2\eta_{t}}{\rho}\|\nabla\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\|^{2}_{\star}\leq\frac{2\eta_{t}}{\rho}Q_{t}^{2}.$
(23)

Combining (21)-(23), we have

$\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\hat{l}_{t}(q^{t};\hat{p}^{t})\leq\frac{D_{\Phi}\left(q^{t}\|\hat{p}^{t}\right)}{\eta_{t}}-\frac{D_{\Phi}\left(q^{t+1}\|\hat{p}^{t+1}\right)}{\eta_{t}}+2H\frac{\|q^{t+1}-q^{t}\|}{\eta_{t}}+\frac{2}{\rho}\eta_{t}Q^{2}_{t}.$

This implies that

$\displaystyle\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})$
$\displaystyle-\sum^{T}_{t=1}\hat{l}_{t}(q^{t};\hat{p}^{t})$
$\displaystyle\leq\frac{D_{\Phi}\left(q^{1}\|\hat{p}^{1}\right)}{\eta_{1}}-\frac{D_{\Phi}\left(q^{T+1}\|\hat{p}^{T+1}\right)}{\eta_{T+1}}+2H\sum^{T}_{t=1}\frac{\|q^{t+1}-q^{t}\|}{\eta_{t}}+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t}$
$\displaystyle\leq\frac{D_{\Phi}\left(q^{1}\|\hat{p}^{1}\right)}{\eta_{1}}+\frac{2H}{\eta_{T}}\sum^{T}_{t=1}\|q^{t+1}-q^{t}\|+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t}$
$\displaystyle\leq\frac{D_{\max}}{\eta_{1}}+\frac{2H}{\eta_{T}}\text{TV}\left(q^{1:T}\right)+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t},$
(24)

since $\hat{p}_{1}$ is the uniform distribution. Finally, note that

$\displaystyle\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\sum^{T}_{t=1}\hat{l}_{t}(q^{t};\hat{p}^{t})\right]$
$\displaystyle=\sum^{T}_{t=1}\mathbb{E}_{S^{1:{t-1}}}\left[\mathbb{E}_{S^{t}}\left[\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})\right]-\mathbb{E}_{S^{t}}\left[\hat{l}_{t}(q^{t};\hat{p}^{t})\right]\right]$
$\displaystyle=\sum^{T}_{t=1}\mathbb{E}_{S^{1:{t-1}}}\left[l_{t}(\hat{p}^{t})\right]-\sum^{T}_{t=1}l_{t}(q^{t})$
$\displaystyle=\text{D-Regret}_{T}(q^{1:T}).$

The conclusion follows by taking expectation
$\mathbb{E}_{S^{1:T}}\left[\cdot\right]$ on both sides of (24). ∎

We are now ready to prove Theorem 1.

###### Proof of Theorem 1.

For any comparator sequence $q^{1:T}$ with $q^{t}\in\mathcal{P}_{M-1}$,
$t\in[T]$, we have

$\text{D-Regret}_{T}(q^{1:T})=\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}l_{t}(\hat{p}^{t})\right]-\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})+\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})-\sum^{T}_{t=1}l_{t}(q^{t}).$
(25)

By Proposition 3, we further have that

$\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}l_{t}(\hat{p}^{t})\right]-\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})\\\
\leq\frac{D_{\max}}{\eta_{1}}+\frac{2H}{\eta_{T}}\text{TV}\left(q^{1:T}\right)+\frac{2}{\rho}\sum^{T}_{t=1}\eta_{t}Q^{2}_{t}+\frac{2H}{\eta_{T}}\left[\text{TV}\left(\tilde{q}^{1:T}\right)-\text{TV}\left(q^{1:T}\right)\right].$
(26)

Therefore, to prove Theorem 1, we design a suitable sequence
$\tilde{q}^{1:T}$, where $\tilde{q}^{t}\in\mathcal{A}$, and bound the terms
$\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})-\sum^{T}_{t=1}l_{t}(q^{t})$ and
$\text{TV}\left(\tilde{q}^{1:T}\right)-\text{TV}\left(q^{1:T}\right)$.

We define $\tilde{q}^{t}$ as

$\tilde{q}^{t}_{m}=\begin{cases}\alpha/M&\text{if }q^{t}_{m}<\alpha/M,\\\
q^{t}_{m}-\omega(q^{t},\alpha)\left(q^{t}_{m}-\frac{\alpha}{M}\right)&\text{if
}q^{t}_{m}\geq\alpha/M,\end{cases}$ (27)

where $\omega(q^{t},\alpha)$ is defined in (9). We now show that
$\tilde{q}^{t}\in\mathcal{A}$, $t\in[T]$, by showing that
$\tilde{q}^{t}_{m}\geq\alpha/M$, $m\in[M]$, and
$\sum_{m\in[M]}\tilde{q}^{t}_{m}=1$. For $m\in[M]$ such that
$q^{t}_{m}<\alpha/M$, we have from (27) that $\tilde{q}^{t}_{m}=\alpha/M$. For
$m\in[M]$ such that $q^{t}_{m}\geq\alpha/M$, by (27), we have
$\tilde{q}^{t}_{m}-{\alpha}/{M}=\left(1-\omega(q^{t},\alpha)\right)\left(q^{t}_{m}-{\alpha}/{M}\right)$.
Thus, we proceed to show that $\omega(q^{t},\alpha)\leq 1$. Since

$1=\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}+\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}\\\
\geq\sum^{M}_{m=1}\frac{\alpha}{M}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}+\sum^{M}_{m=1}\frac{\alpha}{M}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}=\alpha,$

we have

$\sum^{M}_{m=1}\left(q^{t}_{m}-\frac{\alpha}{M}\right)\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}\geq\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}.$

Therefore, $0\leq\omega(q^{t},\alpha)\leq 1$. Furthermore, $\omega(q^{t},0)=1$
and $\omega(q^{t},1)=1$. Finally, we show that
$\sum^{M}_{m=1}\tilde{q}^{t}_{m}=1$. By (27) and the definition of
$\omega(q^{t},\alpha)$ in (9), we have

$\displaystyle\sum^{M}_{m=1}\tilde{q}^{t}_{m}$
$\displaystyle=\sum^{M}_{m=1}\frac{\alpha}{M}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}+\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}-\omega(q^{t},\alpha)\sum^{M}_{m=1}\left(q^{t}_{m}-\frac{\alpha}{M}\right)\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}$
$\displaystyle=\sum^{M}_{m=1}\frac{\alpha}{M}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}+\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}-\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}$
$\displaystyle=\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}+\sum^{M}_{m=1}q^{t}_{m}\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}$
$\displaystyle=1.$

Therefore, $\tilde{q}^{t}\in\mathcal{A}$ for any $t\in[T]$.

We now bound $\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})-\sum^{T}_{t=1}l_{t}(q^{t})$.
When $q^{t}_{m}<\alpha/M$, then $1/\tilde{q}^{t}_{m}-1/q^{t}_{m}<0$; and when
$q^{t}_{m}\geq\alpha/M$, then

$\frac{1}{\tilde{q}^{t}_{m}}-\frac{1}{q^{t}_{m}}=\frac{1}{q^{t}_{m}}\cdot\left[\frac{1}{1-\omega(q^{t},\alpha)\left(1-\frac{\alpha}{Mq^{t}_{m}}\right)}-1\right]=\frac{1}{q^{t}_{m}}\cdot\frac{\omega(q^{t},\alpha)\left(1-\frac{\alpha}{Mq^{t}_{m}}\right)}{1-\omega(q^{t},\alpha)\left(1-\frac{\alpha}{Mq^{t}_{m}}\right)}.$

Since

$\omega(q^{t},\alpha)\left(1-\frac{\alpha}{Mq^{t}_{m}}\right)\leq\omega(q^{t},\alpha)\quad\text{and}\quad
1-\omega(q^{t},\alpha)\left(1-\frac{\alpha}{Mq^{t}_{m}}\right)\geq
1-\omega(q^{t},\alpha)+\frac{\omega(q^{t},\alpha)\alpha}{M}$

as $q^{t}_{m}\leq 1$, we have

$\frac{1}{\tilde{q}^{t}_{m}}-\frac{1}{q^{t}_{m}}\leq\frac{1}{q^{t}_{m}}\cdot\frac{\omega(q^{t},\alpha)}{1-\omega(q^{t},\alpha)\left(1-\frac{\alpha}{M}\right)}=\frac{\phi(q^{t},\alpha)}{q^{t}_{m}}.$

Thus,

$\displaystyle\sum^{T}_{t=1}l_{t}(\tilde{q}^{t})-\sum^{T}_{t=1}l_{t}(q^{t})$
$\displaystyle=\sum^{T}_{t=1}\sum^{M}_{m=1}a^{t}_{m}\left(\frac{1}{\tilde{q}^{t}_{m}}-\frac{1}{q^{t}_{m}}\right)$
$\displaystyle\leq\sum^{T}_{t=1}\sum^{M}_{m=1}a^{t}_{m}\left(\frac{1}{\tilde{q}^{t}_{m}}-\frac{1}{q^{t}_{m}}\right)\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}$
$\displaystyle\leq\sum^{T}_{t=1}\phi(q^{t},\alpha)\sum^{M}_{m=1}\frac{a^{t}_{m}}{q^{t}_{m}}\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}$
$\displaystyle\leq\sum^{T}_{t=1}\phi(q^{t},\alpha)l_{t}(q^{t}).$ (28)

Next, we bound
$\text{TV}\left(\tilde{q}^{1:T}\right)-\text{TV}\left(q^{1:T}\right)$. Note
that

$\displaystyle\text{TV}\left(\tilde{q}^{1:T}\right)$
$\displaystyle=\sum^{T}_{t=2}\left\|\tilde{q}^{t}-\tilde{q}^{t-1}\right\|_{1}$
$\displaystyle=\sum^{T}_{t=2}\left\|\tilde{q}^{t}-q^{t}+q^{t}-q^{t-1}+q^{t-1}-\tilde{q}^{t-1}\right\|_{1}$
$\displaystyle\leq\sum^{T}_{t=2}\left\|\tilde{q}^{t}-q^{t}\right\|_{1}+\sum^{T}_{t=2}\left\|q^{t}-q^{t-1}\right\|_{1}+\sum^{T}_{t=2}\left\|q^{t-1}-\tilde{q}^{t-1}\right\|_{1}$
$\displaystyle\leq\text{TV}\left(q^{1:T}\right)+2\sum^{T}_{t=1}\left\|\tilde{q}^{t}-q^{t}\right\|_{1}.$

We now upper bound $\sum^{T}_{t=1}\left\|\tilde{q}^{t}-q^{t}\right\|_{1}$. If
$q^{t}_{m}<\alpha/M$, then $|\tilde{q}^{t}_{m}-q^{t}_{m}|=\alpha/M-q^{t}_{m}$.
If $q^{t}_{m}\geq\alpha/M$, by (27), we have
$\left|\tilde{q}^{t}_{m}-q^{t}_{m}\right|=\omega(q^{t},\alpha)\left(q^{t}_{m}-{\alpha}/{M}\right)$.
Therefore, recalling the definition of $\psi(q^{t},\alpha)$ in (9), we have

$\displaystyle\left\|\tilde{q}^{t}-q^{t}\right\|_{1}$
$\displaystyle=\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}+\omega(q^{t},\alpha)\sum^{M}_{m=1}\left(q^{t}_{m}-\frac{\alpha}{M}\right)\mathds{1}\left\\{q^{t}_{m}\geq\frac{\alpha}{M}\right\\}$
$\displaystyle=2\sum^{M}_{m=1}\left(\frac{\alpha}{M}-q^{t}_{m}\right)\mathds{1}\left\\{q^{t}_{m}<\frac{\alpha}{M}\right\\}$
$\displaystyle=2\psi(q^{t},\alpha)$

and

$\text{TV}\left(\tilde{q}^{1:T}\right)-\text{TV}\left(q^{1:T}\right)\leq
4\sum^{T}_{t=1}\psi(q^{t},\alpha).$ (29)

Combining (25), (26), (28), and (29), we obtain the result. ∎

###  A.3 Proof of Corollary 1

When $\Phi=\Phi_{e}$, we have $\rho=1$ from Pinsker’s inequality. Furthermore,
$\|\cdot\|_{\star}=\|\cdot\|_{\infty}$, $\nabla\Phi_{e}(p)=(\log
p_{1},\dots,\log p_{M})^{\top}$, and
$\hat{p}^{t}\in\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha/M,\infty)^{M}$,
$t\in[T]$. Then $Q_{t}=M^{3}\bar{a}^{t}/(K^{2}\alpha^{3})$ and
$H=\log(M/\alpha)$ follows by checking the definition. Finally, to show that
$D_{\Phi_{e}}(q\,\|\,p^{\rm unif})\leq\log M$ for all $q\in\mathcal{A}$, we
note that $D_{\Phi_{e}}(q\,\|\,p^{\rm unif})=\log M+\sum^{M}_{m=1}q_{m}\log
q_{m}\leq\log M$.

###  A.4 Proof of Theorem 2

The proof proceeds in two steps. First, we show that there exists an expert
learning rate $\eta_{e}\in\mathcal{E}$ such that the regret bound for
$\hat{p}^{1:T}_{e}$ is close to (11). That is, we show that there exists
$\eta_{e}\in\mathcal{E}$ such that

$\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\right]-\sum^{T}_{t=1}l_{t}(p^{t}_{\star})\leq\frac{3\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\sqrt{T\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]},$
(30)

where $S^{t}\sim\hat{p}^{t}$. Second, we show that the output of meta-
algorithm can track the best expert with small regret. That is, we show that

$\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}l_{t}(\hat{p}^{t})\right]-\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\right]\leq\frac{M}{\alpha_{\star}}\sqrt{\frac{T\bar{a}^{1}}{8K}}(1+2\log
E),\qquad e\in[E].$ (31)

The theorem follows by combining (30) and (31).

We first prove (30). Since $0\leq\text{TV}\left(p^{1:T}_{\star}\right)\leq
2(T-1)$, we have

$\min\mathcal{E}=\frac{K^{2}\alpha^{3}_{\star}}{M^{3}\bar{a}^{1}}\sqrt{\frac{\log
M}{2T}}\leq\eta^{\star}\leq\frac{K^{2}\alpha^{3}_{\star}}{M^{3}\bar{a}^{1}}\sqrt{\frac{\log
M+4\log(M/\alpha_{\star})(T-1)}{2T}}\leq\max\mathcal{E},$

where $\eta^{\star}$ is defined in (10). Thus, there exists
$\eta_{e}\in\mathcal{E}$, such that $\eta_{e}\leq\eta^{\star}\leq 2\eta_{e}$.
Repeating the proof of (24) and proof of Corollary 1, we show that

$\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})-\sum^{T}_{t=1}\hat{l}_{t}(p^{t}_{\star};\hat{p}^{t})\leq\frac{\log
M}{\eta_{e}}+\frac{2\log(M/\alpha_{\star})}{\eta_{e}}\text{TV}\left(p^{1:T}_{\star}\right)+\frac{2\eta_{e}M^{6}}{K^{4}\alpha^{6}_{\star}}\sum^{T}_{t=1}\left(\bar{a}^{t}\right)^{2}.$

Since $\eta^{\star}/2\leq\eta_{e}\leq\eta^{\star}$, we further have

$\displaystyle\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})-\sum^{T}_{t=1}\hat{l}_{t}(p^{t}_{\star};\hat{p}^{t})$
$\displaystyle\leq\frac{2\log
M}{\eta^{\star}}+\frac{4\log(M/\alpha_{\star})}{\eta^{\star}}\text{TV}\left(p^{1:T}_{\star}\right)+\frac{2\eta^{\star}M^{6}}{K^{4}\alpha^{6}_{\star}}\sum^{T}_{t=1}\left(\bar{a}^{t}\right)^{2}$
$\displaystyle=\frac{3\sqrt{2}M^{3}}{K^{2}\alpha_{\star}^{3}}\sqrt{\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]\sum^{T}_{t=1}\left(\bar{a}^{t}\right)^{2}}$
$\displaystyle\leq\frac{3\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\sqrt{T\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]}.$

Now, (30) follows, since

$\mathbb{E}_{S^{1:T}}\left[\sum^{T}_{t=1}\hat{l}_{t}(p^{t}_{\star};\hat{p}^{t})\right]=\sum^{T}_{t=1}\mathbb{E}_{S^{1:(t-1)}}\left[\mathbb{E}_{S^{t}}\left[\hat{l}_{t}(p^{t}_{\star};\hat{p}^{t})\right]\right]=\sum^{T}_{t=1}\mathbb{E}_{S^{1:(t-1)}}\left[l_{t}(p^{t}_{\star})\right]=\sum^{T}_{t=1}l_{t}(p^{t}_{\star}),$

We prove (31) next. Let

$\hat{L}^{e}_{t}=\sum^{t}_{s=1}\hat{l}_{s}(\hat{p}^{s}_{e};\hat{p}^{s})\quad
e\in[E],t\in[T].$

Recall the update for $\theta^{t}_{e}$ in Step 11 of Alg 3. We have

$\theta^{t}_{e}=\frac{\theta^{1}_{e}\exp\left(-\gamma\hat{L}^{e}_{t-1}\right)}{\sum^{E}_{b=1}\theta^{1}_{b}\exp\left(-\gamma\hat{L}^{b}_{t-1}\right)},\qquad
t=2,\dots T.$

Let
$\Theta_{t}=\sum^{E}_{b=1}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{t}\right\\}$.
Then

$\displaystyle\log\Theta_{1}$
$\displaystyle=\log\left(\sum^{E}_{b=1}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{1}\right\\}\right)$
and, for $t\geq 2$,
$\displaystyle\log\left(\frac{\Theta_{t}}{\Theta_{t-1}}\right)$
$\displaystyle=\log\left(\frac{\sum^{E}_{b=1}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{t-1}\right\\}\exp\left\\{-\gamma\hat{l}_{t}(\hat{p}^{t}_{b};\hat{p}^{t})\right\\}}{\sum^{E}_{b=1}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{t-1}\right\\}}\right)$
$\displaystyle=\log\left(\sum^{E}_{b=1}\theta^{t}_{b}\exp\left\\{-\gamma\hat{l}_{t}(\hat{p}^{t}_{b};\hat{p}^{t})\right\\}\right).$

We have

$\displaystyle\log\Theta_{T}$
$\displaystyle=\log\Theta_{1}+\sum^{T}_{t=1}\log\left(\frac{\Theta_{t}}{\Theta_{t-1}}\right)$
$\displaystyle=\sum^{T}_{t=1}\log\left(\sum^{E}_{b=1}\theta^{t}_{b}\exp\left\\{-\gamma\hat{l}_{t}(\hat{p}^{t}_{b};\hat{p}^{t})\right\\}\right)$
$\displaystyle\leq\sum^{T}_{t=1}\left(-\gamma\sum^{E}_{b=1}\theta^{t}_{b}\hat{l}_{t}(\hat{p}^{t}_{b};\hat{p}^{t})+\frac{\gamma^{2}M^{2}\bar{a}^{t}}{8K\alpha^{2}_{\star}}\right)$
(Lemma 3)
$\displaystyle\leq-\gamma\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})+\frac{\gamma^{2}M^{2}\left(\sum^{T}_{t=1}\bar{a}^{t}\right)}{8K\alpha^{2}_{\star}}$
(Jensen’s inequality)

and

$\log\left(\Theta_{T}\right)=\log\left(\sum^{E}_{b=1}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{T}\right\\}\right)\\\
\geq\log\left(\max_{1\leq b\leq
E}\theta^{1}_{b}\exp\left\\{-\gamma\hat{L}^{b}_{T}\right\\}\right)=-\gamma\min_{1\leq
b\leq
E}\left\\{\hat{L}^{b}_{T}+\frac{1}{\gamma}\log\frac{1}{\theta^{1}_{b}}\right\\}.$

Combining the last two displays, we have

$-\gamma\min_{1\leq b\leq
E}\left\\{\hat{L}^{b}_{T}+\frac{1}{\gamma}\log\frac{1}{\theta^{1}_{b}}\right\\}\leq-\gamma\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})+\frac{\gamma^{2}M^{2}\left(\sum^{T}_{t=1}\bar{a}^{t}\right)}{8K\alpha^{2}_{\star}},$

which implies that

$\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\hat{L}^{e}_{T}\leq\frac{\gamma
M^{2}\left(\sum^{T}_{t=1}\bar{a}^{t}\right)}{8K\alpha^{2}_{\star}}+\frac{1}{\gamma}\log\frac{1}{\theta^{1}_{e}}\leq\frac{\gamma
M^{2}T\bar{a}^{1}}{8K\alpha^{2}_{\star}}+\frac{1}{\gamma}\log\frac{1}{\theta^{1}_{e}},\qquad
e\in[E].$

Since $\theta^{1}_{e}\geq\frac{1}{E^{2}}$, $\log 1/\theta^{1}_{e}\leq 2\log
E$. Let $\gamma=\sqrt{8K\alpha^{2}_{\star}/(TM^{2}\bar{a}^{1})}$ to minimize
the right hand side of the above inequality with $\log 1/\theta^{1}_{e}$
substituted by 1. Then

$\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\hat{L}^{e}_{T}=\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t};\hat{p}^{t})-\sum^{T}_{t=1}\hat{l}_{t}(\hat{p}^{t}_{e};\hat{p}^{t})\leq\frac{M}{\alpha_{\star}}\sqrt{\frac{T\bar{a}^{1}}{8K}}\left(1+2\log
E\right),\qquad e\in[E].$

Equation (31) follows by taking expectation $\mathbb{E}_{S^{1:T}}[\cdot]$ on
both sides.

###  A.5 Proof of Theorem 3

Recall that $T_{b}=2^{b-1}$ and $\hat{p}^{T_{b}}$ is reinitialized as the
uniform distribution. Let

$\text{D-Regret}_{b}=\mathbb{E}_{S^{T_{b}:(T_{b+1}-1)}}\left[\sum^{T_{b+1}-1}_{t=T_{b}}l_{t}(\hat{p}^{t})\right]-\sum^{T_{b+1}-1}_{t=T_{b}}l_{t}(p^{t}_{\star}).$

Similar to the proof of (30) and (31), we have

$\displaystyle\text{D-Regret}_{b}$
$\displaystyle\leq\frac{(2C+1)\sqrt{2}M^{3}\bar{a}^{T_{b}}}{K^{2}\alpha_{\star}^{3}}\sqrt{\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{T_{b}:(T_{b+1}-1)}_{\star}\right)\right](T_{b+1}-T_{b})}$
$\displaystyle\qquad\qquad+\frac{M}{\alpha_{\star}}\sqrt{\frac{(T_{b+1}-T_{b})\bar{a}^{T_{b}}}{8K}}(C+2\log
E_{b})$
$\displaystyle\leq\frac{(2C+1)\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\sqrt{\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]}(\sqrt{2})^{b-1}$
$\displaystyle\qquad\qquad+\frac{M}{\alpha_{\star}}\sqrt{\frac{\bar{a}^{1}}{8K}}(C+2\log
E)(\sqrt{2})^{b-1},$

where $E$ is defined in (13). Since
$B=\left\lfloor\log_{2}(T+1)\right\rfloor$, we have $T_{B}\leq T\leq
T_{B+1}-1$, which implies that $1\leq T-T_{B}+1\leq T_{B+1}-T_{B}=2^{B}$.
Thus, we can similarly obtain

$\mathbb{E}_{S^{T_{B}:T}}\left[\sum^{T}_{t=T_{B}}l_{t}(\hat{p}^{t})\right]-\sum^{T}_{t=T_{B}}l_{t}(p^{t}_{\star})\\\
\leq\frac{(2C+1)\sqrt{2}M^{3}\bar{a}^{1}}{K^{2}\alpha_{\star}^{3}}\sqrt{\left[\log
M+2\log\left(M/\alpha_{\star}\right)\text{TV}\left(p^{1:T}_{\star}\right)\right]}(\sqrt{2})^{B}\\\
+\frac{M}{\alpha_{\star}}\sqrt{\frac{\bar{a}^{1}}{8K}}(C+2\log
E)(\sqrt{2})^{B}.$

The result follows combing the last two displays.

##  Appendix B Useful lemmas

######  Lemma 1.

Suppose that $f$ is a differentiable convex function defined on $\text{dom}f$,
and $\mathcal{X}\subseteq\text{dom}f$ is a closed convex set. Then $x$ is the
minimizer of $f$ on $\mathcal{X}$ if and only if

$\nabla f(x)^{\top}(y-x)\geq 0\quad\text{for all }y\in\mathcal{X}.$

###### Proof.

See Section 4.2.3 of [5]. ∎

######  Lemma 2.

For $q\in\mathcal{P}_{M-1}$ we have $D_{\Phi}(q\,\|\,p^{\rm unif})\leq\log M$,
where $\Phi$ is the unnormalized negative entropy.

###### Proof.

Since $\Phi(q)=\sum^{M}_{m=1}q_{m}(\log q_{m}-1)\leq 0$, $\Phi(p^{\rm
unif})=-\log M$, and

$\langle\nabla\Phi(p^{\rm unif}),q-p^{\rm
unif}\rangle=\sum^{M}_{m=1}(q_{m}-\frac{1}{M})\log\frac{1}{M}=0,$

we have $D_{\Phi}(q\,\|\,p)\leq\log M$. ∎

######  Lemma 3 (Hoeffding’s Inequality).

Let $X$ be a random variable with $a\leq X\leq b$ for $a,b\in\mathbb{R}$. Then
for all $s\in\mathbb{R}$, we have

$\log\mathbb{E}\left[e^{sX}\right]\leq
s\mathbb{E}[X]+\frac{s^{2}(b-a)^{2}}{8}.$

###### Proof.

See Section 2 of [38]. ∎

######  Lemma 4 (Based on Exercise 26.12 of [25]).

Let $\alpha\in[0,1]$, $\mathcal{A}=\mathcal{P}_{M-1}\cap[\alpha/M,1]^{M}$,
$\mathcal{D}=[0,\infty)^{M}$, and $\Phi=\Phi_{e}$ is the unnormalised entropy
on $\mathcal{D}$. For $y\in[0,\infty)^{M}$, let
$x=\arg\min_{v\in\mathcal{A}}D_{\Phi}(v\|y)$. Suppose $y_{1}\leq
y_{2}\leq\dots\leq y_{M}$. Let $m^{\star}$ be the smallest value such that

$y_{m^{\star}}\left(1-\frac{m^{\star}-1}{M}\alpha\right)>\frac{\alpha}{M}\sum^{M}_{m=m^{\star}}y_{m}.$

Then

$x_{m}=\begin{cases}\frac{\alpha}{M}&\text{if }m<m^{\star}\\\
\frac{(1-\frac{m^{\star}-1}{M}\alpha)y_{m}}{\sum^{M}_{n=m^{\star}}y_{n}}&\text{otherwise}.\end{cases}$

###### Proof.

Consider the following constrained optimization problem:

$\displaystyle\min_{u\in[0,\infty)^{M}}\,$
$\displaystyle\sum^{M}_{m=1}u_{m}\log\frac{u_{m}}{y_{m}},$ s.t.
$\displaystyle\sum^{M}_{m=1}u_{m}=1,$ $\displaystyle
u_{m}\geq\frac{\alpha}{M},\qquad m\in[M].$

Since $x$ is the solution to this problem, by the optimality condition, there
exists $\lambda,\nu_{1},\dots,\nu_{M}\in\mathbb{R}$ such that

$\displaystyle\log\frac{x_{m}}{y_{m}}+1-\lambda-\nu_{m}$
$\displaystyle=0,\qquad m\in[M],$ (32) $\displaystyle\sum^{M}_{m=1}x_{m}$
$\displaystyle=1,$ (33) $\displaystyle x_{m}-\frac{\alpha}{M}$
$\displaystyle\geq 0,\qquad m\in[M],$ (34) $\displaystyle\nu_{m}$
$\displaystyle\geq 0,\qquad m\in[M],$ (35)
$\displaystyle\nu_{m}\left(x_{m}-\frac{\alpha}{M}\right)$
$\displaystyle=0,\qquad m\in[M].$ (36)

By (32), we have $x_{m}=y_{m}\exp(-1+\lambda+\nu_{m})$. By (35) and (36), when
$x_{m}=\alpha/M$, we have $x_{m}=y_{m}\exp(-1+\lambda+\nu_{m})\geq
y_{m}\exp(-1+\lambda)$; when $x_{m}>\alpha/M$, we have
$x_{m}=y_{m}\exp(-1+\lambda)$. Assume that
$x_{1}=\dots=x_{m^{\star}-1}=\alpha/M<x_{m^{\star}}\leq\dots\leq x_{M}$. Then

$1=\sum^{M}_{m=1}x_{m}=(m^{\star}-1)\frac{\alpha}{M}+\exp(-1+\lambda)\cdot\sum^{M}_{m=m^{\star}}y_{m},$

which implies that

$\exp(-1+\lambda)=\frac{1-(m^{\star}-1)\frac{\alpha}{M}}{\sum^{M}_{m=m^{\star}}y_{m}}.$
(37)

Thus, we have

$x_{m^{\star}}=y_{m^{\star}}\exp(-1+\lambda)=y_{m^{\star}}\frac{1-(m^{\star}-1)\frac{\alpha}{M}}{\sum^{M}_{m=m^{\star}}y_{m}}>\frac{\alpha}{M},$

which implies that

$y_{m^{\star}}\left(1-\frac{m^{\star}-1}{M}\alpha\right)>\frac{\alpha}{M}\sum^{M}_{m=m^{\star}}y_{m}.$
(38)

To complete the proof, we then only need to show that

$y_{m^{\prime}}\left(1-\frac{m^{\prime}-1}{M}\alpha\right)\leq\frac{\alpha}{M}\sum^{M}_{m=m^{\prime}}y_{m}$
(39)

for all $1\leq m^{\prime}\leq m^{\star}-1$. The result then follows from (38)
and (39). To prove(39), recall that for any $1\leq m^{\prime}\leq
m^{\star}-1$, we have
$\alpha/M=y_{m^{\prime}}\exp(-1+\lambda+\nu_{m^{\prime}})$, and because
$y_{1}\leq\cdots\leq y_{M}$, we have $\nu_{1}\geq\cdots\geq\nu_{m^{\star}-1}$.
This way, we have

$\left(m^{\star}-m^{\prime}\right)\frac{\alpha}{M}=\sum^{m^{\star}-1}_{m=m^{\prime}}y_{m}\exp\left(-1+\lambda+\nu_{m}\right)\leq\exp\left(-1+\lambda+\nu_{m^{\prime}}\right)\sum^{m^{\star}-1}_{m=m^{\prime}}y_{m}.$
(40)

On the other hand, by (37), we have

$1-(m^{\star}-1)\frac{\alpha}{M}=\exp(-1+\lambda)\sum^{M}_{m=m^{\star}}y_{m}\leq\exp\left(-1+\lambda+\nu_{m^{\prime}}\right)\sum^{M}_{m=m^{\star}}y_{m}.$
(41)

Combining (40) and (41), we have

$\displaystyle\frac{1-(m^{\prime}-1)\frac{\alpha}{M}}{\sum^{M}_{m=m^{\prime}}y_{m}}$
$\displaystyle=\frac{1-(m^{\star}-1)\frac{\alpha}{M}+(m^{\star}-m^{\prime})\frac{\alpha}{M}}{\sum^{M}_{m=m^{\prime}}y_{m}}$
$\displaystyle\leq\frac{\exp\left(-1+\lambda+\nu_{m^{\prime}}\right)\left(\sum^{m^{\star}-1}_{m=m^{\prime}}y_{m}+\sum^{M}_{m=m^{\star}}y_{m}\right)}{\sum^{M}_{m=m^{\prime}}y_{m}}$
$\displaystyle=\exp\left(-1+\lambda+\nu_{m^{\prime}}\right)$
$\displaystyle=\frac{\frac{\alpha}{M}}{y_{m^{\prime}}},$

which then implies (39). ∎